module MealyRefine where

import Mealy
import Partition (Partition)

import Control.Monad.ST (runST)
import Copar.Algorithm (refine)
import Copar.Functors.Polynomial (Polynomial, PolyF1(..))
import Copar.RefinementInterface (Label, F1)
import Data.Bool (bool)
import Data.CoalgebraEncoding (Encoding(..))
import Data.List.Ordered (nubSort)
import Data.Map.Strict qualified as Map
import Data.Proxy (Proxy(..))
import Data.Vector qualified
import Data.Vector.Unboxed qualified

project :: Ord u => (o -> u) -> MealyMachine s i o -> MealyMachine s i u
project f MealyMachine{..} = MealyMachine
  { outputs = nubSort $ fmap f outputs -- inefficient
  , behaviour = \s i -> case behaviour s i of
      (out, s2) -> (f out, s2)
  , ..
  }

projectToBit :: Ord o => o -> MealyMachine s i o -> MealyMachine s i Bool
projectToBit o = project (o ==)

projectToComponent :: Ord o => (o -> Bool) -> MealyMachine s i o -> MealyMachine s i (Maybe o)
projectToComponent oPred = project oMaybe
  where oMaybe o
          | oPred o = Just o
          | otherwise = Nothing

-- We will project to all (relevant) outputs. Since a lot of data is shared
-- among those mealy machines, I do this in one function. The static parts
-- are converted only once. Only "eStructure" (the state-labels) are different
-- for each projection.
allProjections :: (Ord s, Ord i, Eq o) => MealyMachine s i o -> [o] -> ([Encoding (Label Polynomial) (F1 Polynomial)], Map.Map s Int)
allProjections MealyMachine{..} outs = (fmap mkEncoding outs, state2idx)
  where
    numStates = length states
    numInputs = length inputs
    idx2state = Map.fromList $ zip [0..] states
    state2idx = Map.fromList $ zip states [0..]
    idx2input = Map.fromList $ zip [0..] inputs
    stateInputIndex n = n `divMod` numInputs -- inverse of \(s, i) -> s * numInputs + i
    edgesFrom = Data.Vector.Unboxed.generate (numStates * numInputs) (fst . stateInputIndex)
    edgesLabel = Data.Vector.generate (numStates * numInputs) (snd . stateInputIndex)
    edgesTo = Data.Vector.Unboxed.generate (numStates * numInputs) ((state2idx Map.!) . snd . (\(s, i) -> behaviour (idx2state Map.! s) (idx2input Map.! i)) . stateInputIndex)
    bool2Int = bool 0 1
    structure o = Data.Vector.generate numStates
      (\s -> PolyF1
        { polyF1Summand = 0 -- There is only one summand
        , polyF1Variables = numInputs -- This many transitions per state
        , polyF1Constants = Data.Vector.Unboxed.generate numInputs
          (\i -> bool2Int . (o==) . fst $ (behaviour (idx2state Map.! s) (idx2input Map.! i)))
        }
      )
    mkEncoding o = Encoding
      { eStructure = structure o
      , eInitState = Nothing -- not needed
      , eEdgesFrom = edgesFrom
      , eEdgesLabel = edgesLabel
      , eEdgesTo = edgesTo
      }

-- Refine a encoded mealy machine
refineMealy :: Encoding (Label Polynomial) (F1 Polynomial) -> Partition
refineMealy machine = runST $ refine (Proxy @Polynomial) machine True

mealyMachineToEncoding :: (Ord s, Ord i, Ord o) => MealyMachine s i o -> Encoding (Label Polynomial) (F1 Polynomial)
mealyMachineToEncoding MealyMachine{..} =
  let numStates = length states
      numInputs = length inputs
      idx2state = Map.fromList $ zip [0..] states
      idx2input = Map.fromList $ zip [0..] inputs
      out2idx = Map.fromList $ zip outputs [0..]
      eStructure = Data.Vector.generate numStates
        (\s -> PolyF1
          { polyF1Summand = 0 -- There is only one summand
          , polyF1Variables = numInputs -- This many transitions per state
          , polyF1Constants = Data.Vector.Unboxed.generate numInputs
            (\i -> out2idx Map.! (fst (behaviour (idx2state Map.! s) (idx2input Map.! i))))
          }
        )
      eInitState = Nothing
      state2idx = Map.fromList $ zip states [0..]
      stateInputIndex n = n `divMod` numInputs
      -- stateInputPair (s, i) = s * numInputs + i
      eEdgesFrom = Data.Vector.Unboxed.generate (numStates * numInputs) (fst . stateInputIndex)
      eEdgesLabel = Data.Vector.generate (numStates * numInputs) (snd . stateInputIndex)
      eEdgesTo = Data.Vector.Unboxed.generate (numStates * numInputs) ((state2idx Map.!) . snd . (\(s, i) -> behaviour (idx2state Map.! s) (idx2input Map.! i)) . stateInputIndex)
  in Encoding { .. }