module Control.DFSTTest where

import Control.DFST
import Control.FST hiding (stInitial, stTransition, stAccept)

import Data.Set (Set)
import qualified Data.Set as Set

import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import Data.Sequence (Seq)
import qualified Data.Sequence as Seq

import Data.Maybe (maybeToList)

import Test.Tasty
import Test.Tasty.Hedgehog
import Test.Tasty.HUnit hiding (assert)

import Hedgehog
import qualified Hedgehog.Gen as G
import qualified Hedgehog.Range as R

import Numeric.Natural

import Text.PrettyPrint.Leijen (Pretty(..))


dfstId :: Ord a => Set a -> DFST () a a
dfstId syms = DFST
  { stInitial = ()
  , stTransition = Map.fromList
    [(((), sym), ((), Seq.singleton sym)) | sym <- Set.toList syms]
  , stAccept = Set.singleton ()
  }

dfstDouble :: Ord a => Set a -> DFST () a a
dfstDouble syms = DFST
  { stInitial = ()
  , stTransition = Map.fromList 
    [(((), sym), ((), Seq.fromList [sym, sym])) | sym <- Set.toList syms]
  , stAccept = Set.singleton ()
  }

dfstRunLengthDecode :: Ord a
                    => Set a
                    -> Natural
                    -> DFST (Maybe Natural) (Either Natural a) a
dfstRunLengthDecode syms lim = DFST
  { stInitial = Nothing
  , stTransition = Map.fromList . concat $
    [ [((Nothing, Left n), (Just n, Seq.empty)) | n <- [0..lim]]
    , [((Just n, Right sym), (Nothing, Seq.replicate (fromIntegral n) sym)) | n <- [0..lim], sym <- Set.toList syms]
    ]
  , stAccept = Set.singleton Nothing
  }

genWord :: Gen [Natural]
genWord = G.list (R.linear 0 1000) . G.integral $ R.linear 0 100

genDFST :: (Ord input, Ord output) => Set input -> Set output -> Gen (DFST Natural input output)
genDFST inA outA = do
  states <- G.set (R.linear 1 1000) . G.integral $ R.linear 0 100
  stInitial <- G.element $ Set.toList states
  stAccept <- Set.fromList <$> G.subsequence (Set.toList states)
  stTransition <- fmap Map.fromList . G.list (R.linear 0 1000) . G.small $ do
    st <- G.element $ Set.toList states
    input <- G.element $ Set.toList inA
    st' <- G.element $ Set.toList states
    output <- fmap Seq.fromList . G.list (R.linear 0 20) . G.element $ Set.toList outA
    return ((st, input), (st', output))
  return DFST{..}
    

testDFST :: (Show output, Ord output, Show state, Ord state) => (Set Natural -> DFST state Natural output) -> (Seq Natural -> Seq output) -> Property
testDFST mkDfst f = property $ do
  input <- forAll genWord
  let fst = mkDfst $ Set.fromList input
  Just (f $ Seq.fromList input) === runDFST fst (Seq.fromList input)

hprop_runDFSTId, hprop_runDFSTDouble :: Property
hprop_runDFSTId = testDFST dfstId id
hprop_runDFSTDouble = testDFST dfstDouble double
  where
    double :: Seq a -> Seq a
    double Seq.Empty = Seq.Empty
    double (a Seq.:<| as) = a Seq.:<| a Seq.:<| double as

unit_runLengthDecode :: Assertion
unit_runLengthDecode = runDFST dfst input @?= Just (Seq.fromList "aaacc")
  where
    input = Seq.fromList [Left 3, Right 'a', Left 0, Right 'b', Left 2, Right 'c']
    dfst = dfstRunLengthDecode (Set.fromList "abc") 3

hprop_toFST :: Property
hprop_toFST = property $ do
  input <- forAll genWord
  dfst <- forAllWith (show . pretty . toFST) $ genDFST (Set.fromList $ input ++ [0..20]) (Set.fromList [0..20] :: Set Natural)
  
  runFST (toFST dfst) (Seq.fromList input) === maybeToList (runDFST dfst $ Seq.fromList input)