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

module RandomGen where

import Data.Partition (Block (..))
import SplittingTree

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

data RandomGenOptions = RandomGenOptions
  { numStates :: Int
  , numComponents :: Int
  }
  deriving Show

randomGenOptionsParser :: Parser RandomGenOptions
randomGenOptionsParser =
  RandomGenOptions
    <$> option auto (long "states" <> short 'n' <> help "Number of states per component (max)" <> metavar "NUM" <> showDefault <> value 10)
    <*> option auto (long "components" <> short 'c' <> help "Number of components" <> metavar "COMP" <> showDefault <> value 2)
      <**> helper

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)

mainRandomGen :: RandomGenOptions -> IO ()
mainRandomGen RandomGenOptions{..} = do
  let
    n = numStates
    c = numComponents

  -- 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 to stdout
  let
    toBlock s = getBarePartition 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 (Block s) i o (Block t) = "s" <> show s <> " -> " <> "s" <> show t <> " " <> showLabel i o

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