{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{-
 - Copyright (C) 2019  Koz Ross <koz.ross@retro-freedom.nz>
 -
 - This program is free software: you can redistribute it and/or modify
 - it under the terms of the GNU General Public License as published by
 - the Free Software Foundation, either version 3 of the License, or
 - (at your option) any later version.
 -
 - This program is distributed in the hope that it will be useful,
 - but WITHOUT ANY WARRANTY; without even the implied warranty of
 - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 - GNU General Public License for more details.
 -
 - You should have received a copy of the GNU General Public License
 - along with this program.  If not, see <http://www.gnu.org/licenses/>.
 -}

module Main where

import Data.Finitary (Finitary (..))
import Data.Int (Int16, Int32, Int8)
import Data.Ord (Down (..))
import qualified Data.Vector.Sized as V
import qualified Data.Vector.Storable.Sized as VS
import Data.Vector.Unboxed.Sized (Unbox)
import qualified Data.Vector.Unboxed.Sized as VU
import Data.Word (Word16, Word32, Word8)
import Foreign.Storable (Storable)
import GHC.Generics (Generic)
import Hedgehog ((===), Gen, PropertyT, forAll)
import qualified Hedgehog.Gen as Gen
import Hedgehog.Range (constantBounded)
import Test.Hspec (SpecWith, describe, hspec, it, parallel)
import Test.Hspec.Hedgehog (hedgehog, modifyMaxSize)

main :: IO ()
main = hspec . parallel $ do
  describe "Bijectivity and order preservation" $ do
    checkBijection "Char" Gen.unicode
    checkBijection "Word8" (Gen.enumBounded @_ @Word8)
    modifyMaxSize (const 10000)
      . checkBijection "Word16"
      $ Gen.enumBounded @_ @Word16
    modifyMaxSize (const 10000)
      . checkBijection "Word32"
      $ Gen.enumBounded @_ @Word32
    modifyMaxSize (const 10000)
      . checkBijection "Word64"
      $ Gen.word64 constantBounded
    checkBijection "Int8" (Gen.enumBounded @_ @Int8)
    modifyMaxSize (const 10000)
      . checkBijection "Int16"
      $ Gen.enumBounded @_ @Int16
    modifyMaxSize (const 10000)
      . checkBijection "Int32"
      $ Gen.enumBounded @_ @Int32
    modifyMaxSize (const 10000)
      . checkBijection "Int64"
      $ Gen.int64 constantBounded
    modifyMaxSize (const 10000)
      . checkBijection "Int"
      $ Gen.int constantBounded
    modifyMaxSize (const 10000)
      . checkBijection "Word"
      $ Gen.word constantBounded
  describe "Down" $ do
    checkMonotonic "Bool" Gen.bool
    modifyMaxSize (const 10000)
      . checkMonotonic "Int"
      $ (Gen.enumBounded @_ @Int)
    modifyMaxSize (const 10000)
      . checkMonotonic "(Either Int Bool)"
      $ Gen.choice
        [ Left <$> Gen.enumBounded @_ @Int,
          Right <$> Gen.enumBounded @_ @Bool
        ]
    modifyMaxSize (const 10000)
      . checkMonotonic "(Int, Bool)"
      $ ( (,)
            <$> Gen.enumBounded @_ @Int
            <*> Gen.enumBounded @_ @Bool
        )
    modifyMaxSize (const 10000)
      . checkMonotonic "of a user-defined type"
      $ genFoo
  describe "Fixed-length vectors" $ do
    modifyMaxSize (const 10000)
      . checkStorable "Int8"
      . genStorable
      $ Gen.enumBounded @_ @Int8
    modifyMaxSize (const 10000)
      . checkUnboxed "Int8"
      . genUnboxed
      $ Gen.enumBounded @_ @Int8
    modifyMaxSize (const 10000)
      . checkRegular "Int8"
      . genRegular
      $ Gen.enumBounded @_ @Int8
    modifyMaxSize (const 10000)
      . checkUnboxed "(Int8, Int8)"
      . genUnboxed
      $ ( (,) <$> Gen.enumBounded @_ @Int8
            <*> Gen.enumBounded @_ @Int8
        )
    modifyMaxSize (const 10000)
      . checkRegular "(Int8, Int8)"
      . genRegular
      $ ( (,) <$> Gen.enumBounded @_ @Int8
            <*> Gen.enumBounded @_ @Int8
        )
    modifyMaxSize (const 10000)
      . checkRegular "Either Int8 Bool"
      . genRegular
      . Gen.choice
      $ [ Left <$> Gen.enumBounded @_ @Int8,
          Right <$> Gen.bool
        ]
    modifyMaxSize (const 10000)
      . checkRegular "a user defined type"
      . genRegular
      $ genFoo

-- Helpers

data Foo
  = Bar
  | Baz Int8
  | Quux (Int8, Int8)
  deriving stock (Eq, Ord, Generic, Show)
  deriving anyclass (Finitary)

checkStorable ::
  forall a.
  (Storable a, Finitary a, Show a, Ord a) =>
  String ->
  Gen (VS.Vector 10 a) ->
  SpecWith ()
checkStorable name =
  it ("should biject a Storable Vector of " <> name)
    . hedgehog
    . bicheck @(VS.Vector 10 a)

checkRegular ::
  forall a.
  (Finitary a, Show a, Ord a) =>
  String ->
  Gen (V.Vector 10 a) ->
  SpecWith ()
checkRegular name =
  it ("should biject a Vector of " <> name)
    . hedgehog
    . bicheck @(V.Vector 10 a)

checkUnboxed ::
  forall a.
  (Unbox a, Finitary a, Show a, Ord a) =>
  String ->
  Gen (VU.Vector 10 a) ->
  SpecWith ()
checkUnboxed name =
  it ("should biject an Unboxed Vector of " <> name)
    . hedgehog
    . bicheck @(VU.Vector 10 a)

bicheck :: forall a. (Show a, Finitary a, Ord a) => Gen a -> PropertyT IO ()
bicheck gen = do
  v <- forAll gen
  let iv = toFinite v
  v === (fromFinite . toFinite $ v)
  iv === (toFinite @a . fromFinite $ iv)
  v' <- forAll gen
  let iv' = toFinite v'
  compare v v' === compare iv iv'

genStorable :: (Storable a) => Gen a -> Gen (VS.Vector 10 a)
genStorable = VS.replicateM

genUnboxed :: (Unbox a) => Gen a -> Gen (VU.Vector 10 a)
genUnboxed = VU.replicateM

genRegular :: Gen a -> Gen (V.Vector 10 a)
genRegular = V.replicateM

genFoo :: Gen Foo
genFoo =
  Gen.choice
    [ pure Bar,
      Baz <$> Gen.enumBounded,
      Quux <$> ((,) <$> Gen.enumBounded <*> Gen.enumBounded)
    ]

checkBijection :: forall a. (Show a, Ord a, Finitary a) => String -> Gen a -> SpecWith ()
checkBijection name gen =
  it ("should biject " <> name <> " with fromFinite and toFinite preserving order")
    . hedgehog
    $ go
  where
    go = do
      x <- forAll gen
      let ix = toFinite x
      x === (fromFinite . toFinite $ x)
      ix === (toFinite @a . fromFinite $ ix)
      y <- forAll gen
      let iy = toFinite y
      compare x y === compare ix iy

checkMonotonic :: (Show a, Finitary a) => String -> Gen a -> SpecWith ()
checkMonotonic name gen =
  it ("should be Ord-monotonic on Down " <> name)
    . hedgehog
    $ go
  where
    go = do
      x <- forAll gen
      y <- forAll gen
      let dx = toFinite . Down $ x
      let dy = toFinite . Down $ y
      let ix = toFinite x
      let iy = toFinite y
      case compare ix iy of
        LT -> compare dx dy === GT
        EQ -> compare dx dy === EQ
        GT -> compare dx dy === LT