{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeInType          #-}
{-# LANGUAGE TypeOperators       #-}

module Eigen.SparseMatrix
  ( -- * Types
    SparseMatrix(..)
  , SparseMatrixXf
  , SparseMatrixXd
  , SparseMatrixXcf
  , SparseMatrixXcd

    -- * Matrix internal data
  , elems
  , values
  , innerIndices
  , outerStarts
  , innerNNZs

    -- * Accessors
  , cols
  , rows
  , coeff
  , (!)
  , getRow
  , getCol
  , getRows
  , getCols
  , unsafeGetCol
  , unsafeGetRow
  , unsafeFromRows
  , unsafeFromCols

    -- * Matrix conversions
  , fromList
  , toList
  , fromVector
  , toVector
  , fromVectorSized
  , fromListSized
  , fromDenseList
  , toDenseList
  , fromMatrix
  , toMatrix

    -- * Matrix properties
  , norm
  , squaredNorm
  , blueNorm
  , block
  , unsafeBlock
  , nonZeros
  , innerSize
  , outerSize

    -- * Basic matrix algebra
  , add
  , sub
  , mul

    -- * Matrix tranformations
  , pruned
  , scale
  , transpose
  , adjoint
  , map
  , imap

    -- * Matrix representation
  , compress
  , uncompress
  , compressed

    -- * Matrix serialisation
  , encode
  , decode

    -- * Mutable matrices
  , thaw
  , freeze
  , unsafeThaw
  , unsafeFreeze
  ) where

import Control.Monad (when, guard)
import Control.Monad.Primitive (PrimMonad(..), unsafePrimToPrim)
import qualified Prelude
import Prelude hiding (read, map)
import Data.Binary (Binary(..))
import qualified Data.Binary as Binary
import qualified Data.ByteString.Lazy as BSL
import Data.Complex (Complex)
import Data.Kind (Type)
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import Foreign.C.String (CString)
import Foreign.C.Types (CInt)
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, touchForeignPtr)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Ptr (Ptr)
import Foreign.Storable (Storable(..))
import qualified Foreign.Concurrent as FC
import GHC.TypeLits (Nat, KnownNat, type (<=))
import Eigen.Internal
  ( Elem
  , Cast(..)
  , CSparseMatrix
  , CSparseMatrixPtr
  , natToInt
  , Row(..)
  , Col(..)
  , CTriplet(..)
  )
import qualified Data.List as List
import qualified Eigen.Internal as Internal
import qualified Eigen.Matrix as M
import qualified Eigen.Matrix.Mutable as MM
import qualified Eigen.SparseMatrix.Mutable as SMM

{-| A versatible sparse matrix representation.
SparseMatrix is the main sparse matrix representation of Eigen's sparse module.
It offers high performance and low memory usage.
It implements a more versatile variant of the widely-used Compressed Column (or Row) Storage scheme.
It consists of four compact arrays:
* `values`: stores the coefficient values of the non-zeros.
* `innerIndices`: stores the row (resp. column) indices of the non-zeros.
* `outerStarts`: stores for each column (resp. row) the index of the first non-zero in the previous two arrays.
* `innerNNZs`: stores the number of non-zeros of each column (resp. row). The word inner refers to an inner vector that is a column for a column-major matrix, or a row for a row-major matrix. The word outer refers to the other direction.
This storage scheme is better explained on an example. The following matrix
@
0   3   0   0   0
22  0   0   0   17
7   5   0   1   0
0   0   0   0   0
0   0   14  0   8
@
and one of its possible sparse, __column major__ representation:
@
values:         22  7   _   3   5   14  _   _   1   _   17  8
innerIndices:   1   2   _   0   2   4   _   _   2   _   1   4
outerStarts:    0   3   5   8   10  12
innerNNZs:      2   2   1   1   2
@
Currently the elements of a given inner vector are guaranteed to be always sorted by increasing inner indices.
The "\_" indicates available free space to quickly insert new elements. Assuming no reallocation is needed,
the insertion of a random element is therefore in @O(nnz_j)@ where @nnz_j@ is the number of nonzeros of the
respective inner vector. On the other hand, inserting elements with increasing inner indices in a given inner
vector is much more efficient since this only requires to increase the respective `innerNNZs` entry that is a @O(1)@ operation.
The case where no empty space is available is a special case, and is refered as the compressed mode.
It corresponds to the widely used Compressed Column (or Row) Storage schemes (CCS or CRS).
Any `SparseMatrix` can be turned to this form by calling the `compress` function.
In this case, one can remark that the `innerNNZs` array is redundant with `outerStarts` because we the equality:
@InnerNNZs[j] = OuterStarts[j+1]-OuterStarts[j]@. Therefore, in practice a call to `compress` frees this buffer.
The results of Eigen's operations always produces compressed sparse matrices.
On the other hand, the insertion of a new element into a `SparseMatrix` converts this later to the uncompressed mode.
For more infomration please see Eigen <http://eigen.tuxfamily.org/dox/classEigen_1_1SparseMatrix.html documentation page>.
-}
newtype SparseMatrix :: Nat -> Nat -> Type -> Type where
  SparseMatrix :: ForeignPtr (CSparseMatrix a) -> SparseMatrix n m a

instance forall n m a. (Elem a, Show a, KnownNat n, KnownNat m) => Show (SparseMatrix n m a) where
  show m = concat
    [ "SparseMatrix"
    , show (rows_ m)
    , "x"
    , show (cols_ m)
    , "\n"
    , List.intercalate "\n" $ Prelude.map (List.intercalate "\t" . Prelude.map show) $ toDenseList m
    , "\n"
    ]

instance forall n m a. (Elem a, KnownNat n, KnownNat m) => Binary (SparseMatrix n m a) where
  put mat = do
    put $ Internal.magicCode (undefined :: C a)
    put $ natToInt @n
    put $ natToInt @m
    put $ toVector mat

  get = do
    get >>= (`when` fail "wrong matrix type") . (/= Internal.magicCode (undefined :: C a))
    fromVector <$> get

-- | Encode the sparse matrix as a lazy bytestring
encode :: (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> BSL.ByteString
encode = Binary.encode

-- | Decode the sparse matrix from a lazy bytestring
decode :: (Elem a, KnownNat n, KnownNat m) => BSL.ByteString -> SparseMatrix n m a
decode = Binary.decode

-- | Alias for single precision sparse matrix
type SparseMatrixXf  n m = SparseMatrix n m Float
-- | Alias for double precision sparse matrix
type SparseMatrixXd  n m = SparseMatrix n m Double
-- | Alias for single previsiom sparse matrix of complex numbers
type SparseMatrixXcf n m = SparseMatrix n m (Complex Float)
-- | Alias for double prevision sparse matrix of complex numbers
type SparseMatrixXcd n m = SparseMatrix n m (Complex Double)

-- | Get the coefficient values of the non-zeros.
values :: Elem a => SparseMatrix n m a -> VS.Vector a
values = VS.map fromC . _getvec Internal.sparse_values

-- | Get the row indices of the non-zeros.
innerIndices :: Elem a => SparseMatrix n m a -> VS.Vector Int
innerIndices = VS.map fromC . _getvec Internal.sparse_innerIndices

-- | Gets for each column the index of the first non-zero in the previous two arrays.
outerStarts :: Elem a => SparseMatrix n m a -> VS.Vector Int
outerStarts = VS.map fromC . _getvec Internal.sparse_outerStarts

-- | Gets the number of non-zeros of each column.
-- The word inner refers to an inner vector that is a column for a column-major matrix, or a row for a row-major matrix.
-- The word outer refers to the other direction
innerNNZs :: Elem a => SparseMatrix n m a -> Maybe (VS.Vector Int)
innerNNZs m
  | compressed m = Nothing
  | otherwise    = Just $ VS.map fromC $ _getvec Internal.sparse_innerNNZs m

-- | Number of rows in the sparse matrix
rows :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Row n
rows _ = Row @n

-- | Number of colums in the sparse matrix
cols :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Col m
cols _ = Col @m

-- | Number of rows in the sparse matrix
rows_ :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Int
rows_ _ = natToInt @n

-- | Number of colums in the sparse matrix
cols_ :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Int
cols_ _ = natToInt @m

-- | Sparse matrix coefficient at the given row and column
coeff :: forall n m r c a. (Elem a, KnownNat n, KnownNat m, KnownNat r, KnownNat c, r <= n, c <= m)
  => Row r -> Col c -> SparseMatrix n m a -> a
coeff _ _ (SparseMatrix fp) =
  let !c_row = toC $! natToInt @r
      !c_col = toC $! natToInt @c
  in Internal.performIO $ withForeignPtr fp $ \p -> alloca $ \pq -> do
       Internal.call $ Internal.sparse_coeff p c_row c_col pq
       fromC <$> peek pq

-- | Matrix coefficient at the given row and column
(!) :: forall n m r c a. (Elem a, KnownNat n, KnownNat m, KnownNat r, KnownNat c, r <= n, c <= m)
  => SparseMatrix n m a -> (Row r, Col c) -> a
(!) m (row,col) = coeff row col m

{-| 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 :: Elem a => SparseMatrix n m a -> a
norm = _unop Internal.sparse_norm (pure . fromC)

-- | 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 :: Elem a => SparseMatrix n m a -> a
squaredNorm = _unop Internal.sparse_squaredNorm (pure . fromC)

-- | 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 :: Elem a => SparseMatrix n m a -> a
blueNorm = _unop Internal.sparse_blueNorm (pure . fromC)

unsafeBlock ::
     (Elem a)
  => Int
  -> Int
  -> Int
  -> Int
  -> SparseMatrix n m a
  -> SparseMatrix n1 m1 a
unsafeBlock !sr !sc !br !bc =
  let !c_startRow = toC $! sr
      !c_startCol = toC $! sc
      !c_rows     = toC $! br
      !c_cols     = toC $! bc
  in _unop (\p pq -> Internal.sparse_block p c_startRow c_startCol c_rows c_cols pq) _mk

-- | Extract a rectangular block from the sparse matrix, given a startRow, startCol, blockRows, blockCols
block :: forall sr sc br bc n m a.
     (Elem a, KnownNat sr, KnownNat sc, KnownNat br, KnownNat bc, KnownNat n, KnownNat m)
  => (sr <= n, sc <= m, br <= n, bc <= m)
  => Row sr
  -> Col sc
  -> Row br
  -> Col bc
  -> SparseMatrix n m a
  -> SparseMatrix br bc a
block _ _ _ _ =
  let !c_startRow = toC $! natToInt @sr
      !c_startCol = toC $! natToInt @sc
      !c_rows     = toC $! natToInt @br
      !c_cols     = toC $! natToInt @bc
  in _unop (\p pq -> Internal.sparse_block p c_startRow c_startCol c_rows c_cols pq) _mk

-- | Number of non-zeros elements in the sparse matrix
nonZeros :: Elem a => SparseMatrix n m a -> Int
nonZeros = _unop Internal.sparse_nonZeros (pure . fromC)

-- | Number of elements in the sparse matrix, including zeros
elems :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Int
elems _ = (natToInt @n) * (natToInt @m)

-- | The matrix in the compressed format
compress :: Elem a => SparseMatrix n m a -> SparseMatrix n m a
compress = _unop Internal.sparse_makeCompressed _mk

-- | The matrix in the uncompressed format
uncompress :: Elem a => SparseMatrix n m a -> SparseMatrix n m a
uncompress = _unop Internal.sparse_uncompress _mk

-- | Is the matrix compressed?
compressed :: Elem a => SparseMatrix n m a -> Bool
compressed = _unop Internal.sparse_isCompressed (pure . (/=0))

-- | Minor dimension with respect to the storage order
innerSize :: Elem a => SparseMatrix n m a -> Int
innerSize = _unop Internal.sparse_innerSize (pure . fromC)

-- | Major dimension with respect to the storage order
outerSize :: Elem a => SparseMatrix n m a -> Int
outerSize = _unop Internal.sparse_outerSize (pure . fromC)

-- | Suppresses all nonzeros which are much smaller than the reference under the tolerance @epsilon@
pruned :: Elem a => a -> SparseMatrix n m a -> SparseMatrix n m a
pruned r = _unop (\p pq -> alloca $ \pr -> poke pr (toC r) >> Internal.sparse_prunedRef p pr pq) _mk

-- | Multiply matrix on a given scalar
scale :: Elem a => a -> SparseMatrix n m a -> SparseMatrix n m a
scale x = _unop (\p pq -> alloca $ \px -> poke px (toC x) >> Internal.sparse_scale p px pq) _mk

-- | Transpose of the sparse matrix
transpose :: Elem a => SparseMatrix n m a -> SparseMatrix m n a
transpose = _unop Internal.sparse_transpose _mk

-- | Adjoint of the sparse matrix
adjoint :: Elem a => SparseMatrix n m a -> SparseMatrix m n a
adjoint = _unop Internal.sparse_adjoint _mk

-- | Add two sparse matrices by adding the corresponding entries together.
add :: Elem a => SparseMatrix n m a -> SparseMatrix n m a -> SparseMatrix n m a
add = _binop Internal.sparse_add _mk

-- | Subtract two sparse matrices by subtracting the corresponding entries together.
sub :: Elem a => SparseMatrix n m a -> SparseMatrix n m a -> SparseMatrix n m a
sub = _binop Internal.sparse_sub _mk

-- | Matrix multiplication.
mul :: Elem a => SparseMatrix p q a -> SparseMatrix q r a -> SparseMatrix p r a
mul = _binop Internal.sparse_mul _mk

-- | Map a function over the 'SparseMatrix'.
map :: (Elem a, Elem b, KnownNat n, KnownNat m) => (a -> b) -> SparseMatrix n m a -> SparseMatrix n m b
map f m = fromVector . VS.map g . toVector $ m where
  g (CTriplet r c v) = CTriplet r c $ (toC . f . fromC) v

-- | Map an indexed function over the 'SparseMatrix'.
imap :: (Elem a, Elem b, KnownNat n, KnownNat m) => (Int -> Int -> a -> b) -> SparseMatrix n m a -> SparseMatrix n m b
imap f m = fromVector . VS.map g . toVector $ m where
  g (CTriplet r c v) =
    let !_r = fromC r
        !_c = fromC c
        !_v = fromC v
    in CTriplet r c $ toC $ f _r _c _v

-- | Construct asparse matrix of the given size from the storable vector of triplets (row, col, val)
fromVector :: forall n m a. (Elem a, KnownNat n, KnownNat m)
  => VS.Vector (CTriplet a)
  -> SparseMatrix n m a
fromVector = fromVectorSized (natToInt @n) (natToInt @m)

-- | Convert a sparse matrix to the list of triplets (row, col, val). Compressed elements will not be included.
toVector :: Elem a => SparseMatrix n m a -> VS.Vector (CTriplet a)
toVector m@(SparseMatrix fp) = Internal.performIO $ do
  let !size = nonZeros m
  tris <- VSM.new size
  withForeignPtr fp $ \p ->
    VSM.unsafeWith tris $ \q ->
      Internal.call $ Internal.sparse_toList p q (toC size)
  VS.unsafeFreeze tris

-- | Convert a sparse matrix to the list of triplets (row, col, val). Compressed elements will not be included.
toList :: Elem a => SparseMatrix n m a -> [(Int, Int, a)]
toList = Prelude.map fromC . VS.toList . toVector

-- | Construct a sparse matrix from a list of triples (row, val, col)
--
fromList :: (Elem a, KnownNat n, KnownNat m) => [(Int, Int, a)] -> SparseMatrix n m a
fromList = fromVector . VS.fromList . fmap toC

fromListSized :: (Elem a, KnownNat n, KnownNat m) => Int -> Int -> [(Int, Int, a)] -> SparseMatrix n m a
fromListSized _rows _cols = fromVectorSized _rows _cols . VS.fromList . Prelude.map Internal.toC

fromVectorSized :: forall n m a. (Elem a, KnownNat n, KnownNat m)
  => Int
  -> Int
  -> VS.Vector (CTriplet a)
  -> SparseMatrix n m a
fromVectorSized _rows _cols tris =
  let !c_rs = toC $! _rows
      !c_cs = toC $! _cols
      !len  = toC $! VS.length tris
  in Internal.performIO $ VS.unsafeWith tris $ \p ->
       alloca $ \pq -> do
         Internal.call $ Internal.sparse_fromList c_rs c_cs p len pq
         peek pq >>= _mk

-- | Construct asparse matrix of the given size from the storable vector of triplets (row, col, val)
--fromVector :: forall n m a. (Elem a, KnownNat n, KnownNat m)
--  => VS.Vector (CTriplet a)
--  -> SparseMatrix n m a
--fromVector tris =
--  let !c_rs = toC $! natToInt @n
--      !c_cs = toC $! natToInt @m
--      !len  = toC $! VS.length tris
--  in Internal.performIO $ VS.unsafeWith tris $ \p ->
--       alloca $ \pq -> do
--         Internal.call $ Internal.sparse_fromList c_rs c_cs p len pq
--         peek pq >>= _mk

-- | Convert a sparse matrix to a (n X m) dense list of values.
toDenseList :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> [[a]]
toDenseList mat = [[_unsafeCoeff row col mat | col <- [0 .. _unsafeCols mat - 1]] | row <- [0 .. _unsafeRows mat - 1]]

-- | Construct a sparsematrix from a two-dimensional list of values.
--   If the dimensions of the list do not match that of the list, 'Nothing' is returned.
--   Zero values will be compressed.
fromDenseList :: forall n m a. (Elem a, Eq a, KnownNat n, KnownNat m) => [[a]] -> Maybe (SparseMatrix n m a)
fromDenseList list =
  let _rows = List.length list
      _cols = List.foldl' max 0 $ List.map length list
  in if ((_rows /= (natToInt @n)) || (_cols /= (natToInt @m)))
    then Nothing
    else Just $ fromList $ do
      (row, vals) <- zip [0..] list
      (col, val) <- zip [0..] vals
      guard $ val /= 0
      return (row, col, val)

--  fromList $ do
-- (row, vals) <- zip [0..] list
-- (col, val) <- zip [0..] vals
-- guard $ val /= 0
-- return (row, col, val)
-- where
--   rows = List.length list
--   cols = List.foldl' max 0 $ List.map length list

-- | Construct a dense matrix from a sparse matrix
toMatrix :: (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> M.Matrix n m a
toMatrix (SparseMatrix fp) = Internal.performIO $ do
  m0 :: MM.IOMatrix n m a <- MM.new
  MM.unsafeWith m0 $ \_vals _rows _cols ->
    withForeignPtr fp $ \pm1 ->
      Internal.call $ Internal.sparse_toMatrix pm1 _vals _rows _cols
  M.unsafeFreeze m0

-- | Construct a sparse matrix from a dense matrix. zero-elements will be compressed.
fromMatrix :: (Elem a, KnownNat n, KnownNat m) => M.Matrix n m a -> SparseMatrix n m a
fromMatrix m1 = Internal.performIO $ alloca $ \pm0 ->
  M.unsafeWith m1 $ \_vals _rows _cols -> do
    Internal.call $ Internal.sparse_fromMatrix _vals _rows _cols pm0
    peek pm0 >>= _mk

-- | Yield an immutable copy of the mutable matrix
freeze :: (Elem a, PrimMonad p) => SMM.MSparseMatrix n m (PrimState p) a -> p (SparseMatrix n m a)
freeze (SMM.MSparseMatrix fp) = SparseMatrix <$> _clone fp

-- | Yield a mutable copy of the immutable matrix.
thaw :: (Elem a, PrimMonad p) => SparseMatrix n m a -> p (SMM.MSparseMatrix n m (PrimState p) a)
thaw (SparseMatrix fp) = SMM.MSparseMatrix <$> _clone fp

-- | Unsafely convert a mutable matrix to an immutable one without copying. The mutable matrix may not be used after this operation.
unsafeFreeze :: (Elem a, PrimMonad p) => SMM.MSparseMatrix n m (PrimState p) a -> p (SparseMatrix n m a)
unsafeFreeze (SMM.MSparseMatrix fp) = return $! SparseMatrix fp

-- | Unsafely convert an immutable matrix to a mutable one without copying. The immutable matrix may not be used after this operation.
unsafeThaw :: (Elem a, PrimMonad p) => SparseMatrix n m a -> p (SMM.MSparseMatrix n m (PrimState p) a)
unsafeThaw (SparseMatrix fp) = return $! SMM.MSparseMatrix fp

-- | Return a single row of the sparse matrix.
getRow :: forall n m r a. (Elem a, KnownNat n, KnownNat m, KnownNat r, r <= n, 1 <= n) => Row r -> SparseMatrix n m a -> SparseMatrix 1 m a
getRow row mat = block row (Col @0) (Row @1) (Col @m) mat

-- | Return a single column of the sparse matrix.
getCol :: forall n m c a. (Elem a, KnownNat n, KnownNat m, KnownNat c, c <= m, 1 <= m) => Col c -> SparseMatrix n m a -> SparseMatrix n 1 a
getCol col mat = block (Row @0) col (Row @n) (Col @1) mat

-- | Return all rows of the sparse matrix.
getRows :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> [SparseMatrix 1 m a]
getRows mat = fmap (flip unsafeGetRow mat) [0 .. (natToInt @n) - 1]

-- | Return all columns of a sparse matrix.
getCols :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> [SparseMatrix n 1 a]
getCols mat = fmap (flip unsafeGetCol mat) [0 .. (natToInt @m) - 1]

-- | Returns a single row of the sparse matrix. This is unsafe because it allows out-of-bounds access.
unsafeGetRow :: forall n m a. (Elem a, KnownNat n, KnownNat m) => Int -> SparseMatrix n m a -> SparseMatrix 1 m a
unsafeGetRow row mat = unsafeBlock row 0 1 (natToInt @m) mat

-- | Returns a single column of the sparse matrix. This is unsafe because it allows out-of-bounds access.
unsafeGetCol :: forall n m a. (Elem a, KnownNat n, KnownNat m) => Int -> SparseMatrix n m a -> SparseMatrix n 1 a
unsafeGetCol col mat = unsafeBlock 0 col (natToInt @n) 1 mat

-- | Get a sparse matrix from a list of columns.
--   This operation is unsafe because we cannot unify the length of the list with the type variable 'm'.
unsafeFromCols :: forall n m a. (Elem a, KnownNat n, KnownNat m) => [SparseMatrix n 1 a] -> SparseMatrix n m a
unsafeFromCols [] = fromListSized 0 0 []
unsafeFromCols _cols = fromListSized n m . concatMap toList . zipWith updateColIdx [0..] $ _cols
  where
    n = maximum . fmap _unsafeRows $ _cols
    m = length _cols

-- | Get a sparse matrix from a list of rows.
--   This operation is unsafe because we cannot unify the length of the list with the type variable 'n'.
unsafeFromRows :: forall n m a. (Elem a, KnownNat n, KnownNat m) => [SparseMatrix 1 m a] -> SparseMatrix n m a
unsafeFromRows [] = fromListSized 0 0 []
unsafeFromRows _rows = fromListSized n m . concatMap toList . zipWith updateRowIdx [0..] $ _rows
  where
    n = length _rows
    m = maximum . fmap _unsafeCols $ _rows

-- | Update the row indices of a sparse matrix.
updateRowIdx :: forall n m a. (KnownNat n, KnownNat m, Elem a) => Int -> SparseMatrix n m a -> SparseMatrix n m a
updateRowIdx row mat = fromListSized n m . fmap (\(_, j, v) -> (row, j, v)) . toList $ mat
  where
    n = row + 1
    m = _unsafeCols mat

-- | Update the column indices of a sparse matrix.
updateColIdx :: forall n m a. (KnownNat n, KnownNat m, Elem a) => Int -> SparseMatrix n m a -> SparseMatrix n m a
updateColIdx col mat = fromListSized n m . fmap (\(i, _, v) -> (i, col, v)) . toList $ mat
  where
    n = _unsafeRows mat
    m = col + 1

--fromRows :: forall n m a. (Elem a, KnownNat n, KnownNat m)
--  => [SparseMatrix 1 m a]
--  -> SparseMatrix n m a

_unop :: Storable b => (CSparseMatrixPtr a -> Ptr b -> IO CString) -> (b -> IO c) -> SparseMatrix n m a -> c
_unop f g (SparseMatrix fp) = Internal.performIO $
  withForeignPtr fp $ \p ->
    alloca $ \pq -> do
      Internal.call (f p pq)
      peek pq >>= g

_binop :: Storable b => (CSparseMatrixPtr a -> CSparseMatrixPtr a -> Ptr b -> IO CString) -> (b -> IO c) -> SparseMatrix n m a -> SparseMatrix n1 m1 a -> c
_binop f g (SparseMatrix fp1) (SparseMatrix fp2) = Internal.performIO $
  withForeignPtr fp1 $ \p1 ->
    withForeignPtr fp2 $ \p2 ->
      alloca $ \pq -> do
        Internal.call (f p1 p2 pq)
        peek pq >>= g

_getvec :: (Elem a, Storable b) => (Ptr (CSparseMatrix a) -> Ptr CInt -> Ptr (Ptr b) -> IO CString) -> SparseMatrix n m a -> VS.Vector b
_getvec f (SparseMatrix fm) = Internal.performIO $
  withForeignPtr fm $ \m ->
    alloca $ \ps ->
      alloca $ \pq -> do
        Internal.call $ f m ps pq
        s <- fromIntegral <$> peek ps
        q <- peek pq
        fr <- FC.newForeignPtr q $ touchForeignPtr fm
        pure $! VS.unsafeFromForeignPtr0 fr s

_clone :: (PrimMonad p, Elem a) => ForeignPtr (CSparseMatrix a) -> p (ForeignPtr (CSparseMatrix a))
_clone fp = unsafePrimToPrim $ withForeignPtr fp $ \p -> alloca $ \pq -> do
  Internal.call $ Internal.sparse_clone p pq
  q <- peek pq
  FC.newForeignPtr q $ Internal.call $ Internal.sparse_free q

_mk :: Elem a => Ptr (CSparseMatrix a) -> IO (SparseMatrix n m a)
_mk p = SparseMatrix <$> FC.newForeignPtr p (Internal.call $ Internal.sparse_free p)

-- | Number of rows in the sparse matrix
_unsafeRows :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Int
{-# INLINE _unsafeRows #-}
_unsafeRows _ = natToInt @n

-- | Number of colums in the sparse matrix
_unsafeCols :: forall n m a. (Elem a, KnownNat n, KnownNat m) => SparseMatrix n m a -> Int
{-# INLINE _unsafeCols #-}
_unsafeCols _ = natToInt @m

_unsafeCoeff :: (Elem a, KnownNat n) => Int -> Int -> SparseMatrix n m a -> a
{-# INLINE _unsafeCoeff #-}
_unsafeCoeff !row !col (SparseMatrix fp) =
  let !c_row = toC row
      !c_col = toC col
  in Internal.performIO $ withForeignPtr fp $ \p -> alloca $ \pq -> do
       Internal.call $ Internal.sparse_coeff p c_row c_col pq
       fromC <$> peek pq


{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
module Taiji.Utils.DataFrame
    ( DataFrame(..)
    , DataFrameIndex(..)
    , mkDataFrame
    , fromMatrix
    , dim
    , isEmpty
    , Taiji.Utils.DataFrame.transpose
    , rbind
    , cbind
    , cjoin
    , rowNames
    , colNames
    , ReodrderFn
    , reorderColumns
    , reorderRows
    , Taiji.Utils.DataFrame.map
    , mapRows
    , imapRows
    , mapCols
    , imapCols
    , filterRows
    , filterCols
    , Taiji.Utils.DataFrame.zip
    , Taiji.Utils.DataFrame.unzip
    , readData
    , writeTable
    , readTable
    , readTableWith
    , orderByHClust
    , orderByKMeans
    , orderDataFrame

    , pearson
    , spearman
    ) where

import           Bio.Utils.Functions    (ihs')
import           Bio.Utils.Misc         (readDouble)
import Conduit
import Control.Monad (forM_)
import GHC.Generics (Generic)
import qualified Data.ByteString.Char8  as B
import qualified Data.Text as T
import qualified Data.CaseInsensitive   as CI
import           Data.Function          (on)
import           Data.List as L
import AI.Clustering.Hierarchical hiding (normalize)
import AI.Clustering.KMeans
import Data.Vector.Binary ()
import Data.Maybe
import qualified Data.Matrix            as M
import qualified Data.Matrix.Unboxed    as MU
import qualified Data.Vector            as V
import qualified Data.Map.Strict    as HM
import qualified Data.HashSet as HS
import Statistics.Correlation (pearsonMatByRow, spearmanMatByRow)
import Statistics.Matrix (toRowLists, fromRowLists)
import Data.Binary (Binary)
import Control.DeepSeq (NFData)

data DataFrame a = DataFrame
    { _dataframe_row_names :: V.Vector T.Text
    , _dataframe_row_names_idx :: HM.Map T.Text Int
    , _dataframe_col_names :: V.Vector T.Text
    , _dataframe_col_names_idx :: HM.Map T.Text Int
    , _dataframe_data :: M.Matrix a
    } deriving (Show, Generic)

instance Binary a => Binary (M.Matrix a)
instance Binary a => Binary (DataFrame a)
instance NFData a => NFData (DataFrame a)

mkDataFrame :: [T.Text]     -- ^ row names
            -> [T.Text]     -- ^ col names
            -> [[a]]        -- ^ data
            -> DataFrame a
mkDataFrame r c d = DataFrame
    { _dataframe_row_names = V.fromList r
    , _dataframe_row_names_idx = HM.fromList $ L.zip r [0..]
    , _dataframe_col_names = V.fromList c
    , _dataframe_col_names_idx = HM.fromList $ L.zip c [0..]
    , _dataframe_data = M.fromLists d }

fromMatrix :: [T.Text]     -- ^ row names
           -> [T.Text]     -- ^ col names
           -> M.Matrix a  -- ^ data
           -> DataFrame a
fromMatrix r c mat = DataFrame
    { _dataframe_row_names = V.fromList r
    , _dataframe_row_names_idx = HM.fromList $ L.zip r [0..]
    , _dataframe_col_names = V.fromList c
    , _dataframe_col_names_idx = HM.fromList $ L.zip c [0..]
    , _dataframe_data = mat }

class DataFrameIndex i where
    csub :: DataFrame a -> [i] -> DataFrame a
    cindex :: DataFrame a -> i -> V.Vector a
    rsub :: DataFrame a -> [i] -> DataFrame a
    rindex :: DataFrame a -> i -> V.Vector a
    (!)  :: DataFrame a -> (i,i) -> a
    indexMaybe  :: DataFrame a -> (i,i) -> Maybe a

instance DataFrameIndex Int where
    csub df idx = df
        { _dataframe_col_names = V.fromList col_names
        , _dataframe_col_names_idx = HM.fromList $ L.zip col_names [0..]
        , _dataframe_data = M.fromColumns $ L.map (_dataframe_data df `M.takeColumn`) idx }
      where
        col_names = L.map (_dataframe_col_names df V.!) idx

    cindex df idx = _dataframe_data df `M.takeColumn` idx

    rsub df idx = df
        { _dataframe_row_names = V.fromList row_names
        , _dataframe_row_names_idx = HM.fromList $ L.zip row_names [0..]
        , _dataframe_data = M.fromRows $ L.map (_dataframe_data df `M.takeRow`) idx }
      where
        row_names = L.map (_dataframe_row_names df V.!) idx

    rindex df idx = _dataframe_data df `M.takeRow` idx

    (!) df (i,j) = _dataframe_data df M.! (i,j)
    indexMaybe df (i,j)
        | i < r && j < c = Just $ _dataframe_data df M.! (i,j)
        | otherwise = Nothing
      where
        (r, c) = dim df

instance DataFrameIndex T.Text where
    csub df idx = csub df idx'
      where
        idx' = L.map (\i -> HM.findWithDefault
            (error $ "index doesn't exist: " ++ T.unpack i) i $
            _dataframe_col_names_idx df) idx

    cindex df i = cindex df i'
      where
        i' = HM.findWithDefault (error $ "index doesn't exist: " ++ T.unpack i) i $
            _dataframe_col_names_idx df

    rsub df idx = rsub df idx'
      where
        idx' = L.map (\i -> HM.findWithDefault
            (error $ "index doesn't exist: " ++ T.unpack i) i $
            _dataframe_row_names_idx df) idx

    rindex df i = rindex df i'
      where
        i' = HM.findWithDefault (error $ "index doesn't exist: " ++ T.unpack i) i $
            _dataframe_row_names_idx df

    (!) df (a,b) = (!) df (i,j)
      where
        i = HM.findWithDefault (error $ "index doesn't exist: " ++ T.unpack a) a $
            _dataframe_row_names_idx df
        j = HM.findWithDefault (error $ "index doesn't exist: " ++ T.unpack b) b $
            _dataframe_col_names_idx df

    indexMaybe df (a,b) = case HM.lookup a (_dataframe_row_names_idx df) of
        Nothing -> Nothing
        Just i -> case HM.lookup b (_dataframe_col_names_idx df) of
            Nothing -> Nothing
            Just j -> Just $ (!) df (i,j)

dim :: DataFrame a -> (Int, Int)
dim df = M.dim $ _dataframe_data df

isEmpty :: DataFrame a -> Bool
isEmpty df = r == 0 || c == 0
  where
    (r,c) = M.dim $ _dataframe_data df

rbind :: [DataFrame a] -> DataFrame a
rbind dfs | allTheSame (L.map _dataframe_col_names dfs) = (head dfs)
    { _dataframe_row_names = row_names
    , _dataframe_row_names_idx = HM.fromList $ L.zip (V.toList row_names) [0..]
    , _dataframe_data = M.fromBlocks undefined $ L.map (return . _dataframe_data) dfs }
          | otherwise = error "Columns names differ"
  where
    allTheSame xs = all (== head xs) (tail xs)
    row_names = V.concat $ L.map (_dataframe_row_names) dfs

cjoin :: DataFrame a -> DataFrame a -> DataFrame a
cjoin df1 df2 = mkDataFrame rownames (colNames df1 ++ colNames df2) $
    flip L.map rownames $ \nm ->
        V.toList (df1 `rindex` nm) ++ V.toList (df2 `rindex` nm)
  where
    rownames = HS.toList $ HS.fromList (rowNames df1) `HS.intersection`
        HS.fromList (rowNames df2)

cbind :: [DataFrame a] -> DataFrame a
cbind dfs | allTheSame (L.map (HS.fromList . rowNames) dfs) = (head dfs)
    { _dataframe_col_names = col_names
    , _dataframe_col_names_idx = HM.fromList $ L.zip (V.toList col_names) [0..]
    , _dataframe_data = M.fromBlocks undefined [L.map reorderRow dfs] }
          | otherwise = error "Row names differ"
  where
    reorderRow df | rowNames df == row_names = _dataframe_data df
                  | otherwise = M.fromRows $ L.map (df `rindex`) row_names
    row_names = rowNames $ head dfs
    col_names = V.concat $ L.map (_dataframe_col_names) dfs
    allTheSame xs = all (== head xs) (tail xs)

map :: (a -> b) -> DataFrame a -> DataFrame b
map f df = df{_dataframe_data = M.map f $ _dataframe_data df}

zip :: DataFrame a -> DataFrame b -> DataFrame (a,b)
zip df1 df2
    | _dataframe_col_names df1 == _dataframe_col_names df2 &&
      _dataframe_row_names df1 == _dataframe_row_names df2 = df1
        {_dataframe_data = M.zip (_dataframe_data df1) $ _dataframe_data df2}
    | otherwise = error "names mismatch"

unzip :: DataFrame (a,b) -> (DataFrame a, DataFrame b)
unzip df = (df{_dataframe_data = a}, df{_dataframe_data = b})
  where
    (a,b) = M.unzip $ _dataframe_data df

transpose :: DataFrame a -> DataFrame a
transpose DataFrame{..} = DataFrame _dataframe_col_names _dataframe_col_names_idx
    _dataframe_row_names _dataframe_row_names_idx $ M.tr _dataframe_data

rowNames :: DataFrame a -> [T.Text]
rowNames DataFrame{..} = V.toList _dataframe_row_names

colNames :: DataFrame a -> [T.Text]
colNames DataFrame{..} = V.toList _dataframe_col_names

mapRows :: (V.Vector a -> V.Vector b) -> DataFrame a -> DataFrame b
mapRows fn df = df
    { _dataframe_data = M.fromRows $ L.map fn $ M.toRows $ _dataframe_data df }

imapRows :: (T.Text -> V.Vector a -> V.Vector b) -> DataFrame a -> DataFrame b
imapRows fn df = df{ _dataframe_data = M.fromRows $ L.zipWith fn (rowNames df) $ M.toRows $ _dataframe_data df}

mapCols :: (V.Vector a -> V.Vector b) -> DataFrame a -> DataFrame b
mapCols fn df = df
    { _dataframe_data = M.fromColumns $ L.map fn $ M.toColumns $ _dataframe_data df }

imapCols :: (T.Text -> V.Vector a -> V.Vector b) -> DataFrame a -> DataFrame b
imapCols fn df = df
    { _dataframe_data = M.fromColumns $ L.zipWith fn (colNames df) $ M.toColumns $ _dataframe_data df }

filterRows :: (T.Text -> V.Vector a -> Bool) -> DataFrame a -> DataFrame a
filterRows fn df = df
    { _dataframe_row_names = V.fromList names
    , _dataframe_row_names_idx = HM.fromList $ L.zip names [0..]
    , _dataframe_data = M.fromRows rows }
  where
    (names, rows) = L.unzip $ filter (uncurry fn) $
        L.zip (V.toList $ _dataframe_row_names df) $ M.toRows $ _dataframe_data df

filterCols :: (T.Text -> V.Vector a -> Bool) -> DataFrame a -> DataFrame a
filterCols fn df = df
    { _dataframe_col_names = V.fromList names
    , _dataframe_col_names_idx = HM.fromList $ L.zip names [0..]
    , _dataframe_data = M.fromColumns cols }
  where
    (names, cols) = L.unzip $ filter (uncurry fn) $
        L.zip (V.toList $ _dataframe_col_names df) $ M.toColumns $ _dataframe_data df

type ReodrderFn a = [(T.Text, V.Vector a)] -> [(T.Text, V.Vector a)]

reorderRows :: ReodrderFn a -> DataFrame a -> DataFrame a
reorderRows fn df
    | isEmpty df = df
    | otherwise = df
        { _dataframe_row_names = V.fromList names
        , _dataframe_row_names_idx = HM.fromList $ L.zip names [0..]
        , _dataframe_data = M.fromRows rows }
  where
    (names, rows) = L.unzip $ fn $ L.zip (V.toList $ _dataframe_row_names df) $
        M.toRows $ _dataframe_data df

reorderColumns :: ReodrderFn a -> DataFrame a -> DataFrame a
reorderColumns fn df
    | isEmpty df = df
    | otherwise = df
        { _dataframe_col_names = V.fromList names
        , _dataframe_col_names_idx = HM.fromList $ L.zip names [0..]
        , _dataframe_data = M.fromColumns cols}
  where
    (names, cols) = L.unzip $ fn $ L.zip (V.toList $ _dataframe_col_names df) $
        M.toColumns $ _dataframe_data df

writeTable :: FilePath -> (a -> T.Text) -> DataFrame a -> IO ()
writeTable output f DataFrame{..} = runResourceT $ runConduit $
    source .| unlinesAsciiC .| sinkFile output
  where
    source = do
        yield $ B.pack $ T.unpack $ T.intercalate "\t" $
            "" : V.toList _dataframe_col_names
        forM_ (L.zip (V.toList _dataframe_row_names) $ M.toRows _dataframe_data) $
            \(nm, xs) -> yield $ B.pack $ T.unpack $
                T.intercalate "\t" $ nm : L.map f (V.toList xs)
{-# NOINLINE writeTable #-}

-- | Read data, normalize and calculate p-values.
readData :: FilePath   -- ^ PageRank
         -> FilePath   -- ^ Gene expression
         -> IO (DataFrame (Double, Double))  -- ^ ranks, expression
readData input1 input2 = do
    rank <- readTSV <$> B.readFile input1

    -- Read expression profile and apply "ihs" transformation
    expr <- (fmap ihs' . readTSV) <$> B.readFile input2

    let (labels, xs) = L.unzip $ L.map L.unzip $ groupBy ((==) `on` (fst.fst)) $ sort $
            HM.toList $ HM.intersectionWith (,) rank expr
        rowlab = L.map (T.pack . B.unpack . CI.original) $ fst $ L.unzip $ L.map head labels
        collab = L.map (T.pack . B.unpack . CI.original) $ snd $ L.unzip $ head labels
    return $ mkDataFrame rowlab collab xs

readTSV :: B.ByteString -> HM.Map (CI.CI B.ByteString, CI.CI B.ByteString) Double
readTSV input = HM.fromList $ concatMap (f . B.split '\t') content
  where
    f (x:xs) = L.zipWith (\s v -> ((CI.mk x, CI.mk s), readDouble v)) samples xs
    (header:content) = B.lines input
    samples = tail $ B.split '\t' header

readTable :: FilePath -> IO (DataFrame Double)
readTable = readTableWith readDouble

readTableWith :: (B.ByteString -> a) -> FilePath -> IO (DataFrame a)
readTableWith f input = do
    (header:content) <- B.lines <$> B.readFile input
    let samples = tail $ B.split '\t' header
        (rows, dat) = L.unzip $ L.map ((\(x:xs) ->
            (T.pack $ B.unpack x, L.map f xs)) . B.split '\t') content
    return $ mkDataFrame rows (L.map (T.pack . B.unpack) samples) dat

    {-
orderByName :: [T.Text] -> ReodrderFn a
orderByName prefix = sortBy $ \(a,_) (b,_) ->
    let idx1 = findIdx a
        idx2 = findIdx b
    in case () of
        _ | isJust idx1 && isJust idx2 -> case compare idx1 idx2 of
                LT -> LT
                GT -> GT
                EQ -> compare a b
          | otherwise -> compare a b
  where
    findIdx x = go prefix 0
      where
        go (y:ys) !i | isPrefixOf y x = Just i
                     | otherwise = go ys (i+1)
        go _ _ = Nothing
        -}

orderByHClust :: (a -> Double) -> ReodrderFn a
orderByHClust f xs = flatten $ hclust Ward (V.fromList xs) dist
  where
    dist = euclidean `on` V.map f . snd

orderByKMeans :: Int -> (a -> Double) -> ReodrderFn a
orderByKMeans k f xs = concatMap fst $ flatten $ hclust Ward r dist
  where
    r = V.fromList $ L.zip (fromJust $ clusters kmRes) (MU.toRows $ centers kmRes)
    kmRes = kmeansBy k (V.fromList xs) (V.convert . V.map f . snd) defaultKMeansOpts
    dist = euclidean `on` snd

pearson :: DataFrame Double -> DataFrame Double
pearson DataFrame{..} = DataFrame _dataframe_row_names _dataframe_row_names_idx
    _dataframe_row_names _dataframe_row_names_idx $ M.fromLists $
    toRowLists $ pearsonMatByRow $ fromRowLists $ M.toLists _dataframe_data

spearman :: DataFrame Double -> DataFrame Double
spearman DataFrame{..} = DataFrame _dataframe_row_names _dataframe_row_names_idx
    _dataframe_row_names _dataframe_row_names_idx $ M.fromLists $
    toRowLists $ spearmanMatByRow $ fromRowLists $ M.toLists _dataframe_data

orderDataFrame :: (a -> Double) -> DataFrame a -> DataFrame a
orderDataFrame f = reorderColumns (orderByHClust f) . reorderRows (orderByHClust f)


{-# LANGUAGE TypeFamilies, GeneralizedNewtypeDeriving, DeriveGeneric, TupleSections, RecordWildCards #-}

module BayesStack.Models.Topic.CitationInfluence
  ( -- * Primitives
    NetData(..)
  , netData
  , HyperParams(..)
  , MState(..)
  , CitedUpdateUnit
  , CitingUpdateUnit
  , ItemSource(..)
  , CitedNode(..), CitedNodeItem(..)
  , CitingNode(..), CitingNodeItem(..)
  , Citing(..), Cited(..)
  , Item(..), Topic(..), NodeItem(..), Node(..), Arc(..)
  , setupNodeItems
    -- * Initialization
  , verifyNetData, cleanNetData
  , ModelInit
  , randomInitialize
  , model
  , updateUnits
    -- * Diagnostics
  , modelLikelihood
  , influence
  , arcTopicMixture
  ) where

import qualified Data.Vector as V
import Statistics.Sample (mean)

import           Prelude hiding (mapM_, sum)

import           Data.Set (Set)
import qualified Data.Set as S

import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as M

import           Data.Foldable hiding (product)
import           Control.Applicative ((<$>), (<*>))
import           Control.Monad (when)
import           Control.Monad.Trans.State.Strict
import           Control.Monad.Trans.Writer.Strict

import           Data.Random
import           Data.Random.Lift (lift)
import           Data.Random.Distribution.Categorical (categorical)
import           Numeric.Log hiding (sum)

import           BayesStack.Types
import           BayesStack.Gibbs
import           BayesStack.DirMulti
import           BayesStack.TupleEnum ()
import           BayesStack.Models.Topic.Types

import           GHC.Generics
import           Data.Binary (Binary)
import           Control.DeepSeq

data ItemSource = Shared | Own deriving (Show, Eq, Enum, Ord, Generic)
instance Binary ItemSource
instance NFData ItemSource

newtype Citing a = Citing a deriving (Show, Eq, Enum, Ord, Generic, NFData)
newtype Cited a = Cited a deriving (Show, Eq, Enum, Ord, Generic, NFData)
instance Binary a => Binary (Citing a)
instance Binary a => Binary (Cited a)

type CitingNode = Citing Node
type CitedNode = Cited Node
type CitingNodeItem = Citing NodeItem
type CitedNodeItem = Cited NodeItem

-- ^ A directed edge
data Arc = Arc {citingNode :: !CitingNode, citedNode :: !CitedNode}
            deriving (Show, Eq, Ord, Generic)
instance Binary Arc

data HyperParams = HyperParams
                   { alphaPsi         :: Double
                   , alphaLambda      :: Double
                   , alphaPhi         :: Double
                   , alphaOmega       :: Double
                   , alphaGammaShared :: Double
                   , alphaGammaOwn    :: Double
                   }
                 deriving (Show, Eq, Generic)
instance Binary HyperParams

data NetData = NetData { dHypers             :: !(HyperParams)
                       , dArcs               :: !(Set Arc)
                       , dItems              :: !(Set Item)
                       , dTopics             :: !(Set Topic)
                       , dNodeItems          :: !(Map NodeItem (Node, Item))
                       , dCitingNodes        :: !(Map CitingNode (Set CitedNode))
                         -- ^ Maps each citing node to the set of nodes cited by it
                       , dCitedNodes         :: !(Map CitedNode (Set CitingNode))
                         -- ^ Maps each cited node to the set of nodes citing it
                       }
              deriving (Show, Eq, Generic)
instance Binary NetData

netData :: HyperParams -> Set Arc -> Map NodeItem (Node,Item) -> Set Topic -> NetData
netData hypers arcs nodeItems topics =
    NetData { dHypers       = hypers
            , dArcs         = arcs
            , dItems        = foldMap (S.singleton . snd) $ M.elems nodeItems
            , dTopics       = topics
            , dNodeItems    = nodeItems
            , dCitingNodes  = M.unionsWith S.union
                              $ map (\(Arc a b)->M.singleton a $ S.singleton b)
                              $ S.toList arcs
            , dCitedNodes   = M.unionsWith S.union
                              $ map (\(Arc a b)->M.singleton b $ S.singleton a)
                              $ S.toList arcs
            }

dCitingNodeItems :: NetData -> Map CitingNodeItem (CitingNode, Item)
dCitingNodeItems nd =
    M.mapKeys Citing
    $ M.map (\(n,i)->(Citing n, i))
    $ M.filter (\(n,i)->Citing n `M.member` dCitingNodes nd)
    $ dNodeItems nd

itemsOfCitingNode :: NetData -> CitingNode -> [Item]
itemsOfCitingNode d (Citing u) =
    map snd $ M.elems $ M.filter (\(n,_)->n==u) $ dNodeItems d

connectedNodes :: Set Arc -> Set Node
connectedNodes arcs =
    S.map ((\(Cited n)->n) . citedNode) arcs `S.union` S.map ((\(Citing n)->n) . citingNode) arcs

cleanNetData :: NetData -> NetData
cleanNetData d =
    let nodesWithItems = S.fromList $ map fst $ M.elems $ dNodeItems d
        nodesWithArcs = connectedNodes $ dArcs d
        keptNodes = nodesWithItems `S.intersection` nodesWithArcs
        keepArc (Arc (Citing citing) (Cited cited)) =
            citing `S.member` keptNodes && cited `S.member` keptNodes
    in d { dArcs = S.filter keepArc $ dArcs d
         , dNodeItems = M.filter (\(n,i)->n `S.member` keptNodes) $ dNodeItems d
         }

verifyNetData :: (Node -> String) -> NetData -> [String]
verifyNetData showNode d = execWriter $ do
    let nodesWithItems = S.fromList $ map fst $ M.elems $ dNodeItems d
    forM_ (dArcs d) $ \(Arc (Citing citing) (Cited cited))->do
        when (cited `S.notMember` nodesWithItems)
            $ tell [showNode cited++" has arc yet has no items"]
        when (citing `S.notMember` nodesWithItems)
            $ tell [showNode citing++" has arc yet has no items"]

type CitedModelInit = Map CitedNodeItem (Setting CitedUpdateUnit)
type CitingModelInit = Map CitingNodeItem (Setting CitingUpdateUnit)
data ModelInit = ModelInit CitedModelInit CitingModelInit
               deriving (Show)

randomInitializeCited :: NetData -> CitedModelInit -> RVar CitedModelInit
randomInitializeCited d init = execStateT doInit init
    where doInit = let unset = citedNodeItems `S.difference` M.keysSet init
                       citedNodeItems = S.map Cited (M.keysSet (dNodeItems d))
                   in mapM_ (randomInitCitedUU d) (S.toList unset)

modify' :: Monad m => (a -> a) -> StateT a m ()
modify' f = do x <- get
               put $! f x

randomInitCitedUU :: NetData -> CitedNodeItem -> StateT CitedModelInit RVar ()
randomInitCitedUU d ni = do
    t' <- lift $ randomElement $ toList $ dTopics d
    modify' $ M.insert ni t'

randomInitializeCiting :: NetData -> CitingModelInit -> RVar CitingModelInit
randomInitializeCiting d init = execStateT doInit init
    where doInit :: StateT CitingModelInit RVar ()
          doInit = let unset = M.keysSet (dCitingNodeItems d) `S.difference` M.keysSet init
                   in mapM_ (randomInitCitingUU d) (S.toList unset)

randomInitCitingUU :: NetData -> CitingNodeItem -> StateT CitingModelInit RVar ()
randomInitCitingUU d cni@(Citing ni) =
    let (n,_) = dNodeItems d M.! ni
    in case dCitingNodes d M.! Citing n of
           a | S.null a -> do
               t <- lift $ randomElement $ toList $ dTopics d
               modify' $ M.insert cni $ OwnSetting t

           citedNodes -> do
               s <- lift $ randomElement [Shared, Own]
               c <- lift $ randomElement $ toList citedNodes
               t <- lift $ randomElement $ toList $ dTopics d
               modify' $ M.insert cni $
                   case s of Shared -> SharedSetting t c
                             Own    -> OwnSetting t

randomInitialize :: NetData -> RVar ModelInit
randomInitialize d =
    ModelInit <$> randomInitializeCited d M.empty <*> randomInitializeCiting d M.empty

model :: NetData -> ModelInit -> MState
model d (ModelInit citedInit citingInit) =
    let citingNodes = M.keys $ dCitingNodes d
        HyperParams {..} = dHypers d
        s = MState { -- Citing model
                     stPsis = let dist n = case toList $ dCitingNodes d M.! n of
                                               []    -> M.empty
                                               nodes -> M.singleton n
                                                        $ symDirMulti alphaPsi nodes
                              in foldMap dist citingNodes
                   , stPhis = let dist = symDirMulti alphaPhi (toList $ dItems d)
                              in foldMap (\t->M.singleton t dist) $ dTopics d
                   , stGammas = let dist = multinom [ (Shared, alphaGammaShared)
                                                    , (Own, alphaGammaOwn) ]
                                in foldMap (\t->M.singleton t dist) citingNodes
                   , stOmegas = let dist = symDirMulti alphaOmega (toList $ dTopics d)
                                in foldMap (\t->M.singleton t dist) citingNodes
                   , stCiting = M.empty

                   -- Cited model
                   , stLambdas = let dist = symDirMulti alphaLambda (toList $ dTopics d)
                                 in foldMap (\t->M.singleton t dist) $ M.keys $ dCitedNodes d
                   , stT' = M.empty
                   }

        initCitingUU :: CitingUpdateUnit -> State MState ()
        initCitingUU uu = do
            let err = error $ "CitationInference: Initial value for "++show uu++" not given\n"
                s = maybe err id $ M.lookup (uuNI uu) citingInit
            modify' $ setCitingUU uu (Just s)

        initCitedUU :: CitedUpdateUnit -> State MState ()
        initCitedUU uu = do
            let err = error $ "CitationInference: Initial value for "++show uu++" not given"
                s = maybe err id $ M.lookup (uuNI' uu) citedInit
            modify' $ setCitedUU uu (Just s)

    in execState (do mapM_ initCitingUU $ citingUpdateUnits d
                     mapM_ initCitedUU $ citedUpdateUnits d
                 ) s

updateUnits :: NetData -> [WrappedUpdateUnit MState]
updateUnits d = map WrappedUU (citedUpdateUnits d)
             ++ map WrappedUU (citingUpdateUnits d)

data CitingSetting = OwnSetting !Topic
                   | SharedSetting !Topic !CitedNode
                   deriving (Show, Eq, Generic)
instance Binary CitingSetting
instance NFData CitingSetting where
    rnf (OwnSetting t)      = rnf t `seq` ()
    rnf (SharedSetting t c) = rnf t `seq` rnf c `seq` ()

data MState = MState { -- Citing model state
                       stGammas   :: !(Map CitingNode (Multinom Int ItemSource))
                     , stOmegas   :: !(Map CitingNode (Multinom Int Topic))
                     , stPsis     :: !(Map CitingNode (Multinom Int CitedNode))
                     , stPhis     :: !(Map Topic (Multinom Int Item))

                     , stCiting   :: !(Map CitingNodeItem CitingSetting)

                     -- Cited model state
                     , stLambdas  :: !(Map CitedNode (Multinom Int Topic))

                     , stT'       :: !(Map CitedNodeItem Topic)
                     }
            deriving (Show, Generic)
instance Binary MState

modelLikelihood :: MState -> Probability
modelLikelihood model =
    product $ map likelihood (M.elems $ stGammas model)
           ++ map likelihood (M.elems $ stPhis model)
           ++ map likelihood (M.elems $ stLambdas model)
           ++ map likelihood (M.elems $ stOmegas model)
           ++ map likelihood (M.elems $ stPsis model)

-- | Mixture of the topics of an edge
arcTopicMixture :: NetData -> MState -> Arc -> Topic -> Probability
arcTopicMixture nd m (Arc d c) t =
    let itemObs = zip (itemsOfCitingNode nd d) (repeat 1)
        phi = stPhis m M.! t
        lambda = stLambdas m M.! c
    in realToFrac (sampleProb lambda t) * obsProb phi itemObs

-- | Geometric mean
geomMean :: V.Vector (Log Double) -> Log Double
geomMean = Exp . mean . V.map ln

-- | The geometric mean of the probabilities of a collection of items under a
-- given topic mixture.
topicCompatibility :: MState -> [Item] -> Multinom Int Topic -> Probability
topicCompatibility m items lambda =
    geomMean $ V.fromList
            $ do t <- toList $ dmDomain lambda
                 let phi = stPhis m M.! t
                     itemObs = zip items (repeat 1)
                 return $ realToFrac (sampleProb lambda t) * obsProb phi itemObs

-- | Normalized compatibilities with a set of items
topicCompatibilities :: (Functor f, Foldable f)
                     => MState -> [Item] -> f (Multinom Int Topic) -> f Probability
topicCompatibilities m items topics =
    let scores = fmap (topicCompatibility m items) topics
    in fmap (/sum scores) scores

-- | The influence of cited nodes on a citing node.
influence :: NetData -> MState -> CitingNode -> Map CitedNode Probability
influence d m u =
    let lambdas = foldMap (\f->M.singleton f $ stLambdas m M.! f)
                 $ S.toList $ dCitingNodes d M.! u
    in topicCompatibilities m (itemsOfCitingNode d u) lambdas

-- Cited update unit (LDA-like)
data CitedUpdateUnit = CitedUpdateUnit { uuNI' :: !CitedNodeItem
                                       , uuN'  :: !CitedNode
                                       , uuX' :: !Item
                                       }
                     deriving (Show, Generic)
instance Binary CitedUpdateUnit

instance UpdateUnit CitedUpdateUnit where
    type ModelState CitedUpdateUnit = MState
    type Setting CitedUpdateUnit = Topic
    fetchSetting uu ms = maybe (error $ "Update unit "++show uu++" has no setting") id
                         $ M.lookup (uuNI' uu) (stT' ms)
    evolveSetting ms uu = categorical $ citedFullCond (setCitedUU uu Nothing ms) uu
    updateSetting uu _ s' = setCitedUU uu (Just s') . setCitedUU uu Nothing

citedProb :: MState -> CitedUpdateUnit -> Setting CitedUpdateUnit -> Double
citedProb st (CitedUpdateUnit {uuN'=n', uuX'=x'}) t =
    let lambda = stLambdas st M.! n'
        phi = stPhis st M.! t
    in realToFrac $ sampleProb lambda t * sampleProb phi x'

citedUpdateUnits :: NetData -> [CitedUpdateUnit]
citedUpdateUnits d =
    map (\(ni',(n',x'))->CitedUpdateUnit { uuNI'      = Cited ni'
                                         , uuN'       = Cited n'
                                         , uuX'       = x'
                                         }
        ) $ M.assocs
          $ M.filter (\(n,i)->Cited n `M.member` dCitedNodes d)
          $ dNodeItems d

setCitedUU :: CitedUpdateUnit -> Maybe Topic -> MState -> MState
setCitedUU uu@(CitedUpdateUnit {uuN'=n', uuNI'=ni', uuX'=x'}) setting ms =
    let t' = maybe (fetchSetting uu ms) id setting
        set = maybe Unset (const Set) setting
    in ms { stLambdas = M.adjust (setMultinom set t') n' (stLambdas ms)
          , stPhis = M.adjust (setMultinom set x') t' (stPhis ms)
          , stT' = case setting of Just _  -> M.insert ni' t' $ stT' ms
                                   Nothing -> stT' ms
          }

citedFullCond ::MState -> CitedUpdateUnit -> [(Double, Topic)]
citedFullCond ms uu = do
    t <- M.keys $ stPhis ms
    return (citedProb ms uu t, t)

-- Citing Update unit (Shared Taste-like)
data CitingUpdateUnit = CitingUpdateUnit { uuNI    :: !CitingNodeItem
                                         , uuN     :: !CitingNode
                                         , uuX     :: !Item
                                         , uuCites :: !(Set CitedNode)
                                         }
                      deriving (Show, Generic)
instance Binary CitingUpdateUnit

instance UpdateUnit CitingUpdateUnit where
    type ModelState CitingUpdateUnit = MState
    type Setting CitingUpdateUnit = CitingSetting
    fetchSetting uu ms = stCiting ms M.! uuNI uu
    evolveSetting ms uu = categorical $ citingFullCond (setCitingUU uu Nothing ms) uu
    updateSetting uu _ s' = setCitingUU uu (Just s') . setCitingUU uu Nothing

citingUpdateUnits :: NetData -> [CitingUpdateUnit]
citingUpdateUnits d =
    map (\(ni,(n,x))->CitingUpdateUnit { uuNI      = ni
                                       , uuN       = n
                                       , uuX       = x
                                       , uuCites   = dCitingNodes d M.! n
                                       }
        ) $ M.assocs $ dCitingNodeItems d

citingProb :: MState -> CitingUpdateUnit -> Setting CitingUpdateUnit -> Double
citingProb st (CitingUpdateUnit {uuN=n, uuX=x}) setting =
    let gamma = stGammas st M.! n
        omega = stOmegas st M.! n
        psi = stPsis st M.! n
    in case setting of
        SharedSetting t c -> let phi = stPhis st M.! t
                                 lambda = stLambdas st M.! c
                             in sampleProb gamma Shared
                              * sampleProb psi c
                              * sampleProb lambda t
                              * sampleProb phi x
        OwnSetting t      -> let phi = stPhis st M.! t
                             in sampleProb gamma Own
                              * sampleProb omega t
                              * sampleProb phi x

citingFullCond :: MState -> CitingUpdateUnit -> [(Double, Setting CitingUpdateUnit)]
citingFullCond ms uu = map (\s->(citingProb ms uu s, s)) $ citingDomain ms uu

citingDomain :: MState -> CitingUpdateUnit -> [Setting CitingUpdateUnit]
citingDomain ms uu = do
    s <- [Own, Shared]
    t <- M.keys $ stPhis ms
    case s of
        Shared -> do c <- S.toList $ uuCites uu
                     return $ SharedSetting t c
        Own    -> do return $ OwnSetting t

setCitingUU :: CitingUpdateUnit -> Maybe (Setting CitingUpdateUnit) -> MState -> MState
setCitingUU uu@(CitingUpdateUnit {uuNI=ni, uuN=n, uuX=x}) setting ms =
    let set = maybe Unset (const Set) setting
        ms' = case maybe (fetchSetting uu ms) id  setting of
            SharedSetting t c -> ms { stPsis = M.adjust (setMultinom set c) n $ stPsis ms
                                    , stLambdas = M.adjust (setMultinom set t) c $ stLambdas ms
                                    , stPhis = M.adjust (setMultinom set x) t $ stPhis ms
                                    , stGammas = M.adjust (setMultinom set Shared) n $ stGammas ms
                                    }
            OwnSetting t      -> ms { stOmegas = M.adjust (setMultinom set t) n $ stOmegas ms
                                    , stPhis = M.adjust (setMultinom set x) t $ stPhis ms
                                    , stGammas = M.adjust (setMultinom set Own) n $ stGammas ms
                                    }
    in ms' { stCiting = M.alter (const setting) ni $ stCiting ms' }


{-# LANGUAGE DefaultSignatures      #-}
{-# LANGUAGE DeriveFoldable         #-}
{-# LANGUAGE DeriveFunctor          #-}
{-# LANGUAGE DeriveTraversable      #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TemplateHaskell        #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeOperators          #-}
module Circuits.Analysis.DC where

import           Control.Lens
import           Control.Monad.Primitive
import           Control.Monad.ST
import           Data.Primitive.MutVar
import qualified Data.Vector                  as V
import qualified Data.Vector.Storable         as VS
import qualified Data.Vector.Storable.Mutable as VSM
import           Generics.Deriving.Lens
import           GHC.Generics
import qualified Numeric.LinearAlgebra        as HMatrix

import           Debug.Trace

import           Circuits.Circuit

-- | Describes the time scale of a DC simulation. We can either perform a DC steady-state or a transient analysis.
data TimeScale
  = SteadyState
  -- ^ solving for steady state
  | Step
    { _currentTime :: Double -- ^ the current simulation time in seconds since t0
    , _timeStep    :: Double -- ^ the size of current time step to be taken
    }
    -- ^ solving for next time step

makeLenses ''TimeScale

-- | Type class providing an interface for the component stamps used in DC analysis.
class SimulateDC c where
  -- | Beginning to solve the next time step.
  beginSolve :: Simulation m Double sim => c -> TimeScale -> sim -> m ()
  default beginSolve :: (Generic c, GSimulateDC (Rep c), Simulation m Double sim) => c -> TimeScale -> sim -> m ()
  beginSolve c t s = forOf_ generic c $ \gc -> gbeginSolve gc t s
  -- | Beginning the next iteration of solving the non-linear parts.
  -- This should set up the right hand side based on the current element state.
  beginIteration :: Simulation m Double sim => c -> TimeScale -> sim -> m ()
  default beginIteration :: (Generic c, GSimulateDC (Rep c), Simulation m Double sim) => c -> TimeScale -> sim -> m ()
  beginIteration c t s = forOf_ generic c $ \gc -> gbeginIteration gc t s
  -- | End of the current iteration, after the referenced variables have been updated.
  -- It should update the state of non-linear parts, but not the state of linear parts.
  endIteration :: c -> TimeScale -> c
  default endIteration :: (Generic c, GSimulateDC (Rep c)) => c -> TimeScale -> c
  endIteration c t = over generic (`gendIteration` t) c
  -- | Finished solving the current time step. Here, the state of linear parts should be updated.
  endSolve :: c -> TimeScale -> c
  default endSolve :: (Generic c, GSimulateDC (Rep c)) => c -> TimeScale -> c
  endSolve c t = over generic (`gendSolve` t) c

instance SimulateDC Int where
  beginSolve _ _ _ = pure ()
  beginIteration _ _ _ = pure ()
  endIteration x _ = x
  endSolve x _ = x

instance SimulateDC Double where
  beginSolve _ _ _ = pure ()
  beginIteration _ _ _ = pure ()
  endIteration x _ = x
  endSolve x _ = x

instance SimulateDC Variable
instance SimulateDC Node
instance SimulateDC Branch

instance SimulateDC c => SimulateDC (Circuit c)

-- * Generic deriving of SimulateDC instances

class GSimulateDC f where
  gbeginSolve :: Simulation m Double sim => f p -> TimeScale -> sim -> m ()
  gbeginIteration :: Simulation m Double sim => f p -> TimeScale -> sim -> m ()
  gendIteration :: f p -> TimeScale -> f p
  gendSolve :: f p -> TimeScale -> f p

instance GSimulateDC V1 where
  gbeginSolve _ _ _ = return ()
  gbeginIteration _ _ _ = return ()
  gendIteration x _ = x
  gendSolve x _ = x

instance GSimulateDC U1 where
  gbeginSolve _ _ _ = return ()
  gbeginIteration _ _ _ = return ()
  gendIteration x _ = x
  gendSolve x _ = x

instance SimulateDC a => GSimulateDC (K1 i a) where
  gbeginSolve (K1 x) = beginSolve x
  gbeginIteration (K1 x) = beginIteration x
  gendIteration (K1 x) = K1 . endIteration x
  gendSolve (K1 x) = K1 . endSolve x

instance GSimulateDC f => GSimulateDC (M1 i t f) where
  gbeginSolve (M1 x) = gbeginSolve x
  gbeginIteration (M1 x) = gbeginIteration x
  gendIteration (M1 x) = M1 . gendIteration x
  gendSolve (M1 x) = M1 . gendSolve x

instance (GSimulateDC f, GSimulateDC g) => GSimulateDC (f :+: g) where
  gbeginSolve (L1 x) = gbeginSolve x
  gbeginSolve (R1 x) = gbeginSolve x
  gbeginIteration (L1 x) = gbeginIteration x
  gbeginIteration (R1 x) = gbeginIteration x
  gendIteration (L1 x) = L1 . gendIteration x
  gendIteration (R1 x) = R1 . gendIteration x
  gendSolve (L1 x) = L1 . gendSolve x
  gendSolve (R1 x) = R1 . gendSolve x

instance (GSimulateDC f, GSimulateDC g) => GSimulateDC (f :*: g) where
  gbeginSolve (x :*: y) t s = gbeginSolve x t s >> gbeginSolve y t s
  gbeginIteration (x :*: y) t s = gbeginIteration x t s >> gbeginIteration y t s
  gendIteration (x :*: y) t = gendIteration x t :*: gendIteration y t
  gendSolve (x :*: y) t = gendSolve x t :*: gendSolve y t

-- * DC Simulation

-- | A mutable system of linear equations suitable for representing linear circuits.
data SimulationState s v = SimulationState
  { size                 :: Int -- ^ the number of variables in this system
  , coefficients         :: VSM.MVector s v -- ^ the underlying storage of the equations coefficients in row-major order
  , coefficientsModified :: MutVar s Bool -- ^ flag indicating whether the coefficient matrix has been modified
  , rhs                  :: VSM.MVector s v -- ^ the right hand side of the equations
  , rhsModified          :: MutVar s Bool -- ^ flag indicating whether the RHS has been modified
  }

-- | Creates a new mutable linear circuit equation system with the given number of variables.
newSimulation :: (PrimMonad m, VSM.Storable v, Num v) => Int -> m (SimulationState (PrimState m) v)
newSimulation n =
  SimulationState n <$> VSM.replicate (n * n) 0 <*> newMutVar False <*> VSM.replicate n 0 <*> newMutVar False

cleanCoefficients :: (PrimMonad m) => SimulationState (PrimState m) v -> m ()
cleanCoefficients s = writeMutVar (coefficientsModified s) False

cleanRhs :: (PrimMonad m) => SimulationState (PrimState m) v -> m ()
cleanRhs s = writeMutVar (rhsModified s) False

-- | Computes the index into the flat buffer of matrix variables from a row and column.
flatIndex :: SimulationState s v -- ^ the equation system
          -> Int -- ^ matrix row
          -> Int -- ^ matrix column
          -> Int -- ^ flat buffer index
flatIndex sim row col = row * size sim + col

class Monad m => Simulation m v sys | sys -> v  where
  stampMatrix :: sys -> Variable -> Variable -> v -> m ()
  stampRhs :: sys -> Variable -> v -> m ()

instance (PrimMonad m, PrimState m ~ s, Num v, VS.Storable v) => Simulation m v (SimulationState s v) where
  stampMatrix sys (Index row _) (Index col _) value = do
    VSM.modify (coefficients sys) (+ value) (flatIndex sys row col)
    writeMutVar (coefficientsModified sys) True
  stampMatrix _ _ _ _ = return () -- no stamping when ground is involved

  stampRhs sys (Index row _) value = do
    VSM.modify (rhs sys) (+ value) row
    writeMutVar (rhsModified sys) True
  stampRhs _ Ground _ = return () -- no stamping for ground

initialSolution :: HasVariables c => Circuit c -> VS.Vector Double
initialSolution c = runST $ do
  vec <- VSM.new (c ^. variableCount)
  forOf_ (model . variables . _Index) c $ uncurry (VSM.write vec)
  VS.unsafeFreeze vec

step :: forall c. (HasVariables c, SimulateDC c) => TimeScale -> Circuit c -> Circuit c
step time c = runST simulate where
  simulate :: ST s (Circuit c)
  simulate = do
    sim <- newSimulation (c ^. variableCount)
    -- initialize matrix
    beginSolve c time sim
    mat <- HMatrix.reshape (c ^. variableCount) <$> VS.unsafeFreeze (coefficients sim)
    cleanCoefficients sim
    -- iterate
    writeMutVar (rhsModified sim) True
    c' <- solveNonLinear sim (HMatrix.luPacked mat) (initialSolution c) c
    -- finalize
    return $ endSolve c' time

  solveNonLinear :: SimulationState s Double -> HMatrix.LU Double -> VS.Vector Double -> Circuit c -> ST s (Circuit c)
  solveNonLinear sim lu prevSolution prevCircuit = do
    -- prepare RHS
    beginIteration c time sim
    rhsUpdated <- readMutVar (rhsModified sim)
    cleanRhs sim
    if rhsUpdated
      then do
        -- solve and update components
        rhs' <- VS.unsafeFreeze $ rhs sim
        let solution = HMatrix.flatten $ HMatrix.luSolve lu (HMatrix.asColumn rhs')
            c' = endIteration (updateValues solution c) time
            maxChange = VS.maximum $ VS.zipWith (\x y -> abs (x - y)) solution prevSolution
        if maxChange < 1e-6
          then return c'
          else solveNonLinear sim lu solution c'
      else return prevCircuit

-- | Update all variable references inside a circuit description with its new values.
updateValues :: HasVariables a => VS.Vector Double -> a -> a
updateValues values = over variables replace where
  replace Ground      = Ground
  replace (Index i _) = vars V.! i
  vars = V.imap Index $ VS.convert values


#!/usr/bin/env stack
-- stack --resolver lts-12.4 script
{-# LANGUAGE DeriveDataTypeable, ParallelListComp #-}
module Main ( main ) where

import Hledger

import Control.Exception ( bracket )
import Control.Monad
import System.Exit
import System.IO
import Data.Time.Calendar
import Text.Printf
import Data.Function (on)
import Data.List
import Data.Ord
import Statistics.Math.RootFinding
import Data.Decimal
import qualified Data.Text as T
import System.Console.CmdArgs
import Text.Tabular as Tbl
import qualified Text.Tabular.AsciiArt as Ascii

data Options = Options
  { cashFlow    :: Bool
  , debug       :: Bool
  , file        :: FilePath
  , investments :: String
  , pnl         :: String
  , begin       :: Maybe String
  , end         :: Maybe String
  , daily       :: Bool
  , weekly      :: Bool
  , monthly     :: Bool
  , yearly      :: Bool
  , period      :: Maybe String
  } deriving (Show, Data, Typeable)

queryFromOptions kind day accessor opts =
  -- We ignore QueryOpts with fst as they are irrelevant for this program
  case parseQuery day (T.pack $ accessor opts) of
    (Any, _) -> error $ "you need to specify " ++ kind ++ " query"
    (query, []) -> query
    (_,     opts) -> error $ "this program cannot process query options " ++ show opts

intervalAndSpanFromOptions :: Day -> Options -> (Maybe Interval, DateSpan)
intervalAndSpanFromOptions date opts =
  let begin' = fmap parsedate $ begin opts
      end' = fmap parsedate $ end opts
  in
   case (begin', end', daily opts, weekly opts, monthly opts, yearly opts, period opts) of
     (Nothing, Nothing, False, False, False, False, Nothing) -> (Nothing, nulldatespan)
     (b, e, True,  False, False, False, Nothing) -> (Just (Days 1),   DateSpan b e)
     (b, e, False, True,  False, False, Nothing) -> (Just (Weeks 1),  DateSpan b e)
     (b, e, False, False, True,  False, Nothing) -> (Just (Months 1), DateSpan b e)
     (b, e, False, False, False, True,  Nothing) -> (Just (Years 1),  DateSpan b e)
     (Nothing, Nothing, False, False, False, False, Just pexp) ->
       let (i, s) = parsePeriodExpr' date (T.pack pexp) in (Just i, s)
     _ -> error "Cannot work with this combination of date/period flags"

options =
  Options {
     cashFlow =
      False &= name "c"
      &= help "also show all relevant transactions"
    , debug =
      def
      &= help "print debugging info"
    , file =
      def &= name "f"
      &= typFile
      &= help "input ledger file (pass '-' for stdin)"
    , investments =
      def &= name "i"
      &= typ "QUERY"
      &= help "investments query"
    , pnl =
      def &= name "P"
      &= typ "QUERY"
      &= help "profit-and-loss query"
    , begin =
      Nothing &= name "b"
      &= typ "DATE"
      &= help "calculate interest from this date"
    , end =
      Nothing &= name "e"
      &= typ "DATE"
      &= help "calculate interest until this date"
    , daily =
      def &= name "D"
      &= help "calculate interest for each day"
    , weekly =
      def &= name "W"
      &= help "calculate interest for each week"
    , monthly =
      def &= name "M"
      &= help "calculate interest for each month"
    , yearly =
      def &= name "Y"
      &= help "calculate interest for each year"
    , period =
      Nothing &= name "p"
      &= typ "PERIODEXP"
      &= help "set start date, end date, and/or report interval all at once (overrides the flags above)"
    }
  &= program "hledger-roi"
  &= summary "compute return-on-investment for your portfolio using IRR and TWR"
  &= details [""]

main ::  IO ()
main = bracket (return ()) (\() -> hFlush stdout >> hFlush stderr) $ \() -> do
  opts <-cmdArgs options
  thisDay <- getCurrentDay

  let investmentsQuery = queryFromOptions "investments" thisDay investments opts
  let pnlQuery         = queryFromOptions "pnl" thisDay pnl opts

  jnl <- readJournalFile definputopts (file opts) >>= either fail return

  let trans = jtxns $ filterJournalTransactions investmentsQuery jnl

  when (null trans) $ do
    putStrLn "No relevant transactions found. Check your investments query"
    exitFailure

  when (debug opts) $ do
    putStrLn "DEBUG everything:"
    print $ trans
    putStrLn "DEBUG just non-PnL transactions:"
    print $ filter (matchesTransaction (And [investmentsQuery, Not pnlQuery])) trans
    putStrLn "DEBUG just PnL transactions:"
    print $ filter (matchesTransaction pnlQuery) trans

  let (requestedInterval, requestedSpan) = intervalAndSpanFromOptions thisDay opts
  let existingSpan =
        let dates = map transactionDate2 trans in
        DateSpan (Just $ minimum dates) (Just $ addDays 1 $ maximum dates)
  let wholeSpan = spanDefaultsFrom requestedSpan existingSpan

  let spans = case requestedInterval of
        Nothing -> [wholeSpan]
        Just interval ->
            splitSpan interval $
            spanIntersect existingSpan wholeSpan

  tableBody <- (flip mapM) spans $ \(DateSpan (Just ibegin) (Just iend)) -> do
    -- Spans are [begin,end), and iend here is 1 days after actual end date we are interested in
    let valueBeforeThisPeriod =
          total trans (And [ investmentsQuery
                           , Date (DateSpan Nothing (Just ibegin))])

    let valueAtTheEndOfPeriod  =
          total trans (And [investmentsQuery, Date (DateSpan Nothing (Just iend))])

    let cashFlowInThisPeriod =
          calculateCashFlow trans (And [ Not investmentsQuery
                                       , Not pnlQuery
                                       , Date (DateSpan (Just ibegin) (Just iend)) ] )

    irr <- internalRateOfReturn opts ibegin iend valueBeforeThisPeriod valueAtTheEndOfPeriod cashFlowInThisPeriod
    twr <- timeWeightedReturn opts investmentsQuery trans ibegin iend valueBeforeThisPeriod valueAtTheEndOfPeriod cashFlowInThisPeriod
    let cashFlowAmt = negate $ sum $ map snd cashFlowInThisPeriod
    return [ showDate ibegin
           , showDate (addDays (-1) iend)
           , show valueBeforeThisPeriod
           , show cashFlowAmt
           , show valueAtTheEndOfPeriod
           , show (valueAtTheEndOfPeriod - (valueBeforeThisPeriod + cashFlowAmt))
           , printf "%0.2f%%" irr
           , printf "%0.2f%%" twr ]

  let table = Tbl.Table
              (Group NoLine (map (Header . show) (take (length tableBody) [1..])))
              (Group DoubleLine
               [ Group SingleLine [Header "Begin", Header "End"]
               , Group SingleLine [Header "Value (begin)", Header "Cashflow", Header "Value (end)", Header "PnL"]
               , Group SingleLine [Header "IRR", Header "TWR"]])
              tableBody

  putStrLn $ Ascii.render id id id table

timeWeightedReturn opts investmentsQuery trans ibegin iend valueBefore valueAfter cashFlowInThisPeriod = do
  -- Span is [ibegin,iend), and iend is 1 days after actual end date we are interested in
  let initialUnitPrice = 100
  let initialUnits = valueBefore / initialUnitPrice
  let cashflow =
        -- Aggregate all entries for a single day, assuming that intraday interest is negligible
        map (\date_cash -> let (dates, cash) = unzip date_cash in (head dates, sum cash))
        $ groupBy ((==) `on` fst)
        $ sortBy (comparing fst)
        $ map (\(d,a) -> (d, negate a))
        $ filter ((/=0).snd) cashFlowInThisPeriod

  let units =
        tail $
        (flip scanl)
        (0,0,0,initialUnits)
        (\(_,_,_,unitBalance) (date, amt) ->
          let valueOnDate =
                total trans (And [investmentsQuery, Date (DateSpan Nothing (Just date))])
              unitPrice = if unitBalance == 0.0 then initialUnitPrice else valueOnDate / unitBalance
              unitsBoughtOrSold = amt / unitPrice
          in
           (valueOnDate, unitsBoughtOrSold, unitPrice, unitBalance + unitsBoughtOrSold)
        )
        cashflow

  let finalUnitBalance = if null units then initialUnits else let (_,_,_,u) = last units in u
      finalUnitPrice = valueAfter / finalUnitBalance
      totalTWR = roundTo 2 $ (finalUnitPrice - initialUnitPrice)
      years = (fromIntegral $ diffDays iend ibegin)/365 :: Double
      annualizedTWR = 100*((1+(realToFrac totalTWR/100))**(1/years)-1) :: Double

  let s d = show $ roundTo 2 d
  when (cashFlow opts) $ do
    printf "\nTWR cash flow for %s - %s\n" (showDate ibegin) (showDate (addDays (-1) iend))
    let (dates', amounts') = unzip cashflow
        (valuesOnDate',unitsBoughtOrSold', unitPrices', unitBalances') = unzip4 units
        add x lst = if valueBefore/=0 then x:lst else lst
        dates = add ibegin dates'
        amounts = add valueBefore amounts'
        unitsBoughtOrSold = add initialUnits unitsBoughtOrSold'
        unitPrices = add initialUnitPrice unitPrices'
        unitBalances = add initialUnits unitBalances'
        valuesOnDate = add 0 valuesOnDate'

    putStr $ Ascii.render id id id
      (Tbl.Table
       (Group NoLine (map (Header . showDate) dates))
       (Group DoubleLine [ Group SingleLine [Header "Portfolio value", Header "Unit balance"]
                         , Group SingleLine [Header "Cash", Header "Unit price", Header "Units"]
                         , Group SingleLine [Header "New Unit Balance"]])
       [ [value, oldBalance, amt, prc, udelta, balance]
       | value <- map s valuesOnDate
       | oldBalance <- map s (0:unitBalances)
       | balance <- map s unitBalances
       | amt <- map s amounts
       | prc <- map s unitPrices
       | udelta <- map s unitsBoughtOrSold ])

    printf "Final unit price: %s/%s=%s U.\nTotal TWR: %s%%.\nPeriod: %.2f years.\nAnnualized TWR: %.2f%%\n\n" (s valueAfter) (s finalUnitBalance) (s finalUnitPrice) (s totalTWR) years annualizedTWR

  return annualizedTWR

internalRateOfReturn opts ibegin iend valueBefore valueAfter cashFlowInPeriod = do
  -- Span is [ibegin,iend), and iend is 1 days after actual end date we are interested in
  let prefix = (ibegin, negate valueBefore)

      postfix = (iend, valueAfter)

      totalCF = filter ((/=0) . snd) $ prefix : (sortBy (comparing fst) cashFlowInPeriod) ++ [postfix]

  when (cashFlow opts) $ do
    printf "\nIRR cash flow for %s - %s\n" (showDate ibegin) (showDate (addDays (-1) iend))
    let (dates, amounts) = unzip totalCF
    putStrLn $ Ascii.render id id id
      (Tbl.Table
       (Group NoLine (map (Header . showDate) dates))
       (Group SingleLine [Header "Amount"])
       (map ((:[]) . show) amounts))

  -- 0% is always a solution, so require at least something here
  case ridders 0.00001 (0.000000000001,10000) (interestSum iend totalCF) of
    Root rate -> return ((rate-1)*100)
    NotBracketed -> error "Error: No solution -- not bracketed."
    SearchFailed -> error "Error: Failed to find solution."

type CashFlow = [(Day, Quantity)]

interestSum :: Day -> CashFlow -> Double -> Double
interestSum referenceDay cf rate = sum $ map go cf
    where go (t,m) = (fromRational $ toRational m) * (rate ** (fromIntegral (referenceDay `diffDays` t) / 365))

calculateCashFlow :: [Transaction] -> Query -> CashFlow
calculateCashFlow trans query = map go trans
    where
    go t = (transactionDate2 t, total [t] query)

total :: [Transaction] -> Query -> Quantity
total trans query = unMix $ sumPostings $ filter (matchesPosting query) $ concatMap realPostings trans

unMix :: MixedAmount -> Quantity
unMix a =
  case (normaliseMixedAmount $ costOfMixedAmount a) of
    (Mixed [a]) -> aquantity a
    _ -> error "MixedAmount failed to normalize"


{-# LANGUAGE ConstraintKinds, DataKinds, FlexibleInstances,
             GeneralizedNewtypeDeriving, MultiParamTypeClasses,
             TypeFamilies #-}
-- | Several implementations of the monad in which our dual numbers live
-- and implementations of calculating gradient, derivative and value
-- of an objective function defined on dual numbers.
module HordeAd.Core.Engine
  ( DualMonadValue, primalValueGeneral, primalValue
  , DualMonadGradient, dReverseGeneral, dReverse, dForwardGeneral, dForward, prettyPrintDf
  , DualMonadForward, dFastForwardGeneral, dFastForward
  , generateDeltaVars, initializerFixed
  ) where

import Prelude

import           Control.Monad.Trans.State.Strict
import qualified Data.Array.DynamicS as OT
import           Data.Functor.Identity
import qualified Data.Strict.Vector as Data.Vector
import qualified Data.Vector.Generic as V
import           Numeric.LinearAlgebra (Matrix, Numeric, Vector)
import qualified Numeric.LinearAlgebra as HM

import HordeAd.Core.DualClass
  (Dual, IsPrimal (..), IsPrimalWithScalar, bindInState, dVar)
import HordeAd.Core.DualNumber
import HordeAd.Core.PairOfVectors (DualNumberVariables, makeDualNumberVariables)
import HordeAd.Internal.Delta
  ( CodeOut (..)
  , DeltaState (..)
  , derivativeFromDelta
  , gradientFromDelta
  , ppBindings
  , toDeltaId
  )

-- * The dummy monad implementation that does not collect deltas.
-- It's intended for efficiently calculating the value of the function only.

-- An overcomplicated
-- @type DualMonadValue r = Identity@
-- to avoid @-Wno-orphans@ and @UndecidableInstances@.
newtype DualMonadValue r a = DualMonadValue
  { runDualMonadValue :: Identity a }
  deriving (Monad, Functor, Applicative)

instance IsScalar 'DModeGradient r
         => DualMonad 'DModeGradient r (DualMonadValue r) where
  returnLet (D u _u') = DualMonadValue $ Identity $ D u dZero

-- The general case, needed for old, hacky tests using only scalars.
primalValueGeneral
  :: forall r a. IsScalar 'DModeGradient r
  => (DualNumberVariables 'DModeGradient r
      -> DualMonadValue r a)
  -> Domains r
  -> a
-- Small enough that inline won't hurt.
{-# INLINE primalValueGeneral #-}
primalValueGeneral f (params0, params1, params2, paramsX) =
  let replicateZeros p = V.replicate (V.length p) dZero
      variables = makeDualNumberVariables
                    (params0, params1, params2, paramsX)
                    ( replicateZeros params0  -- dummy
                    , replicateZeros params1
                    , replicateZeros params2
                    , replicateZeros paramsX )
  in runIdentity $ runDualMonadValue $ f variables

primalValue
  :: forall r a. IsScalar 'DModeGradient r
  => (DualNumberVariables 'DModeGradient r
      -> DualMonadValue r (DualNumber 'DModeGradient a))
  -> Domains r
  -> a
-- Small enough that inline won't hurt.
{-# INLINE primalValue #-}
primalValue f parameters =
  let D value _ = primalValueGeneral f parameters
  in value

-- * The fully-fledged monad implementation for gradients
-- and the code that uses it to compute gradients and also derivatives.

newtype DualMonadGradient r a = DualMonadGradient
  { runDualMonadGradient :: State (DeltaState r) a }
  deriving (Monad, Functor, Applicative)

instance IsScalar 'DModeGradient r
         => DualMonad 'DModeGradient
                      r (DualMonadGradient r) where
  returnLet (D u u') = DualMonadGradient $ do
    st <- get
    let (!stNew, !dId) = bindInState u' st
    put stNew
    return $! D u (dVar dId)

initializeState :: forall r. IsScalar 'DModeGradient r
                => Domains r -> DeltaState r
initializeState (params0, params1, params2, paramsX) =
  DeltaState { deltaCounter0 = toDeltaId (V.length params0)
             , deltaCounter1 = toDeltaId (V.length params1)
             , deltaCounter2 = toDeltaId (V.length params2)
             , deltaCounterX = toDeltaId (V.length paramsX)
             , deltaBindings = []
             }

dReverseGeneral
  :: forall r. HasDelta r
  => r
  -> DualNumberVariables 'DModeGradient r
  -> (DualNumberVariables 'DModeGradient r
      -> DualMonadGradient r (DualNumber 'DModeGradient r))
  -> (Domains r, r)
-- The functions in which @dReverseGeneral@ inlines are not inlined themselves
-- in client code, so the bloat is limited.
{-# INLINE dReverseGeneral #-}
dReverseGeneral dt
                variables@(params0, _, params1, _, params2, _, paramsX, _)
                f =
  let dim0 = V.length params0
      dim1 = V.length params1
      dim2 = V.length params2
      dimX = V.length paramsX
      initialState = initializeState (params0, params1, params2, paramsX)
      (D value d, st) = runState (runDualMonadGradient (f variables))
                                 initialState
      inlineDerivative primCode primalArgs dualArgs =
        let D _ u' = outDualNumber primCode primalArgs dualArgs
        in u'
      gradient =
        gradientFromDelta inlineDerivative inlineDerivative inlineDerivative
                          inlineDerivative inlineDerivative
                          dim0 dim1 dim2 dimX st d dt
  in (gradient, value)

dReverse :: HasDelta r
   => r
   -> (DualNumberVariables 'DModeGradient r
       -> DualMonadGradient r (DualNumber 'DModeGradient r))
   -> Domains r
   -> (Domains r, r)
dReverse dt f parameters =
  let varDeltas = generateDeltaVars parameters
      variables = makeDualNumberVariables parameters varDeltas
  in dReverseGeneral dt variables f

-- This function uses @DualMonadGradient@ for an inefficient computation
-- of forward derivaties. See @dFastForwardGeneral@ for an efficient variant.
dForwardGeneral
  :: forall r. HasDelta r
  => DualNumberVariables 'DModeGradient r
  -> (DualNumberVariables 'DModeGradient r
      -> DualMonadGradient r (DualNumber 'DModeGradient r))
  -> Domains r
  -> (r, r)
{-# INLINE dForwardGeneral #-}
dForwardGeneral variables@(params0, _, params1, _, params2, _, paramsX, _)
                f ds =
  let initialState = initializeState (params0, params1, params2, paramsX)
      (D value d, st) = runState (runDualMonadGradient (f variables))
                                 initialState
      inlineDerivative primCode primalArgs dualArgs =
        let D _ u' = outDualNumber primCode primalArgs dualArgs
        in u'
  in ( derivativeFromDelta inlineDerivative inlineDerivative inlineDerivative
                           inlineDerivative inlineDerivative
                           st d ds
     , value )

-- The direction vector ds is taken as an extra argument.
dForward
  :: HasDelta r
  => (DualNumberVariables 'DModeGradient r
      -> DualMonadGradient r (DualNumber 'DModeGradient r))
  -> Domains r
  -> Domains r
  -> (r, r)
dForward f parameters ds =
  let varDeltas = generateDeltaVars parameters
      variables = makeDualNumberVariables parameters varDeltas
  in dForwardGeneral variables f ds

-- * A monad for efficiently computing forward derivatives.

newtype DualMonadForward r a = DualMonadForward
  { runDualMonadForward :: Identity a }
  deriving (Monad, Functor, Applicative)

instance IsScalar 'DModeDerivative r
         => DualMonad 'DModeDerivative
                      r (DualMonadForward r) where
  returnLet (D u u') = DualMonadForward $ Identity $ D u u'

-- This the efficient variant of forward derivative computation.
dFastForwardGeneral
  :: Dual 'DModeDerivative r ~ r
  => DualNumberVariables 'DModeDerivative r
  -> (DualNumberVariables 'DModeDerivative r
      -> DualMonadForward r (DualNumber 'DModeDerivative r))
  -> (r, r)
{-# INLINE dFastForwardGeneral #-}
dFastForwardGeneral variables f =
  let D value d = runIdentity $ runDualMonadForward $ f variables
  in (d, value)

-- The direction vector ds is taken as an extra argument.
dFastForward
  :: forall r. (Numeric r, Dual 'DModeDerivative r ~ r)
  => (DualNumberVariables 'DModeDerivative r
      -> DualMonadForward r (DualNumber 'DModeDerivative r))
  -> Domains r
  -> Domains r
  -> (r, r)
dFastForward f parameters (params0, params1, params2, paramsX) =
  let variables =
        makeDualNumberVariables
          parameters
          (V.convert params0, params1, params2, paramsX)  -- ds
      (derivative, value) = dFastForwardGeneral variables f
  in (derivative, value)

-- * Additional mechanisms

prettyPrintDf
  :: forall r. HasDelta r
  => Bool
  -> (DualNumberVariables 'DModeGradient r
      -> DualMonadGradient r (DualNumber 'DModeGradient r))
  -> Domains r
  -> String
prettyPrintDf reversed f parameters =
  let varDeltas = generateDeltaVars parameters
      variables = makeDualNumberVariables parameters varDeltas
      initialState = initializeState parameters
      (D _ deltaTopLevel, st) = runState (runDualMonadGradient (f variables))
                                         initialState
  in ppBindings reversed st deltaTopLevel

generateDeltaVars
  :: forall r. IsScalar 'DModeGradient r
  => Domains r
  -> ( Data.Vector.Vector (Dual 'DModeGradient r)
     , Data.Vector.Vector (Dual 'DModeGradient (Vector r))
     , Data.Vector.Vector (Dual 'DModeGradient (Matrix r))
     , Data.Vector.Vector (Dual 'DModeGradient (OT.Array r)) )
generateDeltaVars (params0, params1, params2, paramsX) =
  let vVar :: forall a v. (IsPrimalWithScalar 'DModeGradient a r, V.Vector v a)
           => v a -> Data.Vector.Vector (Dual 'DModeGradient a)
      vVar p = V.generate (V.length p) (dVar . toDeltaId)
      !v0 = vVar params0
      !v1 = vVar params1
      !v2 = vVar params2
      !vX = vVar paramsX
  in (v0, v1, v2, vX)

-- | Initialize parameters using a uniform distribution with a fixed range
-- taken from an argument.
--
-- Must be Double, because @randomVector@ only works on that.
--
-- This only works fine for nets with levels that have similar size and use
-- the same activation function. Otherwise, the range should vary per level.
-- A rule of thumb range for weights is @sqrt(6 / (F_in + F_out)@,
-- where @F_in + F_out@ is the sum of inputs and outputs of the largest level.
-- See https://github.com/pytorch/pytorch/issues/15314 and their newer code.
initializerFixed :: Int -> Double -> (Int, [Int], [(Int, Int)], [OT.ShapeL])
                 -> ((Int, Int, Int, Int), Int, Double, Domains Double)
initializerFixed seed range (nParams0, lParams1, lParams2, lParamsX) =
  let vParams1 = V.fromList lParams1
      vParams2 = V.fromList lParams2
      vParamsX = V.fromList lParamsX
      createRandomVector n seedV =
        HM.scale (2 * range)
        $ HM.randomVector seedV HM.Uniform n - HM.scalar 0.5
      params0Init = createRandomVector nParams0 seed
      params1Init =
        V.imap (\i nPV -> createRandomVector nPV (seed + nPV + i)) vParams1
      params2Init =
        V.imap (\i (rows, cols) ->
                 HM.reshape cols
                 $ createRandomVector (rows * cols) (seed + rows + i)) vParams2
      paramsXInit =
        V.imap (\i sh ->
                 let sz = product sh
                 in OT.fromVector sh
                    $ createRandomVector sz (seed + sz + i)) vParamsX
      totalParams = nParams0
                    + V.sum vParams1
                    + V.sum (V.map (uncurry (*)) vParams2)
                    + V.sum (V.map product vParamsX)
  in ( (nParams0, V.length vParams1, V.length vParams2, V.length vParamsX)
     , totalParams
     , range
     , (params0Init, params1Init, params2Init, paramsXInit) )

outDualNumber :: (IsPrimal d a, RealFloat a, Show a, Show (Dual d a))
              => CodeOut -> [a] -> [Dual d a]
              -> DualNumber d a
outDualNumber PlusOut [u, v] [u', v'] = D u u' + D v v'
outDualNumber MinusOut [u, v] [u', v'] = D u u' - D v v'
outDualNumber TimesOut [u, v] [u', v'] = D u u' * D v v'
outDualNumber NegateOut [u] [u'] = negate $ D u u'
outDualNumber AbsOut [u] [u'] = abs $ D u u'
outDualNumber SignumOut [u] [u'] = signum $ D u u'
outDualNumber DivideOut [u, v] [u', v'] = D u u' / D v v'
outDualNumber RecipOut [u] [u'] = recip $ D u u'
outDualNumber ExpOut [u] [u'] = exp $ D u u'
outDualNumber LogOut [u] [u'] = log $ D u u'
outDualNumber SqrtOut [u] [u'] = sqrt $ D u u'
outDualNumber PowerOut [u, v] [u', v'] = D u u' ** D v v'
outDualNumber LogBaseOut [u, v] [u', v'] = logBase (D u u') (D v v')
outDualNumber SinOut [u] [u'] = sin $ D u u'
outDualNumber CosOut [u] [u'] = cos $ D u u'
outDualNumber TanOut [u] [u'] = tan $ D u u'
outDualNumber AsinOut [u] [u'] = asin $ D u u'
outDualNumber AcosOut [u] [u'] = acos $ D u u'
outDualNumber AtanOut [u] [u'] = atan $ D u u'
outDualNumber SinhOut [u] [u'] = sinh $ D u u'
outDualNumber CoshOut [u] [u'] = cosh $ D u u'
outDualNumber TanhOut [u] [u'] = tanh $ D u u'
outDualNumber AsinhOut [u] [u'] = asinh $ D u u'
outDualNumber AcoshOut [u] [u'] = acosh $ D u u'
outDualNumber AtanhOut [u] [u'] = atanh $ D u u'
outDualNumber Atan2Out [u, v] [u', v'] = atan2 (D u u') (D v v')
outDualNumber primCode primalArgs dualArgs =
  error $ "outDualNumber: wrong number of arguments"
          ++ show (primCode, primalArgs, dualArgs)