{-# OPTIONS -fglasgow-exts #-}

module Encoding (Trit(..), Encode(..), Stream(..), encode, decode, decodeStrict, toTrit, fromTrit) where

data Trit = T0 | T1 | T2 deriving (Enum, Eq)

instance Show Trit where
  show = show . fromEnum

instance Read Trit where
  readsPrec _ "0" = [(T0, "")]
  readsPrec _ "1" = [(T1, "")]
  readsPrec _ "2" = [(T2, "")]

instance Num Trit where
  (+) = liftToTrit (+)
  (*) = liftToTrit (*)
  (-) = liftToTrit (-)
  abs = id
  signum = const 0
  fromInteger 0 = T0
  fromInteger 1 = T1
  fromInteger 2 = T2

liftToTrit :: (Int -> Int -> Int) -> Trit -> Trit -> Trit
liftToTrit f x y  = toEnum $ (fromEnum x `f` fromEnum y) `mod` 3

toTrit :: Int -> Trit
toTrit = toEnum
fromTrit :: Integral i => Trit -> i
fromTrit = fromIntegral . fromEnum

type DR a = ([Trit], a)

class Stream a where
	streamread :: a -> [Trit]
	streamwrite :: [Trit] -> a

instance Stream [Trit] where
	streamread = id
	streamwrite = id

instance Stream [Int] where
	streamread = map toEnum
	streamwrite = map fromEnum

instance Stream String where
	streamread = map (toEnum . read . (:[]))
	streamwrite = concat . map (show . fromEnum)

class Encode a where
	sdecode :: [Trit] -> ([Trit], a)
	sencode :: a -> [Trit]

encode :: (Stream s, Encode a) => a -> s
encode = streamwrite . sencode
decode  :: (Stream s, Encode a) => s -> a
decode = snd . sdecode . streamread
decodeStrict :: (Stream s, Encode a) => s -> a
decodeStrict x = case (sdecode $ streamread x) of
	([], a) -> a
	(_, _) -> error "Expected end of input"

instance Encode Bool where
	sdecode xs = case sdecode xs :: DR Int of
		(x0, 0) -> (x0, False)
		(x0, 1) -> (x0, True)
		(x0, i) -> error "Expected 0 or 1"
	sencode True = [1]
	sencode False = [0]

instance Encode Trit where
	sdecode [] = error "Unexpected end of stream"
	sdecode (i:xs) = (xs, i)
	sencode t = [t]

instance Encode a => Encode [a] where
	sdecode [] = error "Unexpected end of stream"
	sdecode (T0:xs) = (xs, [])
	sdecode (T1:xs) = let (x0, a) = sdecode xs in (x0, [a])
	sdecode (T2:T2:xs) = let (x0, a) = sdecode xs :: DR Int in get [] (a+2) x0  where
		get r 0 x = (x, r)
		get r k x | k > 0 = let (xn, e) = sdecode x in get (r ++ e) (k-1) xn
	sencode [] = [T0]
	sencode [x] = T1:sencode x
	sencode xs = (sencode $ length xs - 2) ++ concat (map sencode xs)

instance Encode Int where
	sdecode xs = let (x0, i) = sdecode xs :: DR Integer in (x0, fromIntegral i)
	sencode i = sencode (fromIntegral i :: Integer)

instance Encode Integer where
	sdecode xs = let (x0, digits) = sdecode xs :: DR [Trit] in
		let sum = decodeTern digits in (x0, (3^(length digits - 1) `div` 2) + sum)
	sencode n
		| n < 0 = error "Can't encode negative numbers"
		| n == 0 = [0]
		| n > 0 = let len = log3 n in sencode (encodeTern len (n - base len)) where
			base len = (3^len - 1) `div` 2
			log3 n = (head $ filter (\len -> n < base len) [0..]) - 1

instance (Encode a, Encode z) => Encode (a, z) where
	sdecode xs = let (x0, a) = sdecode xs in let (x1, z) = sdecode x0 in (x1, (a, z))
	sencode (a, z) = sencode a ++ sencode z

instance (Encode (a, b), Encode z) => Encode (a, b, z) where
	sdecode xs = let (x0, (a, b)) = sdecode xs in let (x1, z) = sdecode x0 in (x1, (a, b, z))
	sencode (a, b, z) = sencode (a, b) ++ sencode z

decodeTern :: Integral i => [Trit] -> i
decodeTern digits = dec 0 digits where
	dec s [] = s
	dec s (i:xs) = dec (3*s + (fromTrit i)) xs

encodeTern :: Integral i => Int -> i -> [Trit]
encodeTern 0 0 = []
encodeTern k x = (encodeTern (k-1) (x `div` 3)) ++ [toTrit $ fromIntegral (x `mod` 3)]

-- cartanks :: Car -> Nat
-- cartanks = (1+) . maximum . concat . map (\(upper, _, lower) -> upper ++ lower)