{-# 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.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)
    ]
  ,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)))
    ]
  ]