{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE NumericUnderscores  #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}

module Main (main) where

-- base
import Prelude hiding
  ( Num(..), (^) )
import Data.Foldable
  ( toList )
import Data.List.NonEmpty
  ( NonEmpty(..) )
import Data.Maybe
  ( catMaybes )
import Data.Traversable
  ( for )
import Unsafe.Coerce
  ( unsafeCoerce )

-- brush-strokes
import Math.Algebra.Dual
import Math.Linear
import Math.Module
import Math.Monomial
  ( multiSubsetSum, multiSubsetsSum
  , MonomialBasisQ ( monTabulateQ, monIndexQ )
  )
import Math.Ring

-- hspray
import Math.Algebra.Hspray
  ( Spray )
import qualified Math.Algebra.Hspray as Spray

-- falsify
import Test.Tasty.Falsify
import qualified Test.Falsify.Generator as Falsify
  ( Gen )
import qualified Test.Falsify.Generator as Falsify.Gen
import Test.Falsify.Predicate
  ( (.$) )
import qualified Test.Falsify.Predicate as Falsify.Prop
import qualified Test.Falsify.Property as Falsify
  ( Property
  , assert
  , discard
  , gen, genWith
  )
import qualified Test.Falsify.Range as Falsify

-- tasty
import qualified Test.Tasty as Tasty

-- unordered-containers
import qualified Data.HashMap.Lazy as HashMap

--------------------------------------------------------------------------------


main :: IO ()
main =
  Tasty.defaultMain $
    Tasty.testGroup "brush-strokes property tests"
    [ Tasty.testGroup "Automatic differentiation"
      [ Tasty.testGroup "Monomial basis"
        [ testProperty "Round trip D33" testMonomialBasisQD33
        ]
      , Tasty.testGroup "Monomials"
        [ Tasty.testGroup "multiSubsetSum"
          [ testProperty "multiSubsetSum valid" testMultiSubsetSumValid
          , testProperty "multiSubsetSum exhaustive" testMultiSubsetSumExhaustive
          ]
        -- , Tasty.testGroup "multiSubsetsSum"
        --   [ testProperty "multiSubsetsSum exhaustive" testMultiSubsetsSumExhaustive
        --   ]
        ]
      , Tasty.testGroup "chainRule1NQ"
        [ testProperty "chainRule1NQ_1" testChainRule1NQ_1
        , testProperty "chainRule1NQ_2" testChainRule1NQ_2
        , testProperty "chainRule1NQ_3" testChainRule1NQ_3
        ]
      ]
    ]

-- | Check that the 'multiSubsetSum' function returns valid answers, i.e.
-- all returned multisubsets have the desired size and sum.
testMultiSubsetSumValid :: Falsify.Property ()
testMultiSubsetSumValid = do
  rg  <- Falsify.genWith (\ rg  -> Just $ "range = " ++ show rg ) $ Falsify.Gen.inRange $ Falsify.between ( 1, 6 )
  sz  <- Falsify.genWith (\ sz  -> Just $  "size = " ++ show sz ) $ Falsify.Gen.inRange $ Falsify.between ( 0, 20 )
  tot <- Falsify.genWith (\ tot -> Just $   "tot = " ++ show tot) $ Falsify.Gen.inRange $ Falsify.between ( sz, sz * rg )
  let range = [ 1 .. rg ]
      mss = multiSubsetSum sz tot range
  case mss of
    [] -> Falsify.discard
    r:rs -> do
      ms <- Falsify.gen $ Falsify.Gen.elem ( r :| rs )
      Falsify.assert
         $ Falsify.Prop.eq
        .$ ("(sz, tot)", (sz, tot) )
        .$ ("computed (sz, tot)", (size ms, total ms))
  where
    size, total :: [ ( Word, Word ) ] -> Word
    size [] = 0
    size ((_,n):ins) = n + size ins
    total [] = 0
    total ((i,n):ins) = i * n + total ins

-- | Check that the 'multiSubsetSum' function returns all multisubsets of
-- the given set, by generating a random multisubset, computing its size, and
-- checking it belongs to the output of the 'multiSubsetSum' function.
testMultiSubsetSumExhaustive :: Falsify.Property ()
testMultiSubsetSumExhaustive = do
  rg <- Falsify.genWith (\ rg -> Just $ "range = " ++ show rg) $ Falsify.Gen.inRange $ Falsify.between ( 1, 6 )
  sz <- Falsify.genWith (\ sz -> Just $  "size = " ++ show sz) $ Falsify.Gen.inRange $ Falsify.between ( 0, 10 )
  let range = [ 1 .. rg ]
  (multiSubset, tot) <- Falsify.genWith (\ ms -> Just $ "multisubset = " ++ show ms) $ genMultiSubset range sz
  Falsify.assert
     $ Falsify.Prop.elem
    .$ ("all multisubsets", multiSubsetSum sz tot range )
    .$ ("random multisubset", multiSubset)

genMultiSubset :: [ Word ] -> Word -> Falsify.Gen ( [ ( Word, Word ) ] , Word )
genMultiSubset [i] sz =
  return $
    if sz == 0
    then ( [], 0 )
    else ( [ ( i, sz ) ], i * sz )
genMultiSubset (i:is) sz = do
  nb <- Falsify.Gen.inRange $ Falsify.between ( 0, sz )
  (rest, tot) <- genMultiSubset is ( sz - nb )
  return $ ( if nb == 0 then rest else ( i, nb ) : rest, tot + nb * i )
genMultiSubset [] _ = error "impossible"

coerceVec1 :: [ a ] -> Vec n a
coerceVec1 = unsafeCoerce

coerceVec2 :: Vec n a -> [ a ]
coerceVec2 = toList

-- | Check that the 'multiSubsetSums' function returns all collections of
-- multisubsets of the given set (see 'testMultiSubsetSumExhaustive').
testMultiSubsetsSumExhaustive :: Falsify.Property ()
testMultiSubsetsSumExhaustive = do
  rg <- Falsify.genWith (\ rg -> Just $ "range = " ++ show rg) $ Falsify.Gen.inRange $ Falsify.between ( 1, 5 )
  let range = [ 1 .. rg ]
  n  <- Falsify.genWith (\ n  -> Just $     "n = " ++ show n ) $ Falsify.Gen.inRange $ Falsify.between ( 1, 10 )
  multiSubsets <- for ( [ 0 .. n - 1 ] :: [ Word ] ) \ i -> do
    sz <- Falsify.gen $ Falsify.Gen.inRange $ Falsify.between ( 0, 5 )
    ( ms, tot ) <- Falsify.genWith ( \ ms -> Just $ "ms_" ++ show i ++ " = " ++ show ms ) $ genMultiSubset range sz
    return ( ms, sz, tot )
  let mss = map ( \ (ms, _,_) -> ms ) multiSubsets
      szs = map ( \ (_,sz,_) -> sz) multiSubsets
      tot = sum $ map ( \(_,_,t) -> t) multiSubsets
  Falsify.assert
     $ Falsify.Prop.elem
    .$ ("all multisubsets", map coerceVec2 $ multiSubsetsSum range tot $ coerceVec1 szs )
    .$ ("random multisubset", mss)


testRoundTrip
  :: ( Show a, Eq a )
  => Falsify.Gen a
  -> ( a -> a )
  -> Falsify.Property ()
testRoundTrip g roundTrip = do
  d <- Falsify.gen g
  Falsify.assert
     $ Falsify.Prop.eq
    .$ ("value", d )
    .$ ("round tripped", roundTrip d )

testMonomialBasisQD33 :: Falsify.Property ()
testMonomialBasisQD33 =
  testRoundTrip genD33 \ d -> $$( monTabulateQ \ mon -> monIndexQ [|| d ||] mon )
  where
    genD33 :: Falsify.Gen ( D3𝔸3 Double )
    genD33 =
      D33 <$> (unT <$> g)
          <*> g <*> g <*> g
          <*> g <*> g <*> g <*> g <*> g <*> g
          <*> g <*> g <*> g <*> g <*> g <*> g <*> g <*> g <*> g <*> g

    g :: Falsify.Gen ( T Double )
    g = T . fromIntegral <$> Falsify.Gen.inRange ( Falsify.withOrigin ( -100, 100 ) ( 0 :: Int ) )

-- | Test the Faà di Bruno formula on polynomials, with a composition
-- \( g(f_1(x), f_2(x), .., f_n(x)) \).
testChainRule1NQ_1 :: Falsify.Property ()
testChainRule1NQ_1 = do
  f <- genSpray "f" 1
  g <- genSpray "g" 1
  let gof_spray = Spray.composeSpray g [f]
      gof_chain =
        chain @_ @3 @( ℝ 1 ) ( ℝ1 <$> fromSpray @3 @( ℝ 1 ) f ) ( fromSpray @3 @( ℝ 1 ) g )
  Falsify.assert
     $ Falsify.Prop.eq
    .$ ("direct", fromSpray @3 @( ℝ 1 ) gof_spray )
    .$ ("chain rule", gof_chain )

-- | Test the Faà di Bruno formula on polynomials, with a composition
-- \( g(f_1(x), f_2(x), .., f_n(x)) \).
testChainRule1NQ_2 :: Falsify.Property ()
testChainRule1NQ_2 = do
  f1 <- genSpray "f1" 1
  f2 <- genSpray "f2" 1
  g  <- genSpray "g"  2
  let gof_spray = Spray.composeSpray g [f1, f2]
      f = ℝ2 <$> fromSpray @3 @( ℝ 1 ) f1
             <*> fromSpray @3 @( ℝ 1 ) f2
      gof_chain =
        chain @_ @3 @( ℝ 2 ) f ( fromSpray @3 @( ℝ 2 ) g )
  Falsify.assert
     $ Falsify.Prop.eq
    .$ ("direct", fromSpray @3 @( ℝ 1 ) gof_spray )
    .$ ("chain rule", gof_chain )

-- | Test the Faà di Bruno formula on polynomials, with a composition
-- \( g(f_1(x), f_2(x), .., f_n(x)) \).
testChainRule1NQ_3 :: Falsify.Property ()
testChainRule1NQ_3 = do
  f1 <- genSpray "f1" 1
  f2 <- genSpray "f2" 1
  f3 <- genSpray "f3" 1
  g  <- genSpray "g"  3
  let gof_spray = Spray.composeSpray g [f1, f2, f3]
      f = ℝ3 <$> fromSpray @3 @( ℝ 1 ) f1
             <*> fromSpray @3 @( ℝ 1 ) f2
             <*> fromSpray @3 @( ℝ 1 ) f3
      gof_chain =
        chain @_ @3 @( ℝ 3 ) f ( fromSpray @3 @( ℝ 3 ) g )
  Falsify.assert
     $ Falsify.Prop.eq
    .$ ("direct", fromSpray @3 @( ℝ 1 ) gof_spray )
    .$ ("chain rule", gof_chain )

class FromSpray v where
  varFn :: Int -> v
  linFn :: v -> Int -> Double

instance FromSpray ( ℝ 1 ) where
  varFn = \case
    0 -> ℝ1 1
    i -> error $ "fromSpray in 1d but variable " ++ show i
  linFn ( ℝ1 x ) = \case
    0 -> x
    i -> error $ "fromSpray in 1d but variable " ++ show i
instance FromSpray ( ℝ 2 ) where
  varFn = \case
    0 -> ℝ2 1 0
    1 -> ℝ2 0 1
    i -> error $ "fromSpray in 2d but variable " ++ show i
  linFn ( ℝ2 x y ) = \case
    0 -> x
    1 -> y
    i -> error $ "fromSpray in 2d but variable " ++ show i
instance FromSpray ( ℝ 3 ) where
  varFn = \case
    0 -> ℝ3 1 0 0
    1 -> ℝ3 0 1 0
    2 -> ℝ3 0 0 1
    i -> error $ "fromSpray in 3d but variable " ++ show i
  linFn ( ℝ3 x y z ) = \case
    0 -> x
    1 -> y
    2 -> z
    i -> error $ "fromSpray in 3d but variable " ++ show i

genSpray :: String -> Word -> Falsify.Property ( Spray Double )
genSpray lbl nbVars = Falsify.genWith (\ p -> Just $ lbl ++ " = " ++ Spray.prettySpray show "x" p) $ do
  deg <- Falsify.Gen.inRange $ Falsify.between ( 0, 10 )
  let mons = allMonomials deg nbVars
  coeffs <-
    for mons $ \ mon -> do
      if all (== 0) mon
      then return Nothing
      else do
        nonZero <- Falsify.Gen.bool False
        if nonZero
        then return Nothing
        else do
          -- Just use (small) integral values in tests for now,
          -- to avoid errors arising from rounding.
          c <- Falsify.Gen.inRange $ Falsify.withOrigin ( -100, 100 ) ( 0 :: Int )
          return $ Just ( map fromIntegral mon, fromIntegral c )
  return $ Spray.fromList $ catMaybes coeffs

allMonomials :: Word -> Word -> [ [ Word ] ]
allMonomials k _ | k < 0 = []
allMonomials _ 0 = [ [] ]
allMonomials 0 n = [ replicate ( fromIntegral n ) 0 ]
allMonomials k n = [ i : is | i <- reverse [ 0 .. k ], is <- allMonomials ( k - i ) ( n - 1 ) ]

-- | Convert a multivariate polynomial from the @hspray@ library to the dual algebra.
fromSpray
  :: forall k v
  .  ( HasChainRule Double k v
     , Module Double (T v)
     , Applicative ( D k v )
     , Ring ( D k v Double )
     , FromSpray v
     )
  => Spray Double
  -> D k v Double
fromSpray coeffs = HashMap.foldlWithKey' addMonomial ( konst @Double @k @v $ HashMap.lookupDefault 0 (Spray.Powers mempty 0) coeffs ) coeffs
  where
    addMonomial :: D k v Double -> Spray.Powers -> Double -> D k v Double
    addMonomial a xs c = a + monomial c ( toList $ Spray.exponents xs )
    monomial :: Double -> [ Int ] -> D k v Double
    monomial _ [] = konst @Double @k @v 0
    monomial c is = fmap ( c * ) $ go 0 is
    go :: Int -> [ Int ] -> D k v Double
    go _ [] = konst @Double @k @v 1
    go d (i : is) = pow d i * go ( d + 1 ) is
    pow :: Int -> Int -> D k v Double
    pow _ 0 = konst @Double @k @v 1
    pow d i = linearD @Double @k @v ( \ x -> linFn @v x d ) ( unT origin :: v ) ^ ( fromIntegral i )