{-# LANGUAGE TypeFamilies #-}
module Numeric.LAPACK.Eigen.General (
   values,
   schur,
   decompose,
   ComplexOf,
   ) where

import Numeric.LAPACK.Matrix.Square (Square)

import qualified Numeric.LAPACK.Matrix.Shape.Private as MatrixShape
import Numeric.LAPACK.Matrix.Shape.Private (Order(RowMajor,ColumnMajor))
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Private
         (ComplexOf, RealOf, zero, withAutoWorkspaceInfo,
          copyToTemp, copyToColumnMajor, allocArray)

import qualified Numeric.LAPACK.FFI.Complex as LapackComplex
import qualified Numeric.LAPACK.FFI.Real as LapackReal
import qualified Numeric.BLAS.FFI.Complex as BlasComplex
import qualified Numeric.BLAS.FFI.Real as BlasReal
import qualified Numeric.Netlib.Utility as Call
import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Comfort.Storable.Internal as Array
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable.Internal (Array(Array))

import System.IO.Unsafe (unsafePerformIO)

import Foreign.Marshal.Array (advancePtr, peekArray)
import Foreign.C.Types (CInt, CChar)
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (Ptr, nullPtr, nullFunPtr, castPtr)
import Foreign.Storable (Storable)

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

import Data.Complex (Complex)


values ::
   (Shape.C sh, Class.Floating a) =>
   Square sh a -> Vector sh (ComplexOf a)
values =
   getValues $
   Class.switchFloating
      (Values valuesAux) (Values valuesAux)
      (Values valuesAux) (Values valuesAux)

type Values_ sh a = Square sh a -> Vector sh (ComplexOf a)

newtype Values sh a = Values {getValues :: Values_ sh a}

valuesAux ::
   (Shape.C sh, Class.Floating a, RealOf a ~ ar, Storable ar) =>
   Values_ sh a
valuesAux (Array (MatrixShape.Square _order size) a) =
      Array.unsafeCreateWithSize size $ \n wPtr -> do
   let lda = n
   evalContT $ do
      jobvsPtr <- Call.char 'N'
      sortPtr <- Call.char 'N'
      aPtr <- copyToTemp (n*n) a
      ldaPtr <- Call.cint lda
      sdimPtr <- Call.alloca
      let vsPtr = nullPtr
      ldvsPtr <- Call.cint n
      let bworkPtr = nullPtr
      liftIO $ withAutoWorkspaceInfo "gees" $ \workPtr lworkPtr infoPtr ->
         gees
            jobvsPtr sortPtr n aPtr ldaPtr sdimPtr
            wPtr vsPtr ldvsPtr workPtr lworkPtr bworkPtr infoPtr


{- |
If @(q,r) = schur a@, then @a = q \<#\> r \<#\> adjoint q@,
where @q@ is unitary (orthogonal)
and @r@ is a right-upper triangular matrix for complex @a@
and a 1x1-or-2x2-block upper triangular matrix for real @a@.
With @getDiagonal r@ you get all eigenvalues of @a@ if @a@ is complex
and the real parts of the eigenvalues if @a@ is real.
Complex conjugated eigenvalues of a real matrix @a@
are encoded as 2x2 blocks along the diagonal.
-}
schur ::
   (Shape.C sh, Class.Floating a) =>
   Square sh a -> (Square sh a, Square sh a)
schur =
   getSchur $
   Class.switchFloating
      (Schur schurAux) (Schur schurAux)
      (Schur schurAux) (Schur schurAux)

type Schur_ sh a = Square sh a -> (Square sh a, Square sh a)

newtype Schur sh a = Schur {getSchur :: Schur_ sh a}

schurAux ::
   (Shape.C sh, Class.Floating a, RealOf a ~ ar, Storable ar) =>
   Schur_ sh a
schurAux (Array (MatrixShape.Square order size) a) = unsafePerformIO $ do
   let n = Shape.size size
   let lda = n
   let sh = MatrixShape.Square ColumnMajor size
   evalContT $ do
      jobvsPtr <- Call.char 'V'
      sortPtr <- Call.char 'N'
      aPtr <- ContT $ withForeignPtr a
      (s,sPtr) <- allocArray sh
      liftIO $ copyToColumnMajor order n n aPtr sPtr
      ldaPtr <- Call.cint lda
      sdimPtr <- Call.alloca
      wPtr <- Call.allocaArray n
      (vs,vsPtr) <- allocArray sh
      ldvsPtr <- Call.cint n
      let bworkPtr = nullPtr
      liftIO $ withAutoWorkspaceInfo "gees" $ \workPtr lworkPtr infoPtr ->
         gees
            jobvsPtr sortPtr n sPtr ldaPtr sdimPtr
            wPtr vsPtr ldvsPtr workPtr lworkPtr bworkPtr infoPtr
      return (vs, s)



type GEES_ ar a =
   Ptr CChar -> Ptr CChar -> Int -> Ptr a -> Ptr CInt ->
   Ptr CInt -> Ptr (Complex ar) -> Ptr a -> Ptr CInt ->
   Ptr a -> Ptr CInt -> Ptr Bool -> Ptr CInt -> IO ()

newtype GEES a = GEES {getGEES :: GEES_ (RealOf a) a}

gees :: Class.Floating a => GEES_ (RealOf a) a
gees =
   getGEES $
   Class.switchFloating
      (GEES geesReal) (GEES geesReal) (GEES geesComplex) (GEES geesComplex)

geesReal :: Class.Real a => GEES_ a a
geesReal
      jobvsPtr sortPtr n aPtr ldaPtr sdimPtr
      wPtr vsPtr ldvsPtr workPtr lworkPtr bworkPtr infoPtr =
   evalContT $ do
      let selectPtr = nullFunPtr
      nPtr <- Call.cint n
      wrPtr <- Call.allocaArray n
      wiPtr <- Call.allocaArray n
      liftIO $
         LapackReal.gees
            jobvsPtr sortPtr selectPtr nPtr aPtr ldaPtr sdimPtr
            wrPtr wiPtr vsPtr ldvsPtr workPtr lworkPtr bworkPtr infoPtr
      liftIO $ zipComplex n wrPtr wiPtr wPtr

geesComplex :: Class.Real a => GEES_ a (Complex a)
geesComplex
      jobvsPtr sortPtr n aPtr ldaPtr sdimPtr
      wPtr vsPtr ldvsPtr workPtr lworkPtr bworkPtr infoPtr =
   evalContT $ do
      let selectPtr = nullFunPtr
      nPtr <- Call.cint n
      rworkPtr <- Call.allocaArray n
      liftIO $
         LapackComplex.gees
            jobvsPtr sortPtr selectPtr nPtr aPtr ldaPtr sdimPtr
            wPtr vsPtr ldvsPtr workPtr lworkPtr rworkPtr bworkPtr infoPtr



{- |
@(vr,d,vl) = Eigen.decompose a@

Counterintuitively, @vr@ contains the right eigenvectors
and @vl@ contains the left eigenvectors as columns.
The idea is to provide a decomposition of @a@.
If @a@ is diagonalizable, then @vr@ and @vl@ are almost inverse to each other.
More precisely, @adjoint vl \<#\> vr@ is a diagonal matrix.
This is because all eigenvectors are normalized to Euclidean norm 1.
With the following scaling, the decomposition becomes perfect:

> let scal = Array.map recip $ getDiagonal $ adjoint vl <#> vr
> a == vr <#> diagonal d <#> diagonal scal <#> adjoint vl

If @a@ is non-diagonalizable then some columns of @vr@ and @vl@ are left zero
and the above property does not hold.
-}
decompose ::
   (Shape.C sh, Class.Floating a) =>
   Square sh a ->
   (Square sh (ComplexOf a),
    Vector sh (ComplexOf a),
    Square sh (ComplexOf a))
decompose =
   getDecompose $
   Class.switchFloating
      (Decompose decomposeReal)
      (Decompose decomposeReal)
      (Decompose decomposeComplex)
      (Decompose decomposeComplex)

newtype Decompose sh a =
   Decompose {
      getDecompose ::
         Square sh a ->
         (Square sh (ComplexOf a),
          Vector sh (ComplexOf a),
          Square sh (ComplexOf a))
   }

decomposeReal ::
   (Shape.C sh, Class.Real a) =>
   Square sh a ->
   (Square sh (Complex a), Vector sh (Complex a), Square sh (Complex a))
decomposeReal (Array (MatrixShape.Square order size) a) =
      unsafePerformIO $ do
   let n = Shape.size size
   let lda = n
   evalContT $ do
      jobvlPtr <- Call.char 'V'
      jobvrPtr <- Call.char 'V'
      nPtr <- Call.cint n
      aPtr <- copyToTemp (n*n) a
      ldaPtr <- Call.cint lda
      wrPtr <- Call.allocaArray n
      wiPtr <- Call.allocaArray n
      vlPtr <- Call.allocaArray (n*n)
      ldvlPtr <- Call.cint n
      vrPtr <- Call.allocaArray (n*n)
      ldvrPtr <- Call.cint n
      liftIO $ withAutoWorkspaceInfo "geev" $
         LapackReal.geev
            jobvlPtr jobvrPtr nPtr aPtr ldaPtr
            wrPtr wiPtr vlPtr ldvlPtr vrPtr ldvrPtr
      (w,wPtr) <- allocArray size
      liftIO $ zipComplex n wrPtr wiPtr wPtr
      let sh = MatrixShape.Square ColumnMajor size
      (vlc,vlcPtr) <- allocArray sh
      (vrc,vrcPtr) <- allocArray sh
      liftIO $ eigenvectorsToComplex n wiPtr vlPtr vlcPtr
      liftIO $ eigenvectorsToComplex n wiPtr vrPtr vrcPtr
      return $
         case order of
            RowMajor -> (vlc, w, vrc)
            ColumnMajor -> (vrc, w, vlc)

eigenvectorsToComplex ::
   (Eq a, Class.Real a) =>
   Int -> Ptr a -> Ptr a -> Ptr (Complex a) -> IO ()
eigenvectorsToComplex n wiPtr vPtr vcPtr = evalContT $ do
   nPtr <- Call.cint n
   zeroPtr <- Call.real zero
   inc0Ptr <- Call.cint 0
   inc1Ptr <- Call.cint 1
   inc2Ptr <- Call.cint 2
   liftIO $ do
      let go _ _ [] = return ()
          go xPtr yPtr (False:wi) = do
            let yrPtr = castPtr yPtr
            let yiPtr = advancePtr yrPtr 1
            BlasReal.copy nPtr xPtr    inc1Ptr yrPtr inc2Ptr
            BlasReal.copy nPtr zeroPtr inc0Ptr yiPtr inc2Ptr
            go (advancePtr xPtr n) (advancePtr yPtr n) wi
          go xPtr yPtr (True:True:wi) = do
            let xrPtr = xPtr
            let xiPtr = advancePtr xPtr n
            let yrPtr = castPtr yPtr
            let yiPtr = advancePtr yrPtr 1
            let y1Ptr = advancePtr yPtr n
            BlasReal.copy nPtr xrPtr inc1Ptr yrPtr inc2Ptr
            BlasReal.copy nPtr xiPtr inc1Ptr yiPtr inc2Ptr
            BlasComplex.copy nPtr yPtr inc1Ptr y1Ptr inc1Ptr
            LapackComplex.lacgv nPtr y1Ptr inc1Ptr
            go (advancePtr xPtr (2*n)) (advancePtr yPtr (2*n)) wi
          go _xPtr _yPtr wi =
            error $ "eigenvectorToComplex: invalid non-real pattern " ++ show wi
      go vPtr vcPtr . map (zero/=) =<< peekArray n wiPtr

decomposeComplex ::
   (Shape.C sh, Class.Real a) =>
   Square sh (Complex a) ->
   (Square sh (Complex a), Vector sh (Complex a), Square sh (Complex a))
decomposeComplex (Array (MatrixShape.Square order size) a) =
      unsafePerformIO $ do
   let n = Shape.size size
   let lda = n
   evalContT $ do
      jobvlPtr <- Call.char 'V'
      jobvrPtr <- Call.char 'V'
      nPtr <- Call.cint n
      aPtr <- copyToTemp (n*n) a
      ldaPtr <- Call.cint lda
      (w,wPtr) <- allocArray size
      let sh = MatrixShape.Square ColumnMajor size
      (vl,vlPtr) <- allocArray sh
      ldvlPtr <- Call.cint n
      (vr,vrPtr) <- allocArray sh
      ldvrPtr <- Call.cint n
      rworkPtr <- Call.allocaArray (2*n)
      liftIO $ withAutoWorkspaceInfo "geev" $ \workPtr lworkPtr infoPtr ->
         LapackComplex.geev
            jobvlPtr jobvrPtr nPtr aPtr ldaPtr
            wPtr vlPtr ldvlPtr vrPtr ldvrPtr
            workPtr lworkPtr rworkPtr infoPtr
      return $
         case order of
            RowMajor -> (vl, w, vr)
            ColumnMajor -> (vr, w, vl)


zipComplex ::
   (Class.Real a) => Int -> Ptr a -> Ptr a -> Ptr (Complex a) -> IO ()
zipComplex n vr vi vc =
   evalContT $ do
      nPtr <- Call.cint n
      incxPtr <- Call.cint 1
      incyPtr <- Call.cint 2
      let yPtr = castPtr vc
      liftIO $ BlasReal.copy nPtr vr incxPtr yPtr incyPtr
      liftIO $ BlasReal.copy nPtr vi incxPtr (advancePtr yPtr 1) incyPtr