module Control.FSTTest where

import Control.FST

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 (fromMaybe)
import Data.Foldable (Foldable(..))

import Control.Monad (when)

import Data.Void

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(..))

import Control.DeepSeq (force)


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

fstDouble :: Ord a => Set a -> FST (Maybe a) a a
fstDouble syms = FST
  { stInitial = Set.singleton Nothing
  , stTransition = Map.fromListWith Set.union . concat $
    [ [((Nothing, Just sym), Set.singleton (Just sym, Just sym)) | sym <- Set.toList syms]
    , [((Just sym, Nothing), Set.singleton (Nothing, Just sym)) | sym <- Set.toList syms]
    ]
  , stAccept = Set.singleton Nothing
  }

fstRunLengthDecode :: Ord a
                   => Set a -- ^ Alphabet
                   -> Natural -- ^ Upper limit to run length
                   -> FST (Maybe (Natural, Maybe a)) (Either Natural a) a
fstRunLengthDecode syms lim = FST
  { stInitial = Set.singleton Nothing
  , stTransition = Map.fromListWith Set.union . concat $
    [ [((Nothing, Just (Left n)), Set.singleton (Just (n, Nothing), Nothing)) | n <- [0..lim]]
    , [((Just (n, Nothing), Just (Right sym)), Set.singleton (Just (n, Just sym), Nothing)) | n <- [0..lim], sym <- Set.toList syms]
    , [((Just (n, Just sym), Nothing), Set.singleton (Just (pred n, Just sym), Just sym)) | n <- [1..lim], sym <- Set.toList syms]
    , [((Just (0, Just sym), Nothing), Set.singleton (Nothing, Nothing)) | sym <- Set.toList syms]
    ]
  , stAccept = Set.singleton Nothing
  }

data StRunLengthEncode a = STREInitial
                         | STRECountUp a Natural
                         | STRESwitch (Maybe a) a
                         | STREFinish
  deriving (Show, Eq, Ord)

fstRunLengthEncode :: Ord a
                   => Bool -- ^ Generate /all/ run length encodings instead of the best
                   -> Set a -- ^ Alphabet
                   -> Natural -- ^ Upper limit to run length
                   -> FST (StRunLengthEncode a) a (Either Natural a)
fstRunLengthEncode genAll syms lim = FST
  { stInitial = Set.singleton STREInitial
  , stTransition = Map.fromListWith Set.union . concat $
    [ [((STREInitial, Just sym), Set.singleton (STRECountUp sym 1, Nothing)) | sym <- Set.toList syms]
    , [((STRECountUp sym n, Just sym), Set.singleton (STRECountUp sym (succ n), Nothing)) | sym <- Set.toList syms, n <- [1..pred lim]]
    , [((STRECountUp sym n, Just sym'), Set.singleton (STRESwitch (Just sym') sym, Just $ Left n)) | n <- [1..lim], sym <- Set.toList syms, sym' <- Set.toList syms, sym /= sym']
    , [((STRECountUp sym lim, Just sym), Set.singleton (STRESwitch (Just sym) sym, Just $ Left lim)) | sym <- Set.toList syms]
    , [((STRECountUp sym n, Nothing), Set.singleton (STRESwitch Nothing sym, Just $ Left n)) | sym <- Set.toList syms, n <- [1..lim]]
    , [((STRESwitch (Just sym') sym, Nothing), Set.singleton (STRECountUp sym' 1, Just $ Right sym)) | sym <- Set.toList syms, sym' <- Set.toList syms]
    , [((STRESwitch Nothing sym, Nothing), Set.singleton (STREFinish, Just $ Right sym)) | sym <- Set.toList syms]
    , [((STRECountUp sym n, Just sym), Set.singleton (STRESwitch (Just sym) sym, Just $ Left n)) | n <- [1..lim], sym <- Set.toList syms, genAll]
    ]
  , stAccept = Set.fromList [STREInitial, STREFinish]
  }


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

  
runFSTDet :: (Show output, Ord output, Show state, Ord state) => (Set Natural -> FST state Natural output) -> (Seq Natural -> Seq output) -> Property
runFSTDet mkFst f = property $ do
  input <- forAll genWord
  let fst = mkFst $ Set.fromList input
  annotateShow $ pretty fst
  [f $ Seq.fromList input] === runFST fst (Seq.fromList input)


hprop_runFSTId, hprop_runFSTDouble :: Property
hprop_runFSTId = runFSTDet fstId id
hprop_runFSTDouble = runFSTDet fstDouble double
  where
    double :: Seq a -> Seq a
    double Seq.Empty = Seq.Empty
    double (a Seq.:<| as) = a Seq.:<| a Seq.:<| double as

hprop_runWordFST :: Property
hprop_runWordFST = property $ do
  input <- forAll genWord
  let fst = wordFST $ Seq.fromList input
  annotateShow $ pretty fst
  [Seq.fromList input] === runFST fst (Seq.empty :: Seq Void)

unit_runLengthDecode, unit_runLengthEncode :: Assertion
unit_runLengthDecode = runFST fst input @?= [Seq.fromList "aaacc"]
  where
    input = Seq.fromList [Left 3, Right 'a', Left 0, Right 'b', Left 2, Right 'c']
    fst = fstRunLengthDecode (Set.fromList "abc") 3
unit_runLengthEncode = runFST fst input @?= [Seq.fromList [Left 3, Right 'a', Left 2, Right 'c']]
  where
    input = Seq.fromList "aaacc"
    fst = fstRunLengthEncode False (Set.fromList "abc") 3

hprop_runLength :: Property
hprop_runLength = property $ do
  input <- forAll genWord
  let maxRun = fromMaybe 1 $ countMaxRun input
      alphabet = Set.fromList input
      encode = fstRunLengthEncode False alphabet maxRun
      decode = fstRunLengthDecode alphabet maxRun
  annotateShow $ countMaxRun input
  annotateShow $ pretty encode
  annotateShow $ pretty decode
  [encoded] <- return . runFST encode $ Seq.fromList input
  annotateShow $ toList encoded
  [Seq.fromList input] === runFST decode encoded

hprop_runLength' :: Property
hprop_runLength' = property $ do
  input <- forAll genWord

  let maxRun = fromMaybe 1 $ countMaxRun input
      alphabet = Set.fromList input
      encode = fstRunLengthEncode True alphabet maxRun
      decode = fstRunLengthDecode alphabet maxRun
  annotateShow $ countMaxRun input
  annotateShow $ pretty encode
  annotateShow $ pretty decode

  let encoded = runFST encode $ Seq.fromList input
  annotateShow encoded
  when (maxRun > 1) $
    assert $ encoded `longerThan` 1

  encoded' <- forAll $ G.element encoded
  [Seq.fromList input] === runFST decode encoded'
  where
    longerThan :: [a] -> Natural -> Bool
    longerThan []     _ = False
    longerThan  _     0 = True
    longerThan (_:xs) n = longerThan xs $ pred n

countMaxRun :: (Eq a, Foldable f) => f a -> Maybe Natural
countMaxRun = goMaxRun Nothing Nothing . toList
  where
    goMaxRun :: Eq a
             => Maybe (a, Natural) -- ^ Maximum
             -> Maybe (a, Natural) -- ^ Current
             -> [a]
             -> Maybe Natural
    goMaxRun Nothing Nothing [] = Nothing
    goMaxRun (Just (_, n)) (Just (_, n')) []
      | n' > n    = Just n'
      | otherwise = Just n
    goMaxRun Nothing Nothing (a:as) = goMaxRun (Just (a, 1)) (Just (a, 1)) as
    goMaxRun (Just (a, n)) (Just (a', n')) (x:xs)
      | x == a'   = goMaxRun (Just (a, n)) (Just (a', succ n')) xs
      | n' > n    = goMaxRun (Just (a', n')) (Just (x, 1)) xs
      | otherwise = goMaxRun (Just (a, n)) (Just (x, 1)) xs