{-# 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 RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE ViewPatterns        #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE FlexibleContexts    #-}
--------------------------------------------------------------------------------
-- |
-- Module      : ArrayFire.Data
-- Copyright   : David Johnson (c) 2019-2020
-- License     : BSD 3
-- Maintainer  : David Johnson <djohnson.m@gmail.com>
-- Stability   : Experimental
-- Portability : GHC
--
-- Functions for populating 'Array' with Data.
--
-- @
-- >>> constant @Double [2,2] 2.0
--  ArrayFire Array
-- [2 2 1 1]
--    2.0000     2.0000
--    2.0000     2.0000
-- @
--
--------------------------------------------------------------------------------
module ArrayFire.Data where

import Control.Exception
import Control.Monad
import Data.Complex
import Data.Int
import Data.Proxy
import Data.Word
import Foreign.C.Types
import Foreign.ForeignPtr
import Foreign.Marshal          hiding (void)
import Foreign.Storable
import System.IO.Unsafe
import Unsafe.Coerce

import ArrayFire.Exception
import ArrayFire.FFI
import ArrayFire.Internal.Data
import ArrayFire.Internal.Defines
import ArrayFire.Internal.Types
import ArrayFire.Arith

-- | Creates an 'Array' from a scalar value from given dimensions
--
-- >>> constant @Double [2,2] 2.0
--  ArrayFire Array
-- [2 2 1 1]
--    2.0000     2.0000
--    2.0000     2.0000
constant
  :: forall a . AFType a
  => [Int]
  -- ^ Dimensions
  -> a
  -- ^ Scalar value
  -> Array a
constant dims val =
  case dtyp of
    x | x == c64 ->
        cast $ constantComplex dims (unsafeCoerce val :: Complex Double)
      | x == c32 ->
        cast $ constantComplex dims (unsafeCoerce val :: Complex Float)
      | x == s64 ->
        cast $ constantLong dims (unsafeCoerce val :: Int)
      | x == u64 ->
        cast $ constantULong dims (unsafeCoerce val :: Word64)
      | x == s32 ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: Int32) :: Double)
      | x == s16 ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: Int16) :: Double)
      | x == u32 ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: Word32) :: Double)
      | x == u8 ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: Word8) :: Double)
      | x == u16 ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: Word16) :: Double)
      | x == f64 ->
        cast $ constant' dims (unsafeCoerce val :: Double)
      | x == b8  ->
        cast $ constant' dims (fromIntegral (unsafeCoerce val :: CBool) :: Double)
      | x == f32 ->
        cast $ constant' dims (realToFrac (unsafeCoerce val :: Float))
      | otherwise -> error "constant: Invalid array fire type"
  where
    dtyp = afType (Proxy @ a)

    constant'
      :: [Int]
      -- ^ Dimensions
      -> Double
      -- ^ Scalar value
      -> Array Double
    constant' dims' val' =
      unsafePerformIO . mask_ $ do
        ptr <- alloca $ \ptrPtr -> do
          zeroOutArray ptrPtr
          withArray (fromIntegral <$> dims') $ \dimArray -> do
            throwAFError =<< af_constant ptrPtr val' n dimArray typ
            peek ptrPtr
        Array <$>
          newForeignPtr
            af_release_array_finalizer
              ptr
          where
            n = fromIntegral (length dims')
            typ = afType (Proxy @ Double)

    -- | Creates an 'Array (Complex Double)' from a scalar val'ue
    --
    -- @
    -- >>> constantComplex [2,2] (2.0 :+ 2.0)
    -- @
    --
    constantComplex
      :: forall arr . (Real arr, AFType (Complex arr))
      => [Int]
      -- ^ Dimensions
      -> Complex arr
      -- ^ Scalar val'ue
      -> Array (Complex arr)
    constantComplex dims' ((realToFrac -> x) :+ (realToFrac -> y)) = unsafePerformIO . mask_ $ do
      ptr <- alloca $ \ptrPtr -> do
        zeroOutArray ptrPtr
        withArray (fromIntegral <$> dims') $ \dimArray -> do
          throwAFError =<< af_constant_complex ptrPtr x y n dimArray typ
          peek ptrPtr
      Array <$>
        newForeignPtr
          af_release_array_finalizer
            ptr
          where
            n = fromIntegral (length dims')
            typ = afType (Proxy @ (Complex arr))

    -- | Creates an 'Array Int64' from a scalar val'ue
    --
    -- @
    -- >>> constantLong [2,2] 2.0
    -- @
    --
    constantLong
      :: [Int]
      -- ^ Dimensions
      -> Int
      -- ^ Scalar val'ue
      -> Array Int
    constantLong dims' val' = unsafePerformIO . mask_ $ do
      ptr <- alloca $ \ptrPtr -> do
        zeroOutArray ptrPtr
        withArray (fromIntegral <$> dims') $ \dimArray -> do
          throwAFError =<< af_constant_long ptrPtr (fromIntegral val') n dimArray
          peek ptrPtr
      Array <$>
        newForeignPtr
          af_release_array_finalizer
            ptr
          where
            n = fromIntegral (length dims')

    -- | Creates an 'Array Word64' from a scalar val'ue
    --
    -- @
    -- >>> constantULong [2,2] 2.0
    -- @
    --
    constantULong
      :: [Int]
      -> Word64
      -> Array Word64
    constantULong dims' val' = unsafePerformIO . mask_ $ do
      ptr <- alloca $ \ptrPtr -> do
        zeroOutArray ptrPtr
        withArray (fromIntegral <$> dims') $ \dimArray -> do
          throwAFError =<< af_constant_ulong ptrPtr (fromIntegral val') n dimArray
          peek ptrPtr
      Array <$>
        newForeignPtr
          af_release_array_finalizer
            ptr
          where
            n = fromIntegral (length dims')

-- | Creates a range of values in an Array
--
-- >>> range @Double [10] (-1)
-- ArrayFire Array
-- [10 1 1 1]
--     0.0000
--     1.0000
--     2.0000
--     3.0000
--     4.0000
--     5.0000
--     6.0000
--     7.0000
--     8.0000
--     9.0000
range
  :: forall a
   . AFType a
  => [Int]
  -> Int
  -> Array a
range dims (fromIntegral -> k) = unsafePerformIO $ do
  ptr <- alloca $ \ptrPtr -> mask_ $ do
    withArray (fromIntegral <$> dims) $ \dimArray -> do
      throwAFError =<< af_range ptrPtr n dimArray k typ
      peek ptrPtr
  Array <$>
    newForeignPtr
      af_release_array_finalizer
        ptr
      where
        n = fromIntegral (length dims)
        typ = afType (Proxy @ a)

-- | Create an sequence [0, dims.elements() - 1] and modify to specified dimensions dims and then tile it according to tile_dims.
--
-- <http://arrayfire.org/docs/group__data__func__iota.htm>
--
-- >>> iota @Double [5,3] []
-- ArrayFire Array
-- [5 3 1 1]
--     0.0000     5.0000    10.0000
--     1.0000     6.0000    11.0000
--     2.0000     7.0000    12.0000
--     3.0000     8.0000    13.0000
--     4.0000     9.0000    14.0000
--
-- >>> iota @Double [5,3] [1,2]
-- ArrayFire Array
-- [5 6 1 1]
--     0.0000     5.0000    10.0000     0.0000     5.0000    10.0000
--     1.0000     6.0000    11.0000     1.0000     6.0000    11.0000
--     2.0000     7.0000    12.0000     2.0000     7.0000    12.0000
--     3.0000     8.0000    13.0000     3.0000     8.0000    13.0000
--     4.0000     9.0000    14.0000     4.0000     9.0000    14.0000
iota
  :: forall a . AFType a
  => [Int]
  -- ^ is the array containing sizes of the dimension
  -> [Int]
  -- ^ is array containing the number of repetitions of the unit dimensions
  -> Array a
  -- ^ is the generated array
iota dims tdims = unsafePerformIO $ do
  let dims' = take 4 (dims ++ repeat 1)
      tdims' =  take 4 (tdims ++ repeat 1)
  ptr <- alloca $ \ptrPtr -> mask_ $ do
    zeroOutArray ptrPtr
    withArray (fromIntegral <$> dims') $ \dimArray ->
      withArray (fromIntegral <$> tdims') $ \tdimArray -> do
        throwAFError =<< af_iota ptrPtr 4 dimArray 4 tdimArray typ
        peek ptrPtr
  Array <$>
    newForeignPtr
      af_release_array_finalizer
        ptr
      where
        typ = afType (Proxy @ a)

-- | Creates the identity `Array` from given dimensions
--
-- >>> identity [2,2]
-- ArrayFire Array
-- [2 2 1 1]
--    1.0000     0.0000
--    0.0000     1.0000
identity
  :: forall a . AFType a
  => [Int]
  -- ^ Dimensions
  -> Array a
identity dims = unsafePerformIO . mask_ $ do
  let dims' = take 4 (dims ++ repeat 1)
  ptr <- alloca $ \ptrPtr -> mask_ $ do
    zeroOutArray ptrPtr
    withArray (fromIntegral <$> dims') $ \dimArray -> do
      throwAFError =<< af_identity ptrPtr n dimArray typ
      peek ptrPtr
  Array <$>
    newForeignPtr
      af_release_array_finalizer
        ptr
      where
        n = fromIntegral (length dims)
        typ = afType (Proxy @ a)

-- | Create a diagonal matrix from input array when extract is set to false
--
-- >>> diagCreate (vector @Double 2 [1..]) 0
-- ArrayFire Array
-- [2 2 1 1]
--    1.0000     0.0000
--    0.0000     2.0000
diagCreate
  :: AFType (a :: *)
  => Array a
  -- ^	is the input array which is the diagonal
  -> Int
  -- ^ is the diagonal index
  -> Array a
diagCreate x (fromIntegral -> n) =
  x `op1` (\p a -> af_diag_create p a n)

-- | Create a diagonal matrix from input array when extract is set to false
--
-- >>> diagExtract (matrix @Double (2,2) [[1,2],[3,4]]) 0
-- ArrayFire Array
-- [2 1 1 1]
--     1.0000
--     4.0000
diagExtract
  :: AFType (a :: *)
  => Array a
  -> Int
  -> Array a
diagExtract x (fromIntegral -> n) =
  x `op1` (\p a -> af_diag_extract p a n)

-- | Join two Arrays together along a specified dimension
--
-- >>> join 0 (matrix @Double (2,2) [[1,2],[3,4]]) (matrix @Double (2,2) [[5,6],[7,8]])
-- ArrayFire Array
-- [4 2 1 1]
--     1.0000     3.0000
--     2.0000     4.0000
--     5.0000     7.0000
--     6.0000     8.0000
--
join
  :: Int
  -> Array (a :: *)
  -> Array a
  -> Array a
join (fromIntegral -> n) arr1 arr2 = op2 arr1 arr2 (\p a b -> af_join p n a b)

-- | Join many Arrays together along a specified dimension
--
-- *FIX ME*
--
-- >>> joinMany 0 [1,2,3]
-- ArrayFire Array
-- [3 1 1 1]
--    1.0000     2.0000     3.0000
--
joinMany
  :: Int
  -> [Array a]
  -> Array a
joinMany (fromIntegral -> n) arrays = unsafePerformIO . mask_ $ do
  fptrs <- forM arrays $ \(Array fptr) -> pure fptr
  newPtr <-
    alloca $ \fPtrsPtr -> do
      forM_ fptrs $ \fptr ->
        withForeignPtr fptr (poke fPtrsPtr)
      alloca $ \aPtr -> do
        zeroOutArray aPtr
        throwAFError =<< af_join_many aPtr n nArrays fPtrsPtr
        peek aPtr
  Array <$>
    newForeignPtr af_release_array_finalizer newPtr
  where
    nArrays = fromIntegral (length arrays)

-- | Tiles an Array according to specified dimensions
--
-- >>> tile @Double (scalar 22.0) [5,5]
-- ArrayFire Array
-- [5 5 1 1]
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
--
tile
  :: Array (a :: *)
  -> [Int]
  -> Array a
tile a (take 4 . (++repeat 1) -> [x,y,z,w]) =
  a `op1` (\p k -> af_tile p k (fromIntegral x) (fromIntegral y) (fromIntegral z) (fromIntegral w))
tile _ _ = error "impossible"

-- | Reorders an Array according to newly specified dimensions
--
-- *FIX ME*
--
-- >>> reorder @Double (scalar 22.0) [5,5]
-- ArrayFire Array
-- [5 5 1 1]
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
-- 22.0000    22.0000    22.0000    22.0000    22.0000
--
reorder
  :: Array (a :: *)
  -> [Int]
  -> Array a
reorder a (take 4 . (++ repeat 0) -> [x,y,z,w]) =
  a `op1` (\p k -> af_reorder p k (fromIntegral x) (fromIntegral y) (fromIntegral z) (fromIntegral w))
reorder _ _ = error "impossible"

-- | Shift elements in an Array along a specified dimension (elements will wrap).
--
-- >>> shift (vector @Double 4 [1..]) 2 0 0 0
-- ArrayFire Array
-- [4 1 1 1]
--     3.0000
--     4.0000
--     1.0000
--     2.0000
--
shift
  :: Array (a :: *)
  -> Int
  -> Int
  -> Int
  -> Int
  -> Array a
shift a (fromIntegral -> x) (fromIntegral -> y) (fromIntegral -> z) (fromIntegral -> w) =
  a `op1` (\p k -> af_shift p k x y z w)

-- | Modify dimensions of array
--
-- >>> moddims (vector @Double 3 [1..]) [1,3]
-- ArrayFire Array
-- [1 3 1 1]
--     1.0000     2.0000     3.0000
--
moddims
  :: forall a
   . Array (a :: *)
  -> [Int]
  -> Array a
moddims (Array fptr) dims =
  unsafePerformIO . mask_ . withForeignPtr fptr $ \ptr -> do
    newPtr <- alloca $ \aPtr -> do
      zeroOutArray aPtr
      withArray (fromIntegral <$> dims) $ \dimsPtr -> do
        throwAFError =<< af_moddims aPtr ptr n dimsPtr
        peek aPtr
    Array <$> newForeignPtr af_release_array_finalizer newPtr
  where
    n = fromIntegral (length dims)

-- | Flatten an Array into a single dimension
--
-- >>> flat (matrix @Double (2,2) [[1..],[1..]])
-- ArrayFire Array
-- [4 1 1 1]
--     1.0000
--     2.0000
--     1.0000
--     2.0000
--
-- >>> flat $ cube @Int (2,2,2) [[[1,1],[1,1]],[[1,1],[1,1]]]
-- ArrayFire Array
-- [8 1 1 1]
--          1
--          1
--          1
--          1
--          1
--          1
--          1
--          1
--
flat
  :: Array a
  -> Array a
flat = (`op1` af_flat)

-- | Flip the values of an Array along a specified dimension
--
-- >>> matrix @Double (2,2) [[2,2],[3,3]]
-- ArrayFire Array
-- [2 2 1 1]
--     2.0000     3.0000
--     2.0000     3.0000
--
-- >>> A.flip (matrix @Double (2,2) [[2,2],[3,3]]) 1
-- ArrayFire Array
-- [2 2 1 1]
--     3.0000     2.0000
--     3.0000     2.0000
--
flip
  :: Array a
  -> Int
  -> Array a
flip a (fromIntegral -> dim) =
  a `op1` (\p k -> af_flip p k dim)

-- | Create a lower triangular matrix from input array.
--
-- >>> lower (constant [2,2] 10 :: Array Double) True
-- ArrayFire Array
-- [2 2 1 1]
--     1.0000     0.0000
--    10.0000     1.0000
--
lower
  :: Array a
  -- ^ is the input matrix
  -> Bool
  -- ^ boolean parameter specifying if the diagonal elements should be 1
  -> Array a
lower a (fromIntegral . fromEnum -> b) =
  a `op1` (\p k -> af_lower p k b)

-- | Create an upper triangular matrix from input array.
--
-- >>> upper (constant [2,2] 10 :: Array Double) True
-- ArrayFire Array
-- [2 2 1 1]
--    1.0000     10.0000
--    0.0000     1.0000
--
upper
  :: Array a
  -> Bool
  -> Array a
upper a (fromIntegral . fromEnum -> b) =
  a `op1` (\p k -> af_upper p k b)

-- | Selects elements from two arrays based on the values of a binary conditional array.
--
-- >>> cond = vector @CBool 5 [1,0,1,0,1]
-- >>> arr1 = vector @Double 5 (repeat 1)
-- >>> arr2 = vector @Double 5 (repeat 2)
-- >>> select cond arr1 arr2
-- ArrayFire Array
-- [5 1 1 1]
--     1.0000
--     2.0000
--     1.0000
--     2.0000
--     1.0000
--
select
  :: Array CBool
  -- ^ is the conditional array
  -> Array a
  -- ^ is the array containing elements from the true part of the condition
  -> Array a
  -- ^	is the array containing elements from the false part of the condition
  -> Array a
  -- ^ is the output containing elements of a when cond is true else elements from b
select a b c = op3 a b c af_select

-- | Selects elements from two arrays based on the values of a binary conditional array.
--
-- <http://arrayfire.org/docs/group__data__func__select.htm#gab6886120d0bac4717276910e468bbe88>
--
-- >>> cond = vector @CBool 5 [1,0,1,0,1]
-- >>> arr1 = vector @Double 5 (repeat 1)
-- >>> x = 99
-- >>> selectScalarR cond arr1 x
-- ArrayFire Array
-- [5 1 1 1]
--     1.0000
--    99.0000
--     1.0000
--    99.0000
--     1.0000
--
selectScalarR
  :: Array CBool
  -- ^ is the conditional array
  -> Array a
  -- ^ is the array containing elements from the true part of the condition
  -> Double
  -- ^	is a scalar assigned to out when cond is false
  -> Array a
  -- ^ the output containing elements of a when cond is true else elements from b
selectScalarR a b c = op2 a b (\p w x -> af_select_scalar_r p w x c)

-- | Selects elements from two arrays based on the values of a binary conditional array.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__data__func__select.htm#ga0ccdc05779f88cab5095bce987c2da9d)
--
-- >>> cond = vector @CBool 5 [1,0,1,0,1]
-- >>> arr1 = vector @Double 5 (repeat 1)
-- >>> x = 99
-- >>> selectScalarL cond x arr1
-- ArrayFire Array
-- [5 1 1 1]
--    99.0000
--     1.0000
--    99.0000
--     1.0000
--    99.0000
--
selectScalarL
  :: Array CBool
  -- ^ the conditional array
  -> Double
  -- ^ a scalar assigned to out when cond is true
  -> Array a
  -- ^ the array containing elements from the false part of the condition
  -> Array a
  -- ^ is the output containing elements of a when cond is true else elements from b
selectScalarL a n b = op2 a b (\p w x -> af_select_scalar_l p w n x)

-- af_err af_replace(af_array a, const af_array cond, const af_array b);
-- af_err af_replace_scalar(af_array a, const af_array cond, const double b);


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE DeriveAnyClass, DeriveGeneric #-}
module Data.GameObject where

import qualified Data.Vector.Storable as VS
import           Numeric.LinearAlgebra (Vector, Matrix, flatten, (<>), tr,(#>),(><))
import           Control.Lens      hiding (indices)
import           Number.Quaternion ((+::))
import qualified Graphics.Rendering.OpenGL.GL as GL
import           Graphics.Rendering.OpenGL.GL (($=))
import qualified Graphics.GL as GLRaw
import qualified Data.HashMap               as HM
import           Foreign.Ptr (intPtrToPtr, IntPtr(..))
import           Control.Arrow
import           Data.Maybe
import Data.Vector.Strategies
import GHC.Generics

import           Data.Mesh
import           Data.PhysicsData
import           Utils.Quaternions
import           Data.Shaders
import           Physics.Collision

import           Utils (posToMat,posToMatV,glTexture,pointAt)

data GameObject = GameObject {
  _mesh     :: Mesh,
  _location :: Vector Float,
  _physics  :: Maybe PhysicsData,
  _children :: [GameObject],
  _primitive:: Primitive
  }

instance Show GameObject where
  show (GameObject m l p _ pr) = show l ++ "\n" ++ show p ++ "\n" ++ show (m^.quaternion) ++ "\n" ++ show pr ++ "\n"

makeLenses ''GameObject

modelMat :: GameObject -> Matrix Float
modelMat =
  view location &&& view (mesh.quaternion)  >>>
  posToMat *** quatToMat >>>
  uncurry (<>) --- >>> tr

translatePoint :: GameObject -> Vector Float -> Vector Float
translatePoint obj =
  let rot = (#>) . tr . quatToMat3 $ obj^.mesh.quaternion
      tra = (obj^.location +)
  in tra . rot

prepareCPrim :: GameObject -> PrimitiveComputation
prepareCPrim =
  modelMat &&& view primitive >>>
  uncurry preparePrimitive

initializePhysics :: Float -> GameObject -> GameObject
initializePhysics mInv obj =
  let quat = obj^.mesh.quaternion
  in set physics (Just $ startPhysics mInv quat) obj

turnQuatZ :: Float -> GameObject -> GameObject
turnQuatZ = over (mesh.quaternion) . quatConcat . rotQuatZ
  where
    rotQuatZ t = cos (t/2) +:: (0,sin $ t / 2,0)

getInertiaInv :: GameObject -> Matrix Float
getInertiaInv =
  fromMaybe ((3><3) $ replicate 9 0) .
  preview (physics._Just.inertiaInv)

drawObject2 cam proj obj = do
  GL.currentProgram $= Just (obj^.mesh.shaders.program)
  GL.bindVertexArrayObject $= Just (obj^.mesh.vao)

  passMatrix "camera" $ flatten cam
  passMatrix "proj" proj
  mapM_ (passMatrix "rot" . pointAt . view speed) $ obj^.physics
  passMatrix "pos" . posToMatV $ obj^.location

  passTexture "texSampler" 0 $ obj^.mesh.texture
  passTexture "normSampler" 1 $ obj^.mesh.normalMap

  GL.drawElements GL.Triangles indLen GL.UnsignedInt (intPtrToPtr $ IntPtr 0)

  GL.bindVertexArrayObject $= Nothing
  where
    indLen    = fromIntegral . VS.length $ obj^.mesh.indices
    passMatrix = passMatrixToProg $ obj^.mesh.shaders
    passTexture = passTextureToProg GL.Texture2D $ obj^.mesh.shaders

passTextureToProg tp shaders name pos mtex = do
  GL.activeTexture $= GL.TextureUnit pos
  GL.textureBinding tp $= mtex
  GL.uniform (GL.UniformLocation $ uniformName shaders name) $= GL.TextureUnit pos
  return ()

passMatrixToProg shaders name mat =
  VS.unsafeWith mat $
     GLRaw.glUniformMatrix4fv (uniformName shaders name) 1 0

drawObject mskyTex cam proj obj = do
  GL.blend $= GL.Enabled
  -- GL.depthFunc $= Nothing

  GL.blendFunc $= (GL.SrcAlpha, GL.OneMinusSrcAlpha)
  GL.currentProgram $= Just (obj^.mesh.shaders.program)
  GL.bindVertexArrayObject $= Just (obj^.mesh.vao)

  passMatrix "camera" $ flatten cam
  passMatrix "proj" proj
  passMatrix "rot" . quatToMatV $ obj^.mesh.quaternion
  passMatrix "pos" . posToMatV $ obj^.location

  passTexture "texSampler" 1 $ obj^.mesh.texture
  passTexture "normSampler" 2 $ obj^.mesh.normalMap
  passTexture "skytex" 1 mskyTex

  GL.drawElements GL.Triangles indLen GL.UnsignedInt (intPtrToPtr $ IntPtr 0)
  GL.blend $= GL.Disabled
  -- GL.depthFunc $= Just GL.Less

  GL.bindVertexArrayObject $= Nothing
  where
    indLen    = fromIntegral . VS.length $ obj^.mesh.indices
    passMatrix = passMatrixToProg $ obj^.mesh.shaders
    passTexture = passTextureToProg GL.Texture2D $ obj^.mesh.shaders


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE FlexibleContexts #-}

import CPVO.IO
import CPVO.Numeric
import CPVO.IO.DOS
import CPVO.IO.Fortran
import CPVO.IO.Reader.Ecalj.DOS
import CPVO.IO.Reader.Ecalj.Common
import CPVO.IO.Reader.Ecalj.MMOM
import CPVO.IO.Plot.Gnuplot.DOS
import CPVO.IO.Plot.Gnuplot.Common
import CPVO.IO.Plot.PDOS

import qualified Language.C.Inline as C

import qualified Control.Foldl as Fold
import System.Environment (getArgs)
import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Text.Read as T
import qualified Data.Text.Format as T
import Text.Printf
import Data.List.Split
import Data.List
import Data.Maybe
import Data.Either
-------------------------

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel (readMatrix)
import Numeric.LinearAlgebra.Data hiding (find)
import Data.Char
import System.Process
import System.IO
-------------------------
import Text.PrettyPrint.Boxes hiding ((<>),cols,rows)
import qualified Text.PrettyPrint.Boxes as TB
import Data.List
-------------------------
import Text.Pandoc
import Control.Monad ((<=<),forM)

rydberg=1/13.605

testArgs = concat [ [ "GWtable01.tex"
                   , "Total and partial DOS at *E*~F~ in states/eV/unit-cell, QSGW method."
                   , "|*N*$_\\uparrow$(*E*~F~)|*N*$_\\downarrow$(*E*~F~)|*N*(*E*~F~)"
                     ]
                   , words "@{}lSSS@{} -9:6 25 T o flipSpin nico2o4 extendedNiCo2O4.normal/nico2o4.6G10 Ni:Ni#3d:6:7:9:8:10 Co:Co#3d:6:7:8:9:10 O:O#2p:3:4:5"
                   ]

main = do
  putStrLn "=======end Test========="

main1 = do
--    allArgs <- getArgs
    (invStat, ymax, xmin, xmax, ctrlAtoms, uniqAtoms, ctrlAtomicAOs,jdTable, jdHeads, foldernya, tailer, colAlign, texFile) <- readHeaderData testArgs
    totalDOS <- readTotalDOSText tailer foldernya
    let resTotY0 = map (\i -> getY0 $ totalDOS ¿ [0,i]) $ flipBy invStat [1,2]
    putStrLn "===========NF total DOS========="
    putStrLn $ show resTotY0
    pdosAtAll <- readPDOS invStat tailer foldernya ctrlAtomicAOs
    let y0AtAll = map (\(mp,b) -> (getY0 mp,b)) pdosAtAll
    let rs' = rendertable
              $ (:) jdHeads
              $ (:)  ("Total" : (map (showDouble 3) $ processCols $ map ((* rydberg) . last . toList) $ resTotY0))
              $ map (\a -> (((fst $ head a):(map (showDouble 3) $ processCols $ sumIt [0,0] $ map snd a))))
              $ groupBy (\a b -> fst a == fst b)
              $ (\[us,ds] -> zipWith (\(yu,(_,(j,_))) (yd,_) -> (j, [yu,yd])) us ds )
              $ groupBy (\(_,(s,_)) (_,(s',_)) ->  s == s')
              $ map (\(mp,b) -> ((* rydberg) $ last $ toList $ getY0 mp, b)) pdosAtAll
    let rs = unlines [ rs' , jdTable ]
    res' <- markdownToTex rs
    let res = T.replace "\\}" "}"
            $ T.replace "\\{" "{"
            $ T.pack $ unlines  [ "\\begin{longtable}[]{" ++ colAlign ++ "}"
                                , unlines $ tail $ lines $ T.unpack res'
                                ]
    T.writeFile texFile res
    putStrLn $ texFile ++ "=============done:t1dosEf.hs=========="
      where
        processCols [u,d] = [u,d,u+d]
        sumIt x [] = x
        sumIt [u0,d0] ([up,dn]:xs) = sumIt [u0+up,d0+dn] xs


import Mnist
import SampleWeight
import Numeric.LinearAlgebra

main = do
    sw <- loadSW
    print $ sw !! 3
    putStrLn "Done"


{-|
Module: TerminalProgress
Description: Learn function with progress bar for terminal.
Copyright: (c) Alexander Ignatyev, 2017
License: BSD-3
Stability: experimental
Portability: POSIX

Learn function with progress bar for terminal.
-}

module MachineLearning.TerminalProgress
(
  learnWithProgressBar
  , learnOneVsAllWithProgressBar
)

where

import Data.List (transpose)
import MachineLearning.Types (Vector, Matrix)
import qualified MachineLearning.Classification.Internal as MLC
import Control.Monad (foldM, mapAndUnzipM)
import Control.DeepSeq (deepseq)
import qualified System.Console.AsciiProgress as AP
import qualified Numeric.LinearAlgebra as LA

-- | Learn the given function displaying progress bar in terminal.
-- It takes function, initial theta and number of iterations to call the function.
-- It returns theta and optimization path (see "MachineLearning.Optimization" for details).
learnWithProgressBar :: (Vector -> (Vector, Matrix)) -> Vector -> Int -> IO (Vector, Matrix)
learnWithProgressBar func initialTheta nIterations = AP.displayConsoleRegions $ do
  pg <- newProgressBar nIterations
  (theta, optPaths) <- foldM (doLoop pg func) (initialTheta, []) [1..nIterations]
  return (theta, buildOptPathMatrix $ reverse optPaths)

-- | Learn the given function displaying progress bar in terminal.
-- It takes function, list of outputs and list of initial thetas and number of iterations to call the function.
-- It returns list of thetas and list of optimization paths (see "MachineLearning.Optimization" for details).
learnOneVsAllWithProgressBar :: (Vector -> Vector -> (Vector, Matrix)) -> Vector -> [Vector] -> Int -> IO ([Vector], [Matrix])
learnOneVsAllWithProgressBar func y initialThetaList nIterations = AP.displayConsoleRegions $ do
    let numLabels = length initialThetaList
        ys = MLC.processOutputOneVsAll numLabels y
    pg <- newProgressBar $ nIterations * (length ys)
    mapAndUnzipM (learnOneClass pg func nIterations) $ zip ys initialThetaList

newProgressBar nIterations = AP.newProgressBar AP.def {
  AP.pgTotal = fromIntegral nIterations
  , AP.pgFormat = "Learning :percent [:bar] (for :elapsed, :eta remaining)"
  }

doLoop pg func (theta, optPaths) _ = do
  let (theta', optPath) = func theta
  theta' `deepseq` AP.tick pg
  return (theta', (optPath : optPaths))

learnOneClass pg func nIterations (y, theta) = do
  (theta, optPaths) <- foldM (doLoop pg $ func y) (theta, []) [1..nIterations]
  return (theta, buildOptPathMatrix $ reverse optPaths)

-- | Build a single optimazation path matrix from list of optimization path matrices.
buildOptPathMatrix :: [Matrix] -> Matrix
buildOptPathMatrix matrices = LA.fromBlocks $ map (\m -> [m]) matrices


{-# LANGUAGE AllowAmbiguousTypes    #-}
{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE DeriveDataTypeable     #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE InstanceSigs           #-}
{-# LANGUAGE KindSignatures         #-}
{-# LANGUAGE LambdaCase             #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE PatternSynonyms        #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TupleSections          #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeInType             #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE UndecidableInstances   #-}
{-# LANGUAGE ViewPatterns           #-}

module Backprop.Learn.Model.Combinator (
    Chain(..)
  , (~++)
  , chainParamLength
  , chainStateLength
  , LearnFunc(..), learnFunc
  , LMap(..), RMap(..)
  , Dimap, pattern DM, _dmPre, _dmPost, _dmLearn
  , (.~)
  , nilLF, onlyLF
  , (:.~)(..)
  , (:&&&)(..)
  , KAutoencoder, kAutoencoder, kaeEncoder, kaeDecoder
  , Feedback(..), feedbackId
  , FeedbackTrace(..), feedbackTraceId
  ) where

import           Backprop.Learn.Model.Class
import           Backprop.Learn.Model.Function
import           Control.Applicative
import           Control.Category
import           Control.Monad
import           Control.Monad.Primitive
import           Control.Monad.Trans.State
import           Data.Bifunctor
import           Data.Kind
import           Data.Type.Equality
import           Data.Type.Length
import           Data.Type.Mayb                        as Mayb
import           Data.Type.NonEmpty
import           Data.Type.Tuple hiding                (T2(..), T3(..))
import           Data.Typeable
import           GHC.TypeNats
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Prelude hiding                        ((.), id)
import           Type.Class.Known
import           Type.Class.Witness
import           Type.Family.List                      as List
import qualified Data.Type.Tuple                       as T
import qualified Data.Vector.Sized                     as SV
import qualified System.Random.MWC                     as MWC

-- | Chain components linearly, retaining the ability to deconstruct at
-- a later time.
data Chain :: [Type] -> Type -> Type -> Type where
    CNil  :: Chain '[] a a
    (:~>) :: (Learn a b l, KnownMayb (LParamMaybe l), KnownMayb (LStateMaybe l))
          => l
          -> Chain ls        b c
          -> Chain (l ': ls) a c
  deriving (Typeable)
infixr 5 :~>

instance ( ListC (Backprop List.<$> LParams ls)
         , ListC (Backprop List.<$> LStates ls)
         )
      => Learn a b (Chain ls a b) where
    type LParamMaybe (Chain ls a b) = NETup Mayb.<$> ToNonEmpty (LParams ls)
    type LStateMaybe (Chain ls a b) = NETup Mayb.<$> ToNonEmpty (LStates ls)

    runLearn      = runChainLearn
    runLearnStoch = runChainLearnStoch

runChainLearn
    :: (Reifies s W, ListC (Backprop List.<$> LParams ls), ListC (Backprop List.<$> LStates ls))
    => Chain ls a b
    -> Mayb (BVar s) (NETup Mayb.<$> ToNonEmpty (LParams ls))
    -> BVar s a
    -> Mayb (BVar s) (NETup Mayb.<$> ToNonEmpty (LStates ls))
    -> (BVar s b, Mayb (BVar s) (NETup Mayb.<$> ToNonEmpty (LStates ls)))
runChainLearn = \case
  CNil -> \_ x _ -> (x, N_)
  (l :: l) :~> ls ->
    let lenPs = chainParamLength ls
        lenSs = chainStateLength ls
    in case knownMayb @(LParamMaybe l) of
      N_ -> \ps x -> case knownMayb @(LStateMaybe l) of
        N_ -> \ss -> flip (runChainLearn ls ps) ss
                   . runLearnStateless l N_
                   $ x
        J_ _ -> case lenSs of
          LZ -> \case
            J_ (isoVar (tOnly . netT) (NETT . onlyT)->s) ->
              let (y, J_ s') = runLearn      l  N_ x (J_ s)
                  (z, _    ) = runChainLearn ls ps y N_
              in  (z, J_ $ isoVar (NETT . onlyT) (tOnly . netT) s')
          LS _ -> \case
            J_ ss -> lenSs //
              let (y, J_ s' ) = runLearn      l  N_ x (J_ (ss ^^. netHead))
                  ssTail      = isoVar NETT netT $ ss ^^. netTail
                  (z, J_ ss') = runChainLearn ls ps y (J_ ssTail)
              in  (z, J_ $ isoVar2 NET unNet s' $ isoVar netT NETT ss')
      J_ _ -> case lenPs of
        LZ -> \case
          J_ (isoVar (tOnly . netT) (NETT . onlyT)->p) -> \x -> case knownMayb @(LStateMaybe l) of
            N_ -> \ss -> flip (runChainLearn ls N_) ss
                       . runLearnStateless l (J_ p)
                       $ x
            J_ _ -> case lenSs of
              LZ -> \case
                J_ (isoVar (tOnly . netT) (NETT . onlyT)->s) ->
                  let (y, J_ s') = runLearn      l  (J_ p)  x (J_ s)
                      (z, _    ) = runChainLearn ls N_      y N_
                  in  (z, J_ $ isoVar (NETT . onlyT) (tOnly . netT) s')
              LS _ -> \case
                J_ ss -> lenSs //
                  let (y, J_ s' ) = runLearn      l  (J_ p) x (J_ (ss ^^. netHead))
                      ssTail      = isoVar NETT netT $ ss ^^. netTail
                      (z, J_ ss') = runChainLearn ls N_     y (J_ ssTail)
                  in  (z, J_ $ isoVar2 NET unNet s' $ isoVar netT NETT ss')
        LS _ -> \case
          J_ ps -> \x -> lenPs //
            let psHead = ps ^^. netHead
                psTail = isoVar NETT netT $ ps ^^. netTail
            in  case knownMayb @(LStateMaybe l) of
                  N_ -> \ss -> flip (runChainLearn ls (J_ psTail)) ss
                             . runLearnStateless l (J_ psHead)
                             $ x
                  J_ _ -> case lenSs of
                    LZ -> \case
                      J_ (isoVar (tOnly . netT) (NETT . onlyT)->s) ->
                        let (y, J_ s') = runLearn      l  (J_ psHead) x (J_ s)
                            (z, _    ) = runChainLearn ls (J_ psTail) y N_
                        in  (z, J_ $ isoVar (NETT . onlyT) (tOnly . netT) s')
                    LS _ -> \case
                      J_ ss -> lenSs //
                        let (y, J_ s' ) = runLearn      l  (J_ psHead) x (J_ (ss ^^. netHead))
                            ssTail      = isoVar NETT netT $ ss ^^. netTail
                            (z, J_ ss') = runChainLearn ls (J_ psTail) y (J_ ssTail)
                        in  (z, J_ $ isoVar2 NET unNet s' $ isoVar netT NETT ss')

runChainLearnStoch
    :: ( Reifies s W
       , ListC (Backprop List.<$> LParams ls)
       , ListC (Backprop List.<$> LStates ls)
       , PrimMonad m
       )
    => Chain ls a b
    -> MWC.Gen (PrimState m)
    -> Mayb (BVar s) (NETup Mayb.<$> ToNonEmpty (LParams ls))
    -> BVar s a
    -> Mayb (BVar s) (NETup Mayb.<$> ToNonEmpty (LStates ls))
    -> m (BVar s b, Mayb (BVar s) (NETup Mayb.<$> ToNonEmpty (LStates ls)))
runChainLearnStoch = \case
  CNil -> \_ _ x _ -> pure (x, N_)
  (l :: l) :~> ls -> \g ->
    let lenPs = chainParamLength ls
        lenSs = chainStateLength ls
    in case knownMayb @(LParamMaybe l) of
      N_ -> \ps x -> case knownMayb @(LStateMaybe l) of
        N_ -> \ss -> flip (runChainLearnStoch ls g ps) ss
                 <=< runLearnStochStateless l g N_
                   $ x
        J_ _ -> case lenSs of
          LZ -> \case
            J_ (isoVar (tOnly . netT) (NETT . onlyT)->s) -> do
              (y, s') <- second fromJ_
                     <$> runLearnStoch      l  g N_ x (J_ s)
              (z, _ ) <- runChainLearnStoch ls g ps y N_
              pure (z, J_ $ isoVar (NETT . onlyT) (tOnly . netT) s')
          LS _ -> \case
            J_ ss -> lenSs // do
              (y, s' ) <- second fromJ_
                      <$> runLearnStoch      l  g N_ x (J_ (ss ^^. netHead))
              let ssTail = isoVar NETT netT $ ss ^^. netTail
              (z, ss') <- second fromJ_
                      <$> runChainLearnStoch ls g ps y (J_ ssTail)
              pure  (z, J_ $ isoVar2 NET unNet s' $ isoVar netT NETT ss')
      J_ _ -> case lenPs of
        LZ -> \case
          J_ (isoVar (tOnly . netT) (NETT . onlyT)->p) -> \x -> case knownMayb @(LStateMaybe l) of
            N_ -> \ss -> flip (runChainLearnStoch ls g N_) ss
                     <=< runLearnStochStateless l g (J_ p)
                       $ x
            J_ _ -> case lenSs of
              LZ -> \case
                J_ (isoVar (tOnly . netT) (NETT . onlyT)->s) -> do
                  (y, s') <- second fromJ_
                         <$> runLearnStoch      l  g (J_ p)  x (J_ s)
                  (z, _ ) <- runChainLearnStoch ls g N_      y N_
                  pure (z, J_ $ isoVar (NETT . onlyT) (tOnly . netT) s')
              LS _ -> \case
                J_ ss -> lenSs // do
                  (y, s' ) <- second fromJ_
                          <$> runLearnStoch      l  g (J_ p) x (J_ (ss ^^. netHead))
                  let ssTail = isoVar NETT netT $ ss ^^. netTail
                  (z, ss') <- second fromJ_
                          <$> runChainLearnStoch ls g N_     y (J_ ssTail)
                  pure (z, J_ $ isoVar2 NET unNet s' $ isoVar netT NETT ss')
        LS _ -> \case
          J_ ps -> \x -> lenPs //
            let psHead = ps ^^. netHead
                psTail = isoVar NETT netT $ ps ^^. netTail
            in  case knownMayb @(LStateMaybe l) of
                  N_ -> \ss -> flip (runChainLearnStoch ls g (J_ psTail)) ss
                           <=< runLearnStochStateless l g (J_ psHead)
                             $ x
                  J_ _ -> case lenSs of
                    LZ -> \case
                      J_ (isoVar (tOnly . netT) (NETT . onlyT)->s) -> do
                        (y, s') <- second fromJ_
                               <$> runLearnStoch      l  g (J_ psHead) x (J_ s)
                        (z, _ ) <- runChainLearnStoch ls g (J_ psTail) y N_
                        pure (z, J_ $ isoVar (NETT . onlyT) (tOnly . netT) s')
                    LS _ -> \case
                      J_ ss -> lenSs // do
                        (y, s' ) <- second fromJ_
                                <$> runLearnStoch      l  g (J_ psHead) x (J_ (ss ^^. netHead))
                        let ssTail = isoVar NETT netT $ ss ^^. netTail
                        (z, ss') <- second fromJ_
                                <$> runChainLearnStoch ls g (J_ psTail) y (J_ ssTail)
                        pure (z, J_ $ isoVar2 NET unNet s' $ isoVar netT NETT ss')

-- | Appending 'Chain'
(~++)
    :: forall ls ms a b c. ()
    => Chain ls a b
    -> Chain ms b c
    -> Chain (ls ++ ms) a c
(~++) = \case
    CNil     -> id
    (l :: l) :~> (ls :: Chain ls' a' b) ->
      case assocMaybAppend @(LParamMaybe l) @(LParams ls') @(LParams ms) known of
        Refl -> case assocMaybAppend @(LStateMaybe l) @(LStates ls') @(LStates ls) known of
          Refl -> \ms -> (l :~> (ls ~++ ms))
            \\ appendLength (chainParamLength ls) (chainParamLength ms)
            \\ appendLength (chainStateLength ls) (chainStateLength ms)

chainParamLength
    :: Chain ls a b
    -> Length (LParams ls)
chainParamLength = \case
    CNil -> LZ
    (_ :: l) :~> ls -> case knownMayb @(LParamMaybe l) of
      N_   -> chainParamLength ls
      J_ _ -> LS $ chainParamLength ls

chainStateLength
    :: Chain ls a b
    -> Length (LStates ls)
chainStateLength = \case
    CNil -> LZ
    (_ :: l) :~> ls -> case knownMayb @(LStateMaybe l) of
      N_   -> chainStateLength ls
      J_ _ -> LS $ chainStateLength ls

appendLength
    :: forall as bs. ()
    => Length as
    -> Length bs
    -> Length (as ++ bs)
appendLength LZ     = id
appendLength (LS l) = LS . appendLength l

assocMaybAppend
    :: forall a bs cs. ()
    => Mayb P a
    -> (MaybeToList a ++ (bs ++ cs)) :~: ((MaybeToList a ++ bs) ++ cs)
assocMaybAppend = \case
    N_   -> Refl
    J_ _ -> Refl

-- | Data type representing trainable models.
--
-- Useful for performant composition, but you lose the ability to decompose
-- parts.
data LearnFunc :: Maybe Type -> Maybe Type -> Type -> Type -> Type where
    LF :: { _lfRunLearn
               :: forall q. Reifies q W
               => Mayb (BVar q) p
               -> BVar q a
               -> Mayb (BVar q) s
               -> (BVar q b, Mayb (BVar q) s)
          , _lfRunLearnStoch
               :: forall m q. (PrimMonad m, Reifies q W)
               => MWC.Gen (PrimState m)
               -> Mayb (BVar q) p
               -> BVar q a
               -> Mayb (BVar q) s
               -> m (BVar q b, Mayb (BVar q) s)
          }
       -> LearnFunc p s a b
  deriving (Typeable)

learnFunc
    :: Learn a b l
    => l
    -> LearnFunc (LParamMaybe l) (LStateMaybe l) a b
learnFunc l = LF { _lfRunLearn      = runLearn l
                 , _lfRunLearnStoch = runLearnStoch l
                 }

instance Learn a b (LearnFunc p s a b) where
    type LParamMaybe (LearnFunc p s a b) = p
    type LStateMaybe (LearnFunc p s a b) = s

    runLearn      = _lfRunLearn
    runLearnStoch = _lfRunLearnStoch

instance Category (LearnFunc p s) where
    id = LF { _lfRunLearn      = \_ -> (,)
            , _lfRunLearnStoch = \_ _ x -> pure . (x,)
            }
    f . g = LF { _lfRunLearn  = \p x s0 ->
                    let (y, s1) = _lfRunLearn g p x s0
                    in  _lfRunLearn f p y s1
               , _lfRunLearnStoch = \gen p x s0 -> do
                    (y, s1) <- _lfRunLearnStoch g gen p x s0
                    _lfRunLearnStoch f gen p y s1
               }

-- | Compose two 'LearnFunc' on lists.
(.~)
    :: forall ps qs ss ts a b c.
     ( ListC (Backprop List.<$> ps)
     , ListC (Backprop List.<$> qs)
     , ListC (Backprop List.<$> ss)
     , ListC (Backprop List.<$> ts)
     , ListC (Backprop List.<$> (ss ++ ts))
     , Known Length ps
     , Known Length ss
     , Known Length ts
     )
    => LearnFunc ('Just (T ps        )) ('Just (T ss         )) b c
    -> LearnFunc ('Just (T qs        )) ('Just (T ts         )) a b
    -> LearnFunc ('Just (T (ps ++ qs))) ('Just (T (ss ++ ts ))) a c
f .~ g = LF { _lfRunLearn  = \(J_ psqs) x (J_ ssts) -> appendLength @ss @ts known known //
                let (y, J_ ts) = _lfRunLearn g (J_ (psqs ^^. tDrop @ps @qs known))
                                               x
                                               (J_ (ssts ^^. tDrop @ss @ts known))
                    (z, J_ ss) = _lfRunLearn f (J_ (psqs ^^. tTake @ps @qs known))
                                               y
                                               (J_ (ssts ^^. tTake @ss @ts known))
                in  (z, J_ $ isoVar2 (tAppend @ss @ts) (tSplit @ss @ts known)
                                     ss ts
                    )
            , _lfRunLearnStoch = \gen (J_ psqs) x (J_ ssts) -> appendLength @ss @ts known known // do
                (y, ts) <- second fromJ_
                       <$> _lfRunLearnStoch g gen (J_ (psqs ^^. tDrop @ps @qs known))
                                                  x
                                                  (J_ (ssts ^^. tDrop @ss @ts known))
                (z, ss) <- second fromJ_
                       <$> _lfRunLearnStoch f gen (J_ (psqs ^^. tTake @ps @qs known))
                                                   y
                                                   (J_ (ssts ^^. tTake @ss @ts known))
                pure  (z, J_ $ isoVar2 (tAppend @ss @ts) (tSplit @ss @ts known)
                                       ss ts
                      )
            }

-- | Identity of '.~'
nilLF :: LearnFunc ('Just (T '[])) ('Just (T '[])) a a
nilLF = id

-- | 'LearnFunc' with singleton lists; meant to be used with '.~'
onlyLF
    :: forall p s a b. (KnownMayb p, MaybeC Backprop p, KnownMayb s, MaybeC Backprop s)
    => LearnFunc p s a b
    -> LearnFunc ('Just (T (MaybeToList p))) ('Just (T (MaybeToList s))) a b
onlyLF f = LF
    { _lfRunLearn = \(J_ ps) x ssM@(J_ ss) -> case knownMayb @p of
        N_ -> case knownMayb @s of
          N_ -> (second . const) ssM
              $ _lfRunLearn f N_ x N_
          J_ _ -> second (J_ . isoVar onlyT tOnly . fromJ_)
                $ _lfRunLearn f N_ x (J_ (isoVar tOnly onlyT ss))
        J_ _ ->
          let p = isoVar tOnly onlyT ps
          in  case knownMayb @s of
                N_ -> (second . const) ssM
                    $ _lfRunLearn f (J_ p) x N_
                J_ _ -> second (J_ . isoVar onlyT tOnly . fromJ_)
                      $ _lfRunLearn f (J_ p) x (J_ (isoVar tOnly onlyT ss))
    , _lfRunLearnStoch = \g (J_ ps) x ssM@(J_ ss) -> case knownMayb @p of
        N_ -> case knownMayb @s of
          N_ -> (fmap . second . const) ssM
              $ _lfRunLearnStoch f g N_ x N_
          J_ _ -> (fmap . second) (J_ . isoVar onlyT tOnly . fromJ_)
                $ _lfRunLearnStoch f g N_ x (J_ (isoVar tOnly onlyT ss))
        J_ _ ->
          let p = isoVar tOnly onlyT ps
          in  case knownMayb @s of
                N_ -> (fmap . second . const) ssM
                    $ _lfRunLearnStoch f g (J_ p) x N_
                J_ _ -> (fmap . second) (J_ . isoVar onlyT tOnly . fromJ_)
                      $ _lfRunLearnStoch f g (J_ p) x (J_ (isoVar tOnly onlyT ss))
    }

-- | Compose two layers sequentially
--
-- Note that this composes in the opposite order of '.' and ':.:', for
-- consistency with the rest of the library.
--
-- The basic autoencoder is simply @l ':.~' m@, where @'Learn' a b l@ and
-- @'Learn' b a m@.  However, for sparse autoencoder, look at
-- 'Autoencoder'.
data (:.~) :: Type -> Type -> Type where
    (:.~) :: l -> m -> l :.~ m

infixr 5 :.~

instance ( Learn a b l
         , Learn b c m
         , KnownMayb (LParamMaybe l)
         , KnownMayb (LParamMaybe m)
         , KnownMayb (LStateMaybe l)
         , KnownMayb (LStateMaybe m)
         , MaybeC Backprop (LParamMaybe l)
         , MaybeC Backprop (LParamMaybe m)
         , MaybeC Backprop (LStateMaybe l)
         , MaybeC Backprop (LStateMaybe m)
         )
      => Learn a c (l :.~ m) where
    type LParamMaybe (l :.~ m) = TupMaybe (LParamMaybe l) (LParamMaybe m)
    type LStateMaybe (l :.~ m) = TupMaybe (LStateMaybe l) (LStateMaybe m)

    runLearn (l :.~ m) pq x st = (z, tupMaybe T2B s' t')
      where
        (p, q) = splitTupMaybe @_ @(LParamMaybe l) @(LParamMaybe m)
                  (\(T2B v u) -> (v, u))
                  pq
        (s, t) = splitTupMaybe @_ @(LStateMaybe l) @(LStateMaybe m)
                  (\(T2B v u) -> (v, u))
                  st
        (y, s') = runLearn l p x s
        (z, t') = runLearn m q y t

    runLearnStoch (l :.~ m) g pq x st = do
        (y, s') <- runLearnStoch l g p x s
        (z, t') <- runLearnStoch m g q y t
        pure (z, tupMaybe T2B s' t')
      where
        (p, q) = splitTupMaybe @_ @(LParamMaybe l) @(LParamMaybe m)
                  (\(T2B v u) -> (v, u))
                  pq
        (s, t) = splitTupMaybe @_ @(LStateMaybe l) @(LStateMaybe m)
                  (\(T2B v u) -> (v, u))
                  st

-- | Pre-compose a pure parameterless function to a model.
--
-- An @'LMap' b a@ takes a model from @b@ and turns it into a model from
-- @a@.
data LMap :: Type -> Type -> Type -> Type where
    LM :: { _lmPre   :: forall s. Reifies s W => BVar s a -> BVar s b
          , _lmLearn :: l
          }
       -> LMap b a l

-- | Post-compose a pure parameterless function to a model.
--
-- An @'Rmap' b c@ takes a model returning @b@ and turns it into
-- a model returning @c@.
data RMap :: Type -> Type -> Type -> Type where
    RM :: { _rmPost  :: forall s. Reifies s W => BVar s b -> BVar s c
          , _rmLearn :: l
          }
       -> RMap b c l

-- | Pre- and post-compose pure parameterless functions to a model.
--
-- A @'Dimap' b c a d@ takes a model from @b@ to @c@ and turns it into
-- a model from @a@ to @d@.
--
-- @
-- instance 'Learn' b c => Learn a d ('Dimap' b c a d l) where
--     type 'LParamMaybe' (Dimap b c a d l) = LParamMaybe l
--     type 'LStateMaybe' (Dimap b c a d l) = LStateMaybe l
-- @
type Dimap b c a d l = RMap c d (LMap b a l)

-- | Constructor for 'Dimap'
pattern DM
    :: (forall s. Reifies s W => BVar s a -> BVar s b)
    -> (forall s. Reifies s W => BVar s c -> BVar s d)
    -> l
    -> Dimap b c a d l
pattern DM { _dmPre, _dmPost, _dmLearn } = RM _dmPost (LM _dmPre _dmLearn)

instance Learn b c l => Learn a c (LMap b a l) where
    type LParamMaybe (LMap b a l) = LParamMaybe l
    type LStateMaybe (LMap b a l) = LStateMaybe l

    runLearn (LM f l) p x = runLearn l p (f x)
    runLearnStoch (LM f l) g p x = runLearnStoch l g p (f x)

instance Learn a b l => Learn a c (RMap b c l) where
    type LParamMaybe (RMap b c l) = LParamMaybe l
    type LStateMaybe (RMap b c l) = LStateMaybe l

    runLearn (RM f l) p x = first f . runLearn l p x
    runLearnStoch (RM f l) g p x = (fmap . first) f . runLearnStoch l g p x

-- | Take a model and turn it into a model that runs its output into itself
-- several times, returning the final result.  Parameterized by the number
-- of repeats, and the function to process the output to become the next
-- input.
--
-- I don't know why anyone would ever want this.
--
-- See 'FeedbackTrace' if you want to observe all of the intermediate
-- outputs.
data Feedback :: Type -> Type -> Type -> Type where
    FB :: { _fbTimes :: Int
          , _fbFeed  :: forall s. Reifies s W => BVar s b -> BVar s a
          , _fbLearn :: l
          }
       -> Feedback a b l
  deriving (Typeable)

-- | Construct a 'Feedback' from an endofunction (a function that returns
-- a value fo the same type as its input) by simply providing the output
-- directly as the next intput.
feedbackId :: Int -> l -> Feedback a a l
feedbackId n = FB n id

instance Learn a b l => Learn a b (Feedback a b l) where
    type LParamMaybe (Feedback a b l) = LParamMaybe l
    type LStateMaybe (Feedback a b l) = LStateMaybe l

    runLearn (FB n f l) p = runState
                          . foldr (>=>) go (replicate (n - 1) (fmap f . go))
      where
        go = state . runLearn l p

    runLearnStoch (FB n f l) g p = runStateT
                                 . foldr (>=>) go (replicate (n - 1) (fmap f . go))
      where
        go = StateT . runLearnStoch l g p

-- | Take a model and turn it into a model that runs its output into itself
-- several times, and returns all of the intermediate outputs.
-- Parameterized by the function to process the output to become the next
-- input.
--
-- See 'Feedback' if you only need the final result.
--
-- Compare also to 'Unroll'.
data FeedbackTrace :: Nat -> Type -> Type -> Type -> Type where
    FBT :: { _fbtFeed  :: forall s. Reifies s W => BVar s b -> BVar s a
           , _fbtLearn :: l
           }
        -> FeedbackTrace n a b l
  deriving (Typeable)

-- | Construct a 'FeedbackTrace' from an endofunction (a function that
-- returns a value fo the same type as its input) by simply providing the
-- output directly as the next intput.
feedbackTraceId :: l -> FeedbackTrace n a a l
feedbackTraceId = FBT id

instance (Learn a b l, KnownNat n, Backprop b)
      => Learn a (ABP (SV.Vector n) b) (FeedbackTrace n a b l) where
    type LParamMaybe (FeedbackTrace n a b l) = LParamMaybe l
    type LStateMaybe (FeedbackTrace n a b l) = LStateMaybe l

    runLearn (FBT f l) p x0 =
          second snd
        . runState (collectVar . ABP <$> SV.replicateM (state go))
        . (x0,)
      where
        go (x, s) = (y, (f y, s'))
          where
            (y, s') = runLearn l p x s

    runLearnStoch (FBT f l) g p x0 =
          (fmap . second) snd
        . runStateT (collectVar . ABP <$> SV.replicateM (StateT go))
        . (x0,)
      where
        go (x, s) = do
          (y, s') <- runLearnStoch l g p x s
          pure (y, (f y, s'))

-- | "Fork"/"Fan out" two models from the same input.  Useful for
data (:&&&) :: Type -> Type -> Type where
    (:&&&) :: l -> m -> l :&&& m

infixr 3 :&&&

instance ( Learn a b l
         , Learn a c m
         , KnownMayb (LParamMaybe l)
         , KnownMayb (LParamMaybe m)
         , KnownMayb (LStateMaybe l)
         , KnownMayb (LStateMaybe m)
         , MaybeC Backprop (LParamMaybe l)
         , MaybeC Backprop (LParamMaybe m)
         , MaybeC Backprop (LStateMaybe l)
         , MaybeC Backprop (LStateMaybe m)
         , Backprop b
         , Backprop c
         )
      => Learn a (T.T2 b c) (l :&&& m) where
    type LParamMaybe (l :&&& m) = TupMaybe (LParamMaybe l) (LParamMaybe m)
    type LStateMaybe (l :&&& m) = TupMaybe (LStateMaybe l) (LStateMaybe m)

    runLearn (l :&&& m) pq x st = ( T2B y z
                                  , tupMaybe T2B s' t'
                                  )
      where
        (p, q) = splitTupMaybe @_ @(LParamMaybe l) @(LParamMaybe m)
                  (\(T2B v u) -> (v, u))
                  pq
        (s, t) = splitTupMaybe @_ @(LStateMaybe l) @(LStateMaybe m)
                  (\(T2B v u) -> (v, u))
                  st
        (y, s') = runLearn l p x s
        (z, t') = runLearn m q x t

    runLearnStoch (l :&&& m) g pq x st = do
        (y, s') <- runLearnStoch l g p x s
        (z, t') <- runLearnStoch m g q x t
        pure ( T2B y z
             , tupMaybe T2B s' t'
             )
      where
        (p, q) = splitTupMaybe @_ @(LParamMaybe l) @(LParamMaybe m)
                  (\(T2B v u) -> (v, u))
                  pq
        (s, t) = splitTupMaybe @_ @(LStateMaybe l) @(LStateMaybe m)
                  (\(T2B v u) -> (v, u))
                  st

-- | K-sparse autoencoder.  A normal autoencoder is simply @l ':.~' m@;
-- however a k-sparse autoencoder attempts to ensure that the encoding has
-- about @k@ active components for every input.
--
-- <http://www.ericlwilkinson.com/blog/2014/11/19/deep-learning-sparse-autoencoders>
--
-- @
-- instance ('Learn' a ('R' n) l, Learn (R n) a m) => Learn a a ('KAutoencoder' n l m) where
--     type 'LParamMaybe' ('KAutoencoder' n l m) = 'TupMaybe' (LParamMaybe l) (LParamMaybe m)
--     type 'LStateMaybe' ('KAutoencoder' n l m) = 'TupMaybe' (LStateMaybe l) (LStateMaybe m)
-- @
--
-- To "encode" after training this, get the encoder with 'kaeEncoder'.
type KAutoencoder n l m = RMap (R n) (R n) l :.~ m

-- | Construct a 'KAutoencoder'.
--
-- Note that this only has a 'Learn' instance of @l@ outputs @'R' n@ and
-- @m@ takes @'R' n@.  Also, for this to be an actual autoencoder, the
-- input of @l@ has to be the same as the output of @m@.
kAutoencoder
    :: KnownNat n
    => l
    -> m
    -> Int
    -> KAutoencoder n l m
kAutoencoder l m k = RM (kSparse k) l
                 :.~ m

kaeEncoder
    :: KAutoencoder n l m
    -> l
kaeEncoder (RM _ l :.~ _) = l

kaeDecoder
    :: KAutoencoder n l m
    -> m
kaeDecoder (_ :.~ m) = m

-- TODO: KL-divergence autoencoders, a la
-- <http://ufldl.stanford.edu/tutorial/unsupervised/Autoencoders/>.

-- -- | Simple /sparse/ autoencoder that outputs the average activation of
-- -- a vector encoding as well as the result.
-- --
-- -- The traditional autoencoder is simply @l ':.~' m@.  However, the enforce
-- -- sparsity, you have to also be able to observe the average activation of
-- -- your encoding for your loss function (typically 'klDivergence' for
-- -- Kullback-Leibler divergence)
-- --
-- -- See <http://ufldl.stanford.edu/tutorial/unsupervised/Autoencoders/>.
-- --
-- -- @
-- -- instance ('Learn' a ('R' n) l, Learn (R n) a l, 1 '<=' n) => Learn a ('T2' 'Double' a) where
-- --     type 'LParamMaybe' ('Autoencoder' n l m) = 'TupMaybe' (LParamMaybe l) (LParamMaybe m)
-- --     type 'LStateMaybe' ('Autoencoder' n l m) = 'TupMaybe' (LStateMaybe l) (LStateMaybe m)
-- -- @
-- --
-- -- To /use/ the autoencoder after training it, just pattern match on 'AE'
-- -- or use '_aeEncode' (or '_aeDecode')
-- type Autoencoder n l m = l :.~ (FixedFunc (R n) Double :&&& m)
--
-- -- | Construct an 'Autoencoder' by giving the encoder and decoder.
-- pattern AE
--     :: (Learn a (R n) l, Learn (R n) a m, 1 <= n, KnownNat n)
--     => l                    -- ^ '_aeEncode'
--     -> m                    -- ^ '_aeDecode'
--     -> Autoencoder n l m
-- pattern AE { _aeEncode, _aeDecode } <- _aeEncode :.~ (_ :&&& _aeDecode)
--   where
--     AE e d = e :.~ (FF mean :&&& d)


{-# LANGUAGE FlexibleInstances #-}

import Debug.Trace
-- import Control.Monad.State.Lazy

import System.Environment (getArgs)
import System.IO (readFile)
import Data.Map.Strict (Map, (!), insert, elems, fromList, toList, findWithDefault, size, empty, member, findMin, findMax, singleton, filter)
import qualified Data.Map.Strict as M
import Data.Set (Set)
import qualified Data.Set as S
--import qualified Data.Array as A
import Data.List (find, intercalate, intersperse, permutations, inits, tails, isPrefixOf)
import Data.List.Split (splitOn)
-- import Data.Complex (Complex((:+)), realPart, imagPart) -- define my own complex

-- import UI.NCurses
-- import Data.Char (chr, ord)
-- -- import Data.Complex (Complex((:+)))
import Data.Maybe (fromJust)
type Instructions = Map Integer Integer

data ComputerState = Ready | Running | Blocked | Done
  deriving (Show, Eq)

data Computer = Computer {
  state :: ComputerState,
  memory :: Map Integer Integer,
  iptr :: Integer, -- instruction pointer
  base :: Integer, -- base offset
  input :: [Integer],
  output :: [Integer],
  time :: Integer}

computer0 = Computer { state = Ready, memory = empty, iptr = 0, base = 0, input = [], output = [], time = 0 }

run1 :: Computer -> Computer -- (Integer, Integer) -> Instructions -> [Integer] -> [Integer]
run1 c
  | state c == Ready = run1 $ c {state = Running}
  | state c == Done = c
  | state c == Blocked = if null (input c) then c else run1 $ c {state = Running}
  | otherwise = -- running
    --traceShow (iptr c, base c, memory c) $
    case instr `mod` 100 of
      1 -> -- add
        addTime $ c {iptr = i+4, memory = insert (addr 3) (arg 1 + arg 2) instructions}
      2 -> -- multiply
        addTime $ c {iptr = i+4, memory = insert (addr 3) (arg 1 * arg 2) instructions}
      3 -> -- read input
        if (null $ input c)
        then c {state = Blocked}
        else addTime $ c {iptr = i+2, memory = insert (addr 1) (head $ input c) instructions, input = tail $ input c}
      4 -> -- output
        --traceShow (arg 1) $
        addTime $ c {iptr = i+2, output = output c ++ [arg 1]}
      5 -> -- jump-if-true
        addTime $ c {iptr = if arg 1 == 0 then i+3 else arg 2}
      6 -> -- jump-if-false
        addTime $ c {iptr = if arg 1 == 0 then arg 2 else i+3}
      7 -> -- less than
        addTime $ c {iptr = i+4, memory = insert (addr 3) (if arg 1 < arg 2 then 1 else 0) instructions}
      8 -> -- equals
        addTime $ c {iptr = i+4, memory = insert (addr 3) (if arg 1 == arg 2 then 1 else 0) instructions}
      9 -> -- set relative base
        addTime $ c {iptr = i+2, base = base c + arg 1}
      99 -> -- halt
        addTime $ c {state = Done}
      _ -> error "unknown opcode"
  where instructions = memory c
        i = iptr c
        instr = instructions!i
        ii x = findWithDefault 0 x instructions
        arg :: Integer -> Integer
        arg n = case (instr `mod` (100*10^n)) `div` (10*10^n) of
          0 -> ii $ ii $ i+n
          1 -> ii $ i+n
          2 -> ii ((ii $ i+n) + base c)
          _ -> error "bad argument mode"
        addr n = case (instr `mod` (100*10^n)) `div` (10*10^n) of
          0 -> ii $ i+n
          1 -> error "address in mode 1"
          2 -> --trace "address in mode 2" $
            ii (i+n) + base c
          _ -> error $ "address in unknown mode"
        addTime c = c{time = time c + 1}

run :: Computer -> Computer
run c | state c == Done = c
      | state c == Blocked = if null (input c) then c else run $ run1 c
      | length (output c) == 3 = c
      | otherwise = run $ run1 c

runNetwork :: Map Int Computer -> Set (Integer, Int) -> Integer
runNetwork computers heap
  | isDone = y
  | otherwise = runNetwork computers'' heap''
  where
    ((_, i), heap') = S.deleteFindMin heap -- find the computer with lowest time
    c = computers!i
    hasOutput = 3 <= length (output c)
    [addr_, x, y] = take 3 $ output c -- laziness means this won't be executed unless needed
    addr = fromIntegral addr_
    isDone = hasOutput && addr == 255
    c' = if hasOutput then c{output = drop 3 $ output c} else c
    c'' = if state c' == Blocked && null (input c') then c'{input = [-1]} else c'
    c''' = run c''
    c2 = computers!addr
    computers' = if hasOutput then M.insert addr (c2{input = input c2 ++ [x,y]}) computers else computers
    computers'' = M.insert i c''' computers'
    heap'' = S.insert (time c''', i) heap'

main = do
  -- [instructionFile] <- getArgs
  instructionStrings <- readFile "23.input.txt" -- instructionFile
  let instructions = fromList . zip [0 ..] $ map read $ splitOn "," instructionStrings

  putStrLn "Part 1"
  let addrs = [0 .. 49]
      computers = M.fromList [(i, computer0{memory=instructions, input=[fromIntegral i]}) | i <- addrs]
      heap = S.fromList [(0,i) | i <- addrs]
  print $ runNetwork computers heap


{-# LANGUAGE RecordWildCards #-}

module Elescore.Projection.DisruptionsPerDay
  ( DisruptionsPerDay
  , applyDisruptionPerDayEvent
  ) where

import           ClassyPrelude
import           Data.DateTime                     hiding (toGregorian)

import           Database.SimpleEventStore
import           Elescore.IdTypes
import           Elescore.Integration.Common.Types
import           Statistics.IQR

type Date = Text
type DisruptionsPerDay = (Set FacilityId, Map Date Sample)

applyDisruptionPerDayEvent :: PersistedEvent (DisruptionEvent a) -> DisruptionsPerDay -> DisruptionsPerDay
applyDisruptionPerDayEvent PersistedEvent {..} =
  case evPayload of
    FacilityDisrupted fid _ -> addSample (insertSet fid) (toDate evOccurredOn)
    FacilityRestored fid    -> addSample (deleteSet fid) (toDate evOccurredOn)
    _                       -> id

addSample :: (Set FacilityId -> Set FacilityId) -> Date -> DisruptionsPerDay -> DisruptionsPerDay
addSample f date (set, xs) =
  let newSet = f set
      numberOfDisruptions = fromIntegral (length newSet)
      xs' = insertWith (<>) date (singletonSample numberOfDisruptions) xs
  in (newSet, xs')

toDate :: DateTime -> Date
toDate dt =
  let (year, month, day) = toGregorian (utctDay dt)
  in mconcat [tshow year, "-", prependZero month, "-", prependZero day]

prependZero :: Int -> Text
prependZero a
  | a < 10 = "0" <> tshow a
  | otherwise = tshow a


{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}

module Graphical where

import Control.Monad
import Graphics.X11.Xlib
import System.Exit (exitWith, ExitCode(..))
import Control.Concurrent
import Data.Bits
import Data.Word
import Data.Scientific as Scientific
import Data.Complex
import GHC.Conc (numCapabilities)
import qualified Control.Concurrent.MVar.Strict as SM
import GHC.Generics (Generic)
import Control.DeepSeq
import Mand

data Ret a = Ret [Word64] Int
  deriving (Eq, Generic, NFData)

data Ven a = Ven a Int
  deriving (Eq, Generic, NFData)

startGraphical :: MVar Bool -> IO ()
startGraphical mv = do
  print "graphical started"
  dpy <- openDisplay ""
  let dflt = defaultScreen dpy
  let scr = defaultScreenOfDisplay dpy
  let background = whitePixel dpy dflt
  let border = blackPixel dpy dflt
  let max = 500000
  rootw <- rootWindow dpy dflt
  win  <- createSimpleWindow dpy rootw 0 0 600 600 1 border background
  --win <- mkUnmanagedWindow dpy scr rootw 0 0 500 500
  setTextProperty dpy win "this will not be seen" wM_NAME
  mapWindow dpy win
  --drawKey dpy win max
  --drawManyInWin 10 10 50 ["green","yellow", "blue","black"] dpy win 0 0
  --callNTime 50 10 10 ["green","yellow", "blue","black"] (drawManyInWin 10 10 10 dpy win)
  --drawInWin dpy win "blue"
  col <- initColor dpy "red"
  --drawAPoint 400 400 dpy win col
  drawSetParalell dpy win (0.250001 :: Double) (-0.002 :: Double) 500 500 0.000008 mand_iteration max
  --drawSetParalell dpy win (-2 :: Double) (-2 :: Double) 500 500 0.008 mand_iteration max

  --drawSet dpy win (-1.4) (-0.005) 500 500 0.00002 mand_iteration
  --mapM_ (printAllCol dpy win) [(x,y) | x <- [0..255], y <- [0..255]]
  --drawRGBRec dpy win $ pixelFromRGB 150 200 25
  --drawRGBRec dpy win $ pixelFromRGB 0 0 255
  --allocaXEvent (eventThing dpy win)
  sync dpy False
  return ()
  --threadDelay (10 * 1000000)
  --exitWith ExitSuccess

printAllCol :: Display -> Window -> (Int, Int) -> IO ()
printAllCol dpy win (x, y) = drawAPoint x y dpy win $ scaleUpToPixel (256*255 + 255) $ 256 * x + y

scaleUpToPixel :: Int -> Int -> Word64
scaleUpToPixel max n = floor $ (256*256*256 -1) * ((fromIntegral n) / (fromIntegral max))

mkUnmanagedWindow :: Display -> Screen -> Window -> Position -> Position -> Dimension -> Dimension -> IO Window
mkUnmanagedWindow dpy scr rw x y w h = do
 let visual = defaultVisualOfScreen scr
     attrmask = cWOverrideRedirect .|. cWBorderPixel .|. cWBackPixel
 win <- allocaSetWindowAttributes $ \attributes -> do
     set_override_redirect attributes True
     set_background_pixel attributes $ whitePixel dpy (defaultScreen dpy)
     set_border_pixel attributes $ blackPixel dpy (defaultScreen dpy)
     createWindow dpy rw x y w h 1 (defaultDepthOfScreen scr) inputOutput visual attrmask attributes
 return win

-- without mask:
-- took 1:51 to draw set
-- with mask:
-- took 0:31
drawSet :: RealFloat a => Display -> Window -> a -> a -> Int -> Int -> a -> (Complex a -> Complex a -> Complex a) -> Int ->IO ()
drawSet dpy win startRe startIm numberRe numberIm step nextIteration max = do
  gc <- createGC dpy win
  mapM_ (stuff gc) [(x,y) | x <- [0..(numberRe-1)] , y <- [0..(numberIm -1)] ]
  freeGC dpy gc
    where
      --stuff ::  -> (Int, Int) -> IO ()
      stuff gc (x, y) = do
        let re = startRe + (fromIntegral x) * step
        let im = startIm + (fromIntegral y) * step
        --let pixel = scaleIntToRGB (general nextIteration (re :+ im) 1000) 10
        let pixel = scaleUpToPixel (max -1) $ general nextIteration (Just zzcosMask) (re :+ im) max
        setForeground dpy gc pixel
        drawPoint dpy win gc (fromIntegral x) (fromIntegral y)

--without mask:
--took 56 seconds with 2 processing threads
--took 55 with 3 processing threads
--took 57 with 4 processing threads
--with mask
--took 17 with 2
--took 18 with 3
--took 30 with 4
drawSetParalell ::(NFData a, RealFloat a) => Display -> Window -> a -> a -> Int -> Int -> a -> (Complex a -> Complex a -> Complex a) -> Int ->IO ()
drawSetParalell dpy win startRe startIm numberRe numberIm step nextIteration max = do
  gc <- createGC dpy win
  case numCapabilities of
    1 -> drawSet dpy win startRe startIm numberRe numberIm step nextIteration max
    otherwise -> do

      vendor <- SM.newEmptyMVar
      ret <- SM.newEmptyMVar
      v1 <- forkIO $ vendorThread vendor startRe step max numberRe
      gc <- createGC dpy win
      t1 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "1"
      t2 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "2"
      t3 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "3"
      t4 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t5 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t6 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t7 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t8 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t9 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t10 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t11 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      t12 <- forkOS $ processingThread vendor ret startIm step nextIteration max numberIm "4"
      mapM_ (\bob -> do
        Ret iteration_arr x <- SM.takeMVar ret
        drawLineOfPoints dpy win gc x 0 iteration_arr
        print bob ) [1..max]
      killThread v1
      killThread t1
      killThread t2
      killThread t3
      killThread t4
      killThread t5
      killThread t6
      killThread t7
      killThread t8
      killThread t9
      killThread t10
      killThread t11
      killThread t12
      freeGC dpy gc

processingThread :: (NFData a, RealFloat a) => SM.MVar (Ven a) -> SM.MVar (Ret a) -> a -> a -> (Complex a -> Complex a -> Complex a) -> Int -> Int -> String -> IO ()
processingThread vendor ret im_start step nextIteration max numberIm name = do
  forever (do
    Ven re x <- SM.takeMVar vendor
    --print name
    SM.putMVar ret $ Ret (calcuateLine re im_start step nextIteration max numberIm) x)

vendorThread :: (NFData a, RealFloat a)=> SM.MVar (Ven a) -> a -> a -> Int -> Int -> IO ()
vendorThread vend reStart step max maxn = do mapM_ (\n -> SM.putMVar vend $ Ven (reStart + step * fromIntegral n) n) [0.. maxn]

calcuateLine :: RealFloat a => a -> a -> a -> (Complex a -> Complex a -> Complex a) -> Int -> Int -> [Word64]
calcuateLine re im_start step nextIteration max numberIm =
  map (\ x -> scaleUpToPixel (max-1)  $ general nextIteration (Nothing ) (re :+ (im_start + (fromIntegral x) * step)) max) [0..(numberIm-1)]

drawLineOfPoints :: Display -> Window -> GC -> Int -> Int -> [Word64] -> IO ()
drawLineOfPoints dpy win gc x y [] = return ()
drawLineOfPoints dpy win gc x y (a:as) = do
  setForeground dpy gc a
  drawPoint dpy win gc (fromIntegral x) (fromIntegral y)
  drawLineOfPoints dpy win gc x (y + 1) as

drawKey :: Display -> Window -> Int -> IO ()
drawKey dpy win max = do
  gc <- createGC dpy win
  mapM_ (f gc) [0..500]
  freeGC dpy gc
  where
    f gc p = do
      drawRGBRec' dpy win  (scaleUpToPixel max $ round (fromInteger p * ((fromIntegral max) / 500))) 200 200 200 200

      setForeground dpy gc $ scaleUpToPixel max $ round (fromInteger p * ((fromIntegral max) / 500))
      fillRectangle dpy win gc 510 (fromIntegral p) 20 1
      print $ scaleUpToPixel max $ round (fromInteger p * ((fromIntegral max) / 500))

--the peramiter max should be the third root of the maximum value that n can take
--scaleIntToRGB n max = if n == max * max * max then 0 else pixelFromRGB c1 c2 c3
scaleIntToRGB :: Int -> Int -> Word64
scaleIntToRGB n max = pixelFromRGB c1 c2 c3
  where
    scale x = round $ 255 * ((fromIntegral x) / fromIntegral (max))
    c1 = scale $ div n (max* max)
    c2 = scale $ mod (div n max) (max * max)
    c3 = scale $ mod n max

scaleIntToWord :: Int -> Int -> Word64
scaleIntToWord max i = floor $ 16777215 * fromIntegral i / fromIntegral max

scaleInt :: Int -> Int -> Word64
scaleInt i max = fromIntegral . round $ ((fromIntegral i) / (fromIntegral max)) * (16777215 / fromIntegral max)
{-
scaleInt2 :: Int -> Int -> Word64
scaleInt2 i max
  | i < mrt = fromIntegral i
  | i < mrt * 2 = fromIntegral 256 * i
  | otherwise = fromIntegral 256 * 256 * i
  where
    mrt = max ** (1/3)
-}
drawInWin :: Display -> Window -> String -> IO ()
drawInWin dpy win col = do
  fgcolour <- initColor dpy col
  gc <- createGC  dpy win
  setForeground dpy gc fgcolour
  fillRectangle dpy win gc 0 0 50 50
  freeGC dpy gc

drawAPoint :: Int -> Int -> Display -> Window -> Graphics.X11.Xlib.Pixel -> IO ()
drawAPoint x y dpy win pixel = do
  gc <- createGC dpy win
  setForeground dpy gc pixel
  drawPoint dpy win gc (fromIntegral x) (fromIntegral y)
  freeGC dpy gc

drawRGBRec :: Display -> Window -> Graphics.X11.Xlib.Pixel -> IO ()
drawRGBRec dpy win pixel = do
  gc <- createGC dpy win
  setForeground dpy gc pixel
  fillRectangle dpy win gc 100 100 200 300
  freeGC dpy gc

drawRGBRec' :: Display -> Window -> Graphics.X11.Xlib.Pixel -> Int -> Int -> Int -> Int -> IO ()
drawRGBRec' dpy win pixel x y width height= do
  gc <- createGC dpy win
  setForeground dpy gc pixel
  fillRectangle dpy win gc (fromIntegral x) (fromIntegral y)  (fromIntegral width) (fromIntegral height)
  freeGC dpy gc

drawManyInWin :: Int -> Int -> Int -> Display -> Window -> Int -> Int -> [String] -> IO ()
drawManyInWin width height number dpy win x y (c:cols)= do
  col <- initColor dpy c
  gc <- createGC dpy win
  setForeground dpy gc col
  fillRectangle dpy win gc (fromIntegral x) (fromIntegral y) (fromIntegral width) (fromIntegral height)
  freeGC dpy gc
  if number == 0 then return () else drawManyInWin width height (number -1) dpy win (x + width) y (cols ++ [c])

callNTime :: Int -> Int -> Int -> [String] -> (Int -> Int -> [String] -> IO ()) -> IO ()
callNTime 0 _ _ _ _ = return ()
callNTime number startx starty colar@(c:cols) mon = (mon startx (starty * number) colar) >> callNTime (number -1) startx starty (cols ++ [c]) mon

initColor :: Display -> String -> IO Graphics.X11.Xlib.Pixel
initColor dpy colour = do
  let colourmap = defaultColormap dpy (defaultScreen dpy)
  (apros, real) <- allocNamedColor dpy colourmap colour
  return $ color_pixel apros

printBits :: Word64 -> IO()
printBits word = print "adsadfalkdfdas dadad af" -- $ wordToBinary word --map (\ a -> if a then '1' else '0') (reverse $ wordToBinary word)

wordToBinary :: Word64 -> [Bool]
wordToBinary = work 63
  where
    work word n
      | n >= 0 = if(word `div` (2 ^ n) > 0) then True:(work (word - (2^n)) (n-1)) else False:(work word (n-1))
      | otherwise = []

pixelFromRGB :: Word8 -> Word8 -> Word8 -> Word64
pixelFromRGB r g b = ((shiftL ((shiftL (zeroBits .|. toWord64 r ) 8 ).|. toWord64 g) 8) .|. toWord64 b)
  where
    toWord64 :: Word8 -> Word64
    toWord64 = fromInteger . toInteger
--pixelFromRGB r g b = (fromIntegral (65025 * r) + fromIntegral (256 * g) + fromIntegral b)