{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
module Main where

import qualified CUBLASBatched as Batched
import qualified Foreign.CUDA.Cublas as Cublas

import qualified Data.Array.Accelerate.Arithmetic.LinearAlgebra as ALinAlg
import qualified Data.Array.Accelerate.Utility.Loop as Loop
import qualified Data.Array.Accelerate.CUDA as CUDA
import qualified Data.Array.Accelerate as A
import Data.Array.Accelerate (All(All), Z(Z), (:.)((:.)))

import qualified Control.Concurrent.PooledIO.Independent as Pooled

import qualified Data.Packed.Matrix as Matrix
import qualified Data.Packed.Vector as Vector
import qualified Numeric.Container as Container
import qualified Numeric.LinearAlgebra.Algorithms as HMLinAlg

import Numeric.Container (Container, (<>))
import Data.Packed.Matrix (Matrix)
import Data.Packed.Vector (Vector)

import qualified System.Random as Rnd
import System.TimeIt (timeIt)

import Text.Printf (printf)

import qualified Data.List.HT as ListHT
import Data.Function.HT (nest)
import Data.Tuple.HT (mapPair)



newtonInverseStep ::
   (Num a, Container Vector a, Container.Product a) =>
   Matrix a -> Matrix a -> Matrix a
newtonInverseStep a x =
   Container.sub
      (Container.scale 2 x)
      (x <> a <> x)

newtonInverse ::
   (Num a, Container Vector a, Container.Product a) =>
   Int -> Matrix a -> Matrix a -> Matrix a
newtonInverse count start a =
   nest count (newtonInverseStep a) start


newtonInverseCUBLASStep, newtonInverseCUBLASStepMul ::
   (A.Shape ix, A.Slice ix, Eq ix, Batched.Element a, A.IsNum a, A.Elt a) =>
   Cublas.Handle ->
   ALinAlg.Matrix ix a ->
   ALinAlg.Matrix ix a ->
   ALinAlg.Matrix ix a
newtonInverseCUBLASStep h a x =
   Batched.mac h (-1) x (Batched.mul h 1 a x) 2 x

newtonInverseCUBLASStepMul h a x =
   A.zipWith (-) (A.map (2*) x) $
   Batched.mul h 1 x $ Batched.mul h 1 a x

newtonInverseCUBLAS ::
   (A.Shape ix, A.Slice ix, Eq ix, Batched.Element a, A.IsNum a, A.Elt a) =>
   Cublas.Handle ->
   A.Exp Int ->
   ALinAlg.Matrix ix a ->
   ALinAlg.Matrix ix a ->
   ALinAlg.Matrix ix a
newtonInverseCUBLAS h n seed a =
   Loop.nest n (newtonInverseCUBLASStep h a) seed


randomMatrixInv :: Int -> (Matrix Double, Matrix Double)
randomMatrixInv size =
   let x =
          Matrix.fromLists $ take size $ ListHT.sliceVertical size $
          Rnd.randomRs (-1,1::Double) $ Rnd.mkStdGen 42
   in  (x, HMLinAlg.inv x)


parallel :: [a] -> (Int -> a -> IO ()) -> IO ()
parallel xs f = Pooled.run $ zipWith f [0 ..] xs

disturbedMatrices ::
   (Container Vector a) =>
   Matrix a -> [a] -> [Matrix a]
disturbedMatrices x yelems =
   let size = Matrix.rows x
   in  map (Container.add x . Matrix.fromLists) $
       ListHT.sliceVertical size $
       ListHT.sliceVertical size $
       yelems

mainHMatrixDirect ::
   (Show a, Container Vector a, HMLinAlg.Field a) =>
   String ->
   Int -> (Matrix a, Matrix a) -> [a] -> IO ()
mainHMatrixDirect typ numberOfMatrices (x, _xinv) yelems = do
   let yinvs = map HMLinAlg.inv $ disturbedMatrices x yelems
   putStrLn $ "hmatrix-direct-" ++ typ
   timeIt $ parallel (take numberOfMatrices yinvs) $ \ n y ->
      writeFile (printf "/tmp/hmatrix-direct-%s%03d.txt" typ n) $ show y

mainHMatrix ::
   (Show a, Container Vector a, Container.Product a) =>
   String ->
   Int -> Int -> (Matrix a, Matrix a) -> [a] -> IO ()
mainHMatrix typ numberOfMatrices newtonIts (x, xinv) yelems = do
   let yinvs = map (newtonInverse newtonIts xinv) $ disturbedMatrices x yelems
   putStrLn $ "hmatrix-" ++ typ
   timeIt $ parallel (take numberOfMatrices yinvs) $ \ n y ->
      writeFile (printf "/tmp/hmatrix-%s%03d.txt" typ n) $ show y


mainCUDA ::
   (A.Elt a, A.IsNum a, Container.Element a) =>
   String ->
   Int -> Int -> (Matrix a, Matrix a) -> [a] -> IO ()
mainCUDA typ numberOfMatrices newtonIts (xm, xinvm) yelems = do
   let size = Matrix.rows xm
       matrixAccFromHM =
          A.fromList (Z :. size :. size) .
          Vector.toList . Matrix.flatten
       xarr = matrixAccFromHM xm
       xinvarr = matrixAccFromHM xinvm

   let ysarr =
          A.fromList (Z :. numberOfMatrices :. size :. size) yelems
       rep = A.replicate (A.lift $ Z :. numberOfMatrices :. All :. All)
       yinvs =
          CUDA.run1
             (\args ->
                case A.unlift args of
                   (x, xinv, ys) ->
                      ALinAlg.newtonInverse (A.constant newtonIts) (rep xinv) $
                      A.zipWith (+) ys (rep x))
             (xarr, xinvarr, ysarr)

   putStrLn $ "cuda-" ++ typ
   timeIt $ writeFile ("/tmp/cuda-"++typ++".txt") $ show yinvs


mainCUBLASDirect ::
   (Batched.Element a, Container.Element a, A.IsNum a, A.Elt a) =>
   String ->
   Int -> (Matrix a, Matrix a) -> [a] -> IO ()
mainCUBLASDirect typ numberOfMatrices (xm, _xinvm) yelems = do
   let size = Matrix.rows xm
       matrixAccFromHM =
          A.fromList (Z :. size :. size) .
          Vector.toList . Matrix.flatten
       xarr = matrixAccFromHM xm

   handle <- Cublas.create
   let ysarr =
          A.fromList (Z :. numberOfMatrices :. size :. size) yelems
       rep = A.replicate (A.lift $ Z :. numberOfMatrices :. All :. All)
       yinvs =
          CUDA.run1
             (\args ->
                case A.unlift args of
                   (x, ys) ->
                      fst $ Batched.inv handle $ A.zipWith (+) ys (rep x))
             (xarr, ysarr)

   putStrLn $ "cublas-direct-" ++ typ
   timeIt $ writeFile ("/tmp/cublas-direct-"++typ++".txt") $ show yinvs


mainCUBLAS ::
   (Batched.Element a, Container.Element a, A.IsNum a, A.Elt a) =>
   String ->
   Int -> Int -> (Matrix a, Matrix a) -> [a] -> IO ()
mainCUBLAS typ numberOfMatrices newtonIts (xm, xinvm) yelems = do
   let size = Matrix.rows xm
       matrixAccFromHM =
          A.fromList (Z :. size :. size) .
          Vector.toList . Matrix.flatten
       xarr = matrixAccFromHM xm
       xinvarr = matrixAccFromHM xinvm

   handle <- Cublas.create
   let ysarr =
          A.fromList (Z :. numberOfMatrices :. size :. size) yelems
       rep = A.replicate (A.lift $ Z :. numberOfMatrices :. All :. All)
       yinvs =
          CUDA.run1
             (\args ->
                case A.unlift args of
                   (x, xinv, ys) ->
                      newtonInverseCUBLAS handle (A.constant newtonIts) (rep xinv) $
                      A.zipWith (+) ys (rep x))
             (xarr, xinvarr, ysarr)

   putStrLn $ "cublas-" ++ typ
   timeIt $ writeFile ("/tmp/cublas-"++typ++".txt") $ show yinvs


main :: IO ()
main = do
   let n = 96
   let sz = 50
   let its = 20
   let xmsDouble = randomMatrixInv sz
       ysDouble = Rnd.randomRs (-0.01,0.01::Double) $ Rnd.mkStdGen 23
   let xmsFloat =
          mapPair
             (Container.cmap realToFrac, Container.cmap realToFrac)
             xmsDouble
       ysFloat :: [Float]
       ysFloat = map realToFrac ysDouble
   mainHMatrixDirect "double" n xmsDouble ysDouble
--   mainHMatrixDirect "float" n xmsFloat ysFloat
   mainHMatrix "double" n its xmsDouble ysDouble
   mainHMatrix "float" n its xmsFloat ysFloat
   mainCUBLASDirect "double" n xmsDouble ysDouble
   mainCUBLASDirect "float" n xmsFloat ysFloat
   mainCUBLAS "double" n its xmsDouble ysDouble
   mainCUBLAS "float" n its xmsFloat ysFloat
   mainCUDA "double" n its xmsDouble ysDouble
   mainCUDA "float" n its xmsFloat ysFloat