{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures #-}

import Automata (Word)
import ExampleAutomata
import IO
import Quotient
import OrbitList
import EquivariantMap (EquivariantMap(..), (!))
import qualified EquivariantMap as Map
import qualified EquivariantSet as Set
import Nominal (Nominal, Orbit, Trivially(..))
import Permutable

import Data.List (tails)
import Data.Maybe (catMaybes)
import Control.Monad.State
import System.IO (hFlush, stdout)
import Prelude hiding (filter, null, elem, lookup, product, Word, map, take, init)

newtype PermEquivariantMap k v = PEqMap { unPEqMap :: EquivariantMap k v }
  deriving Nominal via Trivially (EquivariantMap k v)

-- Defined by the join-semilattice structure of EquivariantMap, left biased.
deriving instance Ord (Orbit k) => Monoid (PermEquivariantMap k v)
deriving instance Ord (Orbit k) => Semigroup (PermEquivariantMap k v)

lookupP :: (Permutable k, Nominal k, Nominal v, _) => k -> PermEquivariantMap k v -> Maybe v
lookupP x (PEqMap m) = case catMaybes [Map.lookup (act px) m | px <- allPermuted x] of
  [] -> Nothing
  (v:_) -> Just v -- take first hit, maybe this is wrong? I guess for v ~ Bool it's fine?

insertP :: (Nominal k, Nominal v, _) => k -> v -> PermEquivariantMap k v -> PermEquivariantMap k v
insertP k v = PEqMap . Map.insert k v . unPEqMap

(!~) :: (Permutable k, Nominal k, Nominal v, _) => PermEquivariantMap k v -> k -> v
(!~) m k = case lookupP k m of
  Just v -> v
  Nothing -> error "Key not found (in PermEquivariantMap)"

-- We use Lists, as they provide a bit more laziness
type Rows a    = OrbitList (Word a)
type Columns a = OrbitList (Word a)
type Table a   = PermEquivariantMap (Word a) Bool


-- Utility functions
exists f = not . null . filter f
forAll f = null . filter (not . f)
ext p a = p <> [a]

equalRows :: _ => Word a -> Word a -> Columns a -> Table a -> Bool
equalRows t0 s0 suffs table =
  forAll (\((t, s), e) -> lookupP (s ++ e) table == lookupP (t ++ e) table) $ product (singleOrbit (t0, s0)) suffs

closed :: _ => Word a -> Rows a -> Columns a -> Table a -> Bool
closed t prefs suffs table =
  exists (\(t, s) -> equalRows t s suffs table) (leftProduct (singleOrbit t) prefs)

nonClosedness :: _ => Rows a -> Rows a -> Columns a -> Table a -> Rows a
nonClosedness prefs prefsExt suffs table =
  filter (\t -> not $ closed t prefs suffs table) prefsExt

inconsistencies :: _ => Rows a -> Columns a -> Table a -> OrbitList a -> OrbitList ((Word a, Word a), (a, Word a))
inconsistencies prefs suffs table alph =
  filter (\((s, t), (a, e)) -> lookupP (s ++ (a:e)) table /= lookupP (t ++ (a:e)) table) candidatesExt
  where
    candidates = filter (\(s, t) -> s < t && equalRows s t suffs table) (product prefs prefs)
    candidatesExt = product candidates (product alph suffs)


-- Main state of the L* algorithm
-- invariants: * prefs and prefsExt disjoint, without dups
--             * prefsExt ordered
--             * prefs and (prefs `union` prefsExt) prefix-closed
--             * table defined on (prefs `union` prefsExt) * suffs
data Observations a = Observations
  { alph     :: OrbitList a
  , prefs    :: Rows a
  , prefsExt :: Rows a
  , suffs    :: Columns a
  , table    :: Table a
  }

-- input alphabet, inner monad, return value
type LStar i m a = StateT (Observations i) m a

-- First lookup, then membership query, also update the table
ask mq (p, s) = do
  Observations{..} <- get
  let w = p ++ s
  case lookupP w table of
    Just b -> return (w, b)
    Nothing -> do
      b <- lift (mq w)
      modify $ \o -> o { table = insertP w b table }
      return (w, b)

-- precondition: newPrefs is subset of prefExts
addRows :: _ => Rows a -> (Word a -> m Bool) -> LStar a m ()
addRows newPrefs mq = do
  Observations{..} <- get
  let newPrefsExt = productWith ext newPrefs alph
      rect = product newPrefsExt suffs
  _ <- mapM (ask mq) (OrbitList.toList rect)
  modify $ \o -> o { prefs = prefs <> newPrefs
                   , prefsExt = (prefsExt `minus` newPrefs) `union` newPrefsExt
                   }
  return ()

-- precondition: newSuffs disjoint from suffs
addCols :: _ => Columns a -> (Word a -> m Bool) -> LStar a m ()
addCols newSuffs mq = do
  Observations{..} <- get
  let rect = product (prefs `union` prefsExt) newSuffs
  _ <- mapM (ask mq) (OrbitList.toList rect)
  modify $ \o -> o { suffs = suffs <> newSuffs }
  return ()

fillTable :: _ => (Word a -> m Bool) -> LStar a m ()
fillTable mq = do
  Observations{..} <- get
  let rect = product (prefs `union` prefsExt) suffs
  _ <- mapM (ask mq) (OrbitList.toList rect)
  return ()

-- This could be cleaned up
learn :: _ => (Word a -> m Bool) -> (Automaton _ a -> m (Maybe (Word a))) -> LStar a m (Automaton _ a)
learn mq eq = do
  Observations{..} <- get
  let ncl = nonClosedness prefs prefsExt suffs table
      inc = inconsistencies prefs suffs table alph
  case null ncl of
    False -> do
      -- If not closed, then add 1 orbit of rows. Then start from top
      addRows (take 1 ncl) mq
      learn mq eq
    True -> do
      -- Closed! Now we check consistency
      case null inc of
        False -> do
          -- If not consistent, then add 1 orbit of columns. Then start from top
          addCols (take 1 (map (uncurry (:) . snd) inc)) mq
          learn mq eq
        True -> do
          -- Also consistent! Let's build a minimal automaton!
          let (f, st, _) = quotientf 0 (\s t -> s == t || equalRows s t suffs table) prefs
              trans = Map.fromList . toList . map (\(s, t) -> (s, f ! t)) . filter (\(s, t) -> equalRows s t suffs table) $ product prefsExt prefs
              trans2 pa = if pa `elem` prefsExt then trans ! pa else f ! pa
              hypothesis = Automaton
                { states = map fst st
                , initialState = f ! []
                , acceptance = Map.fromList . toList . map (\p -> (f ! p, table !~ p)) $ prefs
                , transition = Map.fromList . toList . map (\(p, a) -> ((f ! p, a), trans2 (ext p a))) $ product prefs alph
                }
              askCe = do
                ce <- lift (eq hypothesis)
                case ce of
                  Nothing -> return hypothesis
                  Just w -> do
                    let b1 = accepts hypothesis w
                    (_, b2) <- ask mq (w, [])
                    -- Ignore false counterexamples
                    case b1 == b2 of
                      True -> askCe
                      False -> do
                        -- Add all suffixes of a counterexample
                        let allSuffs = Set.fromList $ tails w
                            newSuffs = allSuffs `Set.difference` Set.fromOrbitList suffs
                        addCols (Set.toOrbitList newSuffs) mq
                        learn mq eq
          askCe


-- Here is the teacher: just pose the queries in the terminal
askMember :: _ => Word a -> IO Bool
askMember w = do
  putStrLn (toStr (MQ w))
  hFlush stdout
  a <- getLine
  case a of
    "Y" -> return True
    "N" -> return False
    _   -> askMember w

askEquiv :: _ => Automaton q a -> IO (Maybe (Word a))
askEquiv aut = do
  putStr "EQ \""
  putStr (toStr aut)
  putStrLn "\""
  hFlush stdout
  a <- getLine
  case a of
    "Y"       -> return Nothing
    'N':' ':w -> return $ Just (fst $ fromStr w)
    _         -> askEquiv aut

init alph = Observations
  { alph = alph
  , prefs = singleOrbit []
  , prefsExt = productWith ext (singleOrbit []) alph
  , suffs = singleOrbit[]
  , table = mempty
  }

main :: IO ()
main = do
  putStrLn "ALPHABET"
  hFlush stdout
  alph <- getLine
  case alph of
    "ATOMS" -> do
      aut <- evalStateT (fillTable askMember >> learn askMember askEquiv) (init rationals)
      return ()
    "FIFO" -> do
      let alph = map Put rationals `union` map Get rationals
      aut <- evalStateT (fillTable askMember >> learn askMember askEquiv) (init alph)
      return ()
    al -> do
      putStr "Unknown alphabet "
      putStrLn al