{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE TypeApplications #-}
module Main where

import Control.Exception (bracket)
import Data.Array.RankedS qualified as RS
import Data.Vector.Storable qualified as VS
import Numeric.LinearAlgebra qualified as LA
import Test.Tasty.Bench

import Data.Array.Nested
import Data.Array.Nested.Internal.Mixed (mliftPrim, mliftPrim2)
import Data.Array.Nested.Internal.Ranked (liftRanked1, liftRanked2)
import qualified Data.Array.Strided.Arith.Internal as Arith


enableMisc :: Bool
enableMisc = False

bgroupIf :: Bool -> String -> [Benchmark] -> Benchmark
bgroupIf True = bgroup
bgroupIf False = \name _ -> bgroup name []


main :: IO ()
main = do
  let enable = False
  bracket (Arith.statisticsEnable enable)
          (\() -> do Arith.statisticsEnable False
                     when enable $ Arith.statisticsPrintAll)
          (\() -> main_tests)

main_tests :: IO ()
main_tests = defaultMain
  [bgroup "compare" tests_compare
  ,bgroup "orthotope"
    [bench "normalize [1e6]" $
      let n = 1_000_000
      in nf (\a -> RS.normalize a)
            (RS.rev [0] (RS.iota @Double n))
    ,bench "normalize noop [1e6]" $
      let n = 1_000_000
      in nf (\a -> RS.normalize a)
            (RS.rev [0] (RS.rev [0] (RS.iota @Double n)))
    ]
  ,bgroupIf enableMisc "misc"
    [let n = 1000
         k = 1000
     in bgroup ("fusion [" ++ show k ++ "]*" ++ show n) $
      [bench "sum (concat)" $
        nf (\as -> VS.sum (VS.concat as))
           (replicate n (VS.enumFromTo (1::Int) k))
      ,bench "sum (force (concat))" $
        nf (\as -> VS.sum (VS.force (VS.concat as)))
              (replicate n (VS.enumFromTo (1::Int) k))]
    ,bgroup "concat"
      [bgroup "N"
        [bgroup "hmatrix"
          [bench ("LA.vjoin [500]*1e" ++ show ni) $
            let n = 10 ^ ni
                k = 500
            in nf (\as -> LA.vjoin as)
                  (replicate n (VS.enumFromTo (1::Int) k))
          | ni <- [1::Int ..5]]
        ,bgroup "vectorStorable"
          [bench ("VS.concat [500]*1e" ++ show ni) $
            let n = 10 ^ ni
                k = 500
            in nf (\as -> VS.concat as)
                  (replicate n (VS.enumFromTo (1::Int) k))
          | ni <- [1::Int ..5]]
        ]
      ,bgroup "K"
        [bgroup "hmatrix"
          [bench ("LA.vjoin [1e" ++ show ki ++ "]*500") $
            let n = 500
                k = 10 ^ ki
            in nf (\as -> LA.vjoin as)
                  (replicate n (VS.enumFromTo (1::Int) k))
          | ki <- [1::Int ..5]]
        ,bgroup "vectorStorable"
          [bench ("VS.concat [1e" ++ show ki ++ "]*500") $
            let n = 500
                k = 10 ^ ki
            in nf (\as -> VS.concat as)
                  (replicate n (VS.enumFromTo (1::Int) k))
          | ki <- [1::Int ..5]]
        ]
      ]
    ]
  ]

tests_compare :: [Benchmark]
tests_compare =
  let n = 1_000_000 in
  [bgroup "Num"
    [bench "sum(+) Double [1e6]" $
      nf (\(a, b) -> runScalar (rsumOuter1 (liftRanked2 (mliftPrim2 (+)) a b)))
         (riota @Double n, riota n)
    ,bench "sum(*) Double [1e6]" $
      nf (\(a, b) -> runScalar (rsumOuter1 (liftRanked2 (mliftPrim2 (*)) a b)))
         (riota @Double n, riota n)
    ,bench "sum(/) Double [1e6]" $
      nf (\(a, b) -> runScalar (rsumOuter1 (liftRanked2 (mliftPrim2 (/)) a b)))
         (riota @Double n, riota n)
    ,bench "sum(**) Double [1e6]" $
      nf (\(a, b) -> runScalar (rsumOuter1 (liftRanked2 (mliftPrim2 (**)) a b)))
         (riota @Double n, riota n)
    ,bench "sum(sin) Double [1e6]" $
      nf (\a -> runScalar (rsumOuter1 (liftRanked1 (mliftPrim sin) a)))
         (riota @Double n)
    ,bench "sum Double [1e6]" $
      nf (\a -> runScalar (rsumOuter1 a))
         (riota @Double n)
    ]
  ,bgroup "NumElt"
    [bench "sum(+) Double [1e6]" $
      nf (\(a, b) -> runScalar (rsumOuter1 (a + b)))
         (riota @Double n, riota n)
    ,bench "sum(*) Double [1e6]" $
      nf (\(a, b) -> runScalar (rsumOuter1 (a * b)))
         (riota @Double n, riota n)
    ,bench "sum(/) Double [1e6]" $
      nf (\(a, b) -> runScalar (rsumOuter1 (a / b)))
         (riota @Double n, riota n)
    ,bench "sum(**) Double [1e6]" $
      nf (\(a, b) -> runScalar (rsumOuter1 (a ** b)))
         (riota @Double n, riota n)
    ,bench "sum(sin) Double [1e6]" $
      nf (\a -> runScalar (rsumOuter1 (sin a)))
         (riota @Double n)
    ,bench "sum Double [1e6]" $
      nf (\a -> runScalar (rsumOuter1 a))
         (riota @Double n)
    ,bench "sum(*) Double [1e6] stride 1; -1" $
      nf (\(a, b) -> runScalar (rsumOuter1 (a * b)))
         (riota @Double n, rrev1 (riota n))
    ,bench "dotprod Float [1e6]" $
      nf (\(a, b) -> rdot a b)
         (riota @Float n, riota @Float n)
    ,bench "dotprod Float [1e6] stride 1; -1" $
      nf (\(a, b) -> rdot a b)
         (riota @Float n, rrev1 (riota @Float n))
    ,bench "dotprod Double [1e6]" $
      nf (\(a, b) -> rdot a b)
         (riota @Double n, riota @Double n)
    ,bench "dotprod Double [1e6] stride 1; -1" $
      nf (\(a, b) -> rdot a b)
         (riota @Double n, rrev1 (riota @Double n))
    ]
  ,bgroup "hmatrix"
    [bench "sum(+) Double [1e6]" $
      nf (\(a, b) -> LA.sumElements (a + b))
         (LA.linspace @Double n (0.0, fromIntegral (n - 1))
         ,LA.linspace @Double n (0.0, fromIntegral (n - 1)))
    ,bench "sum(*) Double [1e6]" $
      nf (\(a, b) -> LA.sumElements (a * b))
         (LA.linspace @Double n (0.0, fromIntegral (n - 1))
         ,LA.linspace @Double n (0.0, fromIntegral (n - 1)))
    ,bench "sum(/) Double [1e6]" $
      nf (\(a, b) -> LA.sumElements (a / b))
         (LA.linspace @Double n (0.0, fromIntegral (n - 1))
         ,LA.linspace @Double n (0.0, fromIntegral (n - 1)))
    ,bench "sum(**) Double [1e6]" $
      nf (\(a, b) -> LA.sumElements (a ** b))
         (LA.linspace @Double n (0.0, fromIntegral (n - 1))
         ,LA.linspace @Double n (0.0, fromIntegral (n - 1)))
    ,bench "sum(sin) Double [1e6]" $
      nf (\a -> LA.sumElements (sin a))
         (LA.linspace @Double n (0.0, fromIntegral (n - 1)))
    ,bench "sum Double [1e6]" $
      nf (\a -> LA.sumElements a)
         (LA.linspace @Double n (0.0, fromIntegral (n - 1)))
    ,bench "dotprod Float [1e6]" $
      nf (\(a, b) -> a LA.<.> b)
         (LA.linspace @Double n (0.0, fromIntegral (n - 1))
         ,LA.linspace @Double n (fromIntegral (n - 1), 0.0))
    ,bench "dotprod Double [1e6]" $
      nf (\(a, b) -> a LA.<.> b)
         (LA.linspace @Double n (0.0, fromIntegral (n - 1))
         ,LA.linspace @Double n (fromIntegral (n - 1), 0.0))
    ]
  ]