{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE TypeFamilies #-}

module Tests.Extra.ModInt64 (tests) where

import AtCoder.Extra.Math qualified as ACEM
import AtCoder.Extra.ModInt64 qualified as M
import AtCoder.ModInt qualified as M32
import Data.Semiring (Ring (..), Semiring (..), WrappedNum (..))
import Data.WideWord.Word128 (Word128 (..))
import Data.Word (Word64 (..))
import GHC.Exts (Proxy#, proxy#)
import GHC.TypeNats (KnownNat, natVal')
import Test.QuickCheck.Classes qualified as QCC
import Test.Tasty
import Test.Tasty.HUnit
import Test.Tasty.QuickCheck qualified as QC
import Tests.Util (laws)

deriving via (WrappedNum Word64) instance Semiring (M.ModInt64 a)

deriving via (WrappedNum Word64) instance Ring (M.ModInt64 a)

type M1 = 3

type M2 = 5

type M3 = 998244353

type M4 = 1000000007

type M5 = 4611686018427387847

instance M32.Modulus M1 where
  {-# INLINE isPrimeModulus #-}
  isPrimeModulus _ = True

  -- FIXME: wrong
  primitiveRootModulus _ = (-1)

instance M32.Modulus M2 where
  {-# INLINE isPrimeModulus #-}
  isPrimeModulus _ = True

  -- FIXME: wrong
  primitiveRootModulus _ = (-1)

instance (KnownNat a) => QC.Arbitrary (M.ModInt64 a) where
  arbitrary = M.new <$> QC.arbitrary

to128 :: (Integral a) => a -> Word128
to128 = fromIntegral

mulMod :: Int -> Int -> Int -> Int
mulMod m x y = fromIntegral . word128Lo64 $! (to128 (x `mod` m) * to128 (y `mod` m)) `mod` to128 m

unit_literal :: forall a. (KnownNat a) => Proxy# a -> TestTree
unit_literal proxy = testCase "literal" $ do
  let !m :: Int = fromIntegral $ natVal' proxy
  (@?= (0 `mod` m)) $ fromIntegral (0 :: M.ModInt64 a)
  (@?= ((-1) `mod` m)) $ fromIntegral (-1 :: M.ModInt64 a)
  (@?= (1 `mod` m)) $ fromIntegral (1 :: M.ModInt64 a)
  (@?= (m `mod` m)) $ fromIntegral (M.new @a m)
  (@?= ((m - 1) `mod` m)) $ fromIntegral (M.new @a (m - 1))
  (@?= ((m + 1) `mod` m)) $ fromIntegral (M.new @a (m + 1))

prop_addMod :: forall a. (KnownNat a) => Proxy# a -> Int -> Int -> QC.Property
prop_addMod proxy x y =
  let !m = fromIntegral $ natVal' proxy
      !res = M.val $ M.new @a x + M.new @a y
      !expected = (x + y) `mod` m
   in res QC.=== expected

prop_subMod :: forall a. (KnownNat a) => Proxy# a -> Int -> Int -> QC.Property
prop_subMod proxy x y =
  let !m = fromIntegral $ natVal' proxy
      !res = M.val $ M.new @a x - M.new @a y
      !expected = (x - y) `mod` m
   in res QC.=== expected

prop_mulMod :: forall a. (KnownNat a) => Proxy# a -> Int -> Int -> QC.Property
prop_mulMod proxy x y =
  let !m = fromIntegral $ natVal' proxy
      !res = M.val $ M.new @a x * M.new @a y
      !expected = mulMod m x y
   in res QC.=== expected

prop_powMod :: forall a. (KnownNat a) => Proxy# a -> Int -> QC.Positive Int -> QC.Property
prop_powMod proxy x (QC.Positive n) =
  let !m = fromIntegral $ natVal' proxy
      !res = M.val $ M.pow (M.new @a x) n
      !expected = ACEM.power (mulMod m) n (x `mod` m)
   in res QC.=== expected

prop_inv :: forall a. (KnownNat a) => M.ModInt64 a -> QC.Property
prop_inv x =
  M.val x
    /= 0
    QC.==> QC.counterexample (show x)
    $ QC.conjoin
      [ M.inv x * x QC.=== M.new 1,
        M.new 1 QC.=== M.inv x * x
      ]

prop_quotRem :: forall a. (M32.Modulus a) => Proxy# a -> Int -> QC.NonZero Int -> QC.Property
prop_quotRem _ x (QC.NonZero y) =
  (y `mod` m /= 0) QC.==>
    let (!resQ, !resR) = M.new @a x `quotRem` M.new @a y
        (!expQ, !expR) = M32.new @a x `quotRem` M32.new @a y
     in QC.conjoin
          [ M.val resQ QC.=== M32.val expQ,
            M.val resR QC.=== M32.val expR
          ]
  where
    !m = fromIntegral $ natVal' (proxy# @a)

prop_quotRem2 :: forall a. (KnownNat a) => Proxy# a -> Int -> QC.NonZero Int -> QC.Property
prop_quotRem2 _ x (QC.NonZero y) =
  let !a = M.new @a x
      !d = M.new @a y
   in d /= 0 QC.==>
        let (!q, !r) = a `quotRem` d
         in (d * q + r) QC.=== a

prop_eq :: forall a. (KnownNat a) => Proxy# a -> Word64 -> Word64 -> QC.Property
prop_eq _ x y = lhs QC.=== rhs
  where
    !lhs = M.new64 @a x == M.new64 @a y
    !m = fromIntegral $ natVal' (proxy# @a)
    !rhs = x `mod` m == y `mod` m

-- Cannot create list for unlifted types
--
-- {-# LANGUAGE ImpredicativeTypes #-}
-- modProps :: String -> [forall a. (KnownNat a) => Proxy# a -> QC.Property] -> TestTree
-- modProps title prop =
--   testGroup title $
--     map
--       (\proxy -> QC.testProperty (show (natVal' proxy)) (prop proxy))
--       [proxy# @M1, proxy# @M2, proxy# @M3, proxy# @M4, proxy# @M5]

tests :: [TestTree]
tests =
  [ testGroup
      "literal"
      [ unit_literal (proxy# @M1),
        unit_literal (proxy# @M2),
        unit_literal (proxy# @M3),
        unit_literal (proxy# @M4),
        unit_literal (proxy# @M5)
      ],
    testGroup
      "inv"
      [ QC.testProperty "1" (prop_inv @M1),
        QC.testProperty "2" (prop_inv @M2),
        QC.testProperty "3" (prop_inv @M3),
        QC.testProperty "4" (prop_inv @M4),
        QC.testProperty "5" (prop_inv @M5)
      ],
    testGroup
      "quotRem"
      [ QC.testProperty "1" (prop_quotRem (proxy# @M1)),
        QC.testProperty "2" (prop_quotRem (proxy# @M2)),
        QC.testProperty "3" (prop_quotRem (proxy# @M3)),
        QC.testProperty "4" (prop_quotRem (proxy# @M4))
        -- 64 bit
        -- QC.testProperty "5" (prop_quotRem (proxy# @M5))
      ],
    testGroup
      "quotRem2"
      [ QC.testProperty "1" (prop_quotRem2 (proxy# @M1)),
        QC.testProperty "2" (prop_quotRem2 (proxy# @M2)),
        QC.testProperty "3" (prop_quotRem2 (proxy# @M3)),
        QC.testProperty "4" (prop_quotRem2 (proxy# @M4)),
        QC.testProperty "5" (prop_quotRem2 (proxy# @M5))
      ],
    testGroup
      "eq"
      [ QC.testProperty "1" (prop_eq (proxy# @M1)),
        QC.testProperty "2" (prop_eq (proxy# @M2)),
        QC.testProperty "3" (prop_eq (proxy# @M3)),
        QC.testProperty "4" (prop_eq (proxy# @M4)),
        QC.testProperty "5" (prop_eq (proxy# @M5))
      ],
    testGroup
      "addMod"
      [ QC.testProperty "1" (prop_addMod (proxy# @M1)),
        QC.testProperty "2" (prop_addMod (proxy# @M2)),
        QC.testProperty "3" (prop_addMod (proxy# @M3)),
        QC.testProperty "4" (prop_addMod (proxy# @M4)),
        QC.testProperty "5" (prop_addMod (proxy# @M5))
      ],
    testGroup
      "subMod"
      [ QC.testProperty "1" (prop_subMod (proxy# @M1)),
        QC.testProperty "2" (prop_subMod (proxy# @M2)),
        QC.testProperty "3" (prop_subMod (proxy# @M3)),
        QC.testProperty "4" (prop_subMod (proxy# @M4)),
        QC.testProperty "5" (prop_subMod (proxy# @M5))
      ],
    testGroup
      "mulMod"
      [ QC.testProperty "1" (prop_mulMod (proxy# @M1)),
        QC.testProperty "2" (prop_mulMod (proxy# @M2)),
        QC.testProperty "3" (prop_mulMod (proxy# @M3)),
        QC.testProperty "4" (prop_mulMod (proxy# @M4)),
        QC.testProperty "5" (prop_mulMod (proxy# @M5))
      ],
    testGroup
      "powMod"
      [ QC.testProperty "1" (prop_powMod (proxy# @M1)),
        QC.testProperty "2" (prop_powMod (proxy# @M2)),
        QC.testProperty "3" (prop_powMod (proxy# @M3)),
        QC.testProperty "4" (prop_powMod (proxy# @M4)),
        QC.testProperty "5" (prop_powMod (proxy# @M5))
      ],
    testGroup
      "laws"
      [ laws @(M.ModInt64 M5)
          [ QCC.eqLaws,
            QCC.numLaws,
            QCC.integralLaws,
            QCC.ordLaws,
            QCC.enumLaws,
            QCC.boundedEnumLaws,
            QCC.primLaws,
            QCC.semiringLaws,
            QCC.ringLaws,
            QCC.showReadLaws
          ]
      ]
  ]