{-# LANGUAGE TypeFamilies #-}
module Numeric.LAPACK.Matrix.Banded.Linear (
   solve,
   solveColumnMajor,
   solveTriangular,
   determinant,
   ) where

import qualified Numeric.LAPACK.Matrix.Banded.Basic as Banded
import qualified Numeric.LAPACK.Matrix.Shape.Omni as Omni
import qualified Numeric.LAPACK.Matrix.Layout.Private as Layout
import qualified Numeric.LAPACK.Matrix.Extent.Private as Extent
import qualified Numeric.LAPACK.Permutation.Private as Perm
import qualified Numeric.LAPACK.Private as Private
import Numeric.LAPACK.Linear.Private (solver, withDeterminantInfo, withInfo)
import Numeric.LAPACK.Matrix.Layout.Private
         (Order(RowMajor,ColumnMajor), transposeFromOrder)
import Numeric.LAPACK.Matrix.Private (Full)
import Numeric.LAPACK.Private (copySubMatrix)

import qualified Numeric.LAPACK.FFI.Generic as LapackGen
import qualified Numeric.Netlib.Utility as Call
import qualified Numeric.Netlib.Class as Class

import qualified Type.Data.Num.Unary as Unary
import Type.Data.Num (integralFromProxy)

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

import System.IO.Unsafe (unsafePerformIO)

import Foreign.Marshal.Array (peekArray, advancePtr)
import Foreign.ForeignPtr (withForeignPtr)

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


solve ::
   (Unary.Natural sub, Unary.Natural super,
    Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C sh, Eq sh, Shape.C nrhs, Class.Floating a) =>
   Banded.Square sub super sh a ->
   Full meas vert horiz sh nrhs a -> Full meas vert horiz sh nrhs a
solve (Array (Layout.Banded numOff order extent) a) =
   solver "Banded.solve" (Extent.squareSize extent) $
         \n nPtr nrhsPtr xPtr ldxPtr -> do
      let (kl,ku) = Layout.numOffDiagonals order numOff
      let k = kl+1+ku
      let ldab = kl+k
      transPtr <- Call.char $ transposeFromOrder order
      klPtr <- Call.cint kl
      kuPtr <- Call.cint ku
      aPtr <- ContT $ withForeignPtr a
      abPtr <- Call.allocaArray (n*ldab)
      ldabPtr <- Call.leadingDim ldab
      ipivPtr <- Call.allocaArray n
      liftIO $ do
         copySubMatrix k n k aPtr ldab (advancePtr abPtr kl)
         withInfo "gbtrf" $
            LapackGen.gbtrf nPtr nPtr klPtr kuPtr abPtr ldabPtr ipivPtr
         withInfo "gbtrs" $
            LapackGen.gbtrs transPtr nPtr klPtr kuPtr nrhsPtr
               abPtr ldabPtr ipivPtr xPtr ldxPtr

solveColumnMajor ::
   (Unary.Natural sub, Unary.Natural super,
    Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C sh, Eq sh, Shape.C nrhs, Class.Floating a) =>
   Banded.Square sub super sh a ->
   Full meas vert horiz sh nrhs a -> Full meas vert horiz sh nrhs a
solveColumnMajor
      (Array (Layout.Banded (sub,super) ColumnMajor extent) a) =
   solver "Banded.solve" (Extent.squareSize extent) $
         \n nPtr nrhsPtr xPtr ldxPtr -> do
      let kl = integralFromProxy sub
      let ku = integralFromProxy super
      let k = kl+1+ku
      let ldab = kl+k
      klPtr <- Call.cint kl
      kuPtr <- Call.cint ku
      aPtr <- ContT $ withForeignPtr a
      abPtr <- Call.allocaArray (n*ldab)
      ldabPtr <- Call.leadingDim ldab
      ipivPtr <- Call.allocaArray n
      liftIO $ do
         copySubMatrix k n k aPtr ldab (advancePtr abPtr kl)
         withInfo "gbsv" $
            LapackGen.gbsv nPtr klPtr kuPtr nrhsPtr
               abPtr ldabPtr ipivPtr xPtr ldxPtr
solveColumnMajor (Array (Layout.Banded _ RowMajor _) _) =
   error "Linear.Banded.solveColumnMajor: RowMajor intentionally unimplemented"

solveTriangular ::
   (Omni.BandedTriangular sub super, Omni.TriDiag diag,
    Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C sh, Eq sh, Shape.C nrhs, Class.Floating a) =>
   Omni.DiagSingleton diag ->
   Banded.Square sub super sh a ->
   Full meas vert horiz sh nrhs a -> Full meas vert horiz sh nrhs a
solveTriangular diag (Array (Layout.Banded numOff order extent) a) =
   solver "Banded.solveTriangular" (Extent.squareSize extent) $
         \ _n nPtr nrhsPtr xPtr ldxPtr -> do
      let (kl,ku) = Layout.numOffDiagonals order numOff
      let k = kl+ku
      uploPtr <- Call.char $ if kl==0 then 'U' else 'L'
      transPtr <- Call.char $ Layout.transposeFromOrder order
      diagPtr <- Call.char $ Omni.charFromTriDiag diag
      kPtr <- Call.cint k
      abPtr <- ContT $ withForeignPtr a
      ldabPtr <- Call.leadingDim (k+1)
      liftIO $
         withInfo "tbtrs" $
            LapackGen.tbtrs uploPtr transPtr diagPtr nPtr kPtr nrhsPtr
               abPtr ldabPtr xPtr ldxPtr

determinant ::
   (Unary.Natural sub, Unary.Natural super, Shape.C sh, Class.Floating a) =>
   Banded.Square sub super sh a -> a
determinant (Array (Layout.Banded numOff order extent) a) =
      unsafePerformIO $ do
   let n = Shape.size $ Extent.squareSize extent
   evalContT $ do
      let (kl,ku) = Layout.numOffDiagonals order numOff
      let k = kl+1+ku
      let ldab = kl+k
      nPtr <- Call.cint n
      klPtr <- Call.cint kl
      kuPtr <- Call.cint ku
      aPtr <- ContT $ withForeignPtr a
      abPtr <- Call.allocaArray (n*ldab)
      ldabPtr <- Call.leadingDim ldab
      ipivPtr <- Call.allocaArray n
      liftIO $ do
         copySubMatrix k n k aPtr ldab (advancePtr abPtr kl)
         withDeterminantInfo "gbtrf"
            (LapackGen.gbtrf nPtr nPtr klPtr kuPtr abPtr ldabPtr ipivPtr)
            (do
               det <- Private.product n (advancePtr abPtr (kl+ku)) ldab
               ipiv <- peekArray n ipivPtr
               return $ Perm.condNegate ipiv det)