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

import Data.Array.RankedS qualified as RS
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)


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)))
    ]
  ]