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. Ik heb geprobeerd of Data.Sequence
-- sneller is, maar dat lijkt niet zo te zijn. Dus het blijft [i] voor gemak.
type Word i = [i]
cons = (:)
snoc w a = w <> [a]

-- 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. Dit zorgt ook voor minder werk
-- als we opnieuw closedness moeten uitrekenen.
closednessDefects :: (Ord i, Ord o) => LStarState i o -> [Word i]
closednessDefects table@LStarState{..} = defects
  where
    extendedRowIndices = [ra | r <- Set.toList rowIndices, a <- alphabet, let ra = r `snoc` a, ra `Set.notMember` rowIndices]
    upperRows = Set.map (row table) rowIndices
    defects = [r | r <- 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 `snoc` a)) (row table (r2 `snoc` a))
    differenceOfMaps = MapMerge.merge MapMerge.dropMissing MapMerge.dropMissing (MapMerge.zipWithMaybeMatched (\_ x y -> if x == y then Nothing else Just ()))

-- 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 `snoc` a)) | r <- rowIndicesLs, a <- alphabet]
    outputs = Map.fromList [((row2Int r, a), o) | r <- rowIndicesLs, a <- alphabet, let o = content Map.! (r `snoc` a)]
    initialState = row2Int mempty


-----------------------------------------
-- 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 = initialiseWithA alphabet [mempty] (map pure alphabet)

-- 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
    colIndices = Set.fromList colIdcs
    queries = [p <> m <> s | p <- rowIdcs, m <- []:fmap pure alphabet, s <- colIdcs]
    contentA = Map.traverseWithKey (\ w _ -> mq w) . Map.fromList . map (, ()) $ queries

-- 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 <> Set.fromList newRowIndices }) <$> contentA
  where
    queries = [w | p <- newRowIndices, m <- []:fmap pure alphabet, s <- Set.toList colIndices, let w = p <> m <> s, w `Map.notMember` content]
    contentA = Map.traverseWithKey (\ w _ -> mq w) . Map.fromList . map (, ()) $ queries

-- 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 <> Set.fromList newColIndices }) <$> contentA
  where
    queries = [w | p <- Set.toList rowIndices, m <- []:fmap pure alphabet, s <- newColIndices, let w = p <> m <> s, w `Map.notMember` content]
    contentA = Map.traverseWithKey (\ w _ -> mq w) . Map.fromList . map (, ()) $ queries


------------------
-- 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 `cons` 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]]
  -- Add all suffixes to the table. This ensures that the next hypothesis is bigger
  newSuffixes = filter (not . null) . tails . drop 1 $ shortestSuffix
  -- Alternatively: Only add one suffix. This might result in the same hypothesis (and CE)
  -- But in the end it will progress, possibly with fewer queries.
  --newSuffixes = take 1 . filter (\s -> s `Set.notMember` colIndices) . 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)