module LStar where

-- LStar voor Mealy machines. Input zijn woorden van type i, en output zijn
-- elementen van type o. (We gebruiken alleen de laatste output.)

import Control.Monad.Trans.Class
import Control.Monad.Trans.State.Strict

import Data.Foldable (minimumBy)
import Data.Functor.Identity
import Data.List (tails, stripPrefix)
import Data.Map.Merge.Strict qualified as MapMerge
import Data.Map.Strict qualified as Map
import Data.Ord (comparing)
import Data.Set qualified as Set
import Prelude hiding (Word)

-----------------------------------
-- Datastructuren en basis functies

-- Een woord is niets anders dan een lijst.
-- TODO: Een keertje uitvogelen of Data.Sequence efficienter is of niet.
type Word i = [i]

-- Membership queries. Voor Mealy machines levert dat een o op. De parameter
-- f is voor een monad. Dat is nodig als de membership queries echt moeten
-- praten met een SUL, of als je bijv. een tellertje wilt bijhouden.
type MQ f i o = Word i -> f o

-- Invariant: dom(content) = `rows × columns` ∪ `rows × alph × columns`
-- TODO: misschien hoeven we geen Set datastructuur te gebruiken, maar
-- volstaan lijsten. Het zou ook met een Trie kunnen.
data LStarState i o = LStarState
  { content :: Map.Map (Word i) o
  , rowIndices :: Set.Set (Word i)
  , colIndices :: Set.Set (Word i)
  , alphabet :: [i]
  }
  deriving (Eq, Ord, Show, Read)

-- We maken onderscheid in "rij-index" en "rij". De index is het woord wat bij
-- de rij hoort. De rij zelf is gegeven door de inhoud van de rij, dus een
-- functie van kolom naar o. Dit doen we met een Map data structuur.
type Row i o = Map.Map (Word i) o

-- Kent aan een rij-index de rij (content) toe.
row :: Ord i => LStarState i o -> Word i -> Row i o
row LStarState{..} prefix = Map.fromSet (\suffix -> content Map.! (prefix <> suffix)) colIndices


-----------------
-- LStar functies

-- Geeft alle rijen waarvoor de tabel nog niet gesloten is.
-- TODO: toekenning van "lage rij" -> "hoge rij" ook uitrekenen, zodat we
-- later meteen een automaat kunnen bouwen.
closednessDefects :: (Ord i, Ord o) => LStarState i o -> [Word i]
closednessDefects table@LStarState{..} = defects
  where
    alphabetSingletons = Set.fromList $ map pure alphabet
    allExtendedRowIndices = Set.map (uncurry (<>)) $ Set.cartesianProduct rowIndices alphabetSingletons
    extendedRowIndices = allExtendedRowIndices `Set.difference` rowIndices
    upperRows = Set.map (row table) rowIndices
    defects = [r | r <- Set.toList extendedRowIndices, row table r `Set.notMember` upperRows]

-- Geeft alle paren van symbool en kolom (a, s) die toegevoegd moeten worden
-- om de tabel consistent te maken. Merk op: het paar (a, s) zou vaker voor
-- kunnen komen als er meerdere equivalente rijen zijn.
inconsistencies :: (Ord i, Eq o) => LStarState i o -> [(i, Word i)]
inconsistencies table@LStarState{..} = defects
  where
    equivalentPairs = [(r1, r2) | r1 <- Set.toList rowIndices, r2 <- Set.toList rowIndices, r1 < r2, row table r1 == row table r2]
    defects = [(a, s) | (r1, r2) <- equivalentPairs, a <- alphabet, let d = differenceOfRows r1 r2 a, s <- Map.keys d]
    differenceOfRows r1 r2 a = differenceOfMaps (row table (r1 <> [a])) (row table (r2 ++ [a]))
    differenceOfMaps m1 m2 = MapMerge.merge MapMerge.dropMissing MapMerge.dropMissing (MapMerge.zipWithMaybeMatched (\_ x y -> if x == y then Nothing else Just ())) m1 m2

-- Preconditie: tabel is gesloten en consistent
-- TODO: misschien checken of de automaat uniek is? De constructie van
-- transitions kiest de eerst-mogelijke successor-state, en misschien zijn
-- er meer.
createHypothesis :: (Ord i, Ord o) => LStarState i o -> (Int, Int, Map.Map (Int, i) Int, Map.Map (Int, i) o)
createHypothesis table@LStarState{..} = (initialState, Map.size row2IntMap, transitions, outputs)
  where
    rowIndicesLs = Set.toList rowIndices
    upperRows = map (row table) $ rowIndicesLs
    row2IntMap = Map.fromList $ zip upperRows [0..]
    row2Int = (Map.!) row2IntMap . row table
    transitions = Map.fromList [((row2Int r, a), row2Int (r <> [a])) | r <- rowIndicesLs, a <- alphabet]
    outputs = Map.fromList [((row2Int r, a), o) | r <- rowIndicesLs, a <- alphabet, let o = content Map.! (r <> [a])]
    initialState = row2Int []


-----------------------------------------
-- Initialisatie en observaties toevoegen

-- Een lege tabel heeft altijd een "epsilon-rij" en voor elk symbool een kolom.
-- (Omdat het Mealy machines zijn.)
initialiseA :: (Applicative f, Ord i) => [i] -> MQ f i o -> f (LStarState i o)
initialiseA alphabet mq = initialiseWithA alphabet [[]] (map pure alphabet) mq

-- We kunnen ook een tabel maken met voorgedefinieerde rijen en kolommen.
initialiseWithA :: (Applicative f, Ord i) => [i] -> [Word i] -> [Word i] -> MQ f i o -> f (LStarState i o)
initialiseWithA alphabet rowIdcs colIdcs mq = (\content -> LStarState{..}) <$> contentA
  where
    rowIndices = Set.fromList rowIdcs
    rowIdcsExt = rowIndices <> Set.fromList [r <> [a] | r <- rowIdcs, a <- alphabet]
    colIndices = Set.fromList colIdcs
    domain = Map.fromSet (const ()) . Set.map (uncurry (<>)) $ Set.cartesianProduct rowIdcsExt colIndices
    contentA = Map.traverseWithKey (\w _ -> mq w) domain

-- preconditie: newRowIndices is disjunct van de huidige rowIndices en de
-- vereniging is prefix-gesloten. (Wordt niet gechecked.)
addRowsA :: (Applicative f, Ord i) => [Word i] -> MQ f i o -> LStarState i o -> f (LStarState i o)
addRowsA newRowIndices mq table@LStarState{..} = (\newContent -> table
  { content = content <> newContent
  , rowIndices = rowIndices <> newRowIndicesSet }) <$> contentA
  where
    newRowIndicesExt = Set.fromList [r <> [a] | r <- newRowIndices, a <- alphabet]
    newRowIndicesSet = Set.fromList newRowIndices
    newRowIndices2 = (newRowIndicesSet <> newRowIndicesExt) `Set.difference` rowIndices
    domain = Map.fromSet (const ()) . Set.map (uncurry (<>)) $ Set.cartesianProduct newRowIndices2 colIndices
    contentA = Map.traverseWithKey (\w _ -> mq w) domain

-- preconditie: zie addRows (?)
addColumnsA :: (Applicative f, Ord i) => [Word i] -> MQ f i o -> LStarState i o -> f (LStarState i o)
addColumnsA newColIndices mq table@LStarState{..} = (\newContent -> table
  { content = content <> newContent
  , colIndices = colIndices <> newColIndicesSet }) <$> contentA
  where
    newColIndicesExt = Set.fromList [[a] <> c | c <- newColIndices, a <- alphabet]
    newColIndicesSet = Set.fromList newColIndices
    newColIndices2 = (newColIndicesSet <> newColIndicesExt) `Set.difference` colIndices
    domain = Map.fromSet (const ()) . Set.map (uncurry (<>)) $ Set.cartesianProduct rowIndices newColIndices2
    contentA = Map.traverseWithKey (\w _ -> mq w) domain


------------------
-- LStar algoritme

-- geeft True als de tabel is gewijzigd.
closeA :: (Monad f, Ord i, Ord o) => MQ f i o -> LStarState i o -> f (LStarState i o, Bool)
closeA mq = loop False where
  loop acc table = case closednessDefects table of
    [] -> pure (table, acc)
    w:_ -> addRowsA [w] mq table >>= loop True

-- geeft True als de tabel is gewijzigd.
makeConsistentA :: (Monad f, Ord i, Eq o) => MQ f i o -> LStarState i o -> f (LStarState i o, Bool)
makeConsistentA mq = loop False where
  loop acc table = case inconsistencies table of
    [] -> pure (table, acc)
    (a, w):_ -> addColumnsA [a:w] mq table >>= loop True

makeClosedAndConsistentA :: (Monad f, Ord i, Ord o) => MQ f i o -> LStarState i o -> f (LStarState i o, Bool)
makeClosedAndConsistentA mq = loop False where
  loop acc table1 = do
    (table2, b2) <- closeA mq table1
    (table3, b3) <- makeConsistentA mq table2
    if b3
    then loop True table3
    else pure (table3, acc || b2)

-- We voegen tegenvoorbeelden toe als kolommen. Dat is vaak het meest
-- efficient. We houden de kolommen wel suffix-gesloten, dat is iets
-- eenvoudiger.
processCounterexampleA :: (Applicative f, Ord i) => Word i -> MQ f i o -> LStarState i o -> f (LStarState i o)
processCounterexampleA ce mq table@LStarState{..} = addColumnsA newSuffixes mq table where
  shortestSuffix = minimumBy (comparing length) [suf | r <- Set.toList rowIndices, Just suf <- [stripPrefix r ce]]
  newSuffixes = filter (not . null) . tails . drop 1 $ shortestSuffix


-------------
-- MQ wrapper

-- Bijhouden hoeveel membership queries er zijn gesteld
countingMQ :: Monad f => MQ f i o -> MQ (StateT Int f) i o
countingMQ mq w = modify' succ >> lift (mq w)


-----------------------
-- Minder generieke API

-- Als je geen monad nodig hebt, kun je een simpele membership query omzetten
-- naar het MQ type.
simpleMQ :: (Word i -> o) -> MQ Identity i o
simpleMQ = (.) Identity

initialise :: Ord i => [i] -> (Word i -> o) -> LStarState i o
initialise alphabet = runIdentity . initialiseA alphabet . simpleMQ

initialiseWith :: Ord i => [i] -> [Word i] -> [Word i] -> (Word i -> o) -> LStarState i o
initialiseWith alphabet rowIdcs colIdcs = runIdentity . initialiseWithA alphabet rowIdcs colIdcs . simpleMQ

addRows, addColumns :: Ord i => [Word i] -> (Word i -> o) -> LStarState i o -> LStarState i o
addRows newRowIndices mq = runIdentity . addRowsA newRowIndices (simpleMQ mq)
addColumns newColIndices mq = runIdentity . addColumnsA newColIndices (simpleMQ mq)

close, makeConsistent, makeClosedAndConsistent :: (Ord i, Ord o) => (Word i -> o) -> LStarState i o -> (LStarState i o, Bool)
close mq = runIdentity . closeA (simpleMQ mq)
makeConsistent mq = runIdentity . makeConsistentA (simpleMQ mq)
makeClosedAndConsistent mq = runIdentity . makeClosedAndConsistentA (simpleMQ mq)

processCounterexample :: Ord i => Word i -> (Word i -> o) -> LStarState i o -> LStarState i o
processCounterexample ce mq = runIdentity . processCounterexampleA ce (simpleMQ mq)