{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}

module Test.Core.LinearAlgebra where

import Data.Array.Accelerate hiding (Ord, Eq, map, (<=))
import Data.Array.Accelerate.Interpreter qualified as Naive
import Data.Array.Accelerate qualified as A
import Data.Massiv.Array qualified as Massiv
import Data.Proxy
import Data.Scientific
import Gargantext.Core.LinearAlgebra qualified as LA
import Gargantext.Core.Methods.Matrix.Accelerate.Utils qualified as A
import Gargantext.Core.Methods.Matrix.Accelerate.Utils qualified as Legacy
import Gargantext.Core.Methods.Similarities.Accelerate.Distributional qualified as Legacy
import Gargantext.Orphans.Accelerate (sliceArray)
import Prelude hiding ((^))
import Test.Tasty
import Test.Tasty.QuickCheck


--
-- Utility types and functions
--

newtype SquareMatrix a = SquareMatrix { _SquareMatrix :: Matrix a }
  deriving newtype (Show, Eq)

instance (Elt a, Show a, Prelude.Num a, Ord a, Arbitrary a) => Arbitrary (SquareMatrix a) where
  arbitrary = do
    x <- choose (1,30)
    let sh = Z :. x :. x
    SquareMatrix . A.fromList sh <$> vectorOf (x*x) arbitrary
  shrink = map (SquareMatrix) . sliceArray . _SquareMatrix

type TermDivNanShape = Z :. Int :. Int

twoByTwo :: SquareMatrix Int
twoByTwo = SquareMatrix $ fromList (Z :. 2 :. 2) (Prelude.replicate 4 5)

testMatrix_01 :: SquareMatrix Int
testMatrix_01 = SquareMatrix $ fromList (Z :. 14 :. 14) $
  [  30,  36, -36, -16,   0,   7,  34,  -7,   5,  -4,   0,  21,   6, -35,
      0, -31,  20, -15, -22,  -7, -22, -37, -29, -29,  23, -31, -29, -23,
    -24, -29,  19,  -6,  16,   7,  15, -27, -27, -30,  -9, -33,  18, -23,
      7, -36,  12,  26, -17,  -3,  -2, -15,  -4,  26,  24,   9,  -4,   4,
     32,  28,  -2, -10,  34,  -3,  20,  -9, -22,  20, -26,  34,  18, -21,
      7, -12,  12,  -2,  36,  10,  34, -37,  13,  -9, -28,  34,  33, -18,
     -4, -32,  -1,  29,  29, -28,  24,  28,  35,  19,   8, -18,  25, -35,
    -14,  -4, -24,  -1,   7,  34, -37, -28, -12, -32,  -5, -23,  27,  33,
    -36, -28,  21, -29,  -2, -26,  -4, -31, -26, -21,  33, -11, -33,  20,
     25,  14,   5,  -7,   5,  24,  37,   1,  -3,  23,  25, -16,  17,   5,
    -35,  36,  -2,  -2,   1, -14,  34, -30, -10,  12,  25,  21,   0,  34,
     17,  -1,  20, -19,  15,  20,  -5, -30, -35, -13,   5,  17, -10, -19,
    -34, -11, -18,  26, -29, -28,   0,   3,  23,  -6,  36,   4,  16,  28,
     13, -37, -16,   2,   7, -13,  21, -10, -33, -33, -26, -19,  -1,  29]

testMatrix_02 :: SquareMatrix Int
testMatrix_02 = SquareMatrix $ fromList (Z :. 7 :. 7) $
  [  30,  36, -36, -16,   0,   7,  34,
      0, -31,  20, -15, -22,  -7, -22,
    -24, -29,  19,  -6,  16,   7,  15,
      7, -36,  12,  26, -17,  -3,  -2,
     32,  28,  -2, -10,  34,  -3,  20,
      7, -12,  12,  -2,  36,  10,  34,
     13, -37, -16,   2,   7, -13,  21]

testMatrix_03 :: SquareMatrix Int
testMatrix_03 = SquareMatrix $ fromList (Z :. 11 :. 11) $
  [  1, -1,  1,  0,  1, -1, 0,  1,  1,  0,  0,
     1,  1,  1,  1,  1,  0, 1, -1,  1,  0,  0,
    -1,  1,  0, -1,  0, -1, 0,  1,  0, -1,  0,
     1,  1,  1, -1, -1,  0, 1, -1,  0,  0, -1,
    -1,  1, -1, -1,  0,  1, 1,  1, -1, -1, -1,
     1,  1,  0, -1, -1, -1, 1,  0,  1, -1, -1,
    -1,  1,  0, -1,  1, -1, 0,  1, -1, -1, -1,
     1,  1, -1,  1,  1,  0, 1, -1,  1, -1,  1,
    -1, -1,  0,  1,  1,  0, 1,  1, -1,  1,  0,
     1,  1,  0, -1,  1, -1, 1,  0,  1,  0, -1,
     1,  1, -1,  0, -1, -1, 1,  0,  1,  0, -1]

testMatrix_04 :: SquareMatrix Int
testMatrix_04 = SquareMatrix $ fromList (Z :. 8 :. 8) $
  [  3, -1,  0, 1, -1,  1,  1, -3,
    -2, -2,  2, 1,  1, -2,  1, -1,
    -2, -3, -1, 1,  1, -3, -2, -1,
     1, -2,  2, 0,  1,  0,  2,  0,
    -1, -3, -1, 3, -3,  0, -1,  2,
     0,  0, -3, 3, -1, -2, -1,  1,
    -2,  1, -1, 2,  1, -1, -2,  0,
    -2,  2,  1, 1,  1,  0,  2, -3]
--
-- Main test runner
--

tests :: TestTree
tests = testGroup "LinearAlgebra" [
      testProperty "termDivNan"     compareTermDivNan
    , testProperty "diag"           compareDiag
    , testProperty "sumRows"        compareSumRows
    , testProperty "matMaxMini"     compareMatMaxMini
    , testProperty "sumM_go"        compareSumM_go
    , testProperty "sumMin_go"      compareSumMin_go
    , testProperty "matrixEye"      compareMatrixEye
    , testProperty "diagNull"       compareDiagNull
    , testGroup "distributional" [
        testProperty "reference implementation roundtrips" compareDistributionalImplementations
      , testProperty "2x2" (compareDistributional (Proxy @Double) twoByTwo)
      , testProperty "7x7" (compareDistributional (Proxy @Double) testMatrix_02)
      , testProperty "14x14" (compareDistributional (Proxy @Double) testMatrix_01)
      , testProperty "roundtrips" (compareDistributional (Proxy @Double))
      ]
    , testGroup "logDistributional2" [
       testProperty "2x2" (compareLogDistributional2 (Proxy @Double) twoByTwo)
      , testProperty "7x7" (compareLogDistributional2 (Proxy @Double) testMatrix_02)
      , testProperty "8x8" (compareLogDistributional2 (Proxy @Double) testMatrix_04)
      , testProperty "11x11" (compareLogDistributional2 (Proxy @Double) testMatrix_03)
      , testProperty "14x14" (compareLogDistributional2 (Proxy @Double) testMatrix_01)
      ,testProperty "roundtrips" (compareLogDistributional2 (Proxy @Double))
    ]
    ]

--
-- Tests
--

compareTermDivNan :: (Array TermDivNanShape Double)
                  -> (Array TermDivNanShape Double)
                  -> Property
compareTermDivNan i1 i2
  = let massiv     = LA.termDivNan @Massiv.U (LA.accelerate2MassivMatrix i1) (LA.accelerate2MassivMatrix i2)
        accelerate = Naive.run (Legacy.termDivNan (use i1) (use i2))
    in accelerate === LA.massiv2AccelerateMatrix massiv

compareDiag :: SquareMatrix Int -> Property
compareDiag (SquareMatrix i1)
  = let massiv     = LA.diag (LA.accelerate2MassivMatrix i1)
        accelerate =  Naive.run (Legacy.diag (use i1))
    in accelerate === LA.massiv2AccelerateVector massiv

compareSumRows :: Array (Z :. Int :. Int :. Int) Int -> Property
compareSumRows i1
  = let massiv     = LA.sumRows @_ @Massiv.U (LA.accelerate2Massiv3DMatrix i1)
        massiv'    = LA.sumRowsReferenceImplementation @Massiv.U (LA.accelerate2Massiv3DMatrix i1)
        accelerate =  Naive.run (A.sum (use i1))
    in counterexample "sumRows and reference implementation do not agree" (massiv === massiv') .&&.
       accelerate === LA.massiv2AccelerateMatrix massiv

compareDistributionalImplementations :: SquareMatrix Int -> Property
compareDistributionalImplementations (SquareMatrix i1) =
  let ma = LA.accelerate2MassivMatrix i1
  in LA.distributional @Massiv.U @Double ma === LA.distributionalReferenceImplementation ma

compareDistributional :: forall e.
                      ( Eq e
                      , Show e
                      , FromIntegral Int e
                      , Prelude.RealFloat e
                      , Massiv.Unbox e
                      , A.Ord e
                      , Ord e
                      , Prelude.Fractional (Exp e)
                      , Prelude.Fractional e
                      , Monoid e
                      ) => Proxy e
                      -> SquareMatrix Int
                      -> Property
compareDistributional Proxy (SquareMatrix i1)
  = let massiv     = Massiv.computeAs Massiv.B $ LA.distributional @_ @e (LA.accelerate2MassivMatrix i1)
        accelerate = Legacy.distributionalWith Naive.run i1
        expected   = map conv (A.toList accelerate)
        actual     = map conv (mconcat (Massiv.toLists2 massiv))
    in counterexample "size not equal" (Prelude.length expected === Prelude.length actual) .&&. expected === actual
  where
    conv :: e -> Scientific
    conv = fromFloatDigits

compareLogDistributional2 :: forall e.
                          ( Eq e
                          , Show e
                          , FromIntegral Int e
                          , Prelude.RealFloat e
                          , Massiv.Unbox e
                          , A.Ord e
                          , Ord e
                          , Prelude.Fractional (Exp e)
                          , Prelude.Fractional e
                          , Prelude.Floating (Exp e)
                          , Prelude.Floating e
                          , Monoid e
                          ) => Proxy e
                          -> SquareMatrix Int
                          -> Property
compareLogDistributional2 Proxy (SquareMatrix i1)
  = let massiv     = Massiv.computeAs Massiv.B $ LA.logDistributional2 @_ @e (LA.accelerate2MassivMatrix i1)
        accelerate = Legacy.logDistributional2With Naive.run i1
        expected   = map conv (A.toList accelerate)
        actual     = map conv (mconcat (Massiv.toLists2 massiv))
    in counterexample "size not equal" (Prelude.length expected === Prelude.length actual) .&&. expected === actual
  where
    conv :: e -> Scientific
    conv = fromFloatDigits

compareMatMaxMini :: SquareMatrix Int -> Property
compareMatMaxMini (SquareMatrix i1)
  = let massiv     = LA.matMaxMini (LA.accelerate2MassivMatrix i1)
        accelerate =  Naive.run (A.matMaxMini (use i1))
    in accelerate === LA.massiv2AccelerateMatrix massiv

compareSumMin_go :: SquareMatrix Int -> Property
compareSumMin_go (SquareMatrix i1)
  = let massiv     = LA.sumMin_go (A.dim i1) (LA.accelerate2MassivMatrix i1)
        accelerate =  Naive.run (Legacy.sumMin_go (A.dim i1) (use i1))
    in accelerate === LA.massiv2AccelerateMatrix massiv

compareSumM_go :: SquareMatrix Int -> Property
compareSumM_go (SquareMatrix i1)
  = let massiv     = LA.sumM_go (A.dim i1) (LA.accelerate2MassivMatrix i1)
        accelerate =  Naive.run (Legacy.sumM_go (A.dim i1) (use i1))
    in accelerate === LA.massiv2AccelerateMatrix massiv

compareMatrixEye :: Positive Int -> Property
compareMatrixEye (getPositive -> n)
  = let massiv     = Massiv.compute @Massiv.U $ LA.matrixEye @Int n
        accelerate = Naive.run (Legacy.matrixEye n)
    in accelerate === LA.massiv2AccelerateMatrix massiv

compareDiagNull :: SquareMatrix Int -> Property
compareDiagNull (SquareMatrix i1)
  = let massiv     = Massiv.compute @Massiv.U $ LA.diagNull (A.dim i1) (LA.accelerate2MassivMatrix i1)
        accelerate =  Naive.run (Legacy.diagNull (A.dim i1) (use i1))
    in accelerate === LA.massiv2AccelerateMatrix massiv
