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

import qualified Numeric.LinearAlgebra as LA
import Test.Tasty.Bench

import Data.Array.Nested
import Data.Array.Nested.Internal (mliftPrim2, 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 -> 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 -> runScalar (rsumOuter1 a))
            (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 -> LA.sumElements a)
            (LA.linspace @Double n (0.0, fromIntegral (n - 1)))
    ]
  ]