{-# LANGUAGE PartialTypeSignatures #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures #-}

module Main where

import SplittingTree

import Control.Monad.Random.Strict
import Control.Monad.Trans.State (execStateT)
import Data.Coerce
import Data.List (sortOn)
import Data.Map qualified as Map
import Data.Set qualified as Set
import System.Environment
import System.Random

genTransitions :: _ => Int -> [Char] -> [Char] -> RandT _ _ _
genTransitions size inputs outputs = do
  let
    states = [1 .. size]
    numOutputs = length outputs
    randomOutput = (outputs !!) <$> liftRand (uniformR (0, numOutputs - 1))
    randomTarget = (states !!) <$> liftRand (uniformR (0, size - 1))
    randomTransition s i = (\o t -> (s, i, o, t)) <$> randomOutput <*> randomTarget
  t <- sequence $ randomTransition <$> states <*> inputs
  return (states, t)

-- numC <= 8
genComposition :: _ => Int -> Int -> [Char] -> RandT _ _ _
genComposition size numC inputs = do
  let
    components = [1 .. numC]
    outputSets = [[], "xy", "zw", "uv", "kl", "mn", "op", "qr", "st"]

  allTransitions <- traverse (\c -> genTransitions size inputs (outputSets !! c)) components

  let
    productState = fmap (Map.fromList . zip components) (mapM fst allTransitions)
    allStates = (,) <$> components <*> productState
    compMap = Map.fromList $ zip components ((Map.fromList . fmap (\(s, i, o, t) -> ((s, i), (o, t)))) . snd <$> allTransitions)
    norm c
      | c <= 0 = c + numC
      | c > numC = c - numC
      | otherwise = c
    transition (c, cs) 'L' = (norm (c - 1), cs)
    transition (c, cs) 'R' = (norm (c + 1), cs)
    transition (c, cs) x = (c, Map.adjust (\s -> snd (compMap Map.! c Map.! (s, x))) c cs)
    output _ 'L' = 'L'
    output _ 'R' = 'R'
    output (c, cs) x = fst (compMap Map.! c Map.! (cs Map.! c, x))

  -- initial states, inputs, transition function, outputs
  return (head allStates, 'L' : 'R' : inputs, transition, output)

reachability :: _ => s -> [i] -> (s -> i -> s) -> Map.Map s Int
reachability initialState inputs transitions = go 0 Map.empty [initialState]
 where
  go _ visited [] = visited
  go n visited (s : rest) =
    let newVis = Map.insert s n visited
        newStates = [t | i <- inputs, let t = transitions s i, t `Map.notMember` newVis]
     in go (n + 1) newVis (rest ++ newStates)

main :: IO ()
main = do
  [nStr, cStr] <- getArgs
  let
    n = read nStr
    c = read cStr

  -- create random composition
  (init0, inputs, trans0, outpf0) <- evalRandIO (genComposition n c ['a', 'b'])

  -- reachable part only
  let
    reachableMap = reachability init0 inputs trans0
    inverseMap = Map.fromList . fmap (\(a, b) -> (b, a)) . Map.toList $ reachableMap

    states = Map.elems reachableMap
    init = reachableMap Map.! init0
    trans s i = reachableMap Map.! trans0 (inverseMap Map.! s) i
    outpf s = outpf0 (inverseMap Map.! s)

  -- minimize
  let
    outputFuns = [(i, fun) | i <- inputs, let fun s = outpf s i]
    reverseTransitionMaps i = Map.fromListWith (++) [(t, [s]) | s <- states, let t = trans s i]
    reverseFuns = [(i, fun) | i <- inputs, let m = reverseTransitionMaps i, let fun s = Map.findWithDefault [] s m]

  PRState{..} <- execStateT (refine (const (pure ())) outputFuns reverseFuns) (initialPRState states)

  -- print
  let
    toBlock s = getPartition partition Map.! s
    allTransitions = [(toBlock s, i, o, toBlock t) | s <- states, i <- inputs, let o = outpf s i, let t = trans s i]
    uniqueTransitions = sortOn (\(s, _, _, _) -> s /= toBlock init) . Set.toList . Set.fromList $ allTransitions
    showLabel i o = "[label=\"" <> [i] <> "/" <> [o] <> "\"]"
    showTransition s i o t = "s" <> show (coerce s :: Int) <> " -> " <> "s" <> show (coerce t :: Int) <> " " <> showLabel i o

  putStrLn "digraph g {"
  mapM_ (\(s, i, o, t) -> putStrLn ("  " <> showTransition s i o t)) uniqueTransitions
  putStrLn "}"