module LabB where


import Obsidian

import Prelude hiding (zipWith)
import qualified Prelude as P 

import Data.Word
import Control.Monad

---------------------------------------------------------------------------
-- Array used for inputs 
---------------------------------------------------------------------------
inG :: Pull EWord32 EFloat
inG = namedPull "arr" (variable "x") 


---------------------------------------------------------------------------
-- splitUp
---------------------------------------------------------------------------
splitUp :: Word32 -> Pull EWord32 a -> Pull EWord32 (Pull Word32 a)
splitUp n arr = mkPullArray ((len arr)  `div` fromIntegral n) $
                 \i -> mkPullArray n $ \j -> arr ! (i * fromIntegral n + j) 



---------------------------------------------------------------------------
-- Task 2 (The first task in the Obsidian part of the lab)
---------------------------------------------------------------------------

-- The local computation
vecAddLocal :: SPull EFloat
               -> SPull EFloat
               -> BProgram (SPull EFloat)
vecAddLocal arr1 arr2 = undefined


-- How the computation is distributed over blocks
vecAdd :: DPull EFloat
          -> DPull EFloat
          -> DPush Grid EFloat
vecAdd in1 in2 = undefined

getVecAdd = putStrLn $ genKernel "vecAdd" vecAdd (inG :-> inG) 


---------------------------------------------------------------------------
-- Task 3 Reduction
---------------------------------------------------------------------------

-- Local sum computation
sumLocal :: Int
            -> SPull EFloat
            -> BProgram (SPull EFloat)
sumLocal 0 arr = return arr
sumLocal n arr = undefined 

-- Place many sum computation instances next to eachother. 
sums :: DPull EFloat -> DPush Grid EFloat
sums arr = undefined 

getSums = putStrLn $ genKernel "sums" sums inG 

---------------------------------------------------------------------------
-- Task 4 Dot Product
---------------------------------------------------------------------------

-- the earlier exercises implement a reduction and an element-wise
-- operation. To implement this task you need just that.
-- A reduction with (+) is needed. (already have this part) 
-- Element-wise product is needed.

prodLocal :: Int
            -> SPull EFloat
            -> BProgram (SPull EFloat)
prodLocal = undefined 

prods :: DPull EFloat -> DPush Grid EFloat
prods = undefined

getProds = putStrLn $ genKernel "prods" prods inG 

---------------------------------------------------------------------------
-- Voluntary 1 (single kernel dot product) 
--------------------------------------------------------------------------- 

---------------------------------------------------------------------------
-- Voluntary 2 (Matrix Multiplication)
---------------------------------------------------------------------------

type SMatrix a = SPull (SPull a)

mkMatrix n m f = undefined 

transpose :: SMatrix a -> SMatrix a
transpose mat = undefined

-- possible matMul type
-- Note how the result is a just a 1D array. Obsidia
matMul :: (Num a, MemoryOps a)
          => SMatrix a -> SMatrix a -> SPush Grid a
matMul = undefined 


-- To be able to generate code something like this may be needed
toMatrix :: Word32 -> Word32 -> SPull a -> SMatrix a
toMatrix n m arr = mkMatrix n m $ \i j -> arr ! (i * fromIntegral m + j)


---------------------------------------------------------------------------
-- Voluntary 3 (Experiement of your own)
---------------------------------------------------------------------------