-- This file is auto-generated.  Do not edit directly.
-- |
--   Stability: Experimental
--
--   Generic interface to Blas using safe foreign calls.  Refer to the GHC documentation
--   for more information regarding appropriate use of safe and unsafe foreign calls.
--
--   The functions here are named in a similar fashion to the original Blas interface, with
--   the type-dependent letter(s) removed.  Some functions have been merged with
--   others to allow the interface to work on both real and complex numbers.  If you can't a
--   particular function, try looking for its corresponding complex equivalent (e.g.
--   @symv@ is a special case of 'hemv' applied to real numbers).
--
--   Note: although complex versions of @rot@ and @rotg@ exist in many implementations,
--   they are not part of the official Blas standard and therefore not included here.  If
--   you /really/ need them, submit a ticket so we can try to come up with a solution.
--
--   The documentation here is still incomplete.  Consult the
--   <http://netlib.org/blas/#_blas_routines official documentation> for more
--   information.
--
--   Notation:
--
--     * @⋅@ denotes dot product (without any conjugation).
--     * @*@ denotes complex conjugation.
--     * @⊤@ denotes transpose.
--     * @†@ denotes conjugate transpose (Hermitian conjugate).
--
--   Conventions:
--
--     * All scalars are denoted with lowercase Greek letters
--     * All vectors are denoted with lowercase Latin letters and are
--       assumed to be column vectors (unless transposed).
--     * All matrices are denoted with uppercase Latin letters.
{-# LANGUAGE FlexibleInstances, TypeFamilies #-}
module Blas.Generic.Safe (
    Numeric(..)
  , RealNumeric(..)
  , D.dsdot
  , S.sdsdot
  ) where
import Prelude (Floating, Double, Float, Int, IO)
import Data.Complex (Complex)
import Foreign (Ptr, Storable)
import Blas.Primitive.Types (Order, Transpose, Uplo, Diag, Side)
import qualified Blas.Specialized.Float.Safe as S
import qualified Blas.Specialized.Double.Safe as D
import qualified Blas.Specialized.ComplexFloat.Safe as C
import qualified Blas.Specialized.ComplexDouble.Safe as Z

-- | Blas operations that are applicable to real and complex numbers.
--
--   Instances are defined for the 4 types supported by Blas: the single- and
--   double-precision floating point types and their complex versions.
class (Floating a, Storable a) => Numeric a where

  -- | The corresponding real type of @a@.
  --
  --   In other words, @'RealType' ('Complex' a)@ is an alias for @a@.  For everything
  --   else, @'RealType' a@ is simply @a@.
  type RealType a :: *
  -- | Swap two vectors:
  --
  --   > (x, y) ← (y, x)
  swap :: Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

  -- | Multiply a vector by a scalar.
  --
  --   > x ← α x
  scal :: Int
       -> a
       -> Ptr a
       -> Int
       -> IO ()

  -- | Copy a vector into another vector:
  --
  --   > y ← x
  copy :: Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

  -- | Add a scalar-vector product to a vector.
  --
  --   > y ← α x + y
  axpy :: Int
       -> a
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

  -- | Calculate the bilinear dot product of two vectors:
  --
  --   > x ⋅ y ≡ ∑[i] x[i] y[i]
  dotu :: Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO a

  -- | Calculate the sesquilinear dot product of two vectors.
  --
  --   > x* ⋅ y ≡ ∑[i] x[i]* y[i]
  dotc :: Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO a

  -- | Calculate the Euclidean (L²) norm of a vector:
  --
  --   > ‖x‖₂ ≡ √(∑[i] x[i]²)
  nrm2 :: Int
       -> Ptr a
       -> Int
       -> IO (RealType a)

  -- | Calculate the Manhattan (L¹) norm, equal to the sum of the magnitudes of the elements:
  --
  --   > ‖x‖₁ = ∑[i] |x[i]|
  asum :: Int
       -> Ptr a
       -> Int
       -> IO (RealType a)

  -- | Calculate the index of the element with the maximum magnitude (absolute value).
  iamax :: Int
        -> Ptr a
        -> Int
        -> IO Int

  -- | Perform a general matrix-vector update.
  --
  --   > y ← α T(A) x + β y
  gemv :: Order
       -> Transpose
       -> Int
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> IO ()

  -- | Perform a general banded matrix-vector update.
  --
  --   > y ← α T(A) x + β y
  gbmv :: Order
       -> Transpose
       -> Int
       -> Int
       -> Int
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> IO ()

  -- | Perform a hermitian matrix-vector update.
  --
  --   > y ← α A x + β y
  hemv :: Order
       -> Uplo
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> IO ()

  -- | Perform a hermitian banded matrix-vector update.
  --
  --   > y ← α A x + β y
  hbmv :: Order
       -> Uplo
       -> Int
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> IO ()

  -- | Perform a hermitian packed matrix-vector update.
  --
  --   > y ← α A x + β y
  hpmv :: Order
       -> Uplo
       -> Int
       -> a
       -> Ptr a
       -> Ptr a
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> IO ()

  -- | Multiply a triangular matrix by a vector.
  --
  --   > x ← T(A) x
  trmv :: Order
       -> Uplo
       -> Transpose
       -> Diag
       -> Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

  -- | Multiply a triangular banded matrix by a vector.
  --
  --   > x ← T(A) x
  tbmv :: Order
       -> Uplo
       -> Transpose
       -> Diag
       -> Int
       -> Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

  -- | Multiply a triangular packed matrix by a vector.
  --
  --   > x ← T(A) x
  tpmv :: Order
       -> Uplo
       -> Transpose
       -> Diag
       -> Int
       -> Ptr a
       -> Ptr a
       -> Int
       -> IO ()

  -- | Multiply an inverse triangular matrix by a vector.
  --
  --   > x ← T(A⁻¹) x
  trsv :: Order
       -> Uplo
       -> Transpose
       -> Diag
       -> Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

  -- | Multiply an inverse triangular banded matrix by a vector.
  --
  --   > x ← T(A⁻¹) x
  tbsv :: Order
       -> Uplo
       -> Transpose
       -> Diag
       -> Int
       -> Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

  -- | Multiply an inverse triangular packed matrix by a vector.
  --
  --   > x ← T(A⁻¹) x
  tpsv :: Order
       -> Uplo
       -> Transpose
       -> Diag
       -> Int
       -> Ptr a
       -> Ptr a
       -> Int
       -> IO ()

  -- | Perform an unconjugated rank-1 update of a general matrix.
  --
  --   > A ← α x y⊤ + A
  geru :: Order
       -> Int
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

  -- | Perform a conjugated rank-1 update of a general matrix.
  --
  --   > A ← α x y† + A
  gerc :: Order
       -> Int
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

  -- | Perform a rank-1 update of a Hermitian matrix.
  --
  --   > A ← α x y† + A
  her :: Order
      -> Uplo
      -> Int
      -> RealType a
      -> Ptr a
      -> Int
      -> Ptr a
      -> Int
      -> IO ()

  -- | Perform a rank-1 update of a Hermitian packed matrix.
  --
  --   > A ← α x y† + A
  hpr :: Order
      -> Uplo
      -> Int
      -> RealType a
      -> Ptr a
      -> Int
      -> Ptr a
      -> IO ()

  -- | Perform a rank-2 update of a Hermitian matrix.
  --
  --   > A ← α x y† + y (α x)† + A
  her2 :: Order
       -> Uplo
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

  -- | Perform a rank-2 update of a Hermitian packed matrix.
  --
  --   > A ← α x y† + y (α x)† + A
  hpr2 :: Order
       -> Uplo
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> IO ()

  -- | Perform a general matrix-matrix update.
  --
  --   > C ← α T(A) U(B) + β C
  gemm :: Order -- ^ Layout of all the matrices.
       -> Transpose -- ^ The operation @T@ to be applied to @A@.
       -> Transpose -- ^ The operation @U@ to be applied to @B@.
       -> Int -- ^ Number of rows of @T(A)@ and @C@.
       -> Int -- ^ Number of columns of @U(B)@ and @C@.
       -> Int -- ^ Number of columns of @T(A)@ and number of rows of @U(B)@.
       -> a -- ^ Scaling factor @α@ of the product.
       -> Ptr a -- ^ Pointer to a matrix @A@.
       -> Int -- ^ Stride of the major dimension of @A@.
       -> Ptr a -- ^ Pointer to a matrix @B@.
       -> Int -- ^ Stride of the major dimension of @B@.
       -> a -- ^ Scaling factor @β@ of the original @C@.
       -> Ptr a -- ^ Pointer to a mutable matrix @C@.
       -> Int -- ^ Stride of the major dimension of @C@.
       -> IO ()

  -- | Perform a symmetric matrix-matrix update.
  --
  --   > C ← α A B + β C    or    C ← α B A + β C
  --
  --   where @A@ is symmetric.  The matrix @A@ must be in an unpacked format, although the
  --   routine will only access half of it as specified by the @'Uplo'@ argument.
  symm :: Order -- ^ Layout of all the matrices.
       -> Side -- ^ Side that @A@ appears in the product.
       -> Uplo -- ^ The part of @A@ that is used.
       -> Int -- ^ Number of rows of @C@.
       -> Int -- ^ Number of columns of @C@.
       -> a -- ^ Scaling factor @α@ of the product.
       -> Ptr a -- ^ Pointer to a symmetric matrix @A@.
       -> Int -- ^ Stride of the major dimension of @A@.
       -> Ptr a -- ^ Pointer to a matrix @B@.
       -> Int -- ^ Stride of the major dimension of @B@.
       -> a -- ^ Scaling factor @α@ of the original @C@.
       -> Ptr a -- ^ Pointer to a mutable matrix @C@.
       -> Int -- ^ Stride of the major dimension of @C@.
       -> IO ()

  -- | Perform a Hermitian matrix-matrix update.
  --
  --   > C ← α A B + β C    or    C ← α B A + β C
  --
  --   where @A@ is Hermitian.  The matrix @A@ must be in an unpacked format, although the
  --   routine will only access half of it as specified by the @'Uplo'@ argument.
  hemm :: Order -- ^ Layout of all the matrices.
       -> Side -- ^ Side that @A@ appears in the product.
       -> Uplo -- ^ The part of @A@ that is used.
       -> Int -- ^ Number of rows of @C@.
       -> Int -- ^ Number of columns of @C@.
       -> a -- ^ Scaling factor @α@ of the product.
       -> Ptr a -- ^ Pointer to a Hermitian matrix @A@.
       -> Int -- ^ Stride of the major dimension of @A@.
       -> Ptr a -- ^ Pointer to a matrix @B@.
       -> Int -- ^ Stride of the major dimension of @B@.
       -> a -- ^ Scaling factor @α@ of the original @C@.
       -> Ptr a -- ^ Pointer to a mutable matrix @C@.
       -> Int -- ^ Stride of the major dimension of @C@.
       -> IO ()

  -- | Perform a symmetric rank-k update.
  --
  --   > C ← α A A⊤ + β C    or    C ← α A⊤ A + β C
  syrk :: Order
       -> Uplo
       -> Transpose
       -> Int
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> IO ()

  -- | Perform a Hermitian rank-k update.
  --
  --   > C ← α A A† + β C    or    C ← α A† A + β C
  herk :: Order
       -> Uplo
       -> Transpose
       -> Int
       -> Int
       -> RealType a
       -> Ptr a
       -> Int
       -> RealType a
       -> Ptr a
       -> Int
       -> IO ()

  -- | Perform a symmetric rank-2k update.
  --
  --   > C ← α A B⊤ + α* B A⊤ + β C    or    C ← α A⊤ B + α* B⊤ A + β C
  syr2k :: Order
        -> Uplo
        -> Transpose
        -> Int
        -> Int
        -> a
        -> Ptr a
        -> Int
        -> Ptr a
        -> Int
        -> a
        -> Ptr a
        -> Int
        -> IO ()

  -- | Perform a Hermitian rank-2k update.
  --
  --   > C ← α A B† + α* B A† + β C    or    C ← α A† B + α* B† A + β C
  her2k :: Order
        -> Uplo
        -> Transpose
        -> Int
        -> Int
        -> a
        -> Ptr a
        -> Int
        -> Ptr a
        -> Int
        -> RealType a
        -> Ptr a
        -> Int
        -> IO ()

  -- | Perform a triangular matrix-matrix multiplication.
  --
  --   > B ← α T(A) B    or    B ← α B T(A)
  --
  --   where @A@ is triangular.
  trmm :: Order
       -> Side
       -> Uplo
       -> Transpose
       -> Diag
       -> Int
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

  -- | Perform an inverse triangular matrix-matrix multiplication.
  --
  --   > B ← α T(A⁻¹) B    or    B ← α B T(A⁻¹)
  --
  --   where @A@ is triangular.
  trsm :: Order
       -> Side
       -> Uplo
       -> Transpose
       -> Diag
       -> Int
       -> Int
       -> a
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> IO ()

-- | Blas operations that are only applicable to real numbers.
class Numeric a => RealNumeric a where

  -- | Generate a Givens rotation. (Only available for real floating-point types.)
  rotg :: Ptr a
       -> Ptr a
       -> Ptr a
       -> Ptr a
       -> IO ()

  -- | Generate a modified Givens rotation. (Only available for real floating-point
  --   types.)
  rotmg :: Ptr a
        -> Ptr a
        -> Ptr a
        -> a
        -> Ptr a
        -> IO ()

  -- | Apply a Givens rotation. (Only available for real floating-point types.)
  rot :: Int
      -> Ptr a
      -> Int
      -> Ptr a
      -> Int
      -> a
      -> a
      -> IO ()

  -- | Apply a modified Givens rotation. (Only available for real floating-point types.)
  rotm :: Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> Int
       -> Ptr a
       -> IO ()

instance Numeric Float where
  type RealType Float = Float
  swap = S.swap
  scal = S.scal
  copy = S.copy
  axpy = S.axpy
  dotu = S.dotu
  dotc = S.dotc
  nrm2 = S.nrm2
  asum = S.asum
  iamax = S.iamax
  gemv = S.gemv
  gbmv = S.gbmv
  hemv = S.hemv
  hbmv = S.hbmv
  hpmv = S.hpmv
  trmv = S.trmv
  tbmv = S.tbmv
  tpmv = S.tpmv
  trsv = S.trsv
  tbsv = S.tbsv
  tpsv = S.tpsv
  geru = S.geru
  gerc = S.gerc
  her = S.her
  hpr = S.hpr
  her2 = S.her2
  hpr2 = S.hpr2
  gemm = S.gemm
  symm = S.symm
  hemm = S.hemm
  syrk = S.syrk
  herk = S.herk
  syr2k = S.syr2k
  her2k = S.her2k
  trmm = S.trmm
  trsm = S.trsm

instance Numeric Double where
  type RealType Double = Double
  swap = D.swap
  scal = D.scal
  copy = D.copy
  axpy = D.axpy
  dotu = D.dotu
  dotc = D.dotc
  nrm2 = D.nrm2
  asum = D.asum
  iamax = D.iamax
  gemv = D.gemv
  gbmv = D.gbmv
  hemv = D.hemv
  hbmv = D.hbmv
  hpmv = D.hpmv
  trmv = D.trmv
  tbmv = D.tbmv
  tpmv = D.tpmv
  trsv = D.trsv
  tbsv = D.tbsv
  tpsv = D.tpsv
  geru = D.geru
  gerc = D.gerc
  her = D.her
  hpr = D.hpr
  her2 = D.her2
  hpr2 = D.hpr2
  gemm = D.gemm
  symm = D.symm
  hemm = D.hemm
  syrk = D.syrk
  herk = D.herk
  syr2k = D.syr2k
  her2k = D.her2k
  trmm = D.trmm
  trsm = D.trsm

instance Numeric (Complex Float) where
  type RealType (Complex Float) = Float
  swap = C.swap
  scal = C.scal
  copy = C.copy
  axpy = C.axpy
  dotu = C.dotu
  dotc = C.dotc
  nrm2 = C.nrm2
  asum = C.asum
  iamax = C.iamax
  gemv = C.gemv
  gbmv = C.gbmv
  hemv = C.hemv
  hbmv = C.hbmv
  hpmv = C.hpmv
  trmv = C.trmv
  tbmv = C.tbmv
  tpmv = C.tpmv
  trsv = C.trsv
  tbsv = C.tbsv
  tpsv = C.tpsv
  geru = C.geru
  gerc = C.gerc
  her = C.her
  hpr = C.hpr
  her2 = C.her2
  hpr2 = C.hpr2
  gemm = C.gemm
  symm = C.symm
  hemm = C.hemm
  syrk = C.syrk
  herk = C.herk
  syr2k = C.syr2k
  her2k = C.her2k
  trmm = C.trmm
  trsm = C.trsm

instance Numeric (Complex Double) where
  type RealType (Complex Double) = Double
  swap = Z.swap
  scal = Z.scal
  copy = Z.copy
  axpy = Z.axpy
  dotu = Z.dotu
  dotc = Z.dotc
  nrm2 = Z.nrm2
  asum = Z.asum
  iamax = Z.iamax
  gemv = Z.gemv
  gbmv = Z.gbmv
  hemv = Z.hemv
  hbmv = Z.hbmv
  hpmv = Z.hpmv
  trmv = Z.trmv
  tbmv = Z.tbmv
  tpmv = Z.tpmv
  trsv = Z.trsv
  tbsv = Z.tbsv
  tpsv = Z.tpsv
  geru = Z.geru
  gerc = Z.gerc
  her = Z.her
  hpr = Z.hpr
  her2 = Z.her2
  hpr2 = Z.hpr2
  gemm = Z.gemm
  symm = Z.symm
  hemm = Z.hemm
  syrk = Z.syrk
  herk = Z.herk
  syr2k = Z.syr2k
  her2k = Z.her2k
  trmm = Z.trmm
  trsm = Z.trsm
instance RealNumeric Float where
  rotg = S.rotg
  rotmg = S.rotmg
  rot = S.rot
  rotm = S.rotm

instance RealNumeric Double where
  rotg = D.rotg
  rotmg = D.rotmg
  rot = D.rot
  rotm = D.rotm


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module Main where
import qualified Data.Foldable as Foldable

-- vector
import Data.Vector (Vector)
import qualified Data.Vector as Vector

-- Matrix
import Data.Matrix (Matrix)
import qualified Data.Matrix as Matrix

-- HMatrix
import qualified Numeric.LinearAlgebra.Data as HMatrix
import qualified Numeric.LinearAlgebra.HMatrix as HMatrix

-- bytestring
import Data.ByteString.Lazy (ByteString)
import qualified Data.ByteString.Lazy as ByteString

-- cassava
import Data.Csv
  ( DefaultOrdered(headerOrder)
  , FromField(parseField)
  , FromNamedRecord(parseNamedRecord)
  , Header
  , ToField(toField)
  , ToNamedRecord(toNamedRecord)
  , (.:)
  , (.=)
  )
import qualified Data.Csv as Cassava

-- text
import Data.Text (Text)
import qualified Data.Text.Encoding as Text

-- Exception
import Control.Exception

-- base
import qualified Control.Monad as Monad
import qualified System.Exit as Exit
import qualified Debug.Trace as Trace

data ComponentData =
  ComponentData {
    componentType :: ComponentType,
    nodeK :: Int,
    nodeM :: Int,
    magnitude :: Double,
    param1 :: Double,
    param2 :: Double,
    plot :: Int
     }
  deriving (Eq, Show)

data SimulationData =
  SimulationData {
    nodes :: Int,
    voltageSources :: Int,
    stepSize :: Double,
    tmax :: Double
     }
  deriving (Eq, Show)

data ComponentType = Resistor | Capacitor | Inductor | EAC | EDC | Other Text deriving (Eq, Show)
type SimulationResults = (Vector Double, Matrix Double)

instance FromNamedRecord SimulationData where
  parseNamedRecord m =
    SimulationData
      <$> m .: "Number of Nodes"
      <*> m .: "Number of Voltages Sources"
      <*> m .: "Step Size"
      <*> m .: "Maximum time for simulation"

instance FromNamedRecord ComponentData where
  parseNamedRecord m =
    ComponentData
      <$> m .: "Element Type"
      <*> m .: "Node K"
      <*> m .: "Node M"
      <*> m .: "Value"
      <*> m .: "Source param 1"
      <*> m .: "Source param 2"
      <*> m .: "Plot"

instance FromField ComponentType where
  parseField "R" =
    pure Resistor

  parseField "L" =
    pure Inductor

  parseField "C" =
    pure Capacitor

  parseField "EDC" =
    pure EDC

  parseField "EAC" =
    pure EAC

  parseField otherType =
    Other <$> parseField otherType

decodeItems :: ByteString -> Either String (Vector ComponentData)
decodeItems =
  fmap snd . Cassava.decodeByName

decodeItemsFromFile :: FilePath -> IO (Either String (Vector ComponentData))
decodeItemsFromFile filePath =
  catchShowIO (ByteString.readFile filePath)
    >>= return . either Left decodeItems

decodeSimulation :: ByteString -> Either String (Vector SimulationData)
decodeSimulation =
  fmap snd . Cassava.decodeByName

decodeSimulationFromFile :: FilePath -> IO (Either String (Vector SimulationData))
decodeSimulationFromFile filePath =
  catchShowIO (ByteString.readFile filePath)
    >>= return . either Left decodeSimulation

getSingleSimulationLine :: Vector SimulationData -> SimulationData
getSingleSimulationLine =
    Vector.head

nhComponents :: [ComponentData] -> [ComponentData]
nhComponents =
  filter (\r -> (componentType r == Capacitor) || (componentType r == Inductor))

filterEnergyStorageComponent :: Vector ComponentData -> Vector ComponentData
filterEnergyStorageComponent =
  Vector.filter (\r -> (componentType r == Capacitor) || (componentType r == Inductor))

nh :: Vector ComponentData -> Int
nh components =
  length $ filterEnergyStorageComponent components

-- filter ((== Zaal) . reviewLocation) reviews

filterSources :: Vector ComponentData -> Vector ComponentData
filterSources =
  Vector.filter (\r -> (componentType r == EDC) || (componentType r == EAC))

condutance :: ComponentData -> Double -> Double
condutance component dt =
  case componentType component of
    Resistor -> 1.0 / (magnitude component)
    Capacitor -> (magnitude component) * 0.000001 * 2 / dt
    Inductor -> dt / (2 * 0.001 * (magnitude component))
    _ -> 0.0

gkm :: Vector ComponentData -> Double -> Vector Double
gkm components dt =
  Vector.map (\c -> condutance c dt) components
  -- Matrix.colVector (Vector.map (\c -> condutance c dt) components)
  -- Matrix.fromList (length components) 1 (Vector.toList (Vector.map (\c -> condutance c dt) components))

buildCompactGMatrix :: Double -> [ComponentData] -> Matrix Double -> Matrix Double
buildCompactGMatrix dt [] buffer = buffer
buildCompactGMatrix dt (component:cs) buffer =
  case (nodeK component, nodeM component) of (0, m) -> buildCompactGMatrix dt cs (Matrix.setElem (Matrix.getElem m m buffer + condutance component dt) (m, m) buffer)
                                             (k, 0) -> buildCompactGMatrix dt cs (Matrix.setElem (Matrix.getElem k k buffer + condutance component dt) (k, k) buffer)
                                             (k, m) -> buildCompactGMatrix dt cs (Matrix.setElem (Matrix.getElem k k buffer + condutance component dt) (k, k) (Matrix.setElem (Matrix.getElem m m buffer + condutance component dt) (m, m) (Matrix.setElem (Matrix.getElem k m buffer - condutance component dt) (k, m) (Matrix.setElem (Matrix.getElem m k buffer - condutance component dt) (m, k) buffer))))
                                             (_, _) -> buildCompactGMatrix dt cs buffer

buildGMatrixFromVector :: SimulationData -> Vector ComponentData -> Matrix Double
buildGMatrixFromVector simulation components =
  buildCompactGMatrix (stepSize simulation) (Vector.toList components) (Matrix.zero (nodes simulation) (nodes simulation))

buildIhVector :: [ComponentData] -> Double -> Int -> [Double] -> [Double] -> Matrix Double -> Vector Double
buildIhVector [] _ _ _ ihnew _ = Vector.fromList ihnew
buildIhVector (component:cs) dt n (hold:ihold) ihnew vMatrix =
  case (componentType component, nodeK component, nodeM component) of (Inductor, 0, m) -> buildIhVector cs dt n ihold (ihnew ++ [(2*(condutance component dt)*(Matrix.getElem m n vMatrix) + hold)]) vMatrix
                                                                      (Inductor, k, 0) -> buildIhVector cs dt n ihold (ihnew ++ [(-2*(condutance component dt)*(Matrix.getElem k n vMatrix) + hold)]) vMatrix
                                                                      (Inductor, k, m) -> buildIhVector cs dt n ihold (ihnew ++ [(-2*(condutance component dt)*((Matrix.getElem k n vMatrix) - (Matrix.getElem m n vMatrix)) + hold)]) vMatrix
                                                                      (Capacitor, 0, m) -> buildIhVector cs dt n ihold (ihnew ++ [(-2*(condutance component dt)*(Matrix.getElem m n vMatrix) - hold)]) vMatrix
                                                                      (Capacitor, k, 0) -> buildIhVector cs dt n ihold (ihnew ++ [(2*(condutance component dt)*(Matrix.getElem k n vMatrix) - hold)]) vMatrix
                                                                      (Capacitor, k, m) -> buildIhVector cs dt n ihold (ihnew ++ [(2*(condutance component dt)*((Matrix.getElem k n vMatrix) - (Matrix.getElem m n vMatrix)) - hold)]) vMatrix
                                                                      (_, _, _) -> buildIhVector cs dt n ihold ihnew vMatrix

buildVBVector :: [ComponentData] -> Double -> [Double] -> Vector Double
buildVBVector [] _ buffer = Vector.fromList buffer
buildVBVector (c:components) time buffer =
  case (componentType c) of EDC -> buildVBVector components time ((magnitude c) : buffer)
                            EAC -> buildVBVector components time (((magnitude c * cos (2 * pi * param2 c * time + (param1 c * (pi/180))))) : buffer)
                            _   -> buildVBVector components time buffer

buildIVector :: [ComponentData] -> [Double] -> Vector Double -> Vector Double
buildIVector [] _ iVector = iVector
buildIVector (component:cs) (ihEl:ih) iVector =
  case (componentType component, nodeK component, nodeM component) of (Inductor, k, 0) -> buildIVector cs ih (iVector Vector.// [((k - 1), ((iVector Vector.! (k-1)) + ihEl))])
                                                                      (Inductor, 0, m) -> buildIVector cs ih (iVector Vector.// [((m - 1), ((iVector Vector.! (m-1)) - ihEl))])
                                                                      (Inductor, k, m) -> buildIVector cs ih (iVector Vector.// [((m - 1), ((iVector Vector.! (m-1)) - ihEl)), ((k-1), ((iVector Vector.! (k-1)) + ihEl))])
                                                                      (Capacitor, k, 0) -> buildIVector cs ih (iVector Vector.// [((k - 1), ((iVector Vector.! (k-1)) + ihEl))])
                                                                      (Capacitor, 0, m) -> buildIVector cs ih (iVector Vector.// [((m - 1), ((iVector Vector.! (m-1)) - ihEl))])
                                                                      (Capacitor, k, m) -> buildIVector cs ih (iVector Vector.// [((m - 1), ((iVector Vector.! (m-1)) - ihEl)), ((k-1), ((iVector Vector.! (k-1)) + ihEl))])
                                                                      (_, _, _) -> buildIVector cs ih iVector

thtaControl :: Int -> Double -> Vector Double -> Vector Double -> SimulationData -> (Int, Vector Double, Double)
thtaControl thtactl time ihnew ih simulation
  | thtactl <= 0 = (thtactl, ihnew, (stepSize simulation + time))
  | thtactl < 3 = (thtactl + 1, (Vector.map (\i -> i/2) $ Vector.zipWith (+) ih ihnew), (time + (stepSize simulation/2)))
  | otherwise = (0, ihnew, (stepSize simulation + time))

fromHMatrixTransformer :: HMatrix.Matrix Double -> Matrix Double
fromHMatrixTransformer matrix =
  Matrix.fromLists $ HMatrix.toLists matrix

toHMatrixTransformer :: Matrix Double -> HMatrix.Matrix Double
toHMatrixTransformer matrix =
  HMatrix.fromLists $ Matrix.toLists matrix

fromHMatrixVectorTransformer :: HMatrix.Vector Double -> Vector Double
fromHMatrixVectorTransformer vec =
  Vector.fromList $ HMatrix.toList vec

toHMatrixVectorTransformer :: Vector Double -> HMatrix.Vector Double
toHMatrixVectorTransformer vec =
  HMatrix.fromList $ Vector.toList vec

solver :: HMatrix.Vector Double -> HMatrix.Matrix Double -> HMatrix.Matrix Double -> HMatrix.Matrix Double -> HMatrix.Matrix Double -> HMatrix.Vector Double -> SimulationData -> (Vector Double, Vector Double)
solver iVector gaa gab gba gbb vb simulation =
  let ia = HMatrix.subVector 0 ((nodes simulation) - (voltageSources simulation)) iVector
      rhsa = ia - (gab HMatrix.#> vb)
      va = gaa HMatrix.<\> rhsa
      ib = (gba HMatrix.#> va) + (gbb HMatrix.#> vb)
      iVec = HMatrix.vjoin [ia, ib]
      vVec = HMatrix.vjoin [va, vb]
  in
    ((fromHMatrixVectorTransformer iVec), (fromHMatrixVectorTransformer vVec))

thtaSimulationStep :: [ComponentData] -> Matrix Double -> SimulationData -> Int -> Int -> Double  -> Vector Double -> Matrix Double -> Vector Double -> Vector Double -> SimulationResults
thtaSimulationStep _ _ _ _ 1 _ _ vMatrix _ iVector = (iVector, vMatrix)
thtaSimulationStep components condutances simulation thtactl n time ih vMatrix vbVector iVector =
  let (gaa, gab, gba, gbb) = Matrix.splitBlocks (nodes simulation - voltageSources simulation) (nodes simulation - voltageSources simulation) condutances
      ihBuffer = buildIhVector (nhComponents components) (stepSize simulation) n (Vector.toList ih) [] vMatrix
      (thta, ihThta, timeThta) = thtaControl thtactl time ihBuffer ih simulation
      vbVec = buildVBVector components timeThta []
      iVec = buildIVector (nhComponents components) (Vector.toList ihThta) (Vector.replicate (nodes simulation) 0)
      -- (iVecCalc, vVec) = Trace.trace ("Solver = \n" ++ show (solver (toHMatrixVectorTransformer iVec) (toHMatrixTransformer gaa) (toHMatrixTransformer gab) (toHMatrixTransformer gba) (toHMatrixTransformer gbb) (toHMatrixVectorTransformer vbVec) simulation)) solver (toHMatrixVectorTransformer iVec) (toHMatrixTransformer gaa) (toHMatrixTransformer gab) (toHMatrixTransformer gba) (toHMatrixTransformer gbb) (toHMatrixVectorTransformer vbVec) simulation
      (iVecCalc, vVec) = solver (toHMatrixVectorTransformer iVec) (toHMatrixTransformer gaa) (toHMatrixTransformer gab) (toHMatrixTransformer gba) (toHMatrixTransformer gbb) (toHMatrixVectorTransformer vbVec) simulation
      vMatr = Matrix.mapCol (\r _ -> vVec Vector.! (r - 1)) (n-1) vMatrix
  in
      thtaSimulationStep components condutances simulation thta (n-1) timeThta ihThta vMatr vbVec iVecCalc

thtaSimulation :: Vector ComponentData -> SimulationData -> SimulationResults
thtaSimulation components simulation =
  thtaSimulationStep (Vector.toList components) (buildGMatrixFromVector simulation components) simulation 1 (npoints simulation) 0.0 (Vector.replicate (nh components) 0) (Matrix.zero (nodes simulation) (npoints simulation)) (Vector.replicate (voltageSources simulation) 0) (Vector.replicate (nodes simulation) 0)

npoints :: SimulationData -> Int
npoints sim =
  round ((tmax sim)/(stepSize sim)) + 1

catchShowIO :: IO a -> IO (Either String a)
catchShowIO action =
  fmap Right action
    `catch` handleIOException
  where
    handleIOException :: IOException -> IO (Either String a)
    handleIOException =
      return . Left . show

main :: IO ()
main = do

  eitherSimulation <-
      fmap getSingleSimulationLine
        <$> decodeSimulationFromFile "data/simulation.csv"

  case eitherSimulation of
    Left reason ->
      Exit.die reason

    Right simulation -> do
      components_list <- decodeItemsFromFile "data/components.csv"
      case components_list of
        Left reason -> Exit.die reason
        Right components -> do
          let gmt = buildGMatrixFromVector simulation components
          putStr "GMatrix: \n"
          print (gmt)
          let results = thtaSimulation components simulation
          putStr "Simulation: \n"
          print (results)


{-# LANGUAGE CPP #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Data.Eigen.Matrix (
    -- * Matrix type
    -- | Matrix aliases follows Eigen naming convention
    Matrix(..),
    MatrixXf,
    MatrixXd,
    MatrixXcf,
    MatrixXcd,
    I.Elem,
    I.CComplex,
    valid,
    -- * Matrix conversions
    fromList,
    toList,
    fromFlatList,
    toFlatList,
    generate,
    -- * Standard matrices and special cases
    empty,
    null,
    square,
    zero,
    ones,
    identity,
    constant,
    random,
    -- * Accessing matrix data
    cols,
    rows,
    dims,
    (!),
    coeff,
    unsafeCoeff,
    col,
    row,
    block,
    topRows,
    bottomRows,
    leftCols,
    rightCols,
    -- * Matrix properties
    sum,
    prod,
    mean,
    minCoeff,
    maxCoeff,
    trace,
    norm,
    squaredNorm,
    blueNorm,
    hypotNorm,
    determinant,
    -- * Generic reductions
    fold,
    fold',
    ifold,
    ifold',
    fold1,
    fold1',
    -- * Boolean reductions
    all,
    any,
    count,
    -- * Basic matrix algebra
    add,
    sub,
    mul,
    -- * Mapping over elements
    map,
    imap,
    filter,
    ifilter,
    -- * Matrix transformations
    diagonal,
    transpose,
    inverse,
    adjoint,
    conjugate,
    normalize,
    modify,
    convert,
    TriangularMode(..),
    triangularView,
    lowerTriangle,
    upperTriangle,
    -- * Matrix serialization
    encode,
    decode,
    -- * Mutable matrices
    thaw,
    freeze,
    unsafeThaw,
    unsafeFreeze,
    -- * Raw pointers
    unsafeWith,
) where

import qualified Prelude as P
import qualified Data.List as L
import Prelude hiding (null, sum, all, any, map, filter)
import Data.Tuple
import Data.Complex hiding (conjugate)
import Data.Binary hiding (encode, decode)
import qualified Data.Binary as B
import Foreign.Ptr
import Foreign.C.Types
import Foreign.C.String
import Foreign.Storable
import Foreign.Marshal.Alloc
import Text.Printf
import Control.Monad
import Control.Monad.ST
import Control.Monad.Primitive
#if __GLASGOW_HASKELL__ >= 710
#else
import Control.Applicative hiding (empty)
#endif
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import qualified Data.Eigen.Internal as I
import qualified Data.Eigen.Matrix.Mutable as M
import qualified Data.ByteString.Lazy as BSL

-- | Matrix to be used in pure computations, uses column major memory layout, features copy-free FFI with C++ <http://eigen.tuxfamily.org Eigen> library.

data Matrix a b where
    Matrix :: I.Elem a b => !Int -> !Int -> !(VS.Vector b) -> Matrix a b

-- | Alias for single precision matrix
type MatrixXf = Matrix Float CFloat
-- | Alias for double precision matrix
type MatrixXd = Matrix Double CDouble
-- | Alias for single previsiom matrix of complex numbers
type MatrixXcf = Matrix (Complex Float) (I.CComplex CFloat)
-- | Alias for double prevision matrix of complex numbers
type MatrixXcd = Matrix (Complex Double) (I.CComplex CDouble)

-- | Pretty prints the matrix
instance (I.Elem a b, Show a) => Show (Matrix a b) where
    show m@(Matrix rows cols _) = concat [
        "Matrix ", show rows, "x", show cols,
        "\n", L.intercalate "\n" $ P.map (L.intercalate "\t" . P.map show) $ toList m, "\n"]

-- | Basic matrix math exposed through Num instance: @(*)@, @(+)@, @(-)@, `fromInteger`, `signum`, `abs`, `negate`
instance I.Elem a b => Num (Matrix a b) where
    (*) = mul
    (+) = add
    (-) = sub
    fromInteger = constant 1 1 . fromInteger
    signum = map signum
    abs = map abs
    negate = map negate

-- | Matrix binary serialization
instance I.Elem a b => Binary (Matrix a b) where
    put (Matrix rows cols vals) = do
        put $ I.magicCode (undefined :: b)
        put rows
        put cols
        put vals

    get = do
        get >>= (`when` fail "wrong matrix type") . (/= I.magicCode (undefined :: b))
        Matrix <$> get <*> get <*> get

-- | Encode the matrix as a lazy byte string
encode :: I.Elem a b => Matrix a b -> BSL.ByteString
encode = B.encode

-- | Decode matrix from the lazy byte string
decode :: I.Elem a b => BSL.ByteString -> Matrix a b
decode = B.decode

-- | Empty 0x0 matrix
{-# INLINE empty #-}
empty :: I.Elem a b => Matrix a b
empty = Matrix 0 0 VS.empty

-- | Is matrix empty?
{-# INLINE null #-}
null :: I.Elem a b => Matrix a b -> Bool
null (Matrix rows cols _) = rows == 0 && cols == 0

-- | Is matrix square?
{-# INLINE square #-}
square :: I.Elem a b => Matrix a b -> Bool
square (Matrix rows cols _) = rows == cols

-- | Matrix where all coeffs are filled with given value
{-# INLINE constant #-}
constant :: I.Elem a b => Int -> Int -> a -> Matrix a b
constant rows cols val = Matrix rows cols $ VS.replicate (rows * cols) (I.cast val)

-- | Matrix where all coeff are 0
{-# INLINE zero #-}
zero :: I.Elem a b => Int -> Int -> Matrix a b
zero rows cols = constant rows cols 0

-- | Matrix where all coeff are 1
{-# INLINE ones #-}
ones :: I.Elem a b => Int -> Int -> Matrix a b
ones rows cols = constant rows cols 1

-- | The identity matrix (not necessarily square).
identity :: I.Elem a b => Int -> Int -> Matrix a b
identity rows cols = I.performIO $ do
    m <- M.new rows cols
    I.call $ M.unsafeWith m I.identity
    unsafeFreeze m

-- | The random matrix of a given size
random :: I.Elem a b => Int -> Int -> IO (Matrix a b)
random rows cols = do
    m <- M.new rows cols
    I.call $ M.unsafeWith m I.random
    unsafeFreeze m

-- | Number of rows for the matrix
{-# INLINE rows #-}
rows :: I.Elem a b => Matrix a b -> Int
rows (Matrix rows _ _) = rows

-- | Number of columns for the matrix
{-# INLINE cols #-}
cols :: I.Elem a b => Matrix a b -> Int
cols (Matrix _ cols _) = cols

-- | Mtrix size as (rows, cols) pair
{-# INLINE dims #-}
dims :: I.Elem a b => Matrix a b -> (Int, Int)
dims (Matrix rows cols _) = (rows, cols)

-- | Matrix coefficient at specific row and col
{-# INLINE (!) #-}
(!) :: I.Elem a b => Matrix a b -> (Int, Int) -> a
(!) m (row,col) = coeff row col m

-- | Matrix coefficient at specific row and col
{-# INLINE coeff #-}
coeff :: I.Elem a b => Int -> Int -> Matrix a b -> a
coeff row col m@(Matrix rows cols _)
    | not (valid m) = error "matrix is not valid"
    | row < 0 || row >= rows = error $ printf "Matrix.coeff: row %d is out of bounds [0..%d)" row rows
    | col < 0 || col >= cols = error $ printf "Matrix.coeff: col %d is out of bounds [0..%d)" col cols
    | otherwise = unsafeCoeff row col m

-- | Unsafe version of coeff function. No bounds check performed so SEGFAULT possible
{-# INLINE unsafeCoeff #-}
unsafeCoeff :: I.Elem a b => Int -> Int -> Matrix a b -> a
unsafeCoeff row col (Matrix rows _ vals) = I.cast $ VS.unsafeIndex vals $ col * rows + row

-- | List of coefficients for the given col
{-# INLINE col #-}
col :: I.Elem a b => Int -> Matrix a b -> [a]
col c m@(Matrix rows _ _) = [coeff r c m | r <- [0..pred rows]]

-- | List of coefficients for the given row
{-# INLINE row #-}
row :: I.Elem a b => Int -> Matrix a b -> [a]
row r m@(Matrix _ cols _) = [coeff r c m | c <- [0..pred cols]]

-- | Extract rectangular block from matrix defined by startRow startCol blockRows blockCols
block :: I.Elem a b => Int -> Int -> Int -> Int -> Matrix a b -> Matrix a b
block startRow startCol blockRows blockCols m =
    generate blockRows blockCols $ \row col ->
        coeff (startRow + row) (startCol + col) m

-- | Verify matrix dimensions and memory layout
{-# INLINE valid #-}
valid :: I.Elem a b => Matrix a b -> Bool
valid (Matrix rows cols vals) = rows >= 0 && cols >= 0 && VS.length vals == rows * cols

-- | The maximum coefficient of the matrix
{-# INLINE maxCoeff #-}
maxCoeff :: (I.Elem a b, Ord a) => Matrix a b -> a
maxCoeff = fold1' max

-- | The minimum coefficient of the matrix
{-# INLINE minCoeff #-}
minCoeff :: (I.Elem a b, Ord a) => Matrix a b -> a
minCoeff = fold1' min

-- | Top @N@ rows of matrix
{-# INLINE topRows #-}
topRows :: I.Elem a b => Int -> Matrix a b -> Matrix a b
topRows n m@(Matrix _ cols _) = block 0 0 n cols m

-- | Bottom @N@ rows of matrix
{-# INLINE bottomRows #-}
bottomRows :: I.Elem a b => Int -> Matrix a b -> Matrix a b
bottomRows n m@(Matrix rows cols _) = block (rows - n) 0 n cols m

-- | Left @N@ columns of matrix
{-# INLINE leftCols #-}
leftCols :: I.Elem a b => Int -> Matrix a b -> Matrix a b
leftCols n m@(Matrix rows _ _) = block 0 0 rows n m

-- | Right @N@ columns of matrix
{-# INLINE rightCols #-}
rightCols :: I.Elem a b => Int -> Matrix a b -> Matrix a b
rightCols n m@(Matrix rows cols _) = block 0 (cols - n) rows n m

-- | Construct matrix from a list of rows, column count is detected as maximum row length. Missing values are filled with 0
fromList :: I.Elem a b => [[a]] -> Matrix a b
fromList list = Matrix rows cols vals where
    rows = length list
    cols = L.foldl' max 0 $ P.map length list
    vals = VS.create $ do
        vm <- VSM.replicate (rows * cols) (I.cast (0 `asTypeOf` (head (head list))))
        forM_ (zip [0..] list) $ \(row, vals) ->
            forM_ (zip [0..] vals) $ \(col, val) ->
                VSM.write vm (col * rows + row) (I.cast val)
        return vm

-- | Convert matrix to a list of rows
toList :: I.Elem a b => Matrix a b -> [[a]]
toList m@(Matrix rows cols vals)
    | not (valid m) = error "matrix is not valid"
    | otherwise = [[I.cast $ vals `VS.unsafeIndex` (col * rows + row) | col <- [0..pred cols]] | row <- [0..pred rows]]

-- | Build matrix of given dimensions and values from given list split on rows. Invalid list length results in error.
fromFlatList :: I.Elem a b => Int -> Int -> [a] -> Matrix a b
fromFlatList rows cols list
    | not (rows * cols == (length list)) = error $ concat ["cannot construct ", show rows, "x", show cols, " matrix from ", show $ length list, " values"]
    | otherwise = Matrix rows cols vals where
        vals = VS.create $ do
            vm <- VSM.replicate (rows * cols) (I.cast (0 `asTypeOf` (head list)))
            forM_ (zip [(col * rows + row) | row <- [0..pred rows], col <- [0..pred cols]] list) $ \(idx, val) ->
                VSM.write vm idx (I.cast val)
            return vm

-- | Convert matrix to a list by concatenating rows
toFlatList :: I.Elem a b => Matrix a b -> [a]
toFlatList m@(Matrix rows cols vals)
    | not (valid m) = error "matrix is not valid"
    | otherwise = [I.cast $ vals `VS.unsafeIndex` (col * rows + row) | row <- [0..pred rows], col <- [0..pred cols]]

-- | [generate rows cols (λ row col -> val)]
--
-- Create matrix using generator function @λ row col -> val@
--
generate :: I.Elem a b => Int -> Int -> (Int -> Int -> a) -> Matrix a b
generate rows cols f = Matrix rows cols $ VS.create $ do
    vals <- VSM.new (rows * cols)
    forM_ [0..pred rows] $ \row ->
        forM_ [0..pred cols] $ \col ->
            VSM.write vals (col * rows + row) (I.cast $ f row col)
    return vals

-- | The sum of all coefficients of the matrix
sum :: I.Elem a b => Matrix a b -> a
sum = _prop I.sum

-- | The product of all coefficients of the matrix
prod :: I.Elem a b => Matrix a b -> a
prod = _prop I.prod

-- | The mean of all coefficients of the matrix
mean :: I.Elem a b => Matrix a b -> a
mean = _prop I.mean

-- | The trace of a matrix is the sum of the diagonal coefficients and can also be computed as sum (diagonal m)
trace :: I.Elem a b => Matrix a b -> a
trace = _prop I.trace

-- | Applied to a predicate and a matrix, all determines if all elements of the matrix satisfies the predicate
all :: I.Elem a b => (a -> Bool) -> Matrix a b -> Bool
all f = VS.all (f . I.cast) . _vals

-- | Applied to a predicate and a matrix, any determines if any element of the matrix satisfies the predicate
any :: I.Elem a b => (a -> Bool) -> Matrix a b -> Bool
any f = VS.any (f . I.cast) . _vals

-- | Returns the number of coefficients in a given matrix that evaluate to true
count :: I.Elem a b => (a -> Bool) -> Matrix a b -> Int
count f = VS.foldl' (\n x -> if f (I.cast x) then succ n else n) 0 . _vals

{-| For vectors, the l2 norm, and for matrices the Frobenius norm.
    In both cases, it consists in the square root of the sum of the square of all the matrix entries.
    For vectors, this is also equals to the square root of the dot product of this with itself.
-}
norm :: I.Elem a b => Matrix a b -> a
norm = _prop I.norm

-- | For vectors, the squared l2 norm, and for matrices the Frobenius norm. In both cases, it consists in the sum of the square of all the matrix entries. For vectors, this is also equals to the dot product of this with itself.
squaredNorm :: I.Elem a b => Matrix a b -> a
squaredNorm = _prop I.squaredNorm

-- | The l2 norm of the matrix using the Blue's algorithm. A Portable Fortran Program to Find the Euclidean Norm of a Vector, ACM TOMS, Vol 4, Issue 1, 1978.
blueNorm :: I.Elem a b => Matrix a b -> a
blueNorm = _prop I.blueNorm

-- | The l2 norm of the matrix avoiding undeflow and overflow. This version use a concatenation of hypot calls, and it is very slow.
hypotNorm :: I.Elem a b => Matrix a b -> a
hypotNorm = _prop I.hypotNorm

-- | The determinant of the matrix
determinant :: I.Elem a b => Matrix a b -> a
determinant m
    | square m = _prop I.determinant m
    | otherwise = error "Matrix.determinant: non-square matrix"

-- | Adding two matrices by adding the corresponding entries together. You can use @(+)@ function as well.
add :: I.Elem a b => Matrix a b -> Matrix a b -> Matrix a b
add m1 m2
    | dims m1 == dims m2 = _binop const I.add m1 m2
    | otherwise = error "Matrix.add: matrices should have the same size"

-- | Subtracting two matrices by subtracting the corresponding entries together. You can use @(-)@ function as well.
sub :: I.Elem a b => Matrix a b -> Matrix a b -> Matrix a b
sub m1 m2
    | dims m1 == dims m2 = _binop const I.sub m1 m2
    | otherwise = error "Matrix.add: matrices should have the same size"

-- | Matrix multiplication. You can use @(*)@ function as well.
mul :: I.Elem a b => Matrix a b -> Matrix a b -> Matrix a b
mul m1 m2
    | cols m1 == rows m2 = _binop (\(rows, _) (_, cols) -> (rows, cols)) I.mul m1 m2
    | otherwise = error "Matrix.mul: number of columns for lhs matrix should be the same as number of rows for rhs matrix"

{- | Apply a given function to each element of the matrix.

Here is an example how to implement scalar matrix multiplication:

>>> let a = fromList [[1,2],[3,4]] :: MatrixXf

>>> a
Matrix 2x2
1.0 2.0
3.0 4.0

>>> map (*10) a
Matrix 2x2
10.0    20.0
30.0    40.0

-}
map :: I.Elem a b => (a -> a) -> Matrix a b -> Matrix a b
map f (Matrix rows cols vals) = Matrix rows cols (VS.map (I.cast . f . I.cast) vals)

{- | Apply a given function to each element of the matrix.

Here is an example how upper triangular matrix can be implemented:

>>> let a = fromList [[1,2,3],[4,5,6],[7,8,9]] :: MatrixXf

>>> a
Matrix 3x3
1.0 2.0 3.0
4.0 5.0 6.0
7.0 8.0 9.0

>>> imap (\row col val -> if row <= col then val else 0) a
Matrix 3x3
1.0 2.0 3.0
0.0 5.0 6.0
0.0 0.0 9.0

-}

imap :: I.Elem a b => (Int -> Int -> a -> a) -> Matrix a b -> Matrix a b
imap f (Matrix rows cols vals) = Matrix rows cols (VS.imap (\n -> let (c, r) = divMod n rows in I.cast . f r c . I.cast) vals)

data TriangularMode
    -- | View matrix as a lower triangular matrix.
    = Lower
    -- | View matrix as an upper triangular matrix.
    | Upper
    -- | View matrix as a lower triangular matrix with zeros on the diagonal.
    | StrictlyLower
    -- | View matrix as an upper triangular matrix with zeros on the diagonal.
    | StrictlyUpper
    -- | View matrix as a lower triangular matrix with ones on the diagonal.
    | UnitLower
    -- | View matrix as an upper triangular matrix with ones on the diagonal.
    | UnitUpper deriving (Eq, Enum, Show, Read)

-- | Triangular view extracted from the current matrix
triangularView :: I.Elem a b => TriangularMode -> Matrix a b -> Matrix a b
triangularView Lower         = imap $ \row col val -> case compare row col of { LT -> 0; _ -> val }
triangularView Upper         = imap $ \row col val -> case compare row col of { GT -> 0; _ -> val }
triangularView StrictlyLower = imap $ \row col val -> case compare row col of { GT -> val; _ -> 0 }
triangularView StrictlyUpper = imap $ \row col val -> case compare row col of { LT -> val; _ -> 0 }
triangularView UnitLower     = imap $ \row col val -> case compare row col of { GT -> val; LT -> 0; EQ -> 1 }
triangularView UnitUpper     = imap $ \row col val -> case compare row col of { LT -> val; GT -> 0; EQ -> 1 }

-- | Lower trinagle of the matrix. Shortcut for @triangularView Lower@
lowerTriangle :: I.Elem a b => Matrix a b -> Matrix a b
lowerTriangle = triangularView Lower

-- | Upper trinagle of the matrix. Shortcut for @triangularView Upper@
upperTriangle :: I.Elem a b => Matrix a b -> Matrix a b
upperTriangle = triangularView Upper

-- | Filter elements in the matrix. Filtered elements will be replaced by 0
filter :: I.Elem a b => (a -> Bool) -> Matrix a b -> Matrix a b
filter f = map (\x -> if f x then x else 0)

-- | Filter elements in the matrix. Filtered elements will be replaced by 0
ifilter :: I.Elem a b => (Int -> Int -> a -> Bool) -> Matrix a b -> Matrix a b
ifilter f = imap (\r c x -> if f r c x then x else 0)

-- | Reduce matrix using user provided function applied to each element.
fold :: I.Elem a b => (c -> a -> c) -> c -> Matrix a b -> c
fold f a (Matrix _ _ vals) = VS.foldl (\a x -> f a (I.cast x)) a vals

-- | Reduce matrix using user provided function applied to each element. This is strict version of 'fold'
fold' :: I.Elem a b => (c -> a -> c) -> c -> Matrix a b -> c
fold' f a (Matrix _ _ vals) = VS.foldl' (\a x -> f a (I.cast x)) a vals

-- | Reduce matrix using user provided function applied to each element and it's index
ifold :: I.Elem a b => (Int -> Int -> c -> a -> c) -> c -> Matrix a b -> c
ifold f a (Matrix rows _ vals) = VS.ifoldl (\a n x -> let (c,r) = divMod n rows in f r c a (I.cast x)) a vals

-- | Reduce matrix using user provided function applied to each element and it's index. This is strict version of 'ifold'
ifold' :: I.Elem a b => (Int -> Int -> c -> a -> c) -> c -> Matrix a b -> c
ifold' f a (Matrix rows _ vals) = VS.ifoldl' (\a n x -> let (c,r) = divMod n rows in f r c a (I.cast x)) a vals

-- | Reduce matrix using user provided function applied to each element.
fold1 :: I.Elem a b => (a -> a -> a) -> Matrix a b -> a
fold1 f = foldl1 f . P.map I.cast . VS.toList . _vals

-- | Reduce matrix using user provided function applied to each element. This is strict version of 'fold'
fold1' :: I.Elem a b => (a -> a -> a) -> Matrix a b -> a
fold1' f = L.foldl1' f . P.map I.cast . VS.toList . _vals

-- | Diagonal of the matrix
diagonal :: I.Elem a b => Matrix a b -> Matrix a b
diagonal = _unop (\(rows, cols) -> (min rows cols, 1)) I.diagonal

{- | Inverse of the matrix

For small fixed sizes up to 4x4, this method uses cofactors. In the general case, this method uses PartialPivLU decomposition
-}
inverse :: I.Elem a b => Matrix a b -> Matrix a b
inverse m
    | square m = _unop id I.inverse m
    | otherwise = error "Matrix.inverse: non-square matrix"

-- | Adjoint of the matrix
adjoint :: I.Elem a b => Matrix a b -> Matrix a b
adjoint = _unop swap I.adjoint

-- | Transpose of the matrix
transpose :: I.Elem a b => Matrix a b -> Matrix a b
transpose = _unop swap I.transpose

-- | Conjugate of the matrix
conjugate :: I.Elem a b => Matrix a b -> Matrix a b
conjugate = _unop id I.conjugate

-- | Nomalize the matrix by deviding it on its 'norm'
normalize :: I.Elem a b => Matrix a b -> Matrix a b
normalize (Matrix rows cols vals) = I.performIO $ do
    vals <- VS.thaw vals
    VSM.unsafeWith vals $ \p ->
        I.call $ I.normalize p (I.cast rows) (I.cast cols)
    Matrix rows cols <$> VS.unsafeFreeze vals

-- | Apply a destructive operation to a matrix. The operation will be performed in place if it is safe to do so and will modify a copy of the matrix otherwise.
modify :: I.Elem a b => (forall s. M.MMatrix a b s -> ST s ()) -> Matrix a b -> Matrix a b
modify f (Matrix rows cols vals) = Matrix rows cols (VS.modify (f . M.MMatrix rows cols) vals)

-- | Convert matrix to different type using user provided element converter
convert :: (I.Elem a b, I.Elem c d) => (a -> c) -> Matrix a b -> Matrix c d
convert f (Matrix rows cols vals) = Matrix rows cols $ VS.map (I.cast . f . I.cast) vals

-- | Yield an immutable copy of the mutable matrix
freeze :: I.Elem a b => PrimMonad m => M.MMatrix a b (PrimState m) -> m (Matrix a b)
freeze (M.MMatrix mrows mcols mvals) = VS.freeze mvals >>= return . Matrix mrows mcols

-- | Yield a mutable copy of the immutable matrix
thaw :: I.Elem a b => PrimMonad m => Matrix a b -> m (M.MMatrix a b (PrimState m))
thaw (Matrix rows cols vals) = VS.thaw vals >>= return . M.MMatrix rows cols

-- | Unsafe convert a mutable matrix to an immutable one without copying. The mutable matrix may not be used after this operation.
unsafeFreeze :: I.Elem a b => PrimMonad m => M.MMatrix a b (PrimState m) -> m (Matrix a b)
unsafeFreeze (M.MMatrix mrows mcols mvals) = VS.unsafeFreeze mvals >>= return . Matrix mrows mcols

-- | Unsafely convert an immutable matrix to a mutable one without copying. The immutable matrix may not be used after this operation.
unsafeThaw :: I.Elem a b => PrimMonad m => Matrix a b -> m (M.MMatrix a b (PrimState m))
unsafeThaw (Matrix rows cols vals) = VS.unsafeThaw vals >>= return . M.MMatrix rows cols

-- | Pass a pointer to the matrix's data to the IO action. The data may not be modified through the pointer.
unsafeWith :: I.Elem a b => Matrix a b -> (Ptr b -> CInt -> CInt -> IO c) -> IO c
unsafeWith m@(Matrix rows cols vals) f
    | not (valid m) = fail "Matrix.unsafeWith: matrix layout is invalid"
    | otherwise = VS.unsafeWith vals $ \p -> f p (I.cast rows) (I.cast cols)

{-# INLINE _prop #-}
_prop :: I.Elem a b => (Ptr b -> Ptr b -> CInt -> CInt -> IO CString) -> Matrix a b -> a
_prop f m = I.cast $ I.performIO $ alloca $ \p -> do
    I.call $ unsafeWith m (f p)
    peek p

{-# INLINE _binop #-}
_binop :: I.Elem a b => ((Int, Int) -> (Int, Int) -> (Int, Int)) -> (Ptr b -> CInt -> CInt -> Ptr b -> CInt -> CInt -> Ptr b -> CInt -> CInt -> IO CString) -> Matrix a b -> Matrix a b -> Matrix a b
_binop f g m1 m2 = I.performIO $ do
    m0 <- uncurry M.new $ f (dims m1) (dims m2)
    M.unsafeWith m0 $ \vals0 rows0 cols0 ->
        unsafeWith m1 $ \vals1 rows1 cols1 ->
            unsafeWith m2 $ \vals2 rows2 cols2 ->
                I.call $ g
                    vals0 rows0 cols0
                    vals1 rows1 cols1
                    vals2 rows2 cols2
    unsafeFreeze m0

{-# INLINE _unop #-}
_unop :: I.Elem a b => ((Int,Int) -> (Int,Int)) -> (Ptr b -> CInt -> CInt -> Ptr b -> CInt -> CInt -> IO CString) -> Matrix a b -> Matrix a b
_unop f g m1 = I.performIO $ do
    m0 <- uncurry M.new $ f (dims m1)
    M.unsafeWith m0 $ \vals0 rows0 cols0 ->
        unsafeWith m1 $ \vals1 rows1 cols1 ->
            I.call $ g
                vals0 rows0 cols0
                vals1 rows1 cols1
    unsafeFreeze m0

{-# INLINE _vals #-}
_vals :: I.Elem a b => Matrix a b -> VS.Vector b
_vals (Matrix _ _ vals) = vals


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
module Math.HiddenMarkovModel.Private where

import qualified Math.HiddenMarkovModel.Distribution as Distr
import qualified Math.HiddenMarkovModel.CSV as HMMCSV
import Math.HiddenMarkovModel.Distribution (State(State))

import qualified Numeric.LinearAlgebra.Algorithms as Algo
import qualified Numeric.LinearAlgebra.Util as LinAlg
import qualified Numeric.Container as NC
import qualified Data.Packed.Development as Dev
import qualified Data.Packed.Matrix as Matrix
import qualified Data.Packed.Vector as Vector
import Numeric.Container ((<>))
import Data.Packed.Matrix (Matrix)
import Data.Packed.Vector (Vector)

import Control.DeepSeq (NFData, rnf)
import Foreign.Storable (Storable)

import qualified Data.NonEmpty.Class as NonEmptyC
import qualified Data.NonEmpty as NonEmpty
import qualified Data.Semigroup as Sg
import qualified Data.List as List
import Data.Traversable (Traversable, mapAccumL)
import Data.Tuple.HT (mapPair, mapFst, mapSnd, swap)

{- |
A Hidden Markov model consists of a number of (hidden) states
and a set of emissions.
There is a vector for the initial probability of each state
and a matrix containing the probability for switching
from one state to another one.
The 'distribution' field points to probability distributions
that associate every state with emissions of different probability.
Famous distribution instances are discrete and Gaussian distributions.
See "Math.HiddenMarkovModel.Distribution" for details.

The transition matrix is transposed
with respect to popular HMM descriptions.
But I think this is the natural orientation, because this way
you can write \"transition matrix times probability column vector\".

The type has two type parameters,
although the one for the distribution would be enough.
However, replacing @prob@ by @Distr.Probability distr@
would prohibit the derived Show and Read instances.
-}
data T distr prob =
   Cons {
      initial :: Vector prob,
      transition :: Matrix prob,
      distribution :: distr
   }
   deriving (Show, Read)

instance
   (NFData distr, NFData prob, Storable prob) =>
      NFData (T distr prob) where
   rnf hmm = rnf (initial hmm, transition hmm, distribution hmm)

emission ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   T distr prob -> emission -> Vector prob
emission  =  Distr.emissionProb . distribution

forward ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob -> NonEmpty.T f emission -> prob
forward hmm = NC.sumElements . NonEmpty.last . alpha hmm

alpha ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob ->
   NonEmpty.T f emission -> NonEmpty.T f (Vector prob)
alpha hmm (NonEmpty.Cons x xs) =
   NonEmpty.scanl
      (\alphai xi -> NC.mul (emission hmm xi) (transition hmm <> alphai))
      (NC.mul (emission hmm x) (initial hmm))
      xs

backward ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob -> NonEmpty.T f emission -> prob
backward hmm (NonEmpty.Cons x xs) =
   NC.sumElements $
   NC.mul (initial hmm) $
   NC.mul (emission hmm x) $
   NonEmpty.head $ beta hmm xs

beta ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob ->
   f emission -> NonEmpty.T f (Vector prob)
beta hmm =
   NonEmpty.scanr
      (\xi betai -> NC.mul (emission hmm xi) betai <> transition hmm)
      (NC.constant 1 (NC.dim $ initial hmm))

alphaBeta ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob ->
   NonEmpty.T f emission ->
   (prob, NonEmpty.T f (Vector prob), NonEmpty.T f (Vector prob))
alphaBeta hmm xs =
   let alphas = alpha hmm xs
       betas = beta hmm $ NonEmpty.tail xs
       recipLikelihood = recip $ NC.sumElements $ NonEmpty.last alphas
   in  (recipLikelihood, alphas, betas)

xiFromAlphaBeta ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   T distr prob -> prob ->
   NonEmpty.T [] emission ->
   NonEmpty.T [] (Vector prob) ->
   NonEmpty.T [] (Vector prob) ->
   [Matrix prob]
xiFromAlphaBeta hmm recipLikelihood xs alphas betas =
   zipWith3
      (\x alpha0 beta1 ->
         NC.scale recipLikelihood $
         NC.mul
            (NC.outer (NC.mul (emission hmm x) beta1) alpha0)
            (transition hmm))
      (NonEmpty.tail xs)
      (NonEmpty.init alphas)
      (NonEmpty.tail betas)

zetaFromXi ::
   (NC.Product prob) => T distr prob -> [Matrix prob] -> [Vector prob]
zetaFromXi hmm xis =
   map (NC.constant 1 (Matrix.rows $ transition hmm) <>) xis

zetaFromAlphaBeta ::
   (NC.Container Vector prob) =>
   prob ->
   NonEmpty.T [] (Vector prob) ->
   NonEmpty.T [] (Vector prob) ->
   NonEmpty.T [] (Vector prob)
zetaFromAlphaBeta recipLikelihood alphas betas =
   fmap (NC.scale recipLikelihood) $
   NonEmptyC.zipWith NC.mul alphas betas

{- |
In constrast to Math.HiddenMarkovModel.reveal
this does not normalize the vector.
This is slightly simpler but for long sequences
the product of probabilities might be smaller
than the smallest representable number.
-}
reveal ::
   (Distr.EmissionProb distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission,
    Traversable f) =>
   T distr prob -> NonEmpty.T f emission -> NonEmpty.T f State
reveal hmm (NonEmpty.Cons x xs) =
   fmap State $
   uncurry (NonEmpty.scanr Dev.at') $
   mapFst NC.maxIndex $
   mapAccumL
      (\alphai xi ->
         swap $ mapSnd (NC.mul (emission hmm xi)) $
         matrixMaxMul (transition hmm) alphai)
      (NC.mul (emission hmm x) (initial hmm)) xs

matrixMaxMul ::
   (NC.Container Vector a) =>
   Matrix a -> Vector a -> (Vector Int, Vector a)
matrixMaxMul m v =
   mapPair (Vector.fromList, Vector.fromList) $ unzip $
   map ((\x -> (NC.maxIndex x, NC.maxElement x)) . NC.mul v) $
   Matrix.toRows m

{- |
A trained model is a temporary form of a Hidden Markov model
that we need during the training on multiple training sequences.
It allows to collect knowledge over many sequences with 'mergeTrained',
even with mixed supervised and unsupervised training.
You finish the training by converting the trained model
back to a plain modul using 'finishTraining'.

You can create a trained model in three ways:

* supervised training using an emission sequence with associated states,

* unsupervised training using an emission sequence and an existing Hidden Markov Model,

* derive it from state sequence patterns, cf. "Math.HiddenMarkovModel.Pattern".
-}
data Trained distr prob =
   Trained {
      trainedInitial :: Vector prob,
      trainedTransition :: Matrix prob,
      trainedDistribution :: distr
   }
   deriving (Show, Read)

instance
   (NFData distr, NFData prob, Storable prob) =>
      NFData (Trained distr prob) where
   rnf hmm =
      rnf (trainedInitial hmm, trainedTransition hmm, trainedDistribution hmm)

sumTransitions ::
   (NC.Container Vector e) =>
   T distr e -> [Matrix e] -> Matrix e
sumTransitions hmm =
   List.foldl' NC.add (NC.konst 0 $ LinAlg.size $ transition hmm)

{- |
Baum-Welch algorithm
-}
trainUnsupervised ::
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob, Distr.Emission distr ~ emission) =>
   T distr prob -> NonEmpty.T [] emission -> Trained tdistr prob
trainUnsupervised hmm xs =
   let (recipLikelihood, alphas, betas) = alphaBeta hmm xs
       zetas = zetaFromAlphaBeta recipLikelihood alphas betas

   in  Trained {
          trainedInitial = NonEmpty.head zetas,
          trainedTransition =
             sumTransitions hmm $
             xiFromAlphaBeta hmm recipLikelihood xs alphas betas,
          trainedDistribution =
             Distr.accumulateEmissions $ map (zip (NonEmpty.flatten xs)) $
             List.transpose $ map Vector.toList $ NonEmpty.flatten zetas
       }

mergeTrained ::
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob) =>
   Trained tdistr prob -> Trained tdistr prob -> Trained tdistr prob
mergeTrained hmm0 hmm1 =
   Trained {
      trainedInitial = NC.add (trainedInitial hmm0) (trainedInitial hmm1),
      trainedTransition =
         NC.add (trainedTransition hmm0) (trainedTransition hmm1),
      trainedDistribution =
         Distr.combine
            (trainedDistribution hmm0) (trainedDistribution hmm1)
   }

instance
   (Distr.Estimate tdistr, Distr.Distribution tdistr ~ distr,
    Distr.Probability distr ~ prob) =>
      Sg.Semigroup (Trained tdistr prob) where
   (<>) = mergeTrained

toCells ::
   (Distr.CSV distr, Algo.Field prob, Show prob) =>
   T distr prob -> [[String]]
toCells hmm =
   (HMMCSV.cellsFromVector $ initial hmm) :
   (HMMCSV.cellsFromMatrix $ transition hmm) ++
   [] :
   (Distr.toCells $ distribution hmm)

parseCSV ::
   (Distr.CSV distr, Algo.Field prob, Read prob) =>
   HMMCSV.CSVParser (T distr prob)
parseCSV = do
   v <- HMMCSV.parseNonEmptyVectorCells
   m <- HMMCSV.parseSquareMatrixCells $ Vector.dim v
   HMMCSV.skipEmptyRow
   distr <- Distr.parseCells $ Vector.dim v
   return $ Cons {
      initial = v,
      transition = m,
      distribution = distr
   }


-- |
-- Module      : Cartesian.Internal.Types
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created October 31 2015

-- TODO | - Use TemplateHaskell (?)
--        - Strictness
--        - Performance, inlining

-- SPEC | -
--        -

------------------------------------------------------------------------------------------------------------------------------------------------------
-- GHC Pragmas
------------------------------------------------------------------------------------------------------------------------------------------------------
{-# LANGUAGE TemplateHaskell        #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE FlexibleInstances      #-}

------------------------------------------------------------------------------------------------------------------------------------------------------
-- API
------------------------------------------------------------------------------------------------------------------------------------------------------
module Cartesian.Internal.Types where

------------------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
------------------------------------------------------------------------------------------------------------------------------------------------------
import Control.Lens (Simple, Lens, lens)

import Data.Complex (Complex(..))

import Linear.V1
import Linear.V2
import Linear.V3
import Linear.V4

------------------------------------------------------------------------------------------------------------------------------------------------------
-- Types
------------------------------------------------------------------------------------------------------------------------------------------------------

-- Synonyms ------------------------------------------------------------------------------------------------------------------------------------------

-- TODO: Add some aliased lenses for these aliased types (?)

-- | A lens focusing on a single [vector-]component in a BoundingBox
type BoxLens v v' f f' = Lens (BoundingBox (v f)) (BoundingBox (v' f')) f f'

-- | An axis represented as (begin, length)
type Axis a = (a, a)

-- | A vector where each component represents a single axis (cf. 'Axis')
type Axes v a = v (Axis a)

-- |
-- type Domain

-- |
-- TODO: Rename (eg. 'Shape') (?)
type Polygon m v f = m (v f)

-- | Coordinate system wrappers
newtype Normalised v = Normalised { absolute   :: v } --
newtype Absolute   v = Absoloute  { normalised :: v } --

-- Types ---------------------------------------------------------------------------------------------------------------------------------------------

-- |
-- TODO: Anchors (eg. C, N, S, E W and combinations thereof, perhaps represented as relative Vectors)
-- TODO: Define some standard instances (eg. Functor, Applicative)
data BoundingBox v = BoundingBox { cornerOf :: v, sizeOf :: v } deriving (Show, Eq)

-- |
-- TODO: Use record (eg. from, to) (?)
data Line v = Line v v deriving (Show, Eq)

-- |
data Linear f = Linear { interceptOf :: f, slopeOf :: f } deriving (Show, Eq)

-- |
data Inclusivity r = Inclusive r | Exclusive r -- TODO: Rename (?)
data Interval r    = Interval (Inclusivity r) (Inclusivity r)

-- |
-- TODO: Use existing type instead (?)
-- data Side = SideLeft | SideRight | SideTop | SideBottom

-- Classes -------------------------------------------------------------------------------------------------------------------------------------------

-- TODO: How do you generate lenses for non-record types (?)
class HasX a f | a -> f where { x :: Simple Lens a f }
class HasY a f | a -> f where { y :: Simple Lens a f }
class HasZ a f | a -> f where { z :: Simple Lens a f }

-- Instances -----------------------------------------------------------------------------------------------------------------------------------------

-- TODO: Generate these instances with TemplateHaskell (?)

instance HasX (V1 f) f where
  x = lens (\(V1 x') -> x') (\_ x' -> V1 x')

instance HasX (V2 f) f where
  x = lens (\(V2 x' _) -> x') (\(V2 _ y') x' -> V2 x' y')

instance HasY (V2 f) f where
  y = lens (\(V2 _ y') -> y') (\(V2 x' _) y' -> V2 x' y')

instance HasX (V3 f) f where
  x = lens (\(V3 x' _ _) -> x') (\(V3 _ y' z') x' -> V3 x' y' z')

instance HasY (V3 f) f where
  y = lens (\(V3 _ y' _) -> y') (\(V3 x' _ z') y' -> V3 x' y' z')

instance HasZ (V3 f) f where
  z = lens (\(V3 _ _ z') -> z') (\(V3 x' y' _) z' -> V3 x' y' z')

instance HasX (V4 f) f where
  x = lens (\(V4 x' _ _ _) -> x') (\(V4 _ y' z' w') x' -> V4 x' y' z' w')

instance HasY (V4 f) f where
  y = lens (\(V4 _ y' _ _) -> y') (\(V4 x' _ z' w') y' -> V4 x' y' z' w')

instance HasZ (V4 f) f where
  z = lens (\(V4 _ _ z' _) -> z') (\(V4 x' y' _ w') z' -> V4 x' y' z' w')

instance HasX (Complex f) f where
  x = lens (\(x':+_) -> x') (\(_:+y') x' -> x':+y')

instance HasY (Complex f) f where
  y = lens (\(_':+y') -> y') (\(x':+_) y' -> x':+y')


{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE GADTs                    #-}
module Grenade.Layers.Internal.Update (
    descendMatrix
  , descendVector
  , MatrixInputValues (..)
  , MatrixResult (..)
  , VectorInputValues (..)
  , VectorResult (..)
  ) where

import           Data.Maybe                   (fromJust)
import qualified Data.Vector.Storable         as U (unsafeFromForeignPtr0,
                                                    unsafeToForeignPtr0)
import           Foreign                      (mallocForeignPtrArray, withForeignPtr)
import           Foreign.Ptr                  (Ptr)
import           GHC.TypeLits
import           Numeric.LinearAlgebra        (Vector, flatten)
import qualified Numeric.LinearAlgebra.Devel  as U
import           Numeric.LinearAlgebra.Static
import           System.IO.Unsafe             (unsafePerformIO)

import           Grenade.Core.Optimizer
import           Grenade.Types

data MatrixInputValues rows columns
  = MatrixValuesSGD
      !(L rows columns) -- ^ current weights
      !(L rows columns) -- ^ gradients
      !(L rows columns) -- ^ last update (old momentum)
  | MatrixValuesAdam
      !Int   -- ^ Step
      !(L rows columns) -- ^ current weights
      !(L rows columns) -- ^ gradients
      !(L rows columns) -- ^ current m
      !(L rows columns) -- ^ current v

data MatrixResult rows columns
  = MatrixResultSGD
      { matrixActivations :: !(L rows columns) -- ^ new activations (weights)
      , matrixMomentum    :: !(L rows columns) -- ^ new momentum
      }
  | MatrixResultAdam
      { matrixActivations :: !(L rows columns) -- ^ new activations (weights)
      , matrixM           :: !(L rows columns) -- ^ new m
      , matrixV           :: !(L rows columns) -- ^ new v
      }

data VectorInputValues r
  = VectorValuesSGD
      !(R r) -- ^ current weights
      !(R r) -- ^ gradients
      !(R r) -- ^ last update (old momentum)
  | VectorValuesAdam
      !Int   -- ^ Step
      !(R r) -- ^ current weights
      !(R r) -- ^ current gradients
      !(R r) -- ^ current m
      !(R r) -- ^ current v

data VectorResult r
  = VectorResultSGD
      { vectorBias     :: !(R r) -- ^ new activations (bias)
      , vectorMomentum :: !(R r) -- ^ new momentum
      }
  | VectorResultAdam
      { vectorBias :: !(R r) -- ^ new activations (bias)
      , vectorM    :: !(R r) -- ^ new m
      , vectorV    :: !(R r) -- ^ new v
      }

descendMatrix :: (KnownNat rows, KnownNat columns) => Optimizer o -> MatrixInputValues rows columns -> MatrixResult rows columns
descendMatrix (OptSGD rate momentum regulariser) (MatrixValuesSGD weights gradient lastUpdate) =
  let (rows, cols) = size weights
      len          = rows * cols
      -- Most gradients come in in ColumnMajor,
      -- so we'll transpose here before flattening them
      -- into a vector to prevent a copy.
      --
      -- This gives ~15% speed improvement for LSTMs.
      weights'     = flatten . tr . extract $ weights
      gradient'    = flatten . tr . extract $ gradient
      lastUpdate'  = flatten . tr . extract $ lastUpdate
      (vw, vm)     = descendUnsafeSGD len rate momentum regulariser weights' gradient' lastUpdate'

      -- Note that it's ColumnMajor, as we did a transpose before
      -- using the internal vectors.
      mw           = U.matrixFromVector U.ColumnMajor rows cols vw
      mm           = U.matrixFromVector U.ColumnMajor rows cols vm
  in  MatrixResultSGD (fromJust . create $ mw) (fromJust . create $ mm)
descendMatrix (OptAdam alpha beta1 beta2 epsilon lambda) (MatrixValuesAdam step weights gradient m v) =
  let (rows, cols) = size weights
      len          = rows * cols
      -- Most gradients come in in ColumnMajor,
      -- so we'll transpose here before flattening them
      -- into a vector to prevent a copy.
      --
      -- This gives ~15% speed improvement for LSTMs.
      weights'  = flatten . tr . extract $ weights
      gradient' = flatten . tr . extract $ gradient
      m'        = flatten . tr . extract $ m
      v'        = flatten . tr . extract $ v
      (vw, vm, vv)     = descendUnsafeAdam len step alpha beta1 beta2 epsilon lambda weights' gradient' m' v'

      -- Note that it's ColumnMajor, as we did a transpose before
      -- using the internal vectors.
      mw           = U.matrixFromVector U.ColumnMajor rows cols vw
      mm           = U.matrixFromVector U.ColumnMajor rows cols vm
      mv           = U.matrixFromVector U.ColumnMajor rows cols vv
  in  MatrixResultAdam (fromJust . create $ mw) (fromJust . create $ mm) (fromJust . create $ mv)
descendMatrix opt _ = error $ "optimzer does not match to MatrixInputValues in implementation! Optimizer: " ++ show opt

descendVector :: (KnownNat r) => Optimizer o -> VectorInputValues r -> VectorResult r
descendVector (OptSGD rate momentum regulariser) (VectorValuesSGD weights gradient lastUpdate) =
  let len          = size weights
      weights'     = extract weights
      gradient'    = extract gradient
      lastUpdate'  = extract lastUpdate
      (vw, vm)     = descendUnsafeSGD len rate momentum regulariser weights' gradient' lastUpdate'
  in  VectorResultSGD (fromJust $ create vw) (fromJust $ create vm)
descendVector (OptAdam alpha beta1 beta2 epsilon lambda) (VectorValuesAdam step weights gradient m v) =
  let len       = size weights
      weights'  = extract weights
      gradient' = extract gradient
      m'        = extract m
      v'        = extract v
      (vw, vm, vv)     = descendUnsafeAdam len step alpha beta1 beta2 epsilon lambda weights' gradient' m' v'
  in  VectorResultAdam (fromJust $ create vw) (fromJust $ create vm) (fromJust $ create vv)
descendVector opt _ = error $ "optimzer does not match to VectorInputValues in implementation! Optimizer: " ++ show opt

descendUnsafeSGD :: Int -> RealNum -> RealNum -> RealNum -> Vector RealNum -> Vector RealNum -> Vector RealNum -> (Vector RealNum, Vector RealNum)
descendUnsafeSGD len rate momentum regulariser weights gradient lastUpdate =
  unsafePerformIO $ do
    outWPtr <- mallocForeignPtrArray len
    outMPtr <- mallocForeignPtrArray len
    let (wPtr, _) = U.unsafeToForeignPtr0 weights
    let (gPtr, _) = U.unsafeToForeignPtr0 gradient
    let (lPtr, _) = U.unsafeToForeignPtr0 lastUpdate

    withForeignPtr wPtr $ \wPtr' ->
      withForeignPtr gPtr $ \gPtr' ->
        withForeignPtr lPtr $ \lPtr' ->
          withForeignPtr outWPtr $ \outWPtr' ->
            withForeignPtr outMPtr $ \outMPtr' ->
              descend_sgd_cpu len rate momentum regulariser wPtr' gPtr' lPtr' outWPtr' outMPtr'

    return (U.unsafeFromForeignPtr0 outWPtr len, U.unsafeFromForeignPtr0 outMPtr len)

descendUnsafeAdam ::
     Int -- Len
  -> Int -- Step
  -> RealNum -- Alpha
  -> RealNum -- Beta1
  -> RealNum -- Beta2
  -> RealNum -- Epsilon
  -> RealNum -- Lambda
  -> Vector RealNum -- Weights
  -> Vector RealNum -- Gradient
  -> Vector RealNum -- M
  -> Vector RealNum -- V
  -> (Vector RealNum, Vector RealNum, Vector RealNum)
descendUnsafeAdam len step alpha beta1 beta2 epsilon lambda weights gradient m v =
  unsafePerformIO $ do
    outWPtr <- mallocForeignPtrArray len
    outMPtr <- mallocForeignPtrArray len
    outVPtr <- mallocForeignPtrArray len
    let (wPtr, _) = U.unsafeToForeignPtr0 weights
    let (gPtr, _) = U.unsafeToForeignPtr0 gradient
    let (mPtr, _) = U.unsafeToForeignPtr0 m
    let (vPtr, _) = U.unsafeToForeignPtr0 v
    withForeignPtr wPtr $ \wPtr' ->
      withForeignPtr gPtr $ \gPtr' ->
        withForeignPtr mPtr $ \mPtr' ->
          withForeignPtr vPtr $ \vPtr' ->
            withForeignPtr outWPtr $ \outWPtr' ->
              withForeignPtr outMPtr $ \outMPtr' ->
                withForeignPtr outVPtr $ \outVPtr' -> descend_adam_cpu len step alpha beta1 beta2 epsilon lambda wPtr' gPtr' mPtr' vPtr' outWPtr' outMPtr' outVPtr'
    return (U.unsafeFromForeignPtr0 outWPtr len, U.unsafeFromForeignPtr0 outMPtr len, U.unsafeFromForeignPtr0 outVPtr len)

foreign import ccall unsafe
    descend_sgd_cpu
      :: Int -> RealNum -> RealNum -> RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> IO ()

foreign import ccall unsafe
    descend_adam_cpu
      :: Int -> Int -> RealNum -> RealNum -> RealNum -> RealNum -> RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> Ptr RealNum -> IO ()


module Numerical.HBLAS.Lapack.FFI where
import Foreign.Ptr
import Foreign()
import Foreign.C.Types
import Data.Complex
import Data.Int

{-

stylenote: we will not use the LAPACKE_* operations, only the
LAPACKE_*_work variants that require an explicitly provided work buffer.

This is to ensure that solver routine allocation behavior is transparent

-}

{-
void LAPACK_dgesvx( char* fact, char* trans, lapack_int* n, lapack_int* nrhs,
                    double* a, lapack_int* lda, double* af, lapack_int* ldaf,
                    lapack_int* ipiv, char* equed, double* r, double* c,
                    double* b, lapack_int* ldb, double* x, lapack_int* ldx,
                    double* rcond, double* ferr, double* berr, double* work,
                    lapack_int* iwork, lapack_int *info );

-}

{-
    fortran FFI conventions!
-}

--type Stride_C =

newtype Fact_C = Fact_C CChar
newtype Trans_C = Trans_C CChar
newtype Stride_C = Stride_C Int32
newtype Equilib_C = Equilib_C CChar

type Fun_FFI_GESVX el = Ptr Fact_C  {- fact -}-> Ptr Trans_C {- trans -}
    -> Ptr Int32  {-n -}-> Ptr Int32 {- NRHS -}->
    Ptr el {- a -} -> Ptr Stride_C {- lda -} -> Ptr Double {- af -} -> Ptr Stride_C  {- ldf-}->
    Ptr Int32 -> Ptr Equilib_C {- equed -} -> Ptr el {- r -} -> Ptr el  ->
    Ptr el {- b -} -> Ptr Stride_C {- ld b   -} -> Ptr el {- x -} -> Ptr Stride_C {- ldx -}->
    Ptr el {-rcond -}-> Ptr el {- ferr-} -> Ptr el {-berr-} -> Ptr el {-work-}->
    Ptr Int32 {-iwork -}-> Ptr Int32 {-info  -} -> IO ()

{-

the prefixes mean s=single,d=double,c=complex float,d=complex double

fact will be a 1 character C string
either
    "F", then the inputs af and ipiv already contain the permuted LU factorization
        (act as input rather than result params)
    "E", Matrix input A will be equilibriated if needed, then copied to AF and Factored
    "N", matrix input A will be copied to AF

-}

{-
Xgesvx  is the s -sing

im assuming for now that any real use of *gesvx routines, or any other
n^3 complexity algs from LAPACK, are on inputs typically  n>=15, which means > 1000 flops,
which is > 1µs, and thus ok to
-}

--need to get around to wrapping these, but thats for another day
foreign import ccall  "sgesvx_"  sgesvx :: Fun_FFI_GESVX Float
foreign import ccall  "dgesvx_"  dgesvx :: Fun_FFI_GESVX Double
foreign import ccall  "cgesvx_"  cgesvx :: Fun_FFI_GESVX (Complex Float)
foreign import ccall  "zgesvx_"  zgesvx :: Fun_FFI_GESVX (Complex Double)

--lapack_int ?syev_(  char *jobz, char *uplo, lapack_int *n, ?* a, lapack_int * lda, ?* w );
-- ? is Double or Float

newtype JobTy = JBT CChar
newtype UploTy = UPLT CChar
newtype Info = Info Int32

--basic symmetric eigen value solvers
type SYEV_FUN_FFI elem = Ptr JobTy -> Ptr UploTy -> Ptr Int32  -> Ptr elem -> Ptr Int32 -> Ptr elem -> Ptr Info-> IO ()
foreign import ccall "ssyev_" ssyev_ffi :: SYEV_FUN_FFI Float
foreign import ccall "dsyev_" dsyev_ffi :: SYEV_FUN_FFI Double

{-unsafe versions of lapack routines are meant to ONLY be used for workspace queries-}
foreign import ccall unsafe "ssyev_" ssyev_ffi_unsafe :: SYEV_FUN_FFI Float
foreign import ccall unsafe "dsyev_" dsyev_ffi_unsafe :: SYEV_FUN_FFI Double

--lapack_int LAPACKE_<?>gesv( int matrix_order, lapack_int n, lapack_int nrhs, <datatype>* a, lapack_int lda, lapack_int* ipiv, <datatype>* b, lapack_int ldb );
--call sgesv( n, nrhs, a, lda, ipiv, b, ldb, info )

type GESV_FUN_FFI elem = Ptr Int32  {- n -} -> Ptr Int32 {- nrhs -} -> Ptr elem {- a -}
            -> Ptr Int32 {-  lda -} -> Ptr Int32 {- permutation vector -}
            -> Ptr elem  {- b -} -> Ptr Int32 {- ldb -} -> Ptr Info -> IO ()
-- | basic Linear system solvers. they act inplace
foreign import ccall  "sgesv_"  sgesv_ffi  ::GESV_FUN_FFI Float
foreign import ccall  "dgesv_"  dgesv_ffi :: GESV_FUN_FFI Double
foreign import ccall  "cgesv_"  cgesv_ffi :: GESV_FUN_FFI (Complex Float)
foreign import ccall  "zgesv_"  zgesv_ffi :: GESV_FUN_FFI (Complex Double)

-- / not sure if linear solvers ever are run in < 1 microsecond size instances in practice
foreign import ccall unsafe "sgesv_"  sgesv_ffi_unsafe  ::GESV_FUN_FFI Float
foreign import ccall unsafe "dgesv_"  dgesv_ffi_unsafe :: GESV_FUN_FFI Double
foreign import ccall unsafe "cgesv_"  cgesv_ffi_unsafe :: GESV_FUN_FFI (Complex Float)
foreign import ccall unsafe "zgesv_"  zgesv_ffi_unsafe :: GESV_FUN_FFI (Complex Double)

{- only provide  -}


{-# LANGUAGE ForeignFunctionInterface #-}
module Data.Matrix.Static.Internal
    ( c_dd_mul
    , c_ds_mul
    , c_sd_mul
    , c_ss_mul
    , c_ss_cmul
    , c_sd_plus
    , c_ss_plus
    , c_inverse
    , c_cholesky
    , c_eig
    , c_eigs
    , c_eigsh
    , c_seigs
    , c_seigsh
    , c_geigsh
    , c_bdcsvd

    , computationInfo
    ) where

import Data.Complex (Complex)
import Foreign
import Foreign.C.Types
import Foreign.C.String

-------------------------------------------------------------------------------
-- Arithmetic
-------------------------------------------------------------------------------
foreign import ccall "eigen_dd_mul"
    c_dd_mul :: CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_ds_mul"
    c_ds_mul :: CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_sd_mul"
    c_sd_mul :: CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> Ptr a -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_sd_plus"
    c_sd_plus :: CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> Ptr a -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_ss_mul"
    c_ss_mul :: CInt
          -> Ptr (Ptr a) -> Ptr CInt -> Ptr (Ptr CInt) -> CInt -> CInt -> Ptr CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_ss_cmul"
    c_ss_cmul :: CInt
          -> Ptr (Ptr a) -> Ptr CInt -> Ptr (Ptr CInt) -> CInt -> CInt -> Ptr CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_ss_plus"
    c_ss_plus :: CInt
          -> Ptr (Ptr a) -> Ptr CInt -> Ptr (Ptr CInt) -> CInt -> CInt -> Ptr CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_inverse"
    c_inverse :: CInt
              -> Ptr a -> CInt -> CInt
              -> Ptr a -> CInt -> CInt
              -> IO CString

foreign import ccall "eigen_cholesky"
    c_cholesky :: CInt
               -> Ptr a -> Ptr a ->  CInt -> IO CString

foreign import ccall "eigen_eig"
    c_eig :: Ptr (Complex Double) -> Ptr (Complex Double)
          -> Ptr Double -> CInt -> IO CString

foreign import ccall "spectral_eigs"
    c_eigs :: CInt      -- ^ The number of eigenvectors to search
           -> Ptr (Complex Double)   -- ^
           -> Ptr (Complex Double)   -- ^
           -> Ptr Double   -- ^ Input matrix
           -> CInt         -- ^ Matrix size
           -> CInt   -- ^ ncv
           -> CInt   -- ^ max iterations
           -> Double  -- ^ tolerance
           -> CInt
           -> Double
           -> CInt
           -> IO CInt

foreign import ccall "spectral_eigsh"
    c_eigsh :: CInt
            -> Ptr Double -> Ptr Double
            -> Ptr Double -> CInt
            -> CInt -> CInt -> Double -> CInt -> Double -> CInt
            -> IO CInt

foreign import ccall "spectral_seigs"
    c_seigs :: CInt
            -> Ptr (Complex Double) -> Ptr (Complex Double)
            -> Ptr Double -> Ptr CInt -> Ptr CInt -> CInt -> CInt
            -> CInt -> CInt -> Double -> CInt -> Double -> CInt
            -> IO CInt

foreign import ccall "spectral_seigsh"
    c_seigsh :: CInt
             -> Ptr Double -> Ptr Double
             -> Ptr Double -> Ptr CInt -> Ptr CInt -> CInt -> CInt
             -> CInt -> CInt -> Double -> CInt -> Double -> CInt
             -> IO CInt

foreign import ccall "spectral_geigsh"
    c_geigsh :: CInt
             -> Ptr Double -> Ptr Double
             -> Ptr Double -> CInt
             -> Ptr Double -> Ptr CInt -> Ptr CInt -> CInt
             -> CInt -> CInt -> Double -> CInt
             -> IO CInt

foreign import ccall "eigen_bdcsvd"
    c_bdcsvd :: CInt -> Ptr a -> Ptr b -> Ptr a
             -> Ptr a -> CInt -> CInt -> IO CString

computationInfo :: CInt -> Maybe String
computationInfo 0 = Nothing
computationInfo 1 = Just "NOT_COMPUTED"
computationInfo 2 = Just "NOT_CONVERGING"
computationInfo 3 = Just "NUMERICAL_ISSUE"
computationInfo _ = Just "UNKNOWN ERROR"
{-# INLINE computationInfo #-}