{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE GADTs #-}
module Numeric.LAPACK.Matrix.Mosaic.Packed (
   Mosaic,
   Triangular,
   identity,
   diagonal,
   takeDiagonal,
   forceOrder,

   stackLower,
   stackUpper,
   takeTopLeft,
   Mos.takeTopRight,
   takeBottomLeft,
   takeBottomRight,
   ) where

import qualified Numeric.LAPACK.Matrix.Mosaic.Private as Mos
import qualified Numeric.LAPACK.Matrix.Mosaic.Unpacked as Unpacked
import qualified Numeric.LAPACK.Matrix.Layout.Private as Layout
import qualified Numeric.LAPACK.Matrix.Basic as Basic
import qualified Numeric.LAPACK.Matrix.Private as Matrix
import Numeric.LAPACK.Matrix.Mosaic.Private
         (MosaicLower, MosaicUpper, diagonalPointers, diagonalPointerPairs)
import Numeric.LAPACK.Matrix.Mosaic.Generic (transpose, unpackDirty, pack)
import Numeric.LAPACK.Matrix.Layout.Private (Order, uploOrder)
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Scalar (zero, one)
import Numeric.LAPACK.Private (fill)

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 Data.Array.Comfort.Shape ((::+))

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

import Data.Foldable (forM_)



type Mosaic mirror uplo sh =
         Array (Layout.Mosaic Layout.Packed mirror uplo sh)
type Triangular uplo sh =
         Array (Layout.TriangularP Layout.Packed uplo sh)


identity ::
   (Layout.Mirror mirror, Layout.UpLo uplo,
    Shape.C sh, Class.Floating a) =>
   Order -> sh -> Mosaic mirror uplo sh a
identity order sh =
   let (realOrder, uplo) = autoUploOrder order
   in Array.unsafeCreateWithSize
         (Layout.Mosaic
            Layout.Packed Layout.autoMirror uplo order sh) $
         \size aPtr -> do
      let n = Shape.size sh
      fill zero size aPtr
      mapM_ (flip poke one) (diagonalPointers realOrder n aPtr)

diagonal ::
   (Layout.Mirror mirror, Layout.UpLo uplo,
    Shape.C sh, Class.Floating a) =>
   Order -> Vector sh a -> Mosaic mirror uplo sh a
diagonal order (Array sh x) = do
   let (realOrder, uplo) = autoUploOrder order
   Array.unsafeCreateWithSize
      (Layout.Mosaic Layout.Packed Layout.autoMirror
         uplo order sh) $
            \size aPtr -> do
      let n = Shape.size sh
      fill zero size aPtr
      withForeignPtr x $ \xPtr ->
         forM_ (diagonalPointerPairs realOrder n xPtr aPtr) $
            \(srcPtr,dstPtr) -> poke dstPtr =<< peek srcPtr

takeDiagonal ::
   (Layout.UpLo uplo, Shape.C sh, Class.Floating a) =>
   Mosaic mirror uplo sh a -> Vector sh a
takeDiagonal (Array (Layout.Mosaic _pack _mirror uplo order sh) a) =
   Array.unsafeCreate sh $ \xPtr ->
   withForeignPtr a $ \aPtr ->
      mapM_
         (\(dstPtr,srcPtr) -> poke dstPtr =<< peek srcPtr)
         (diagonalPointerPairs (uploOrder uplo order) (Shape.size sh) xPtr aPtr)


{-
This is not maximally efficient.
It fills up a whole square.
This wastes memory but enables more regular memory access patterns.
-}
forceOrder ::
   (Layout.UpLo uplo, Shape.C sh, Class.Floating a) =>
   Order -> Mosaic mirror uplo sh a -> Mosaic mirror uplo sh a
forceOrder newOrder a =
   let shape = Array.shape a
   in if Layout.mosaicOrder shape == newOrder
         then a
         else pack $ Unpacked.forceOrder newOrder $ unpackDirty a


autoUploOrder ::
   (Layout.UpLo uplo) =>
   Order -> (Order, Layout.UpLoSingleton uplo)
autoUploOrder order =
   case Layout.autoUplo of
      uplo -> (uploOrder uplo order, uplo)


{-
It does not make much sense to put
'stackLower', 'stackUpper', 'stackSymmetric' in one function
because in 'stackLower' and 'stackUpper'
the height and width of matrix b is swapped.
-}
stackLower ::
   (Shape.C height, Eq height, Shape.C width, Eq width, Class.Floating a) =>
   MosaicLower mirror width a ->
   Matrix.General height width a ->
   MosaicLower mirror height a ->
   MosaicLower mirror (width::+height) a
stackLower a b c =
   transpose $ stackUpper (transpose a) (Basic.transpose b) (transpose c)

stackUpper ::
   (Shape.C height, Eq height, Shape.C width, Eq width, Class.Floating a) =>
   MosaicUpper mirror height a ->
   Matrix.General height width a ->
   MosaicUpper mirror width a ->
   MosaicUpper mirror (height::+width) a
stackUpper a b c =
   let order = Layout.fullOrder $ Array.shape b
   in Mos.stack (forceOrder order a) b (forceOrder order c)

takeTopLeft ::
   (Layout.UpLo uplo, Shape.C height, Shape.C width, Class.Floating a) =>
   Mosaic mirror uplo (height::+width) a ->
   Mosaic mirror uplo height a
takeTopLeft =
   Mos.getMap $
   Layout.switchUpLo
      (Mos.Map $ Mos.takeTopLeft)
      (Mos.Map $ transpose . Mos.takeTopLeft . transpose)

takeBottomLeft ::
   (Shape.C height, Shape.C width, Class.Floating a) =>
   MosaicLower mirror (width::+height) a -> Matrix.General height width a
takeBottomLeft = Basic.transpose . Mos.takeTopRight . transpose

takeBottomRight ::
   (Layout.UpLo uplo, Shape.C height, Shape.C width, Class.Floating a) =>
   Mosaic mirror uplo (height::+width) a ->
   Mosaic mirror uplo width a
takeBottomRight =
   Mos.getMap $
   Layout.switchUpLo
      (Mos.Map $ Mos.takeBottomRight)
      (Mos.Map $ transpose . Mos.takeBottomRight . transpose)