{-# LANGUAGE TypeFamilies #-}
module Numeric.LAPACK.LinearSystem (
   leastSquares,
   minimumNorm,
   leastSquaresMinimumNorm,
   pseudoInverseRCond,

   Householder,
   householder,
   householderDecompose,
   householderDeterminant,
   determinant,
   householderExtractQ,
   householderExtractR,
   orthogonalComplement,
   ) where

import Numeric.LAPACK.Matrix
         (General, ZeroInt, zeroInt, transpose, identity, dropColumns)

import qualified Numeric.LAPACK.Matrix.Shape.Private as MatrixShape
import Numeric.LAPACK.Matrix.Shape.Private
         (Order(RowMajor, ColumnMajor), charFromOrder)
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Private
         (RealOf, zero, one, fill, pointerSeq,
          copyTransposed, copySubMatrix, copyBlock)

import qualified Numeric.LAPACK.FFI.Generic as LapackGen
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.Array (allocaArray, advancePtr)
import Foreign.Marshal.Alloc (alloca)
import Foreign.C.Types (CInt)
import Foreign.ForeignPtr (withForeignPtr, mallocForeignPtrArray)
import Foreign.Ptr (Ptr)
import Foreign.Storable (Storable, poke, peek)

import Text.Printf (printf)

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

import qualified Data.Complex as Complex
import Data.Complex (Complex)
import Data.Tuple.HT (mapSnd)


{- |
If @x = leastSquares a b@
then @x@ minimizes @Vector.norm2 (multiply a x `sub` b)@.

Precondition: @a@ must have full rank and @height a >= width a@.
-}
leastSquares ::
   (Shape.C height, Eq height, Shape.C width, Shape.C nrhs,
    Storable a, Class.Floating a) =>
   General height width a -> General height nrhs a -> General width nrhs a
leastSquares
   (Array shapeA@(MatrixShape.General orderA heightA widthA) a)
   (Array        (MatrixShape.General orderB heightB widthB) b) =
      Array.unsafeCreate (MatrixShape.General ColumnMajor widthA widthB) $
         \xPtr -> do
   Call.assert "leastSquares: height shapes mismatch" (heightA == heightB)
   Call.assert "leastSquares: height of 'a' must be at least the width"
      (Shape.size heightA >= Shape.size widthA)
   let (m,n) = MatrixShape.dimensions shapeA
   let lda = m
   let nrhs = Shape.size widthB
   let ldb = Shape.size heightB
   let ldx = Shape.size widthA
   evalContT $ do
      transPtr <- Call.char $ charFromOrder orderA
      mPtr <- Call.cint m
      nPtr <- Call.cint n
      nrhsPtr <- Call.cint nrhs
      aPtr <- ContT $ withForeignPtr a
      ldaPtr <- Call.cint lda
      let aSize = Shape.size (heightA,widthA)
      atmpPtr <- Call.allocaArray aSize
      liftIO $ copyBlock aSize aPtr atmpPtr
      bPtr <- ContT $ withForeignPtr b
      ldbPtr <- Call.cint ldb
      let bSize = Shape.size (heightB,widthB)
      btmpPtr <- Call.allocaArray bSize
      liftIO $ copyToColumnMajor orderB ldb nrhs bPtr btmpPtr
      liftIO $ withAutoWorkspaceInfo "gels" $
         LapackGen.gels transPtr
            mPtr nPtr nrhsPtr atmpPtr ldaPtr btmpPtr ldbPtr
      liftIO $ copySubMatrix ldx nrhs ldb btmpPtr ldx xPtr

{- |
The vector @x@ with @x = minimumNorm a b@
is the vector with minimal @Vector.norm2 x@
that satisfies @multiply a x == b@.

Precondition: @a@ must have full rank and @height a <= width a@.
-}
minimumNorm ::
   (Shape.C height, Eq height, Shape.C width, Shape.C nrhs,
    Storable a, Class.Floating a) =>
   General height width a -> General height nrhs a -> General width nrhs a
minimumNorm
   (Array shapeA@(MatrixShape.General orderA heightA widthA) a)
   (Array        (MatrixShape.General orderB heightB widthB) b) =
      Array.unsafeCreate (MatrixShape.General ColumnMajor widthA widthB) $
         \xPtr -> do
   Call.assert "minimumNorm: height shapes mismatch" (heightA == heightB)
   Call.assert "minimumNorm: width of 'a' must be at least the height"
      (Shape.size widthA >= Shape.size heightA)
   let (m,n) = MatrixShape.dimensions shapeA
   let lda = m
   let nrhs = Shape.size widthB
   let ldb = Shape.size heightB
   let ldx = Shape.size widthA
   evalContT $ do
      transPtr <- Call.char $ charFromOrder orderA
      mPtr <- Call.cint m
      nPtr <- Call.cint n
      nrhsPtr <- Call.cint nrhs
      aPtr <- ContT $ withForeignPtr a
      ldaPtr <- Call.cint lda
      let aSize = Shape.size (heightA,widthA)
      atmpPtr <- Call.allocaArray aSize
      liftIO $ copyBlock aSize aPtr atmpPtr
      bPtr <- ContT $ withForeignPtr b
      ldxPtr <- Call.cint ldx
      liftIO $ copyToSubColumnMajor orderB ldb nrhs bPtr ldx xPtr
      liftIO $ withAutoWorkspaceInfo "gels" $
         LapackGen.gels transPtr
            mPtr nPtr nrhsPtr atmpPtr ldaPtr xPtr ldxPtr

{- |
If @x = leastSquaresMinimumNorm a b@
then @x@ is the vector with minimum @Vector.norm2 x@
that minimizes @Vector.norm2 (multiply a x `sub` b)@.

Matrix @a@ can have any rank
but you must specify the reciprocal condition of the rank-truncated matrix.
-}
leastSquaresMinimumNorm ::
   (Shape.C height, Eq height, Shape.C width, Shape.C nrhs,
    Storable a, Class.Floating a) =>
   RealOf a ->
   General height width a -> General height nrhs a ->
   (Int, General width nrhs a)
leastSquaresMinimumNorm rcond
   (Array (MatrixShape.General orderA heightA widthA) a)
   (Array (MatrixShape.General orderB heightB widthB) b) =
      unsafePerformIO $ do
   Call.assert "minimumNorm: height shapes mismatch" (heightA == heightB)
   let shapeX = MatrixShape.General ColumnMajor widthA widthB
   let m = Shape.size heightA
   let n = Shape.size widthA
   let nrhs = Shape.size widthB
   let aSize = m*n
   let lda = m
   let ldtmp = max m n
   let tmpSize = ldtmp*nrhs
   evalContT $ do
      aPtr <- ContT $ withForeignPtr a
      atmpPtr <- Call.allocaArray aSize
      liftIO $ copyToColumnMajor orderA m n aPtr atmpPtr
      ldaPtr <- Call.cint lda
      bPtr <- ContT $ withForeignPtr b
      let needTmp = m>n
      x <- liftIO $ mallocForeignPtrArray $ Shape.size shapeX
      tmpPtr <-
         ContT $ if needTmp then allocaArray tmpSize else withForeignPtr x
      ldtmpPtr <- Call.cint ldtmp
      liftIO $ copyToSubColumnMajor orderB m nrhs bPtr ldtmp tmpPtr
      jpvtPtr <- Call.allocaArray n
      rankPtr <- Call.alloca
      gelsy m n nrhs atmpPtr ldaPtr tmpPtr ldtmpPtr jpvtPtr rcond rankPtr
      when needTmp $ liftIO $
         withForeignPtr x $ copySubMatrix n nrhs ldtmp tmpPtr n
      rank <- liftIO $ fromIntegral <$> peek rankPtr
      return (rank, Array shapeX x)


newtype GELSY r a =
   GELSY {
      getGELSY ::
         Int -> Int -> Int -> Ptr a -> Ptr CInt -> Ptr a -> Ptr CInt ->
         Ptr CInt -> RealOf a -> Ptr CInt -> ContT r IO ()
   }

gelsy ::
   (Class.Floating a) =>
   Int -> Int -> Int ->
   Ptr a -> Ptr CInt -> Ptr a -> Ptr CInt ->
   Ptr CInt -> RealOf a -> Ptr CInt -> ContT r IO ()
gelsy =
   getGELSY $
   Class.switchFloating
      (GELSY gelsyReal)
      (GELSY gelsyReal)
      (GELSY gelsyComplex)
      (GELSY gelsyComplex)

gelsyReal ::
   (Class.Real a, Class.Floating a) =>
   Int -> Int -> Int ->
   Ptr a -> Ptr CInt -> Ptr a -> Ptr CInt ->
   Ptr CInt -> a -> Ptr CInt -> ContT r IO ()
gelsyReal m n nrhs aPtr ldaPtr bPtr ldbPtr jpvtPtr rcond rankPtr = do
   mPtr <- Call.cint m
   nPtr <- Call.cint n
   nrhsPtr <- Call.cint nrhs
   rcondPtr <- Call.real rcond
   liftIO $ withAutoWorkspaceInfo "gelsy" $
      LapackReal.gelsy mPtr nPtr nrhsPtr
         aPtr ldaPtr bPtr ldbPtr jpvtPtr rcondPtr rankPtr

gelsyComplex ::
   (Class.Real a) =>
   Int -> Int -> Int ->
   Ptr (Complex a) -> Ptr CInt -> Ptr (Complex a) -> Ptr CInt ->
   Ptr CInt -> a -> Ptr CInt -> ContT r IO ()
gelsyComplex m n nrhs aPtr ldaPtr bPtr ldbPtr jpvtPtr rcond rankPtr = do
   mPtr <- Call.cint m
   nPtr <- Call.cint n
   nrhsPtr <- Call.cint nrhs
   rcondPtr <- Call.real rcond
   rworkPtr <- Call.allocaArray (2*n)
   liftIO $ withAutoWorkspaceInfo "gelsy" $ \workPtr lworkPtr infoPtr ->
      LapackComplex.gelsy mPtr nPtr nrhsPtr
         aPtr ldaPtr bPtr ldbPtr jpvtPtr rcondPtr rankPtr
         workPtr lworkPtr rworkPtr infoPtr


pseudoInverseRCond ::
   (Shape.C height, Eq height, Shape.C width, Eq width,
    Storable a, Class.Floating a) =>
   RealOf a -> General height width a -> (Int, General width height a)
pseudoInverseRCond rcond a =
   let (MatrixShape.General _ height width) = Array.shape a
   in if Shape.size height < Shape.size width
         then leastSquaresMinimumNorm rcond a $ identity height
         else mapSnd transpose $
              leastSquaresMinimumNorm rcond (transpose a) $
              identity width


type Householder height width = Array (MatrixShape.Householder height width)

{-
@(q,r) = householder a@
means that @q@ is unitary and @r@ is upper triangular and @a = multiply q r@.
-}
householder ::
   (Shape.C height, Shape.C width, Eq width, Storable a, Class.Floating a) =>
   General height width a ->
   (General height height a, General height width a)
householder a =
   let hh = householderDecompose a
   in  (householderExtractQ hh, householderExtractR $ snd hh)

householderDecompose ::
   (Shape.C height, Shape.C width, Storable a, Class.Floating a) =>
   General height width a -> (Vector width a, Householder height width a)
householderDecompose (Array (MatrixShape.General order height width) a) =
   unsafePerformIO $ do

   let (m,n) =
         case order of
            RowMajor -> (Shape.size width, Shape.size height)
            ColumnMajor -> (Shape.size height, Shape.size width)
   let lda = m
   let mn = min m n
   evalContT $ do
      mPtr <- Call.cint m
      nPtr <- Call.cint n
      aPtr <- ContT $ withForeignPtr a
      ldaPtr <- Call.cint lda
      qr <- liftIO $ mallocForeignPtrArray (m*n)
      qrPtr <- ContT $ withForeignPtr qr
      liftIO $ copyBlock (m*n) aPtr qrPtr
      tau <- liftIO $ mallocForeignPtrArray n
      tauPtr <- ContT $ withForeignPtr tau
      liftIO $ fill zero (n-mn) (advancePtr tauPtr mn)
      liftIO $
         case order of
            RowMajor ->
               withAutoWorkspaceInfo "gelqf" $
                  LapackGen.gelqf mPtr nPtr qrPtr ldaPtr tauPtr
            ColumnMajor ->
               withAutoWorkspaceInfo "geqrf" $
                  LapackGen.geqrf mPtr nPtr qrPtr ldaPtr tauPtr
      return (Array width tau,
              Array (MatrixShape.Householder order height width) qr)

householderDeterminant ::
   (Shape.C height, Shape.C width, Storable a, Class.Floating a) =>
   Householder height width a -> a
householderDeterminant
      (Array (MatrixShape.Householder order height width) a) =
   let m = Shape.size height
       n = Shape.size width
       k = case order of RowMajor -> n; ColumnMajor -> m
   in unsafePerformIO $
      withForeignPtr a $ \aPtr ->
         foldM (\x ptr -> do y <- peek ptr; return $! mul x y) one $
         take (min m n) $ pointerSeq (k+1) aPtr

newtype Mul a = Mul {getMul :: a -> a -> a}

mul :: (Class.Floating a) => a -> a -> a
mul = getMul $ Class.switchFloating (Mul (*)) (Mul (*)) (Mul (*)) (Mul (*))


{-|
Generalized determinant - works also for non-square matrices.
In contrast to the square root of the Gramian determinant
it has the proper sign.
-}
determinant ::
   (Shape.C height, Shape.C width, Eq a, Storable a, Class.Floating a) =>
   General height width a -> a
determinant a =
   let (tau,hh) = householderDecompose a
   in  foldl (\x _ -> neg x)
         (householderDeterminant hh)
         (takeWhile (/=zero) $ Array.toList tau)

newtype Neg a = Neg {getNeg :: a -> a}

neg :: (Class.Floating a) => a -> a
neg =
   getNeg $
   Class.switchFloating (Neg negate) (Neg negate) (Neg negate) (Neg negate)


householderExtractQ ::
   (Shape.C height, Shape.C width, Eq width, Storable a, Class.Floating a) =>
   (Vector width a, Householder height width a) -> General height height a
householderExtractQ
   (Array widthTau tau,
    Array (MatrixShape.Householder order height width) qr) =

   Array.unsafeCreate (MatrixShape.General order height height) $ \qPtr -> do

   Call.assert "householderExtractQ: width shapes mismatch" (widthTau == width)

   let m = Shape.size height
   let k = min m $ Shape.size width
   let lda = m
   evalContT $ do
      mPtr <- Call.cint m
      kPtr <- Call.cint k
      qrPtr <- ContT $ withForeignPtr qr
      ldaPtr <- Call.cint lda
      tauPtr <- ContT $ withForeignPtr tau
      liftIO $
         case order of
            RowMajor -> do
               copySubMatrix k m k qrPtr lda qPtr
               withAutoWorkspaceInfo "unglq" $
                  LapackGen.unglq mPtr mPtr kPtr qPtr ldaPtr tauPtr
            ColumnMajor -> do
               copyBlock (m*k) qrPtr qPtr
               withAutoWorkspaceInfo "ungqr" $
                  LapackGen.ungqr mPtr mPtr kPtr qPtr ldaPtr tauPtr

householderExtractR ::
   (Shape.C height, Shape.C width, Eq width, Storable a, Class.Floating a) =>
   Householder height width a -> General height width a
householderExtractR
      (Array (MatrixShape.Householder order height width) qr) =

   Array.unsafeCreate (MatrixShape.General order height width) $
      \rPtr -> do

   let (uplo, (m,n)) =
         case order of
            RowMajor -> ('L', (Shape.size width, Shape.size height))
            ColumnMajor -> ('U', (Shape.size height, Shape.size width))
   fill zero (m*n) rPtr
   evalContT $ do
      uploPtr <- Call.char uplo
      mPtr <- Call.cint m
      nPtr <- Call.cint n
      qrPtr <- ContT $ withForeignPtr qr
      ldqrPtr <- Call.cint m
      ldrPtr <- Call.cint m
      liftIO $ LapackGen.lacpy uploPtr mPtr nPtr qrPtr ldqrPtr rPtr ldrPtr

{- |
For an m-by-n-matrix @a@ with m>=n
this function computes an m-by-(m-n)-matrix @b@
such that @Matrix.multiply (transpose b) a@ is a zero matrix.
The function does not try to compensate a rank deficiency of @a@.
That is, @a|||b@ has full rank if and only if @a@ has full rank.

For full-rank matrices you might also call this @kernel@ or @nullspace@.
-}
orthogonalComplement ::
   (Shape.C height, Shape.C width, Eq width, Storable a, Class.Floating a) =>
   General height width a -> General height ZeroInt a
orthogonalComplement a =
   dropColumns (Shape.size $ MatrixShape.generalWidth $ Array.shape a) $
   Array.mapShape zeroIntWidth $ householderExtractQ $ householderDecompose a

zeroIntWidth ::
   (Shape.C width) =>
   MatrixShape.General height width -> MatrixShape.General height ZeroInt
zeroIntWidth (MatrixShape.General order height width) =
   MatrixShape.General order height (zeroInt $ Shape.size width)



withAutoWorkspaceInfo ::
   (Storable a, Class.Floating a) =>
   String -> (Ptr a -> Ptr CInt -> Ptr CInt -> IO ()) -> IO ()
withAutoWorkspaceInfo name computation = evalContT $ do
   infoPtr <- Call.alloca
   liftIO $ withAutoWorkspace $ \workPtr lworkPtr ->
      computation workPtr lworkPtr infoPtr
   info <- liftIO $ 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: deficient rank %d" name info

withAutoWorkspace ::
   (Storable a, Class.Floating a) =>
   (Ptr a -> Ptr CInt -> IO ()) -> IO ()
withAutoWorkspace computation = evalContT $ do
   lworkPtr <- Call.cint (-1)
   lwork <- liftIO $ alloca $ \workPtr -> do
      computation workPtr lworkPtr
      ceilingSize <$> peek workPtr
   workPtr <- Call.allocaArray lwork
   liftIO $ poke lworkPtr $ fromIntegral lwork
   liftIO $ computation workPtr lworkPtr


copyToColumnMajor ::
   (Storable a, Class.Floating a) =>
   Order -> Int -> Int -> Ptr a -> Ptr a -> IO ()
copyToColumnMajor order m n aPtr bPtr =
   case order of
      RowMajor -> copyTransposed m n aPtr m bPtr
      ColumnMajor -> copyBlock (m*n) aPtr bPtr

copyToSubColumnMajor ::
   (Storable a, Class.Floating a) =>
   Order -> Int -> Int -> Ptr a -> Int -> Ptr a -> IO ()
copyToSubColumnMajor order m n aPtr ldb bPtr =
   case order of
      RowMajor -> copyTransposed m n aPtr ldb bPtr
      ColumnMajor ->
         if m==ldb
           then copyBlock (m*n) aPtr bPtr
           else copySubMatrix m n m aPtr ldb bPtr


newtype FuncArg b a = FuncArg {runFuncArg :: a -> b}

ceilingSize :: (Class.Floating a) => a -> Int
ceilingSize =
   runFuncArg $
   Class.switchFloating
      (FuncArg ceiling)
      (FuncArg ceiling)
      (FuncArg $ ceiling . Complex.realPart)
      (FuncArg $ ceiling . Complex.realPart)