{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GADTs #-}
module Numeric.LAPACK.Matrix.Triangular.Basic (
   Triangular, TriangularP,
   Layout.UpLo,
   Upper,
   Lower,

   toSquare,
   takeLower,
   takeUpper,

   fromLowerRowMajor, toLowerRowMajor,
   fromUpperRowMajor, toUpperRowMajor,
   forceOrder,

   multiplyVector,
   multiply,
   multiplyFull,
   ) where

import qualified Numeric.LAPACK.Matrix.Mosaic.Unpacked as Unpacked
import qualified Numeric.LAPACK.Matrix.Mosaic.Private as Mos
import qualified Numeric.LAPACK.Matrix.Layout.Private as Layout
import qualified Numeric.LAPACK.Matrix.Extent.Private as Extent
import Numeric.LAPACK.Matrix.Mosaic.Private
         (Triangular, Lower, Upper,
          withPacking, noLabel, applyFuncPair, triArg)
import Numeric.LAPACK.Matrix.Mosaic.Packed (forceOrder)
import Numeric.LAPACK.Matrix.Mosaic.Generic (unpackDirty, repack)
import Numeric.LAPACK.Matrix.Shape.Omni (TriDiag, DiagSingleton, charFromTriDiag)
import Numeric.LAPACK.Matrix.Layout.Private
         (Order(RowMajor,ColumnMajor),
          transposeFromOrder, uploFromOrder, uploOrder)
import Numeric.LAPACK.Matrix.Private (Full, Square)
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Private (copyBlock)

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.Shape as Shape
import Data.Array.Comfort.Storable.Unchecked (Array(Array))

import Foreign.ForeignPtr (withForeignPtr)

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


type TriangularP pack uplo sh = Array (Layout.TriangularP pack uplo sh)


toSquare ::
   (Layout.UpLo uplo, Shape.C sh, Class.Floating a) =>
   Triangular uplo sh a -> Square sh a
toSquare
   (Array
      (Layout.Mosaic Layout.Packed Layout.NoMirror uplo order sh)
      a) =
   Array.unsafeCreate (Layout.square order sh) $ \bPtr ->
      withForeignPtr a $ \aPtr ->
         Mos.unpackZero (uploOrder uplo order) (Shape.size sh) aPtr bPtr

unpack ::
   (Layout.UpLo uplo, Shape.C sh, Class.Floating a) =>
   Triangular uplo sh a -> Unpacked.Triangular uplo sh a
unpack (Array (Layout.Mosaic Layout.Packed mirror uplo order sh) a) =
   Array.unsafeCreate
      (Layout.Mosaic Layout.Unpacked mirror uplo order sh) $
   \bPtr ->
      withForeignPtr a $ \aPtr ->
         Mos.unpackZero (uploOrder uplo order) (Shape.size sh) aPtr bPtr

unpackGen ::
   (Layout.UpLo uplo, Shape.C sh, Class.Floating a) =>
   TriangularP pack uplo sh a -> Unpacked.Triangular uplo sh a
unpackGen a =
   case Layout.mosaicPack $ Array.shape a of
      Layout.Unpacked -> a
      Layout.Packed -> unpack a


takeLower ::
   (Extent.Measure meas, Extent.C horiz,
    Shape.C height, Shape.C width, Class.Floating a) =>
   Full meas Extent.Small horiz height width a -> Lower height a
takeLower = Mos.fromLowerPart Mos.leaveDiagonal Layout.NoMirror

takeUpper ::
   (Extent.Measure meas, Extent.C vert,
    Shape.C height, Shape.C width, Class.Floating a) =>
   Full meas vert Extent.Small height width a -> Upper width a
takeUpper (Array (Layout.Full order extent) a) =
   let (height,width) = Extent.dimensions extent
       m = Shape.size height
       n = Shape.size width
       k = case order of RowMajor -> n; ColumnMajor -> m
   in Array.unsafeCreate
         (Layout.Mosaic Layout.Packed
            Layout.NoMirror Layout.Upper order width) $ \bPtr ->
      withForeignPtr a $ \aPtr -> Mos.packRect order n k aPtr bPtr


fromLowerRowMajor ::
   (Shape.C sh, Class.Floating a) =>
   Array (Shape.Triangular Shape.Lower sh) a -> Lower sh a
fromLowerRowMajor =
   Array.mapShape (Layout.lowerTriangular RowMajor . Shape.triangularSize)

fromUpperRowMajor ::
   (Shape.C sh, Class.Floating a) =>
   Array (Shape.Triangular Shape.Upper sh) a -> Upper sh a
fromUpperRowMajor =
   Array.mapShape (Layout.upperTriangular RowMajor . Shape.triangularSize)

toLowerRowMajor ::
   (Shape.C sh, Class.Floating a) =>
   Lower sh a -> Array (Shape.Triangular Shape.Lower sh) a
toLowerRowMajor =
   Array.mapShape (Shape.Triangular Shape.Lower . Layout.mosaicSize)
   .
   forceOrder Layout.RowMajor

toUpperRowMajor ::
   (Shape.C sh, Class.Floating a) =>
   Upper sh a -> Array (Shape.Triangular Shape.Upper sh) a
toUpperRowMajor =
   Array.mapShape (Shape.Triangular Shape.Upper . Layout.mosaicSize)
   .
   forceOrder Layout.RowMajor


multiplyVector ::
   (Layout.UpLo uplo, TriDiag diag,
    Shape.C sh, Eq sh, Class.Floating a) =>
   DiagSingleton diag ->
   TriangularP pack uplo sh a -> Vector sh a -> Vector sh a
multiplyVector diag
   (Array (Layout.Mosaic pack_ Layout.NoMirror uplo order shA) a)
   (Array shX x) =
      Array.unsafeCreate shX $ \yPtr -> do
   Call.assert "Triangular.multiplyVector: width shapes mismatch" (shA == shX)
   let n = Shape.size shA
   evalContT $ do
      uploPtr <- Call.char $ uploFromOrder $ uploOrder uplo order
      transPtr <- Call.char $ transposeFromOrder order
      diagPtr <- Call.char $ charFromTriDiag diag
      nPtr <- Call.cint n
      aPtr <- ContT $ withForeignPtr a
      xPtr <- ContT $ withForeignPtr x
      incyPtr <- Call.cint 1
      liftIO $ copyBlock n xPtr yPtr
      withPacking pack_ $
         applyFuncPair (noLabel BlasGen.tpmv) (noLabel BlasGen.trmv)
            uploPtr transPtr diagPtr nPtr (triArg aPtr n) yPtr incyPtr


multiply ::
   (Layout.Packing pack, Layout.UpLo uplo, TriDiag diag,
    Shape.C sh, Eq sh, Class.Floating a) =>
   DiagSingleton diag ->
   TriangularP pack uplo sh a ->
   TriangularP pack uplo sh a ->
   TriangularP pack uplo sh a
multiply diag a b =
   repack $ Unpacked.multiplyCompatible diag (unpackDirty a) (unpackGen b)

multiplyFull ::
   (Layout.UpLo uplo, TriDiag diag,
    Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Eq height, Shape.C width,
    Class.Floating a) =>
   DiagSingleton diag ->
   TriangularP pack uplo height a ->
   Full meas vert horiz height width a ->
   Full meas vert horiz height width a
multiplyFull diag = Unpacked.multiplyFull diag . unpackDirty