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

import qualified Numeric.LAPACK.Matrix.Triangular as Triangular
import qualified Numeric.LAPACK.Matrix.Shape.Private as MatrixShape
import Numeric.LAPACK.Matrix.Triangular.Private
         (unpackZero, pack, unpackToTemp, fillTriangle,
          forPointers, rowMajorPointers)
import Numeric.LAPACK.Matrix.Triangular (Triangular)
import Numeric.LAPACK.Matrix.Shape.Private
         (Order(ColumnMajor,RowMajor), caseUplo, uploOrder, triangleSize)
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Private (zero, lacgv, allocArray)

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

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

import System.IO.Unsafe (unsafePerformIO)

import Foreign.Marshal.Alloc (alloca)
import Foreign.C.Types (CInt, CChar)
import Foreign.Ptr (Ptr, nullPtr)
import Foreign.Storable (peek)

import Control.Monad.Trans.Cont (evalContT)
import Control.Monad.IO.Class (liftIO)
import Control.Applicative ((<$>))

import Text.Printf (printf)

import Data.Complex (Complex)


values ::
   (MatrixShape.Uplo uplo, Shape.C sh, Class.Floating a) =>
   Triangular uplo sh a -> Vector sh a
values = Triangular.getDiagonal


{- |
@(vr,d,vlAdj) = TriEigen.decompose a@

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

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

If @a@ is non-diagonalizable
then some columns of @vr@ and corresponding rows of @vlAdj@ are left zero
and the above property does not hold.
-}
decompose ::
   (MatrixShape.Uplo uplo, Shape.C sh, Class.Floating a) =>
   Triangular uplo sh a ->
   (Triangular uplo sh a, Vector sh a, Triangular uplo sh a)
decompose a =
   let (vr,vl) = decomposePlain a
   in  (vr, values a, vl)

decomposePlain ::
   (MatrixShape.Uplo uplo, Shape.C sh, Class.Floating a) =>
   Triangular uplo sh a -> (Triangular uplo sh a, Triangular uplo sh a)
decomposePlain (Array (MatrixShape.Triangular uplo order sh) a) =
      unsafePerformIO $ do
   let n = Shape.size sh
   let n2 = n*n
   let triSize = triangleSize n
   evalContT $ do
      sidePtr <- Call.char 'B'
      howManyPtr <- Call.char 'A'
      let selectPtr = nullPtr
      let unpk =
            case uploOrder uplo order of
               ColumnMajor -> unpackZero ColumnMajor
               RowMajor -> unpackZeroRowMajor
      aPtr <- unpackToTemp unpk n a
      ldaPtr <- Call.cint n
      vlPtr <- Call.allocaArray n2
      vrPtr <- Call.allocaArray n2
      mmPtr <- Call.cint n
      mPtr <- Call.alloca
      liftIO $ withInfo "trevc" $
         trevc sidePtr howManyPtr selectPtr n
            aPtr ldaPtr vlPtr ldaPtr vrPtr ldaPtr mmPtr mPtr
      (vl,vlpPtr) <-
         allocArray $
         MatrixShape.Triangular uplo (uploOrder uplo RowMajor) sh
      (vr,vrpPtr) <-
         allocArray $
         MatrixShape.Triangular uplo (uploOrder uplo ColumnMajor) sh
      sizePtr <- Call.cint triSize
      incPtr <- Call.cint 1
      liftIO $ do
         pack ColumnMajor n vrPtr vrpPtr
         pack RowMajor n vlPtr vlpPtr
         lacgv sizePtr vlpPtr incPtr
      return $ caseUplo uplo (vl,vr) (vr,vl)


unpackZeroRowMajor :: Class.Floating a => Int -> Ptr a -> Ptr a -> IO ()
unpackZeroRowMajor n packedPtr fullPtr = do
   fillTriangle zero RowMajor n fullPtr
   unpackRowMajor n packedPtr fullPtr

unpackRowMajor :: Class.Floating a => Int -> Ptr a -> Ptr a -> IO ()
unpackRowMajor n packedPtr fullPtr = evalContT $ do
   incxPtr <- Call.cint 1
   incyPtr <- Call.cint n
   liftIO $
      forPointers (rowMajorPointers n fullPtr packedPtr) $
            \nPtr (dstPtr,srcPtr) ->
         BlasGen.copy nPtr srcPtr incxPtr dstPtr incyPtr


withInfo :: String -> (Ptr CInt -> IO ()) -> IO ()
withInfo name computation = alloca $ \infoPtr -> do
   computation infoPtr
   info <- fromIntegral <$> peek infoPtr
   case compare info (0::Int) of
      EQ -> return ()
      LT -> error $ printf "%s: illegal value in %d-th argument" name (-info)
      GT -> error $
         printf "%s: %d off-diagonal elements not converging" name info


type TREVC_ a =
   Ptr CChar -> Ptr CChar -> Ptr Bool ->
   Int -> Ptr a -> Ptr CInt -> Ptr a -> Ptr CInt -> Ptr a -> Ptr CInt ->
   Ptr CInt -> Ptr CInt -> Ptr CInt -> IO ()

newtype TREVC a = TREVC {getTREVC :: TREVC_ a}

trevc :: Class.Floating a => TREVC_ a
trevc =
   getTREVC $
   Class.switchFloating
      (TREVC trevcReal) (TREVC trevcReal)
      (TREVC trevcComplex) (TREVC trevcComplex)

trevcReal :: Class.Real a => TREVC_ a
trevcReal sidePtr howmnyPtr selectPtr n
      tPtr ldtPtr vlPtr ldvlPtr vrPtr ldvrPtr mmPtr mPtr infoPtr =
   evalContT $ do
      nPtr <- Call.cint n
      workPtr <- Call.allocaArray (3*n)
      liftIO $
         LapackReal.trevc sidePtr howmnyPtr selectPtr nPtr
            tPtr ldtPtr vlPtr ldvlPtr vrPtr ldvrPtr mmPtr mPtr workPtr infoPtr

trevcComplex :: Class.Real a => TREVC_ (Complex a)
trevcComplex sidePtr howmnyPtr selectPtr n
      tPtr ldtPtr vlPtr ldvlPtr vrPtr ldvrPtr mmPtr mPtr infoPtr =
   evalContT $ do
      nPtr <- Call.cint n
      workPtr <- Call.allocaArray (2*n)
      rworkPtr <- Call.allocaArray n
      liftIO $
         LapackComplex.trevc sidePtr howmnyPtr selectPtr nPtr
            tPtr ldtPtr vlPtr ldvlPtr vrPtr ldvrPtr mmPtr mPtr
            workPtr rworkPtr infoPtr