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

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

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

import qualified Numeric.LAPACK.Matrix.Shape.Private as MatrixShape
import qualified Numeric.LAPACK.Private as Private
import Numeric.LAPACK.Matrix.Shape.Private
         (Order(RowMajor, ColumnMajor), transposeFromOrder)
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Private
         (RealOf, zero, fill,
          copySubMatrix, copyBlock, copyToTemp,
          copyToColumnMajor, copyToSubColumnMajor,
          withAutoWorkspaceInfo, allocArray, allocHigherArray)

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 (advancePtr)
import Foreign.C.Types (CInt)
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (Ptr)
import Foreign.Storable (peek)

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

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, 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 $ transposeFromOrder orderA
      mPtr <- Call.cint m
      nPtr <- Call.cint n
      nrhsPtr <- Call.cint nrhs
      aPtr <- copyToTemp (Shape.size shapeA) a
      ldaPtr <- Call.cint lda
      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 aPtr 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, 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 $ transposeFromOrder orderA
      mPtr <- Call.cint m
      nPtr <- Call.cint n
      nrhsPtr <- Call.cint nrhs
      aPtr <- copyToTemp (Shape.size shapeA) a
      ldaPtr <- Call.cint lda
      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 aPtr 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, 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 "leastSquaresMinimumNorm: 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
   evalContT $ do
      aPtr <- ContT $ withForeignPtr a
      atmpPtr <- Call.allocaArray aSize
      liftIO $ copyToColumnMajor orderA m n aPtr atmpPtr
      ldaPtr <- Call.cint lda
      (x,(tmpPtr,ldtmp)) <- allocHigherArray shapeX m n nrhs
      ldtmpPtr <- Call.cint ldtmp
      bPtr <- ContT $ withForeignPtr b
      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
      rank <- liftIO $ fromIntegral <$> peek rankPtr
      return (rank, x)


type GELSY_ r ar a =
   Int -> Int -> Int -> Ptr a -> Ptr CInt -> Ptr a -> Ptr CInt ->
   Ptr CInt -> ar -> Ptr CInt -> ContT r IO ()

newtype GELSY r a = GELSY {getGELSY :: GELSY_ r (RealOf a) a}

gelsy :: (Class.Floating a) => GELSY_ r (RealOf a) a
gelsy =
   getGELSY $
   Class.switchFloating
      (GELSY gelsyReal)
      (GELSY gelsyReal)
      (GELSY gelsyComplex)
      (GELSY gelsyComplex)

gelsyReal :: (Class.Real a) => GELSY_ r a a
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) => GELSY_ r a (Complex a)
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, 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, Class.Floating a) =>
   General height width a ->
   (Square 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, Class.Floating a) =>
   General height width a -> (Vector width a, Householder height width a)
householderDecompose (Array shape@(MatrixShape.General order height width) a) =
   unsafePerformIO $ do

   let (m,n) = MatrixShape.dimensions shape
   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,qrPtr) <- allocArray $ MatrixShape.Householder order height width
      liftIO $ copyBlock (m*n) aPtr qrPtr
      (tau,tauPtr) <- allocArray width
      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 (tau, qr)

householderDeterminant ::
   (Shape.C height, Shape.C width, 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 ->
      Private.product (min m n) aPtr (k+1)


{-|
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, Class.Floating a) =>
   General height width a -> a
determinant a =
   let (tau,hh) = householderDecompose a
   in  foldl (\x _ -> negate x)
         (householderDeterminant hh)
         (takeWhile (/=zero) $ Array.toList tau)


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

   Array.unsafeCreate (MatrixShape.Square order 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, 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, Class.Floating a) =>
   General height width a -> General height ZeroInt a
orthogonalComplement a =
   dropColumns (Shape.size $ MatrixShape.generalWidth $ Array.shape a) $
   Array.mapShape zeroIntWidth $
   Square.toGeneral $ 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)