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

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 (arithPromoteRanked, arithPromoteRanked2)


enableMisc :: Bool
enableMisc = False

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


main :: IO ()
main = defaultMain
  [bgroup "Num"
    [bench "sum(+) Double [1e6]" $
      let n = 1_000_000
      in nf (\(a, b) -> runScalar (rsumOuter1 (arithPromoteRanked2 (mliftPrim2 (+)) a b)))
            (riota @Double n, riota n)
    ,bench "sum(*) Double [1e6]" $
      let n = 1_000_000
      in nf (\(a, b) -> runScalar (rsumOuter1 (arithPromoteRanked2 (mliftPrim2 (*)) a b)))
            (riota @Double n, riota n)
    ,bench "sum(/) Double [1e6]" $
      let n = 1_000_000
      in nf (\(a, b) -> runScalar (rsumOuter1 (arithPromoteRanked2 (mliftPrim2 (/)) a b)))
            (riota @Double n, riota n)
    ,bench "sum(**) Double [1e6]" $
      let n = 1_000_000
      in nf (\(a, b) -> runScalar (rsumOuter1 (arithPromoteRanked2 (mliftPrim2 (**)) a b)))
            (riota @Double n, riota n)
    ,bench "sum(sin) Double [1e6]" $
      let n = 1_000_000
      in nf (\a -> runScalar (rsumOuter1 (arithPromoteRanked (mliftPrim sin) a)))
            (riota @Double n)
    ,bench "sum Double [1e6]" $
      let n = 1_000_000
      in nf (\a -> runScalar (rsumOuter1 a))
            (riota @Double n)
    ]
  ,bgroup "NumElt"
    [bench "sum(+) Double [1e6]" $
      let n = 1_000_000
      in nf (\(a, b) -> runScalar (rsumOuter1 (a + b)))
            (riota @Double n, riota n)
    ,bench "sum(*) Double [1e6]" $
      let n = 1_000_000
      in nf (\(a, b) -> runScalar (rsumOuter1 (a * b)))
            (riota @Double n, riota n)
    ,bench "sum(/) Double [1e6]" $
      let n = 1_000_000
      in nf (\(a, b) -> runScalar (rsumOuter1 (a / b)))
            (riota @Double n, riota n)
    ,bench "sum(**) Double [1e6]" $
      let n = 1_000_000
      in nf (\(a, b) -> runScalar (rsumOuter1 (a ** b)))
            (riota @Double n, riota n)
    ,bench "sum(sin) Double [1e6]" $
      let n = 1_000_000
      in nf (\a -> runScalar (rsumOuter1 (sin a)))
            (riota @Double n)
    ,bench "sum Double [1e6]" $
      let n = 1_000_000
      in nf (\a -> runScalar (rsumOuter1 a))
            (riota @Double n)
    ,bench "dotprod Float [1e6]" $
      let n = 1_000_000
      in nf (\(a, b) -> rdot a b)
            (riota @Float n, riota @Float n)
    ,bench "dotprod Float [1e6] stride 1; -1" $
      let n = 1_000_000
      in nf (\(a, b) -> rdot a b)
            (riota @Float n, rrev1 (riota @Float n))
    ,bench "dotprod Double [1e6]" $
      let n = 1_000_000
      in nf (\(a, b) -> rdot a b)
            (riota @Double n, riota @Double n)
    ,bench "dotprod Double [1e6] stride 1; -1" $
      let n = 1_000_000
      in nf (\(a, b) -> rdot a b)
            (riota @Double n, rrev1 (riota @Double n))
    ]
  ,bgroup "hmatrix"
    [bench "sum(+) Double [1e6]" $
      let n = 1_000_000
      in 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]" $
      let n = 1_000_000
      in 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]" $
      let n = 1_000_000
      in 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]" $
      let n = 1_000_000
      in 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]" $
      let n = 1_000_000
      in nf (\a -> LA.sumElements (sin a))
            (LA.linspace @Double n (0.0, fromIntegral (n - 1)))
    ,bench "sum Double [1e6]" $
      let n = 1_000_000
      in nf (\a -> LA.sumElements a)
            (LA.linspace @Double n (0.0, fromIntegral (n - 1)))
    ,bench "dotprod Float [1e6]" $
      let n = 1_000_000
      in 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]" $
      let n = 1_000_000
      in nf (\(a, b) -> a LA.<.> b)
            (LA.linspace @Double n (0.0, fromIntegral (n - 1))
            ,LA.linspace @Double n (fromIntegral (n - 1), 0.0))
    ]
  ,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]]
        ]
      ]
    ]
  ]