{-# OPTIONS -fglasgow-exts #-}

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

import Control.Monad

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

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)

data Decoder i o = Decoder { doDecode :: [i] -> ([i], o) }
type Encoder o i = i -> [o]
type DR a = Decoder Trit a

instance Monad (Decoder i) where
	f >>= g = Decoder $ \x -> let (x0, r0) = doDecode f x in doDecode (g r0) x0
	f >> g = Decoder $ \x -> let (x0, r0) = doDecode f x in doDecode g x0
	return a = Decoder $ \x -> (x, a)
	fail = error

class Encode a where
	sdecode :: DR  a
	sencode :: Encoder Trit a

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

getTrit :: Decoder Trit Trit
getTrit = Decoder $ \xs -> case xs of
	[] -> error "Unexpected end of stream"
	(i:xs) -> (xs, i)

instance Encode Bool where
	sdecode = do
		i <- sdecode :: DR Int
		case i of
			0 -> return False
			1 -> return True
			_ -> fail "Expected 0 or 1"
	sencode True = [1]
	sencode False = [0]

instance Encode Trit where
	sdecode = getTrit
	sencode t = [t]

instance Encode a => Encode [a] where
	sdecode = do
		i <- getTrit
		case i of
			T0 -> return []
			T1 -> sdecode >>= (\x -> return [x])
			T2 -> do
				i <- getTrit
				case i of
					T2 -> do
						n <- sdecode :: DR Int
						replicateM (n+2) sdecode
					_ -> fail "Unexpected end of stream"
	sencode [] = [T0]
	sencode [x] = T1:sencode x
	sencode xs = (sencode $ length xs - 2) ++ concat (map sencode xs)

instance Encode Int where
	sdecode = do
		n <- sdecode :: DR Integer
		return $ fromIntegral n
	sencode i = sencode (fromIntegral i :: Integer)

instance Encode Integer where
	sdecode = do
		digits <- sdecode :: DR [Trit]
		return $ decodeTern digits + (3^(length digits - 1) `div` 2)
	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 = do
		a <- sdecode
		z <- sdecode
		return (a, z)
	sencode (a, z) = sencode a ++ sencode z

instance (Encode (a, b), Encode z) => Encode (a, b, z) where
	sdecode = do
		(a, b) <- sdecode
		z <- sdecode
		return (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)]