{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GADTs #-}
module Numeric.LAPACK.Matrix.Mosaic.Basic (
   Mosaic.fromList,
   Mosaic.autoFromList,

   Mosaic.transpose,
   Mosaic.adjoint,

   Mosaic.repack,
   pack,
   unpack,
   reunpack,
   Mosaic.unpackDirty,

   Mos.takeUpper,
   Mos.fromUpper,

   takeDiagonal,
   toSquare,
   square,
   power,
   powers1,
   ) where

import qualified Numeric.LAPACK.Matrix.Symmetric.Unified as Symmetric
import qualified Numeric.LAPACK.Matrix.Mosaic.Unpacked as Unpacked
import qualified Numeric.LAPACK.Matrix.Mosaic.Packed as Packed
import qualified Numeric.LAPACK.Matrix.Mosaic.Generic as Mosaic
import qualified Numeric.LAPACK.Matrix.Mosaic.Private as Mos
import qualified Numeric.LAPACK.Matrix.Triangular.Basic as Triangular
import qualified Numeric.LAPACK.Matrix.Square.Basic as Square
import qualified Numeric.LAPACK.Matrix.Basic as Basic
import qualified Numeric.LAPACK.Matrix.Layout.Private as Layout
import Numeric.LAPACK.Matrix.Mosaic.Private (Mosaic)
import Numeric.LAPACK.Matrix.Shape.Omni (TriDiag, DiagSingleton)
import Numeric.LAPACK.Matrix.Private (Square)
import Numeric.LAPACK.Vector (Vector)

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.Stream (Stream)


takeDiagonal ::
   (Layout.UpLo uplo, Shape.C sh, Class.Floating a) =>
   Mosaic pack mirror uplo sh a -> Vector sh a
takeDiagonal a =
   case Layout.mosaicPack $ Array.shape a of
      Layout.Unpacked -> Square.takeDiagonal $ Unpacked.toSquare a
      Layout.Packed -> Packed.takeDiagonal a

toSquare ::
   (Layout.UpLo uplo, Shape.C sh, Class.Floating a) =>
   Packed.Mosaic mirror uplo sh a -> Square sh a
toSquare a =
   let shape = Array.shape a
   in case Layout.mosaicMirror shape of
         Layout.NoMirror -> Triangular.toSquare a
         Layout.SimpleMirror ->
            case Layout.mosaicUplo shape of
               Layout.Upper -> Symmetric.toSquare a
               Layout.Lower ->
                  Basic.transpose $ Symmetric.toSquare $ Mosaic.transpose a
         Layout.ConjugateMirror ->
            case Layout.mosaicUplo shape of
               Layout.Upper -> Symmetric.toSquare a
               Layout.Lower ->
                  Basic.transpose $ Symmetric.toSquare $ Mosaic.transpose a

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

unpack ::
   (Layout.UpLo uplo, Shape.C sh, Class.Floating a) =>
   Mosaic pack mirror uplo sh a -> Unpacked.Mosaic mirror uplo sh a
unpack a =
   let shape = Array.shape a in
   case Layout.mosaicPack shape of
      Layout.Unpacked -> a
      Layout.Packed ->
         Unpacked.fromSquare (Layout.mosaicMirror shape) $ toSquare a

reunpack ::
   (Layout.Packing pack, Layout.UpLo uplo, Shape.C sh, Class.Floating a) =>
   Packed.Mosaic mirror uplo sh a -> Mosaic pack mirror uplo sh a
reunpack a = Mosaic.withPack $ \packing ->
   case packing of
      Layout.Unpacked -> unpack a
      Layout.Packed -> a


square ::
   (Layout.Packing pack, TriDiag diag, Layout.UpLo uplo,
    Shape.C sh, Class.Floating a) =>
   DiagSingleton diag ->
   Mosaic pack mirror uplo sh a -> Mosaic pack mirror uplo sh a
square diag = Mosaic.repack . Unpacked.square diag . unpack

power ::
   (Layout.Packing pack, TriDiag diag, Layout.UpLo uplo,
    Shape.C sh, Class.Floating a) =>
   DiagSingleton diag ->
   Integer -> Mosaic pack mirror uplo sh a -> Mosaic pack mirror uplo sh a
power diag n = Mosaic.repack . Unpacked.power diag n . unpack

powers1 ::
   (Layout.Packing pack, TriDiag diag, Layout.UpLo uplo,
    Shape.C sh, Class.Floating a) =>
   DiagSingleton diag ->
   Mosaic pack mirror uplo sh a -> Stream (Mosaic pack mirror uplo sh a)
powers1 diag = fmap Mosaic.repack . Unpacked.powers1 diag . unpack