{-# LANGUAGE TypeOperators #-}
module Numeric.LAPACK.Matrix.Square.Basic (
   Square,
   LiberalSquare,
   mapSize,
   toFull,
   fromFull,
   fromFullUnchecked,
   liberalFromFull,
   fromScalar,
   toScalar,
   fromList,
   autoFromList,

   transpose,
   adjoint,

   identity,
   identityOrder,
   identityFrom,
   identityFromWidth,
   identityFromHeight,
   diagonal,
   diagonalOrder,
   takeDiagonal,
   trace,

   stack,

   square,
   power,
   congruence,
   congruenceAdjoint,
   ) where


import qualified Numeric.LAPACK.Matrix.Layout.Private as Layout
import qualified Numeric.LAPACK.Matrix.Extent.Private as ExtentPriv
import qualified Numeric.LAPACK.Matrix.Extent as Extent
import qualified Numeric.LAPACK.Matrix.Basic as Basic
import qualified Numeric.LAPACK.Matrix.Private as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import qualified Numeric.LAPACK.Private as Private
import Numeric.LAPACK.Matrix.Layout.Private
         (Order(RowMajor, ColumnMajor), swapOnRowMajor)
import Numeric.LAPACK.Matrix.Private
         (General, argGeneral, LiberalSquare, Square, argSquare,
          Full, mapExtent, ShapeInt, shapeInt)
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Scalar (zero, one)
import Numeric.LAPACK.Shape.Private (Unchecked(Unchecked))
import Numeric.LAPACK.Private (pokeCInt)

import qualified Numeric.LAPACK.FFI.Generic as LapackGen
import qualified Numeric.BLAS.FFI.Generic as BlasGen
import qualified Numeric.Netlib.Utility as Call
import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Comfort.Storable.Unchecked as Array
import qualified Data.Array.Comfort.Storable as CheckedArray
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable.Unchecked (Array(Array))
import Data.Array.Comfort.Shape ((::+))

import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Storable (Storable, peek, poke)

import System.IO.Unsafe (unsafePerformIO)

import Control.Monad.Trans.Cont (ContT(ContT), evalContT)
import Control.Monad.IO.Class (liftIO)

import Data.Function.HT (powerAssociative)


mapSize :: (sh0 -> sh1) -> Square sh0 a -> Square sh1 a
mapSize = Basic.mapExtent . ExtentPriv.mapSquareSize

toFull ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz) =>
   Square sh a -> Full meas vert horiz sh sh a
toFull = mapExtent ExtentPriv.fromSquare

fromFull ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz, Eq sh) =>
   Full meas vert horiz sh sh a -> Square sh a
fromFull = mapExtent ExtentPriv.squareFromFull

fromFullUnchecked ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz) =>
   Full meas vert horiz sh sh a -> Square sh a
fromFullUnchecked =
   Basic.recheck . mapExtent ExtentPriv.squareFromFull . Basic.uncheck

liberalFromFull ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Shape.C width) =>
   Full meas vert horiz height width a -> LiberalSquare height width a
liberalFromFull = mapExtent ExtentPriv.liberalSquareFromFull


fromScalar :: (Storable a) => a -> Square () a
fromScalar a =
   Array.unsafeCreate (Layout.square RowMajor ()) $ flip poke a

toScalar :: (Storable a) => Square () a -> a
toScalar = argSquare $ \_ () a ->
   unsafePerformIO $ withForeignPtr a peek

fromList :: (Shape.C sh, Storable a) => sh -> [a] -> Square sh a
fromList sh =
   CheckedArray.fromList (Layout.square RowMajor sh)

autoFromList :: (Storable a) => [a] -> Square ShapeInt a
autoFromList xs =
   let n = length xs
       m = round $ sqrt (fromIntegral n :: Double)
   in if n == m*m
        then fromList (shapeInt m) xs
        else error "Square.autoFromList: no quadratic number of elements"


transpose :: Square sh a -> Square sh a
transpose = Array.mapShape Layout.transpose

adjoint :: (Shape.C sh, Class.Floating a) => Square sh a -> Square sh a
adjoint = transpose . Vector.conjugate


identity :: (Shape.C sh, Class.Floating a) => sh -> Square sh a
identity = identityOrder RowMajor

identityFrom :: (Shape.C sh, Class.Floating a) => Square sh a -> Square sh a
identityFrom = argSquare $ \order sh _ -> identityOrder order sh

identityFromWidth ::
   (Shape.C height, Shape.C width, Class.Floating a) =>
   General height width a -> Square width a
identityFromWidth =
   argGeneral $ \order _ width _ -> identityOrder order width

identityFromHeight ::
   (Shape.C height, Shape.C width, Class.Floating a) =>
   General height width a -> Square height a
identityFromHeight =
   argGeneral $ \order height _ _ -> identityOrder order height

identityOrder, _identityOrder ::
   (Shape.C sh, Class.Floating a) => Order -> sh -> Square sh a
identityOrder order sh =
   Array.unsafeCreate (Layout.square order sh) $ \aPtr ->
   evalContT $ do
      uploPtr <- Call.char 'A'
      nPtr <- Call.cint $ Shape.size sh
      alphaPtr <- Call.number zero
      betaPtr <- Call.number one
      liftIO $ LapackGen.laset uploPtr nPtr nPtr alphaPtr betaPtr aPtr nPtr

_identityOrder order sh =
   Array.unsafeCreateWithSize (Layout.square order sh) $ \blockSize yPtr ->
   evalContT $ do
      nPtr <- Call.alloca
      xPtr <- Call.number zero
      incxPtr <- Call.cint 0
      incyPtr <- Call.cint 1
      liftIO $ do
         pokeCInt nPtr blockSize
         BlasGen.copy nPtr xPtr incxPtr yPtr incyPtr
         let n = fromIntegral $ Shape.size sh
         poke nPtr n
         poke xPtr one
         poke incyPtr (n+1)
         BlasGen.copy nPtr xPtr incxPtr yPtr incyPtr

diagonal :: (Shape.C sh, Class.Floating a) => Vector sh a -> Square sh a
diagonal (Array sh x) =
   Array.unsafeCreateWithSize (Layout.square ColumnMajor sh) $
      \blockSize yPtr ->
   evalContT $ do
      nPtr <- Call.alloca
      xPtr <- ContT $ withForeignPtr x
      zPtr <- Call.number zero
      incxPtr <- Call.cint 1
      incyPtr <- Call.cint 1
      inczPtr <- Call.cint 0
      liftIO $ do
         pokeCInt nPtr blockSize
         BlasGen.copy nPtr zPtr inczPtr yPtr incyPtr
         let n = fromIntegral $ Shape.size sh
         poke nPtr n
         poke incyPtr (n+1)
         BlasGen.copy nPtr xPtr incxPtr yPtr incyPtr

diagonalOrder ::
   (Shape.C sh, Class.Floating a) => Order -> Vector sh a -> Square sh a
diagonalOrder order =
   (case order of ColumnMajor -> id; RowMajor -> transpose) . diagonal


takeDiagonal :: (Shape.C sh, Class.Floating a) => Square sh a -> Vector sh a
takeDiagonal = argSquare $ \_ sh x ->
   Array.unsafeCreateWithSize sh $ \n yPtr -> evalContT $ do
      nPtr <- Call.cint n
      xPtr <- ContT $ withForeignPtr x
      incxPtr <- Call.cint (n+1)
      incyPtr <- Call.cint 1
      liftIO $ BlasGen.copy nPtr xPtr incxPtr yPtr incyPtr

trace :: (Shape.C sh, Class.Floating a) => Square sh a -> a
trace = argSquare $ \_ sh x -> unsafePerformIO $ do
   let n = Shape.size sh
   withForeignPtr x $ \xPtr -> Private.sum n xPtr (n+1)


stack ::
   (Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C sizeA, Eq sizeA, Shape.C sizeB, Eq sizeB, Class.Floating a) =>
   Square sizeA a -> Full meas vert horiz sizeA sizeB a ->
   Full meas horiz vert sizeB sizeA a -> Square sizeB a ->
   Square (sizeA::+sizeB) a
stack a b c d = Basic.stack a (Matrix.fromFull b) (Matrix.fromFull c) d


square :: (Shape.C sh, Class.Floating a) => Square sh a -> Square sh a
square a = multiplyCommutativeUnchecked a a

power ::
   (Shape.C sh, Class.Floating a) =>
   Integer -> Square sh a -> Square sh a
power n a =
   powerAssociative multiplyCommutativeUnchecked (identityFrom a) a n

congruence ::
   (Shape.C height, Eq height, Shape.C width, Class.Floating a) =>
   Square height a -> General height width a -> Square width a
congruence b a0 =
   let a = Basic.mapWidth Unchecked a0
   in mapSize (\(Unchecked sh) -> sh) $ fromFull $
      Basic.multiply (Basic.adjoint a) $ Basic.multiply (toFull b) a

congruenceAdjoint ::
   (Shape.C height, Shape.C width, Eq width, Class.Floating a) =>
   General height width a -> Square width a -> Square height a
congruenceAdjoint a0 b =
   let a = Basic.mapHeight Unchecked a0
   in mapSize (\(Unchecked sh) -> sh) $ fromFull $
      Basic.multiply a $ Basic.multiply (toFull b) $ Basic.adjoint a

{-
orderA and orderB must be equal but this is not checked.
-}
multiplyCommutativeUnchecked ::
   (Shape.C sh, Class.Floating a) =>
   Square sh a -> Square sh a -> Square sh a
multiplyCommutativeUnchecked
   (Array shape@(Layout.Full order extent) a)
   (Array _ b) =
      Array.unsafeCreate shape $ \cPtr ->
   let n = Shape.size $ Extent.height extent
       (at,bt) = swapOnRowMajor order (a,b)
   in  Private.multiplyMatrix ColumnMajor ColumnMajor n n n at bt cPtr