{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Numeric.LAPACK.Matrix.BandedHermitian (
   BandedHermitian,
   Transposition(..),

   Hermitian.Semidefinite,
   Hermitian.assureFullRank,
   Hermitian.assureAnyRank,
   Hermitian.relaxSemidefinite,
   Hermitian.relaxIndefinite,
   Hermitian.assurePositiveDefiniteness,
   Hermitian.relaxDefiniteness,

   size,
   fromList,
   identity,
   diagonal,
   takeDiagonal,
   toHermitian,
   toBanded,
   forceOrder,
   takeTopLeft,
   takeBottomRight,
   negate,
   multiplyVector,
   multiplyFull,
   gramian,
   sumRank1,

   eigenvalues,
   eigensystem,
   ) where

import qualified Numeric.LAPACK.Matrix.BandedHermitian.Eigen as Eigen
import qualified Numeric.LAPACK.Matrix.BandedHermitian.Basic as Basic

import qualified Numeric.LAPACK.Matrix.Array.Hermitian as Hermitian
import qualified Numeric.LAPACK.Matrix.Array.Banded as Banded
import qualified Numeric.LAPACK.Matrix.Array.Private as ArrMatrix
import qualified Numeric.LAPACK.Matrix.Shape.Omni as Omni
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.Extent as Extent
import qualified Numeric.LAPACK.Vector as Vector
import qualified Numeric.LAPACK.Shape as ExtShape
import Numeric.LAPACK.Matrix.Array.Banded (Square)
import Numeric.LAPACK.Matrix.Array.Mosaic (FlexHermitian)
import Numeric.LAPACK.Matrix.Array.Private (Full)
import Numeric.LAPACK.Matrix.Layout.Private (Order, UnaryProxy, natFromProxy)
import Numeric.LAPACK.Matrix.Modifier (Transposition(NonTransposed, Transposed))
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Scalar (RealOf)

import qualified Numeric.Netlib.Class as Class

import qualified Type.Data.Num.Unary.Literal as TypeNum
import qualified Type.Data.Num.Unary.Proof as Proof
import qualified Type.Data.Num.Unary as Unary
import qualified Type.Data.Bool as TBool
import Type.Data.Num.Unary ((:+:))

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

import Data.Tuple.HT (mapPair, mapFst)

import Prelude hiding (negate)


type BandedHermitian offDiag sh = Banded.Hermitian offDiag sh
type Diagonal size = BandedHermitian TypeNum.U0 size


size :: BandedHermitian offDiag sh a -> sh
size = Omni.height . ArrMatrix.shape


fromList ::
   (Unary.Natural offDiag, Shape.C size, Class.Floating a) =>
   UnaryProxy offDiag -> Order -> size -> [a] ->
   BandedHermitian offDiag size a
fromList numOff order size_ =
   ArrMatrix.fromVector . Basic.fromList numOff order size_

identity ::
   (Shape.C sh, Class.Floating a) =>
   sh -> Banded.HermitianPosDef TypeNum.U0 sh a
identity = ArrMatrix.lift0 . Basic.identity

diagonal ::
   (Shape.C sh, Class.Floating a) => Vector sh (RealOf a) -> Diagonal sh a
diagonal = ArrMatrix.lift0 . Basic.diagonal

takeDiagonal ::
   (TBool.C neg, TBool.C zero, TBool.C pos) =>
   (Unary.Natural offDiag, Shape.C size, Class.Floating a) =>
   Banded.FlexHermitian neg zero pos offDiag size a ->
   Vector size (RealOf a)
takeDiagonal = Basic.takeDiagonal . ArrMatrix.toVector

toHermitian ::
   (TBool.C neg, TBool.C zero, TBool.C pos) =>
   (Unary.Natural offDiag, Shape.C size, Class.Floating a) =>
   Banded.FlexHermitian neg zero pos offDiag size a ->
   FlexHermitian neg zero pos size a
toHermitian = ArrMatrix.lift1 Basic.toHermitian

toBanded ::
   (TBool.C neg, TBool.C zero, TBool.C pos,
    Unary.Natural offDiag, Shape.C size, Class.Floating a) =>
   Banded.FlexHermitian neg zero pos offDiag size a ->
   Square offDiag offDiag size a
toBanded = ArrMatrix.lift1 Basic.toBanded


forceOrder ::
   (TBool.C neg, TBool.C zero, TBool.C pos,
    Unary.Natural offDiag, Shape.C size, Class.Floating a) =>
   Order ->
   Banded.FlexHermitian neg zero pos offDiag size a ->
   Banded.FlexHermitian neg zero pos offDiag size a
forceOrder = ArrMatrix.lift1 . Basic.forceOrder


takeTopLeft ::
   (TBool.C neg, TBool.C zero, TBool.C pos,
    Unary.Natural offDiag, Shape.C sh0, Shape.C sh1, Class.Floating a) =>
   Banded.FlexHermitian neg zero pos offDiag (sh0 ::+ sh1) a ->
   Banded.FlexHermitian neg zero pos offDiag sh0 a
takeTopLeft = ArrMatrix.lift1 Basic.takeTopLeft

takeBottomRight ::
   (TBool.C neg, TBool.C zero, TBool.C pos,
    Unary.Natural offDiag, Shape.C sh0, Shape.C sh1, Class.Floating a) =>
   Banded.FlexHermitian neg zero pos offDiag (sh0 ::+ sh1) a ->
   Banded.FlexHermitian neg zero pos offDiag sh1 a
takeBottomRight = ArrMatrix.lift1 Basic.takeBottomRight


negate ::
   (TBool.C neg, TBool.C zero, TBool.C pos,
    Unary.Natural offDiag, Shape.C sh, Class.Floating a) =>
   Banded.FlexHermitian neg zero pos offDiag sh a ->
   Banded.FlexHermitian pos zero neg offDiag sh a
negate = ArrMatrix.lift1 Vector.negate


multiplyVector ::
   (TBool.C neg, TBool.C zero, TBool.C pos,
    Unary.Natural offDiag, Shape.C size, Eq size, Class.Floating a) =>
   Transposition ->
   Banded.FlexHermitian neg zero pos offDiag size a ->
   Vector size a -> Vector size a
multiplyVector transposed =
   Basic.multiplyVector transposed .  ArrMatrix.toVector

gramian ::
   (Shape.C size, Eq size, Class.Floating a,
    Unary.Natural sub, Unary.Natural super) =>
   Square sub super size a ->
   Banded.HermitianPosSemidef (sub :+: super) size a
gramian a =
   case mapPair (natFromProxy,natFromProxy) $
        MatrixShape.bandedOffDiagonals $ ArrMatrix.shape a of
      (sub,super) ->
         case Proof.addNat sub super of
            Proof.Nat -> ArrMatrix.lift1 Basic.gramian a

multiplyFull ::
   (TBool.C neg, TBool.C zero, TBool.C pos,
    Unary.Natural offDiag, Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Eq height, Shape.C width, Class.Floating a) =>
   Transposition ->
   Banded.FlexHermitian neg zero pos offDiag height a ->
   Full meas vert horiz height width a ->
   Full meas vert horiz height width a
multiplyFull = ArrMatrix.lift2 . Basic.multiplyFull

{- |
The list represents ragged rows of a sparse matrix.
-}
sumRank1 ::
   (Unary.Natural k, Shape.Indexed sh, Class.Floating a) =>
   Order -> sh ->
   [(RealOf a, (Shape.Index sh, Basic.StaticVector (Unary.Succ k) a))] ->
   Banded.HermitianPosSemidef k sh a
sumRank1 order sh = ArrMatrix.lift0 . Basic.sumRank1 order sh


eigenvalues ::
   (TBool.C neg, TBool.C zero, TBool.C pos, Unary.Natural offDiag) =>
   (ExtShape.Permutable sh, Class.Floating a) =>
   Banded.FlexHermitian neg zero pos offDiag sh a -> Vector sh (RealOf a)
eigenvalues = Eigen.values . ArrMatrix.toVector

{- |
For symmetric eigenvalue problems, @eigensystem@ and @schur@ coincide.
-}
eigensystem ::
   (TBool.C neg, TBool.C zero, TBool.C pos, Unary.Natural offDiag) =>
   (ExtShape.Permutable sh, Class.Floating a) =>
   Banded.FlexHermitian neg zero pos offDiag sh a ->
   (ArrMatrix.Square sh a, Vector sh (RealOf a))
eigensystem = mapFst ArrMatrix.lift0 . Eigen.decompose . ArrMatrix.toVector