{-# language DeriveGeneric #-}
{-# language DerivingVia #-}
{-# language FlexibleContexts #-}
{-# language RecordWildCards #-}
{-# language UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-missing-signatures #-}

module ExampleAutomata
  ( module ExampleAutomata
  , module Automata
  ) where

import Nominal hiding (product)
import Automata
import IO
import OrbitList
import qualified EquivariantMap as Map
import qualified EquivariantSet as Set

import Data.Foldable (fold)
import GHC.Generics (Generic)
import Prelude as P hiding (map, product, words, filter, foldr)


atoms = rationals

-- words of length <= m
words m = fold $ go (m+1) (singleOrbit []) where
  go 0 _   = []
  go k acc = acc : go (k-1) (productWith (:) atoms acc)

fromKeys f = Map.fromSet f . Set.fromOrbitList

ltPair = filter (\(a, b) -> a < b) $ product atoms atoms


data DoubleWord = Store [Atom] | Check [Atom] | Accept | Reject
  deriving (Eq, Ord, Generic)
  deriving Nominal via Generically DoubleWord

instance ToStr DoubleWord where
  toStr (Store w) = "S " ++ toStr w
  toStr (Check w) = "C " ++ toStr w
  toStr Accept    = "A"
  toStr Reject    = "R"

doubleWordAut 0 = Automaton {..} where
  states = fromList [Accept, Reject]
  initialState = Accept
  acceptance = fromKeys (Accept ==) states
  transition = fromKeys (const Reject) $ product states rationals
doubleWordAut n = Automaton {..} where
  states = fromList [Accept, Reject] <> map Store (words (n-1)) <> map Check (productWith (:) rationals (words (n-1)))
  initialState = Store []
  acceptance = fromKeys (Accept ==) states
  trans Accept _ = Reject
  trans Reject _ = Reject
  trans (Store l) a
    | length l + 1 < n = Store (a:l)
    | otherwise        = Check (reverse (a:l))
  trans (Check (a:as)) b
    | a == b    = if (P.null as) then Accept else Check as
    | otherwise = Reject
  transition = fromKeys (uncurry trans) $ product states rationals


-- alphetbet for the Fifo queue example
data FifoA = Put Atom | Get Atom
  deriving (Eq, Ord, Show, Generic)
  deriving Nominal via Generically FifoA

instance ToStr FifoA where
  toStr (Put a) = "Put " ++ toStr a
  toStr (Get a) = "Get " ++ toStr a

instance FromStr FifoA where
  fromStr ('P':'u':'t':' ':a) = let (x, r) = fromStr a in (Put x, r)
  fromStr ('G':'e':'t':' ':a) = let (x, r) = fromStr a in (Get x, r)
  fromStr _ = error "Cannot parse Fifo"

fifoAlph = map Put rationals <> map Get rationals

data FifoS = FifoS [Atom] [Atom]
  deriving (Eq, Ord, Generic)
  deriving Nominal via Generically FifoS

instance ToStr FifoS where
  toStr (FifoS l1 l2) = "F " ++ toStr l1 ++ " - " ++ toStr l2

fifoAut n = Automaton {..} where
  states0 = filter (\(FifoS l1 l2) -> length l1 + length l2 <= n) $ productWith (\l1 l2 -> FifoS l1 l2) (words n) (words n)
  states = fromList [Nothing] <> map Just states0
  initialState = Just (FifoS [] [])
  acceptance = fromKeys (Nothing /=) states
  trans Nothing _ = Nothing
  trans (Just (FifoS l1 l2)) (Put a)
    | length l1 + length l2 >= n = Nothing
    | otherwise                  = Just (FifoS (a:l1) l2)
  trans (Just (FifoS [] [])) (Get _) = Nothing
  trans (Just (FifoS l1 [])) (Get b) = trans (Just (FifoS [] (reverse l1))) (Get b)
  trans (Just (FifoS l1 (a:l2))) (Get b)
    | a == b    = Just (FifoS l1 l2)
    | otherwise = Nothing
  transition = fromKeys (uncurry trans) $ product states fifoAlph


data Lint a = Lint_start | Lint_single a | Lint_full a a | Lint_semi a a | Lint_error
  deriving (Eq, Ord, Show, Generic)
  deriving Nominal via Generically (Lint a)

lintExample ::Automaton (Lint Atom) Atom
lintExample = Automaton {..} where
    states = fromList [Lint_start, Lint_error]
        <> map Lint_single atoms
        <> map (uncurry Lint_full) ltPair
        <> map (uncurry Lint_semi) ltPair
    initialState = Lint_start
    acc (Lint_full _ _) = True
    acc _ = False
    acceptance = fromKeys acc states
    trans Lint_start a = Lint_single a
    trans (Lint_single a) b
      | a < b = Lint_full a b
      | otherwise = Lint_error
    trans (Lint_full a b) c
      | a < c && c < b = Lint_semi c b
      | otherwise = Lint_error
    trans (Lint_semi a b) c
      | a < c && c < b = Lint_full a c
      | otherwise = Lint_error
    trans Lint_error _ = Lint_error
    transition = fromKeys (uncurry trans) $ product states atoms


data Lmax a = Lmax_start | Lmax_single a | Lmax_double a a
  deriving (Eq, Ord, Show, Generic)
  deriving Nominal via Generically (Lmax a)

lmaxExample :: Automaton (Lmax Atom) Atom
lmaxExample = Automaton {..} where
    states = singleOrbit Lmax_start
        <> map Lmax_single atoms
        <> productWith Lmax_double atoms atoms
    initialState = Lmax_start
    acc (Lmax_double a b) = a == b
    acc _ = False
    acceptance = fromKeys acc states
    trans Lmax_start a = Lmax_single a
    trans (Lmax_single a) b = Lmax_double a b
    trans (Lmax_double b c) a
      | b >= c = Lmax_double b a
      | otherwise = Lmax_double c a
    transition = fromKeys (uncurry trans) $ product states atoms