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

import Numeric.LAPACK.Matrix.Hermitian (Hermitian)
import Numeric.LAPACK.Matrix.Square (Square)

import qualified Numeric.LAPACK.Matrix.Shape.Private as MatrixShape
import Numeric.LAPACK.Matrix.Shape.Private
         (Order(ColumnMajor), uploFromOrder, triangleSize)
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Private (RealOf, copyToTemp, allocArray)

import qualified Numeric.LAPACK.FFI.Complex as LapackComplex
import qualified Numeric.LAPACK.FFI.Real as LapackReal
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.Alloc (alloca)
import Foreign.C.Types (CInt, CChar)
import Foreign.Ptr (Ptr, nullPtr)
import Foreign.Storable (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 ::
   (Shape.C sh, Class.Floating a) =>
   Hermitian sh a -> Vector sh (RealOf a)
values =
   getValues $
   Class.switchFloating
      (Values valuesAux) (Values valuesAux)
      (Values valuesAux) (Values valuesAux)

newtype Values sh a =
   Values {getValues :: Hermitian sh a -> Vector sh (RealOf a)}

valuesAux ::
   (Shape.C sh, Class.Floating a, RealOf a ~ ar, Storable ar) =>
   Hermitian sh a -> Vector sh ar
valuesAux (Array (MatrixShape.Hermitian order size) a) =
      Array.unsafeCreateWithSize size $ \n wPtr -> do
   evalContT $ do
      jobzPtr <- Call.char 'N'
      uploPtr <- Call.char $ uploFromOrder order
      aPtr <- copyToTemp (triangleSize n) a
      let zPtr = nullPtr
      ldzPtr <- Call.cint (max 1 n)
      liftIO $ withInfo "hpev" $
         hpev jobzPtr uploPtr n aPtr wPtr zPtr ldzPtr


{- |
For symmetric eigenvalue problems, @Eigen.decompose@ and @schur@ coincide.
-}
decompose ::
   (Shape.C sh, Class.Floating a) =>
   Hermitian sh a -> (Square sh a, Vector sh (RealOf a))
decompose =
   getDecompose $
   Class.switchFloating
      (Decompose decomposeAux) (Decompose decomposeAux)
      (Decompose decomposeAux) (Decompose decomposeAux)

type Decompose_ sh a = Hermitian sh a -> (Square sh a, Vector sh (RealOf a))

newtype Decompose sh a = Decompose {getDecompose :: Decompose_ sh a}

decomposeAux ::
   (Shape.C sh, Class.Floating a, RealOf a ~ ar, Storable ar) =>
   Decompose_ sh a
decomposeAux (Array (MatrixShape.Hermitian order size) a) = unsafePerformIO $ do
   let n = Shape.size size
   let shZ = MatrixShape.Square ColumnMajor size
   evalContT $ do
      jobzPtr <- Call.char 'V'
      uploPtr <- Call.char $ uploFromOrder order
      aPtr <- copyToTemp (triangleSize n) a
      (w,wPtr) <- allocArray size
      (z,zPtr) <- allocArray shZ
      ldzPtr <- Call.cint n
      liftIO $ withInfo "hpev" $
         hpev jobzPtr uploPtr n aPtr wPtr zPtr ldzPtr
      return (z, w)


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 HPEV_ ar a =
   Ptr CChar -> Ptr CChar -> Int -> Ptr a -> Ptr ar ->
   Ptr a -> Ptr CInt -> Ptr CInt -> IO ()

newtype HPEV a = HPEV {getHPEV :: HPEV_ (RealOf a) a}

hpev :: Class.Floating a => HPEV_ (RealOf a) a
hpev =
   getHPEV $
   Class.switchFloating
      (HPEV spevReal) (HPEV spevReal) (HPEV hpevComplex) (HPEV hpevComplex)

spevReal :: Class.Real a => HPEV_ a a
spevReal jobzPtr uploPtr n apPtr wPtr zPtr ldzPtr infoPtr =
   evalContT $ do
      nPtr <- Call.cint n
      workPtr <- Call.allocaArray (3*n)
      liftIO $
         LapackReal.spev
            jobzPtr uploPtr nPtr apPtr wPtr zPtr ldzPtr workPtr infoPtr

hpevComplex :: Class.Real a => HPEV_ a (Complex a)
hpevComplex jobzPtr uploPtr n apPtr wPtr zPtr ldzPtr infoPtr =
   evalContT $ do
      nPtr <- Call.cint n
      workPtr <- Call.allocaArray (max 1 (2*n-1))
      rworkPtr <- Call.allocaArray (max 1 (3*n-2))
      liftIO $
         LapackComplex.hpev
            jobzPtr uploPtr nPtr apPtr wPtr zPtr ldzPtr workPtr rworkPtr infoPtr