{-# LANGUAGE TypeOperators #-}
module Knead where

import qualified Data.Array.Knead.Symbolic.Render as Render
import qualified Data.Array.Knead.Symbolic as Symb
-- import qualified Data.Array.Knead.Symbolic.Slice as Slice
import qualified Data.Array.Knead.Expression as Expr
import qualified Data.Array.Knead.Shape as Shape
-- import qualified Data.Array.Comfort.Storable as Array
-- import qualified Data.Array.Comfort.Shape as ComfortShape
import Data.Array.Comfort.Storable (Array)
-- import Data.Array.Comfort.Shape ((::+))

-- import qualified LLVM.Extra.Nice.Value.Storable as Storable
-- import qualified LLVM.Extra.Nice.Value.Marshal as Marshal
import qualified LLVM.Extra.Nice.Value as NiceValue

import qualified LLVM.Core as LLVM

import qualified Type.Data.Num.Decimal as TypeNum
import Data.Int (Int64)

import System.IO.Unsafe (unsafePerformIO)



type Packed = LLVM.Vector TypeNum.D8
type Large = LLVM.Vector TypeNum.D64

sumSymb ::
   (Shape.C sh, NiceValue.Additive a) =>
   Symb.Array sh a -> Expr.Exp a
sumSymb = Symb.fold1All Expr.add

runSum :: IO (Array (Shape.ZeroBased Int64) (Packed Float) -> Packed Float)
runSum = do
   proc <- Render.run sumSymb
   return $ \xs -> unsafePerformIO $ proc xs

{-
Try to benefit from out-of-order sum.
Seems to make things worse.
-}
runSumLarge ::
   IO (Array (Shape.ZeroBased Int64) (Large Float) -> Large Float)
runSumLarge = do
   proc <- Render.run sumSymb
   return $ \xs -> unsafePerformIO $ proc xs


rowSumSymb ::
   (Shape.C sh0, Shape.C sh1, NiceValue.Additive a) =>
   Symb.Array (sh0,sh1) a -> Symb.Array sh0 a
rowSumSymb = Symb.fold1 Expr.add


sumByRowsSymb ::
   (Shape.C sh0, Shape.C sh1, NiceValue.Additive a) =>
   Symb.Array (sh0,sh1) a -> Expr.Exp a
sumByRowsSymb = Symb.fold1All Expr.add . Symb.fold1 Expr.add

runSumByRows ::
   IO (Array (Shape.ZeroBased Int64, ()) (Packed Float) ->
       Packed Float)
{-
   IO (Array (Shape.ZeroBased Int64, Shape.Enumeration Bool) (Packed Float) ->
       Packed Float)
-}
runSumByRows = do
   proc <- Render.run sumByRowsSymb
   return $ \xs -> unsafePerformIO $ proc xs