module AFRPLawsSpec (spec) where

import AFRP
import Control.Arrow (arr, first, left, (***), (+++))
import Control.Category ((>>>))
import qualified Control.Category as Cat (id)
import Data.Functor.Identity (Identity (..))
import Hedgehog (Gen, PropertyT)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import Support (fakeRequest)
import Test.Hspec (Spec, describe, it)
import Test.Hspec.Hedgehog (forAll, forAllWith, hedgehog, (===))

-- | Law tests for the Mealy instances. Two machines count as equal when
-- they emit equal outputs on every input sequence, so each law runs both
-- sides on generated inputs.

runPure :: Mealy Identity a b -> [a] -> [b]
runPure m = go (runMealy m id)
  where
    go _ [] = []
    go w (a : as) = case runIdentity (stepAuto w fakeRequest a) of
      (b, w') -> b : go w' as

-- | Machines wrap functions and have no Show; name them for forAll instead.
forAllMealy :: Gen (Mealy Identity a b) -> PropertyT IO (Mealy Identity a b)
forAllMealy = forAllWith (const "<mealy>")

intGen :: Gen Int
intGen = Gen.int (Range.linear (-5) 5)

ints :: Gen [Int]
ints = Gen.list (Range.linear 0 30) intGen

intPairs :: Gen [(Int, Int)]
intPairs = Gen.list (Range.linear 0 30) ((,) <$> intGen <*> intGen)

intEithers :: Gen [Either Int Int]
intEithers = Gen.list (Range.linear 0 30) $
  Gen.choice [Left <$> intGen, Right <$> intGen]

nestedPairs :: Gen [((Int, Int), Int)]
nestedPairs = Gen.list (Range.linear 0 30) ((,) <$> ((,) <$> intGen <*> intGen) <*> intGen)

nestedEithers :: Gen [Either (Either Int Int) Int]
nestedEithers = Gen.list (Range.linear 0 30) $
  Gen.choice
    [ Left <$> Gen.choice [Left <$> intGen, Right <$> intGen]
    , Right <$> intGen
    ]

-- | Stateful Int machines: the arrow variables of the laws.
statefulGen :: Gen (Mealy Identity Int Int)
statefulGen = Gen.choice
  [ (\k -> mapAccum (+) k id) <$> intGen
  , (\k -> preMapAccum (+) k id) <$> intGen
  , (\k -> mapAccum (*) 1 (+ k)) <$> intGen
  ]

eventArrowGen :: Gen (Mealy Identity Int (Event Int))
eventArrowGen = Gen.choice
  [ pure changes
  , (\k -> mapAccum (+) k Event) <$> intGen
  , (\k -> preMapAccum (+) k (Event . (* 2))) <$> intGen
  ]

funArrowGen :: Gen (Mealy Identity Int (Int -> Int))
funArrowGen = Gen.choice
  [ (\k -> mapAccum (+) k (*)) <$> intGen
  , pure (preMapAccum (*) 1 (+))
  ]

spec :: Spec
spec = describe "Mealy laws" $ do
  semigroupSpec
  monoidSpec
  categorySpec
  arrowSpec
  arrowChoiceSpec
  functorSpec
  applicativeSpec

semigroupSpec :: Spec
semigroupSpec = describe "Semigroup (<>)" $ do
  it "(a <> b) <> c = a <> (b <> c)" $ hedgehog $ do
    a <- forAllMealy eventArrowGen
    b <- forAllMealy eventArrowGen
    c <- forAllMealy eventArrowGen
    xs <- forAll ints
    runPure ((a <> b) <> c) xs === runPure (a <> (b <> c)) xs

monoidSpec :: Spec
monoidSpec = describe "Monoid" $ do
  it "mempty <> a = a" $ hedgehog $ do
    a <- forAllMealy eventArrowGen
    xs <- forAll ints
    runPure (mempty <> a) xs === runPure a xs

  it "a <> mempty = a" $ hedgehog $ do
    a <- forAllMealy eventArrowGen
    xs <- forAll ints
    runPure (a <> mempty) xs === runPure a xs

categorySpec :: Spec
categorySpec = describe "Category" $ do
  it "id >>> f = f" $ hedgehog $ do
    f <- forAllMealy statefulGen
    xs <- forAll ints
    runPure (Cat.id >>> f) xs === runPure f xs

  it "f >>> id = f" $ hedgehog $ do
    f <- forAllMealy statefulGen
    xs <- forAll ints
    runPure (f >>> Cat.id) xs === runPure f xs

  it "(f >>> g) >>> h = f >>> (g >>> h)" $ hedgehog $ do
    f <- forAllMealy statefulGen
    g <- forAllMealy statefulGen
    h <- forAllMealy statefulGen
    xs <- forAll ints
    runPure ((f >>> g) >>> h) xs === runPure (f >>> (g >>> h)) xs

arrowSpec :: Spec
arrowSpec = describe "Arrow" $ do
  it "arr id = id" $ hedgehog $ do
    xs <- forAll ints
    runPure (arr id :: Mealy Identity Int Int) xs === runPure Cat.id xs

  it "arr (f >>> g) = arr f >>> arr g" $ hedgehog $ do
    p <- forAll intGen
    q <- forAll intGen
    xs <- forAll ints
    let f = (+ p)
        g = (* q)
    runPure (arr (f >>> g)) xs === runPure (arr f >>> arr g) xs

  it "first (arr f) = arr (first f)" $ hedgehog $ do
    p <- forAll intGen
    ps <- forAll intPairs
    let f = (+ p)
    runPure (first (arr f)) ps === runPure (arr (first f)) ps

  it "first (f >>> g) = first f >>> first g" $ hedgehog $ do
    f <- forAllMealy statefulGen
    g <- forAllMealy statefulGen
    ps <- forAll intPairs
    runPure (first (f >>> g)) ps === runPure (first f >>> first g) ps

  it "first f >>> arr fst = arr fst >>> f" $ hedgehog $ do
    f <- forAllMealy statefulGen
    ps <- forAll intPairs
    runPure (first f >>> arr fst) ps === runPure (arr fst >>> f) ps

  it "first f >>> arr (id *** g) = arr (id *** g) >>> first f" $ hedgehog $ do
    f <- forAllMealy statefulGen
    p <- forAll intGen
    ps <- forAll intPairs
    let g = (* p)
    runPure (first f >>> arr (id *** g)) ps
      === runPure (arr (id *** g) >>> first f) ps

  it "first (first f) >>> arr assoc = arr assoc >>> first f" $ hedgehog $ do
    f <- forAllMealy statefulGen
    ts <- forAll nestedPairs
    let assoc ((a, b), c) = (a, (b, c))
    runPure (first (first f) >>> arr assoc) ts
      === runPure (arr assoc >>> first f) ts

arrowChoiceSpec :: Spec
arrowChoiceSpec = describe "ArrowChoice" $ do
  it "left (arr f) = arr (left f)" $ hedgehog $ do
    p <- forAll intGen
    es <- forAll intEithers
    let f = (+ p)
    runPure (left (arr f)) es === runPure (arr (left f)) es

  it "left (f >>> g) = left f >>> left g" $ hedgehog $ do
    f <- forAllMealy statefulGen
    g <- forAllMealy statefulGen
    es <- forAll intEithers
    runPure (left (f >>> g)) es === runPure (left f >>> left g) es

  it "f >>> arr Left = arr Left >>> left f" $ hedgehog $ do
    f <- forAllMealy statefulGen
    xs <- forAll ints
    runPure (f >>> arr (Left @Int @Int)) xs
      === runPure (arr (Left @Int @Int) >>> left f) xs

  it "left f >>> arr (id +++ g) = arr (id +++ g) >>> left f" $ hedgehog $ do
    f <- forAllMealy statefulGen
    p <- forAll intGen
    es <- forAll intEithers
    let g = (* p)
    runPure (left f >>> arr (id +++ g)) es
      === runPure (arr (id +++ g) >>> left f) es

  it "left (left f) >>> arr assocsum = arr assocsum >>> left f" $ hedgehog $ do
    f <- forAllMealy statefulGen
    es <- forAll nestedEithers
    let assocsum (Left (Left x)) = Left x
        assocsum (Left (Right y)) = Right (Left y)
        assocsum (Right z) = Right (Right z)
    runPure (left (left f) >>> arr assocsum) es
      === runPure (arr assocsum >>> left f) es

functorSpec :: Spec
functorSpec = describe "Functor" $ do
  it "fmap id = id" $ hedgehog $ do
    m <- forAllMealy statefulGen
    xs <- forAll ints
    runPure (fmap id m) xs === runPure m xs

  it "fmap (f . g) = fmap f . fmap g" $ hedgehog $ do
    m <- forAllMealy statefulGen
    p <- forAll intGen
    q <- forAll intGen
    xs <- forAll ints
    let f = (+ p)
        g = (* q)
    runPure (fmap (f . g) m) xs === runPure (fmap f (fmap g m)) xs

applicativeSpec :: Spec
applicativeSpec = describe "Applicative" $ do
  it "pure id <*> v = v" $ hedgehog $ do
    v <- forAllMealy statefulGen
    xs <- forAll ints
    runPure (pure id <*> v) xs === runPure v xs

  it "pure f <*> pure x = pure (f x)" $ hedgehog $ do
    p <- forAll intGen
    x <- forAll intGen
    xs <- forAll ints
    let f = (+ p)
    runPure (pure f <*> pure x) xs === runPure (pure (f x)) xs

  it "u <*> pure y = pure ($ y) <*> u" $ hedgehog $ do
    u <- forAllMealy funArrowGen
    y <- forAll intGen
    xs <- forAll ints
    runPure (u <*> pure y) xs === runPure (pure ($ y) <*> u) xs

  it "pure (.) <*> u <*> v <*> w = u <*> (v <*> w)" $ hedgehog $ do
    u <- forAllMealy funArrowGen
    v <- forAllMealy funArrowGen
    w <- forAllMealy statefulGen
    xs <- forAll ints
    runPure (pure (.) <*> u <*> v <*> w) xs
      === runPure (u <*> (v <*> w)) xs

  it "fmap f x = pure f <*> x" $ hedgehog $ do
    x <- forAllMealy statefulGen
    p <- forAll intGen
    xs <- forAll ints
    let f = (+ p)
    runPure (fmap f x) xs === runPure (pure f <*> x) xs