{-# LANGUAGE OverloadedLists #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Props.Mat where

import Data.Kind (Type)
import Data.Type.Nat (Nat (..), Nat0, Nat1, Nat3, SNat (..), SNatI, snat, snatToNat)
import Data.Vec.Lazy (Vec (..), repeat)
import Test.Falsify.Generator (elem)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.Falsify (testProperty)
import Prelude hiding (elem, repeat)

import Proarrow.Category.Instance.Mat (App, Mat (..), MatK (..))
import Proarrow.Category.Monoidal.CompactClosed (dimension)
import Proarrow.Core (CAT, type (+->))
import Proarrow.Profunctor.Representable (Rep)

import Proarrow.Testing
  ( GenTotal (..)
  , Testable (..)
  , TestableProfunctor
  , TestableType (..)
  , TestingEqShow (..)
  , expect
  , genSomeDef
  , invmap
  , oneElem
  , pattern GenNonEmpty
  )
import Proarrow.Testing.Laws
import Props.Hask ()

-- | The one entry of a @1x1@ matrix, which is what an endo-arrow on 'Unit' is.
scalar :: Mat (M Nat1 :: MatK Int) (M Nat1) -> Int
scalar (Mat ((x ::: VNil) ::: VNil)) = x

test :: TestTree
test =
  testGroup
    "Matrix"
    [ testCategory @(MatK Int)
    , testDagger @(MatK Int)
    , testTerminalObject @(MatK Int)
    , testInitialObject @(MatK Int)
    , testBinaryProducts_ @(MatK Int)
    , testBinaryCoproducts_ @(MatK Int)
    , testHypergraph_ @(MatK Int)
    , testMonoidal_ @(MatK Int)
    , testSymMonoidal_ @(MatK Int)
    , testDistributive_ @(MatK Int)
    , testClosed_ @(MatK Int)
    , testStarAutonomous_ @(MatK Int)
    , testCompactClosed_ @(MatK Int)
    , testTraced_ @(MatK Int)
    , testCopyDiscard_ @(MatK Int)
    , testGroup "App functor" [testProfunctor @(Rep App :: MatK Int +-> Type)]
    , -- the trace of an identity, and the one place the traced object could silently be dropped
      testProperty "dimension counts the object" $
        expect "dimensions 0, 1, 3" [0, 1, 3] (map scalar [dimension @(M Nat0), dimension @(M Nat1), dimension @(M Nat3)])
    , testEqualizers_ @(MatK Rational)
    , testCoequalizers_ @(MatK Rational)
    , testEpiMonoFactorization_ @(MatK Rational)
    , testPullbacks_ @(MatK Rational)
    , testPushouts_ @(MatK Rational)
    ]

type TestableNum n = (Num n, Eq n, Show n, TestableType n)

instance (TestableNum n) => Testable (MatK n) where
  showOb @(M a) = show $ snatToNat $ snat @a
  genSome = genSomeDef @'[M Z, M (S Z), M (S (S Z)), M (S (S (S Z)))]

instance (TestOb (a :: MatK n), TestOb b, TestableNum n) => TestableType (Mat a b) where
  gen = invmap Mat unMat gen
instance (TestOb (a :: MatK n), TestOb b, TestableNum n) => TestingEqShow (Mat a b) where
  eqP (Mat l) (Mat r) = pure $ l == r
  showP (Mat m) = show m
instance (TestableNum n) => TestableProfunctor (Mat :: CAT (MatK n))

instance (Eq a, Show a) => TestingEqShow (Vec n a)
instance (Eq a, Show a, TestableType a, SNatI n) => TestableType (Vec n a) where
  gen = case gen of
    GenEmpty absurd -> case snat @n of
      SZ -> oneElem VNil
      SS -> GenEmpty \(a ::: _) -> absurd a
    GenNonEmpty g -> GenNonEmpty $ sequence (repeat @n g)

instance TestingEqShow Int
instance TestableType Int where
  gen = GenNonEmpty $ liftA2 (*) (elem [1, -1]) (elem [0 .. 9])

instance TestingEqShow Rational
instance TestableType Rational where
  gen = GenNonEmpty $ liftA2 (*) (elem [1, -1]) (fromInteger <$> elem [0 .. 9])

instance TestableProfunctor (Rep App :: MatK Int +-> Type)