{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Numeric.LAPACK.Matrix.BandedHermitian (
   BandedHermitian,
   Transposition(..),
   size,
   fromList,
   identity,
   diagonal,
   takeDiagonal,
   toHermitian,
   toBanded,
   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.Banded as Banded
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix.Shape.Private as MatrixShape
import qualified Numeric.LAPACK.Matrix.Extent as Extent
import Numeric.LAPACK.Matrix.Array.Banded (Square)
import Numeric.LAPACK.Matrix.Array.Triangular (Hermitian)
import Numeric.LAPACK.Matrix.Array (Full)
import Numeric.LAPACK.Matrix.Shape.Private (Order, UnaryProxy)
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 as Unary
import Type.Data.Num.Unary ((:+:))

import qualified Data.Array.Comfort.Shape as Shape

import Foreign.Storable (Storable)

import Data.Tuple.HT (mapFst)


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


size :: BandedHermitian offDiag sh a -> sh
size = MatrixShape.bandedHermitianSize . ArrMatrix.shape


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

identity :: (Shape.C sh, Class.Floating a) => sh -> Diagonal 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 ::
   (Unary.Natural offDiag, Shape.C size, Class.Floating a) =>
   BandedHermitian offDiag size a -> Vector size (RealOf a)
takeDiagonal = Basic.takeDiagonal . ArrMatrix.toVector

toHermitian ::
   (Unary.Natural offDiag, Shape.C size, Class.Floating a) =>
   BandedHermitian offDiag size a -> Hermitian size a
toHermitian = ArrMatrix.lift1 Basic.toHermitian

toBanded ::
   (Unary.Natural offDiag, Shape.C size, Class.Floating a) =>
   BandedHermitian offDiag size a ->
   Square offDiag offDiag size a
toBanded = ArrMatrix.lift1 Basic.toBanded

multiplyVector ::
   (Unary.Natural offDiag, Shape.C size, Eq size, Class.Floating a) =>
   Transposition -> BandedHermitian 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 ->
   BandedHermitian (sub :+: super) size a
gramian = ArrMatrix.lift1 Basic.gramian

multiplyFull ::
   (Unary.Natural offDiag, Extent.C vert, Extent.C horiz,
    Shape.C height, Eq height, Shape.C width, Class.Floating a) =>
   Transposition -> BandedHermitian offDiag height a ->
   Full vert horiz height width a ->
   Full 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))] ->
   BandedHermitian k sh a
sumRank1 order sh = ArrMatrix.lift0 . Basic.sumRank1 order sh


eigenvalues ::
   (Unary.Natural offDiag, Shape.C sh, Class.Floating a) =>
   BandedHermitian offDiag sh a -> Vector sh (RealOf a)
eigenvalues = Eigen.values . ArrMatrix.toVector

{- |
For symmetric eigenvalue problems, @eigensystem@ and @schur@ coincide.
-}
eigensystem ::
   (Unary.Natural offDiag, Shape.C sh, Class.Floating a) =>
   BandedHermitian offDiag sh a -> (ArrMatrix.Square sh a, Vector sh (RealOf a))
eigensystem = mapFst ArrMatrix.lift0 . Eigen.decompose . ArrMatrix.toVector