{-# LANGUAGE LambdaCase #-}

import Data.Fix
import Data.List
import Data.Maybe


-- Lib
type Coalgebra f x = x -> f x

behaviour :: (Functor f) => Coalgebra f x -> x -> Fix f
behaviour = ana -- = (~> Fix) = Fix (fmap (behaviour phi) (phi x))

data MooreF i o x = MooreF o (i -> x)
type Moore i o x = Coalgebra (MooreF i o) x

getO :: MooreF i o x -> o
getO (MooreF o _) = o

apply :: MooreF i o x -> i -> x
apply (MooreF _ f) = f

instance Functor (MooreF i o) where
	fmap f (MooreF o phi) = MooreF o (f . phi)


type DFA i x = Moore i Bool x
type Language i = Fix (MooreF i Bool)
type Word i = [i]

singleton :: (Eq i) => DFA i (Maybe (Word i))
singleton Nothing       = MooreF False (const Nothing)
singleton (Just [])     = MooreF True  (const Nothing)
singleton (Just (y:ys)) = MooreF False (\i -> if i == y then Just ys else Nothing)

isMember :: Word i -> Language i -> Bool
isMember [] l = getO . unFix $ l
isMember (x:xs) l = isMember xs $ apply (unFix l) x

test :: x -> [Word i] -> DFA i x -> ([Word i], [Word i])
test s ws m = partition (\w -> isMember w (behaviour m s)) ws

enumerate :: [x] -> [[x]] -> [[x]]
enumerate xs [] = []
enumerate xs (a:as) = map (:a) xs ++ enumerate xs as

enumerateAll :: [x] -> [[x]]
enumerateAll xs = concat $ iterate (enumerate xs) [[]]

type Color    = Int
type NewColor = Int

algebra :: [i] -> (o -> Int) -> [Color] -> MooreF i o Color -> NewColor
algebra is mp cs (MooreF o phi) = (mp o) * totalWidth + foldl (\acc d -> width * acc + d) 0 digits
	where
		width = length is
		totalWidth = length cs ^ width
		digits = map phi is

allColours :: [x] -> (x -> Int) -> [Int]
allColours xs p = map head . group . sort . map p $ xs

-- maps [x1 ... xn] to [0 ... n]
-- Can be implemented a lot better ;D
reduce :: [Int] -> (Int -> Int)
reduce xs x = fromJust $ lookup x indexed
	where
		sorted = map head . group $ sort xs
		indexed = zip sorted [0..]


base :: [x] -> [i] -> (o -> Int) -> Moore i o x -> (x -> Int)
base _ _ _ _ = \x -> 0

step :: [x] -> [i] -> (o -> Int) -> Moore i o x -> (x -> Int) -> (x -> Int)
step states inputs outputs machine p = reducedP
	where
		colors = allColours states p
		newP = algebra inputs outputs colors . fmap p . machine
		reducedP = reduce (allColours states newP) . newP

-- minimize :: [x] -> [i] -> (o -> Int) -> Moore i o x -> (x -> Int)
minimize' states inputs outputs machine = stabilize $ iterate s b
	where
		b = base states inputs outputs machine
		s = step states inputs outputs machine
		size p = length (allColours states p)
		conv l = zip l (tail l)
		stabilize = snd . last . takeWhile (\(a, b) -> size a < size b) . conv

-- Example
data States = A | B | C | D | E | F
data Input = Zero | One
	deriving Show

machine :: DFA Input States
machine A = MooreF False (\case Zero -> B; One -> C)
machine B = MooreF False (\case Zero -> A; One -> D)
machine C = MooreF True  (\case Zero -> E; One -> F)
machine D = MooreF True  (\case Zero -> E; One -> F)
machine E = MooreF True  (\case Zero -> E; One -> F)
machine F = MooreF False (\case Zero -> F; One -> F)

states = [A, B, C, D, E, F]
inputs = [Zero, One]
outputs = (\o -> if o then 1 else 0)

b = base states inputs outputs machine
s = step states inputs outputs machine

main :: IO ()
main = print . limit 10 . onBoth evaluate . limit 200000 $ test startingState tests machine
	where
		startingState = A
		tests = enumerateAll [Zero, One]
		onBoth f (a, b) = (f a, f b)
		limit n = onBoth (take n)
		evaluate xs = seq (sum (map length xs)) xs