{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
module Main (main) where

import Control.Monad ((<=<))
import Control.Monad.Trans.State (runState, state)
import Data.Hashable (Hashable)
import Data.Coerce (coerce)
import Data.Function (on)
import Data.List (nub, nubBy)
import Data.Proxy (Proxy (..))
import Data.Typeable (Typeable, typeRep)
import Test.QuickCheck (Arbitrary (..), Fun, Property, applyFun, Function (..), functionMap, CoArbitrary, (===))
import Test.QuickCheck.Instances ()
import Test.Tasty (defaultMain, testGroup, TestTree)
import Test.Tasty.QuickCheck (testProperty)

import qualified Data.HashMap.Lazy as HashMap
import qualified Data.IntMap as IntMap
import qualified Data.Map.Lazy as Map
import qualified Data.Vector as V
import qualified Data.Sequence as Seq

import Witherable
import Prelude hiding (filter)

main :: IO ()
main = defaultMain $ testGroup "witherable"
  [ testGroup "Filterable"
    [ filterableLaws (Proxy @[])
    , filterableLaws (Proxy @Maybe)
    , filterableLaws (Proxy @(Either String))
    , filterableLaws (Proxy @V.Vector)
    , filterableLaws (Proxy @Seq.Seq)
    , filterableLaws (Proxy @IntMap.IntMap)
    , filterableLaws (Proxy @(Map.Map K))
    , filterableLaws (Proxy @(HashMap.HashMap K))
    ]

  , testGroup "Witherable"
    [ witherableLaws (Proxy @[])
    , witherableLaws (Proxy @Maybe)
    , witherableLaws (Proxy @(Either String))
    , witherableLaws (Proxy @V.Vector)
    , witherableLaws (Proxy @Seq.Seq)
    , witherableLaws (Proxy @IntMap.IntMap)
    , witherableLaws (Proxy @(Map.Map K))
    , witherableLaws (Proxy @(HashMap.HashMap K))
    ]

  , nubProperties
  ]

-------------------------------------------------------------------------------
-- Filterable laws
-------------------------------------------------------------------------------

filterableLaws
  :: forall f.
     ( Filterable f, Typeable f
     , Arbitrary (f A), Show (f A), Eq (f A)
     , Arbitrary (f (Maybe A)), Show (f (Maybe A))
     , Show (f B), Eq (f B), Show (f C), Eq (f C)
     )
  => Proxy f
  -> TestTree
filterableLaws p = testGroup (show (typeRep p))
  [ testProperty "conservation" prop_conservation
  , testProperty "composition" prop_composition
  , testProperty "default filter" prop_default_filter
  , testProperty "default mapMaybe" prop_default_mapMaybe
  , testProperty "default catMaybes" prop_default_catMaybes
  ]
  where
    prop_conservation :: Fun A B -> f A -> Property
    prop_conservation f' xs =
        mapMaybe (Just . f) xs === fmap f xs
      where
        f = applyFun f'

    prop_composition :: Fun B (Maybe C) -> Fun A (Maybe B) -> f A -> Property
    prop_composition f' g' xs =
        mapMaybe f (mapMaybe g xs) === mapMaybe (f <=< g) xs
      where
        f = applyFun f'
        g = applyFun g'

    prop_default_filter :: Fun A Bool -> f A -> Property
    prop_default_filter f' xs =
        filter f xs === mapMaybe (\a -> if f a then Just a else Nothing) xs
      where
        f = applyFun f'

    prop_default_mapMaybe :: Fun A (Maybe B) -> f A -> Property
    prop_default_mapMaybe f' xs =
        mapMaybe f xs === catMaybes (fmap f xs)
      where
        f = applyFun f'

    prop_default_catMaybes :: f (Maybe A) -> Property
    prop_default_catMaybes xs = catMaybes xs === mapMaybe id xs

-------------------------------------------------------------------------------
-- Witherable laws
-------------------------------------------------------------------------------

witherableLaws
  :: forall f.
     ( Witherable f, Typeable f
     , Arbitrary (f A), Show (f A), Eq (f A)
     , Arbitrary (f (Maybe A)), Show (f (Maybe A))
     , Show (f B), Eq (f B), Show (f C), Eq (f C)
     )
  => Proxy f
  -> TestTree
witherableLaws p = testGroup (show (typeRep p))
  [ testProperty "default wither" prop_default_wither
  , testProperty "default witherM" prop_default_witherM
  , testProperty "default filterA" prop_default_filterA
  ]
  where
    prop_default_wither :: S -> Fun (A, S) (Maybe B, S) -> f A -> Property
    prop_default_wither s0 f' xs =
        runState (wither f xs) s0 === runState (fmap catMaybes (traverse f xs)) s0
      where
        f a = state $ \s -> applyFun f' (a, s)

    prop_default_witherM :: S -> Fun (A, S) (Maybe B, S) -> f A -> Property
    prop_default_witherM s0 f' xs =
        runState (witherM f xs) s0 === runState (wither f xs) s0
      where
        f a = state $ \s -> applyFun f' (a, s)

    prop_default_filterA :: S -> Fun (A, S) (Bool, S) -> f A -> Property
    prop_default_filterA s0 f' xs =
        runState (filterA f xs) s0 === runState (wither (\a -> (\b -> if b then Just a else Nothing) <$> f a) xs) s0
      where
        f a = state $ \s -> applyFun f' (a, s)

-------------------------------------------------------------------------------
-- Nub "laws"
-------------------------------------------------------------------------------

nubProperties :: TestTree
nubProperties = testGroup "nub"
  [ testProperty "ordNub" prop_ordNub
  , testProperty "ordNubOn" prop_ordNubOn
  , testProperty "hashNub" prop_hashNub
  , testProperty "hashNubOn" prop_hashNubOn
  , testProperty "ordNub is lazy" prop_lazy_ordNub
  , testProperty "hashNub is lazy" prop_lazy_hashNub
  ]
  where
    prop_ordNub :: [A] -> Property
    prop_ordNub xs = nub xs === ordNub xs

    prop_hashNub :: [A] -> Property
    prop_hashNub xs = nub xs === hashNub xs

    prop_ordNubOn :: Fun A B -> [A] -> Property
    prop_ordNubOn f' xs = nubBy ((==) `on` f) xs === ordNubOn f xs
      where
        f = applyFun f'

    prop_hashNubOn :: Fun A B -> [A] -> Property
    prop_hashNubOn f' xs = nubBy ((==) `on` f) xs === hashNubOn f xs
      where
        f = applyFun f'

    prop_lazy_ordNub :: Property
    prop_lazy_ordNub = take 3 (ordNub ('x' : 'y' : 'z' : 'z' : error "bottom")) === "xyz"

    prop_lazy_hashNub :: Property
    prop_lazy_hashNub = take 3 (hashNub ('x' : 'y' : 'z' : 'z' : error "bottom")) === "xyz"

-------------------------------------------------------------------------------
-- "Poly"
-------------------------------------------------------------------------------

newtype A = A Int
  deriving (Eq, Ord, Show, Hashable, Arbitrary, CoArbitrary)

instance Function A where
  function = functionMap coerce A

newtype B = B Int
  deriving (Eq, Ord, Show, Hashable, Arbitrary, CoArbitrary)

instance Function B where
  function = functionMap coerce B

newtype C = C Int
  deriving (Eq, Ord, Show, Hashable, Arbitrary, CoArbitrary)

instance Function C where
  function = functionMap coerce C

newtype K = K Int
  deriving (Eq, Ord, Show, Hashable, Arbitrary, CoArbitrary)

instance Function K where
  function = functionMap coerce K

newtype S = S Int
  deriving (Eq, Ord, Show, Hashable, Arbitrary, CoArbitrary)

instance Function S where
  function = functionMap coerce S