{-# LANGUAGE NoMonomorphismRestriction #-}

import Data.Discrimination
import Data.Maybe (fromJust)
import Prelude hiding (lookup)
import System.Environment (getArgs)

-- Straight forward record type
data Machine s i o = Machine
	{ states :: [s]
	, inputs :: [i]
	, output :: s -> o
	, delta :: s -> i -> s
	}

-- Split states on their output. O(n) calls to output
-- and O(n) running time. Can work with any discrimination,
-- so eventually I might want :: Discriminating f => f o -> [s] -> (s -> o) -> [[s]]
-- just really is just the type of disc.
partitionOnOutput :: Grouping o => Machine s i o -> [[s]]
partitionOnOutput m = groupWith (output m) (states m)

-- Split states on their transition, given a symbol
-- The input could be generalised to any function s -> s
-- Currently O(n^2) because of a wrong datastructure
-- Should be O(n) in future
tryRefine :: Eq s => Machine s i o -> i -> [[s]] -> [[s]]
tryRefine m i partition = concat $ map (groupWith d) partition
	where
		-- TODO: define efficient data structure for this
		-- we want (amortized constant time?) State -> Int lookup
		-- where the integer is determined by the block
		-- I tried Map s Int and [(Set s, Int)], both were much slower :(
		d s = lookup (delta m s i) indexedPartition
		indexedPartition = zip partition [0 :: Int ..]
		lookup x [] = undefined
		lookup x ((ss,y):ys) = if elem x ss then y else lookup x ys

-- Refine with all inputs once
refine :: Eq s => Machine s i o -> [[s]] -> [[s]]
refine m partition = foldl1 (.) (map (tryRefine m) (inputs m)) partition

-- Refine until stable. In this case we stop depending on
-- the size of the machine. Ultimately we want to do this
-- by comparing the partitions (just by counting)
-- n*p calls to tryRefine, hence (in the future) O(pn^2)
moore :: (Eq s, Grouping o) => Machine s i o -> [[s]]
moore m = foldl1 (.) (replicate n r) acceptablePartition
	where
		r = refine m
		n = length (states m)
		acceptablePartition = partitionOnOutput m

testMachine :: Int -> Machine Int Bool Bool
testMachine n = Machine [0..n] [False, True] (== 0) (\s i -> if i then s `div` 2 else s `div` 3)

testMachine2 :: Int -> Machine Int Bool Bool
testMachine2 n = Machine (reverse [0..n]) [False, True] (== 0) (\s i -> if i then s `div` 2 else s `div` 3)

hopcroftA :: Int -> Machine Int () Bool
hopcroftA n = Machine [1..n] [()] (==1) (\s _ -> if s > 1 then s-1 else s)

hopcroftA2 :: Int -> Machine Int () Bool
hopcroftA2 n = Machine (reverse [1..n]) [()] (==1) (\s _ -> if s > 1 then s-1 else s)

main = do
	[n, machine] <- getArgs
	case machine of
		"hopcroftA" -> action $ hopcroftA (read n)
		"hopcroftA2" -> action $ hopcroftA2 (read n)
		"testMachine" -> action $ testMachine (read n)
		"testMachine2" -> action $ testMachine2 (read n)
		_ -> putStrLn $ "Unknown machine identifier " ++ machine
	where
		printAll = putStrLn . unlines . map show
		printUneqs = print . length
		action = printUneqs . moore