{-# LANGUAGE Rank2Types #-}
module Test.Format where

import qualified Numeric.LAPACK.Orthogonal.Householder as Hh
import qualified Numeric.LAPACK.Linear.LowerUpper as LU
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.BandedHermitian as BandedHermitian
import qualified Numeric.LAPACK.Matrix.Banded as Banded
import qualified Numeric.LAPACK.Matrix.Hermitian as Hermitian
import qualified Numeric.LAPACK.Matrix.Triangular as Triangular
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import qualified Numeric.LAPACK.Permutation as Perm
import Numeric.LAPACK.Matrix.Shape (Order(RowMajor, ColumnMajor), UnaryProxy)
import Numeric.LAPACK.Matrix.Array (ArrayMatrix)
import Numeric.LAPACK.Matrix (ShapeInt, shapeInt)
import Numeric.LAPACK.Format (Format, (##))

import qualified Numeric.Netlib.Class as Class

import qualified Type.Data.Num.Unary.Literal as TypeNum
import qualified Type.Data.Num.Unary as Unary
import Type.Data.Num.Unary (unary)

import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable (Array)

import Data.Complex as Cplx (Complex((:+)))
import Data.Word (Word64)


randomMatrix ::
   (Shape.C sh, Class.Floating a) => sh -> Word64 -> ArrayMatrix sh a
randomMatrix sh = ArrMatrix.lift0 . Vector.random Vector.UniformBoxPM1 sh


vector :: (Class.Floating a) => Vector.Vector ShapeInt a
vector = Vector.random Vector.UniformBoxPM1 (shapeInt 4) 419

general :: (Class.Floating a) => Order -> Matrix.General ShapeInt ShapeInt a
general order =
   randomMatrix (MatrixShape.general order (shapeInt 3) (shapeInt 4)) 420

split ::
   (Eq lower, Shape.C height, Shape.C width, Class.Floating a) =>
   lower -> height -> width -> Order ->
   Array (MatrixShape.SplitGeneral lower height width) a
split lowerPart height width order =
   Vector.random Vector.UniformBoxPM1
      (MatrixShape.splitGeneral lowerPart order height width) 420

hermitian :: (Class.Floating a) => Order -> Hermitian.Hermitian ShapeInt a
hermitian order =
   randomMatrix (MatrixShape.hermitian order (shapeInt 4)) 421

diagonal :: (Class.Floating a) => Order -> Triangular.Diagonal ShapeInt a
diagonal order =
   randomMatrix (MatrixShape.diagonal order (shapeInt 4)) 422

lowerTriangular ::
   (Class.Floating a) => Order -> Triangular.Lower ShapeInt a
lowerTriangular order =
   randomMatrix (MatrixShape.lowerTriangular order (shapeInt 4)) 423

upperTriangular ::
   (Class.Floating a) => Order -> Triangular.Upper ShapeInt a
upperTriangular order =
   randomMatrix (MatrixShape.upperTriangular order (shapeInt 4)) 424

symmetric :: (Class.Floating a) => Order -> Triangular.Symmetric ShapeInt a
symmetric order =
   randomMatrix (MatrixShape.symmetric order (shapeInt 4)) 425


bandedHermitian ::
   (Unary.Natural offDiag, Class.Floating a) =>
   UnaryProxy offDiag -> Order ->
   BandedHermitian.BandedHermitian offDiag ShapeInt a
bandedHermitian numOff order =
   randomMatrix (MatrixShape.bandedHermitian numOff order (shapeInt 4)) 426

banded ::
   (Unary.Natural sub, Unary.Natural super,
    Shape.C height, Shape.C width, Class.Floating a) =>
   (UnaryProxy sub, UnaryProxy super) -> height -> width -> Order ->
   Banded.General sub super height width a
banded offDiag height width order =
   randomMatrix (MatrixShape.bandedGeneral offDiag order height width) 427


permutation :: Perm.Permutation ShapeInt
permutation =
   Perm.fromPivots Perm.NonInverted $
   Vector.fromList (Perm.Shape $ shapeInt 5) $ map Perm.Element [3,2,4,5,5]


lu :: (Class.Floating a) => Order -> LU.Tall ShapeInt ShapeInt a
lu order =
   LU.fromMatrix $
   randomMatrix (MatrixShape.tall order (shapeInt 4) (shapeInt 3)) 428

qr :: (Class.Floating a) => Order -> Hh.Tall ShapeInt ShapeInt a
qr order =
   Hh.fromMatrix $
   randomMatrix (MatrixShape.tall order (shapeInt 4) (shapeInt 3)) 429


fmt :: String
fmt = "%.4g"

printFormatted :: Format a => a -> IO ()
printFormatted x = putStrLn "" >> (x ## fmt)

printVectorFloat :: (Format (f Float)) => f Float -> IO ()
printVectorFloat = printFormatted

printVectorComplex ::
   (Format (f (Complex Float))) => f (Complex Float) -> IO ()
printVectorComplex = printFormatted

printVectorWithOrder ::
   Format (f Float) =>
   Format (f (Complex Float)) =>
   (forall a. (Class.Floating a) => Order -> f a) -> IO ()
printVectorWithOrder f = do
   printFormatted $ floatVector $ f RowMajor
   printFormatted $ floatVector $ f ColumnMajor
   printFormatted $ complexVector $ f RowMajor
   printFormatted $ complexVector $ f ColumnMajor

floatVector :: f Float -> f Float
floatVector = id

complexVector :: f (Complex Float) -> f (Complex Float)
complexVector = id

main :: IO ()
main = do
   printFormatted (pi :: Float)
   printFormatted permutation
   printVectorFloat $ sin (1:+1)
   printVectorFloat vector
   printVectorComplex vector
   printVectorWithOrder general
   printVectorWithOrder $ split MatrixShape.Reflector (shapeInt 4) (shapeInt 3)
   printVectorWithOrder $ split MatrixShape.Reflector (shapeInt 3) (shapeInt 4)
   printVectorWithOrder $ split MatrixShape.Triangle (shapeInt 4) (shapeInt 3)
   printVectorWithOrder hermitian
   printVectorWithOrder diagonal
   printVectorWithOrder lowerTriangular
   printVectorWithOrder upperTriangular
   printVectorWithOrder symmetric
   printVectorWithOrder $ bandedHermitian $ unary TypeNum.u0
   printVectorWithOrder $ bandedHermitian $ unary TypeNum.u1
   printVectorWithOrder $ bandedHermitian $ unary TypeNum.u2
   printVectorWithOrder $
      banded (unary TypeNum.u0, unary TypeNum.u0) (shapeInt 4) (shapeInt 3)
   printVectorWithOrder $
      banded (unary TypeNum.u0, unary TypeNum.u2) (shapeInt 4) (shapeInt 3)
   printVectorWithOrder $
      banded (unary TypeNum.u2, unary TypeNum.u0) (shapeInt 4) (shapeInt 3)
   printVectorWithOrder $
      banded (unary TypeNum.u1, unary TypeNum.u2) (shapeInt 4) (shapeInt 3)
   printVectorWithOrder $
      banded (unary TypeNum.u1, unary TypeNum.u2) (shapeInt 3) (shapeInt 4)
   printVectorWithOrder lu
   printVectorWithOrder qr