{-# LANGUAGE CPP, DataKinds, GADTs, KindSignatures, RankNTypes,
             StandaloneDeriving, TypeOperators #-}
#if !MIN_VERSION_base(4,16,0)
{-# LANGUAGE IncoherentInstances #-}
#endif
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | The second component of dual numbers, @Delta@, with its semantics.
-- Neel Krishnaswami calls it \"sparse vector expressions\",
-- and indeed even in the simplest case of an objective function
-- defined on scalars only, the codomain of the function that computes
-- gradients from such delta expressions is a set of vectors, because
-- the gradient of an @R^n@ to @R@ function is an @R^n@ vector.
--
-- The \'sparsity\' is less obvious when the domain of the function consists
-- of multiple vectors, matrices and tensors and when the expressions themselves
-- contain vectors, matrices and tensors. However, a single tiny delta
-- expression (e.g., a sum of two variables) may denote a vector of matrices.
-- Even a delta expression containing a big matrix denotes something much
-- bigger: a whole vector of such matrices and more.
--
-- The algebraic structure here is an extension of vector space.
-- The crucial extra constructor of a variable is used both to represent
-- sharing in order to avoid exponential blowup and to replace the one-hot
-- access to parameters with something cheaper and more uniform.
-- A lot of the remaining additional structure is for introducing
-- and reducing dimensions (ranks).
--
-- This is an internal API now, superseded by "HordeAd.Core.DualClass"
-- that permits other kinds of second component of dual numbers,
-- e.g., the same as primal component, for fast computation
-- of forward derivatives (because @derivativeFromDelta@ below,
-- computing derivatives from delta-expressions, is slow once the expressions
-- grow large enough to affect cache misses).
module HordeAd.Internal.Delta
  ( -- * Abstract syntax trees of the delta expressions
    Delta0 (..), Delta1 (..), Delta2 (..), DeltaX (..), DeltaS (..)
  , -- * Delta expression identifiers
    DeltaId, toDeltaId, convertDeltaId
  , -- * Evaluation of the delta expressions
    DeltaBinding
  , DeltaState (..)
  , Domain0, Domain1, Domain2, DomainX, Domains
  , gradientFromDelta, derivativeFromDelta, ppBindings
  , bindInState0, bindInState1, bindInState2, bindInStateX
  , isTensorDummy
  , CodeOut(..)
  ) where

import Prelude

import           Control.Exception (assert)
import           Control.Monad (unless, zipWithM_)
import           Control.Monad.ST.Strict (ST, runST)
import qualified Data.Array.Convert
import qualified Data.Array.Dynamic as OTB
import qualified Data.Array.DynamicS as OT
import qualified Data.Array.Internal
import qualified Data.Array.Internal.DynamicG
import qualified Data.Array.Internal.DynamicS
import qualified Data.Array.Shaped as OSB
import qualified Data.Array.ShapedS as OS
import           Data.Kind (Type)
import           Data.List (foldl')
import           Data.Proxy (Proxy)
import qualified Data.Strict.Vector as Data.Vector
import qualified Data.Strict.Vector.Autogen.Mutable as Data.Vector.Mutable
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Generic.Mutable as VM
import qualified Data.Vector.Storable.Mutable
import           GHC.TypeLits (KnownNat, Nat, natVal, type (+))
import           Numeric.LinearAlgebra (Matrix, Numeric, Vector, (#>), (<.>))
import qualified Numeric.LinearAlgebra as HM
import           Text.Show.Pretty (ppShow)

import qualified HordeAd.Internal.MatrixOuter as MO
import           HordeAd.Internal.OrthotopeOrphanInstances (liftVS2, liftVT2)

-- * Abstract syntax trees of the delta expressions

-- | This is the grammar of delta-expressions at tensor rank 0, that is,
-- at scalar level. The first few operations have analogues
-- at the level of vectors, matrices and arbitrary tensors.
--
-- For each choice of the underlying scalar type @r@,
-- we have several primitive differentiable types based on the scalar:
-- the scalar type @r@ itself, @Vector r@, @Matrix r@ and tensors.
-- Many operations span the ranks and so span the datatypes, which makes
-- the datatypes mutually recursive.
--
-- The @Outline@ constructors represent primitive numeric function applications
-- for which we delay computing and forgo inlining of the derivative.
data Delta0 r =
    Zero0
  | Scale0 r (Delta0 r)
  | Add0 (Delta0 r) (Delta0 r)
  | Var0 (DeltaId r)

  | SumElements0 (Delta1 r) Int  -- ^ see Note [SumElements0]
  | Index0 (Delta1 r) Int Int  -- ^ second integer is the length of the vector

  | Dot0 (Vector r) (Delta1 r)  -- ^ Dot0 v vd == SumElements0 (Scale1 v vd) n

  | FromX0 (DeltaX r)  -- ^ one of many conversions
  | FromS0 (DeltaS '[] r)

  | Outline0 CodeOut [r] [Delta0 r]
  | Delay0 ~(Delta0 r)

deriving instance (Show r, Numeric r) => Show (Delta0 r)

-- | This is the grammar of delta-expressions at tensor rank 1, that is,
-- at vector level.
data Delta1 r =
    Zero1
  | Scale1 (Vector r) (Delta1 r)
  | Add1 (Delta1 r) (Delta1 r)
  | Var1 (DeltaId (Vector r))

  | Seq1 (Data.Vector.Vector (Delta0 r))  -- ^ "unboxing" conversion
  | Konst1 (Delta0 r) Int  -- ^ length; needed only for forward derivative
  | Append1 (Delta1 r) Int (Delta1 r)  -- ^ the length of the first argument
  | Slice1 Int Int (Delta1 r) Int  -- ^ last integer is the length of argument
  | SumRows1 (Delta2 r) Int  -- ^ the integer is the number of columns
  | SumColumns1 (Delta2 r) Int  -- ^ the integer is the number of rows

  | M_VD1 (Matrix r)
          (Delta1 r)  -- ^ M_VD1 m vd == SumRows1 (M_MD2 m (AsRow2 vd))
  | MD_V1 (Delta2 r)
          (Vector r)  -- ^ MD_V1 md v == SumRows1 (MD_M2 md (asRow v))

  | FromX1 (DeltaX r)
  | forall len. KnownNat len
    => FromS1 (DeltaS '[len] r)

    -- unsorted and undocumented yet
  | Reverse1 (Delta1 r)
  | Flatten1 Int Int (Delta2 r)
  | FlattenX1 OT.ShapeL (DeltaX r)
  | forall sh. OS.Shape sh
    => FlattenS1 (DeltaS sh r)

  | Outline1 CodeOut [Vector r] [Delta1 r]
  | Delay1 ~(Delta1 r)

deriving instance (Show r, Numeric r) => Show (Delta1 r)

-- | This is the grammar of delta-expressions at tensor rank 2, that is,
-- at matrix level.
data Delta2 r =
    Zero2
  | Scale2 (Matrix r) (Delta2 r)
  | Add2 (Delta2 r) (Delta2 r)
  | Var2 (DeltaId (Matrix r))

  | FromRows2 (Data.Vector.Vector (Delta1 r))  -- ^ "unboxing" conversion again
  | FromColumns2 (Data.Vector.Vector (Delta1 r))
  | Konst2 (Delta0 r) (Int, Int)  -- ^ size; needed only for forward derivative
  | Transpose2 (Delta2 r)
  | M_MD2 (Matrix r) (Delta2 r)  -- ^ matrix-(matrix-expression) multiplication
  | MD_M2 (Delta2 r) (Matrix r)  -- ^ (matrix-expression)-matrix multiplication
  | RowAppend2 (Delta2 r) Int (Delta2 r)  -- ^ row-length of first argument
  | ColumnAppend2 (Delta2 r) Int (Delta2 r)  -- ^ col-length of first argument
  | RowSlice2 Int Int (Delta2 r) Int  -- ^ last arg is row-length of the matrix
  | ColumnSlice2 Int Int (Delta2 r) Int  -- ^ column-length of the matrix

  | AsRow2 (Delta1 r)  -- ^ AsRow2 vd == FromRows2 (V.replicate n vd)
  | AsColumn2 (Delta1 r)  -- ^ AsColumn2 vd == FromColumns2 (V.replicate n vd)

  | FromX2 (DeltaX r)
  | forall rows cols. (KnownNat rows, KnownNat cols)
    => FromS2 (DeltaS '[rows, cols] r)

    -- unsorted and undocumented yet
  | Flipud2 (Delta2 r)
  | Fliprl2 (Delta2 r)
  | Reshape2 Int (Delta1 r)
  | Conv2 (Matrix r) (Delta2 r)

  | Outline2 CodeOut [Matrix r] [Delta2 r]
  | Delay2 ~(Delta2 r)

deriving instance (Show r, Numeric r) => Show (Delta2 r)

-- | This is the grammar of delta-expressions at arbitrary tensor rank.
--
-- Warning: not tested enough nor benchmarked.
data DeltaX r =
    ZeroX
  | ScaleX (OT.Array r) (DeltaX r)
  | AddX (DeltaX r) (DeltaX r)
  | VarX (DeltaId (OT.Array r))

  | KonstX (Delta0 r) OT.ShapeL  -- ^ size; needed only for forward derivative
  | AppendX (DeltaX r) Int (DeltaX r)
      -- ^ Append two arrays along the outermost dimension.
      -- All dimensions, except the outermost, must be the same.
      -- The integer argument is the outermost size of the first array.
  | SliceX Int Int (DeltaX r) Int
      -- ^ Extract a slice of an array along the outermost dimension.
      -- The extracted slice must fall within the dimension.
      -- The last argument is the outermost size of the argument array.
  | IndexX (DeltaX r) Int Int
      -- ^ The sub-tensors at the given index of the outermost dimension.
      -- The second integer is the length of the dimension.
  | RavelFromListX [DeltaX r]
      -- ^ Create a tensor from a list treated as the outermost dimension.
  | ReshapeX OT.ShapeL OT.ShapeL (DeltaX r)
      -- ^ Change the shape of the tensor from the first to the second.

  | From0X (Delta0 r)
  | From1X (Delta1 r)
  | From2X (Delta2 r) Int
  | forall sh. OS.Shape sh
    => FromSX (DeltaS sh r)

  | OutlineX CodeOut [OT.Array r] [DeltaX r]
  | DelayX ~(DeltaX r)

deriving instance (Show r, Numeric r) => Show (DeltaX r)

-- | This is the grammar of delta-expressions at arbitrary tensor rank,
-- the fully typed Shaped version.
--
-- Warning: not tested enough nor benchmarked.
data DeltaS :: [Nat] -> Type -> Type where
  ZeroS :: DeltaS sh r
  ScaleS :: OS.Array sh r -> DeltaS sh r -> DeltaS sh r
  AddS :: DeltaS sh r -> DeltaS sh r -> DeltaS sh r
  VarS :: DeltaId (OS.Array sh r) -> DeltaS sh r

  KonstS :: Delta0 r -> DeltaS sh r
  AppendS :: (OS.Shape sh, KnownNat m, KnownNat n)
          => DeltaS (m ': sh) r -> DeltaS (n ': sh) r
          -> DeltaS ((m + n) ': sh) r
    -- ^ Append two arrays along the outermost dimension.
  SliceS :: (KnownNat i, KnownNat n, KnownNat k, OS.Shape rest)
         => Proxy i -> Proxy n -> DeltaS (i + n + k ': rest) r
         -> DeltaS (n ': rest) r
    -- ^ Extract a slice of an array along the outermost dimension.
  IndexS :: (KnownNat ix, KnownNat k, OS.Shape rest)
         => DeltaS (ix + 1 + k ': rest) r -> Proxy ix -> DeltaS rest r
    -- ^ The sub-tensors at the given index of the outermost dimension.
  RavelFromListS :: (KnownNat k, OS.Shape rest)
                 => [DeltaS rest r] -> DeltaS (k : rest) r
    -- ^ Create a tensor from a list treated as the outermost dimension.
  ReshapeS :: (OS.Shape sh, OS.Shape sh', OS.Size sh ~ OS.Size sh')
           => DeltaS sh r -> DeltaS sh' r
    -- ^ Change the shape of the tensor.

  From0S :: Delta0 r -> DeltaS '[] r
  From1S :: Delta1 r -> DeltaS '[n] r
  From2S :: KnownNat cols
         => Proxy cols -> Delta2 r -> DeltaS '[rows, cols] r
  FromXS :: DeltaX r -> DeltaS sh r

  OutlineS :: CodeOut -> [OS.Array sh r] -> [DeltaS sh r] -> DeltaS sh r
  DelayS :: ~(DeltaS sh r) -> DeltaS sh r

instance Show (DeltaS sh r) where
  show _ = "a DeltaS delta expression"

-- * Delta expression identifiers

newtype DeltaId a = DeltaId Int
  deriving Show

toDeltaId :: Int -> DeltaId a
toDeltaId = DeltaId

convertDeltaId :: DeltaId (OT.Array r) -> DeltaId (OS.Array sh r)
convertDeltaId (DeltaId i) = DeltaId i

-- The key property is that it preserves the phantom type.
succDeltaId :: DeltaId a -> DeltaId a
succDeltaId (DeltaId i) = DeltaId (succ i)

-- * Evaluation of the delta expressions

-- | Binding at one of the ranks, with a given underlying scalar.
--
-- The 'DeltaId' components could be re-computed on the fly in 'buildFinMaps',
-- but it costs more (they are boxed, so re-allocation is expensive)
-- than storing them here at the time of binding creation and accessing
-- in `buildFinMaps`.
data DeltaBinding r =
    DeltaBinding0 (DeltaId r) (Delta0 r)
  | DeltaBinding1 (DeltaId (Vector r)) (Delta1 r)
  | DeltaBinding2 (DeltaId (Matrix r)) (Delta2 r)
  | DeltaBindingX (DeltaId (OT.Array r)) (DeltaX r)

data DeltaState r = DeltaState
  { deltaCounter0 :: DeltaId r
  , deltaCounter1 :: DeltaId (Vector r)
  , deltaCounter2 :: DeltaId (Matrix r)
  , deltaCounterX :: DeltaId (OT.Array r)
  , deltaBindings :: [DeltaBinding r]
  }

-- | Helper definitions to shorten type signatures. Note that these
-- differ from their counterparts in all other modules, because the type
-- argument here is the underlying scalar (e.g., @Double),
-- while elsewhere it's the dual component of dual numbers from
-- rank 0 (scalar) level (e.g., @Delta0 Double@).
-- By chance, these definitions and definitions from other modules
-- coincide in case of "forward derivatives computed on the spot"
-- where @r@ is @Double@ and @Double@ is also the dual component.
--
-- More generally, @r@ in this module tends to refer to the underlying
-- scalar type, while in all other modules it refers to the rank 0 dual
-- component type.
type Domain0 r = Vector r

type Domain1 r = Data.Vector.Vector (Vector r)

type Domain2 r = Data.Vector.Vector (Matrix r)

type DomainX r = Data.Vector.Vector (OT.Array r)

type Domains r = (Domain0 r, Domain1 r, Domain2 r, DomainX r)

-- | Delta expressions naturally denote forward derivatives,
-- as encoded in function 'derivativeFromDelta'. However, we are more
-- interested in computing gradients, which is what @gradientFromDelta@ does.
-- The two functions are bound by the equation from Lemma 5 from the paper
-- "Provably correct, asymptotically efficient, higher-order reverse-mode
-- automatic differentiation":
--
-- > dt <.> derivativeFromDelta st d ds = gradientFromDelta st d dt <.> ds
--
-- where @\<.\>@ denotes generalized dot product (multiplying
-- all tensors element-wise and summing the results),
-- @st@ contains bindings of delta variables and @d@ is the top level
-- delta expression from translation of the objective function @f@ to dual
-- numbers, @ds@ belongs to the domain of @f@ and @dt@ to the codomain.
-- We omitted for clarity the @dim0@, @dim1@, @dim2@ and @dimX@ arguments
-- that are the lengths of vectors of the tensors in the domain of @f@.
--
-- Intuitively, @ds@ is a tiny perturbation of the arguments of @f@,
-- for which we compute the derivative, that is, the induced change
-- in the result of @f@. Similarly, @dt@ is a tiny perturbation of the
-- result of @f@, for which we compute the gradient, that is, the change
-- of arguments of @f@ sufficient to cause the perturbation.
-- Note that the scaling factor @r@ in functions @eval*@ in @gradientFromDelta@
-- locally plays the role of @dt@, just as the argument @parameters@
-- in @eval*@ in @derivativeFromDelta@ corresponds to @ds@.
--
-- Let's first discuss in detail the semantics of delta-expressions
-- in terms of forward derivatives, since it's more straightforward.
-- Let @r@ be the type of underlying scalars. Let @f@ be a mathematical
-- differentiable function that takes a collection of type @C@
-- of arguments and produces a single result of type @r@.
-- Let a dual number counterpart of @f@ applied to a collection
-- of parameters @P@ of type @C@ be represented as a Haskell value @b@.
-- Let @d :: Delta0 r@ be the closed delta expression that is the second
-- component of @b@, let @ds@ belong to @C@. The semantics of @d@ is a linear
-- function from @C@ to @r@ that is the derivative of @f@ at point @P@
-- with respect to the perturbation @ds@. The mathematical formula
-- for the derivative follows straightforwardly the syntactic form
-- of the delta expression @d@ (see 'derivativeFromDelta').
--
-- Let's now describe the semantics of closed delta expression @d@
-- as the gradient of @f@ at point @P@ with respect to a @dt@ that belongs
-- to @r@. Here the semantics of @d@ is a collection of four finite maps
-- (vectors) @v0@, @v1@, @v2@, @vX@, corresponding to @C@,
-- each map @vi@ taking indexes of type @DeltaId ai@ to values of type @ai@,
-- where @a0@ is @r@, @a1@ is @Vector r@, @a2@ is @Matrix r@
-- and @aX@ is the type of tensors of @r@.
-- The value of @vi@ at index @DeltaId k@ is the partial derivative
-- of function @f@ at @P@ with respect to its parameter of type @ai@.
-- The position of the @ai@ parameter is represented by @DeltaId k@
-- (in other words, the partial derivative is with respect to a variable
-- quantity tagged with @DeltaId k@) and its value comes from @dt@.
--
-- The semantics of a delta expression that is not closed but contains
-- occurrences of variables that do not correspond to parameters of @f@ is only
-- defined in the context of four vectors that contain values associated
-- to its free variables or, alternatively, of bindings from which the values
-- can be computed, or of a mixture of both. This semantics does not change
-- if a bound expression is substituted for a variable instead of being used
-- to compute a value. (Note however that a computed value can't be
-- substituted for all occurrences of the variable in an expression,
-- because the "computing backwards" trick, needed to get gradients
-- from derivative expressions, computes a value for each occurrence
-- of a variable separately and sums over all occurrences instead
-- of substituting a single value into each occurrence.)
--
-- Function @gradientFromDelta@ computes the four vectors described above.
-- Requested lengths of the vectors are given in the first few arguments.
-- The delta state contains a list of mutually-referencing delta bindings
-- that are to be evaluated, in the given order, starting with the top-level
-- binding of a scalar type provided in the next argument and with respect
-- to perturbation @dt@ (usually set to @1@) in the last argument.
gradientFromDelta :: (Eq r, Numeric r, Num (Vector r))
                  => (CodeOut -> [r] -> [Delta0 r] -> Delta0 r)
                  -> (CodeOut -> [Vector r] -> [Delta1 r] -> Delta1 r)
                  -> (CodeOut -> [Matrix r] -> [Delta2 r] -> Delta2 r)
                  -> (CodeOut -> [OT.Array r] -> [DeltaX r] -> DeltaX r)
                  -> (forall sh. OS.Shape sh
                      => CodeOut -> [OS.Array sh r] -> [DeltaS sh r]
                      -> DeltaS sh r)
                  -> Int -> Int -> Int -> Int -> DeltaState r -> Delta0 r -> r
                  -> Domains r
gradientFromDelta inlineDerivative0 inlineDerivative1 inlineDerivative2
                  inlineDerivativeX inlineDerivativeS
                  dim0 dim1 dim2 dimX st deltaTopLevel dt =
  -- This is morally @V.create@ and so totally safe,
  -- but we can't just call @V.create@ thrice, because it would run
  -- the @ST@ action thrice, so we inline and extend @V.create@ here.
  runST $ do
    (finMap0, finMap1, finMap2, finMapX) <-
      buildFinMaps inlineDerivative0 inlineDerivative1 inlineDerivative2
                   inlineDerivativeX inlineDerivativeS
                   st deltaTopLevel dt
    v0 <- V.unsafeFreeze $ VM.take dim0 finMap0
    v1 <- V.unsafeFreeze $ VM.take dim1 finMap1
    v2 <- V.unsafeFreeze $ VM.take dim2 finMap2
    vX <- V.unsafeFreeze $ VM.take dimX finMapX
    -- Convert to normal matrices, but only the portion of vector
    -- that is not discarded.
    return (v0, v1, V.map MO.convertMatrixOuterOrNull v2, vX)

-- | Create vectors (representing finite maps) that hold delta-variable
-- values. They are initialized with dummy values so that it's cheap to check
-- if any update has already been performed to a cell (allocating big matrices
-- filled with zeros is too costly, especially if never used in an iteration,
-- and adding to such matrices and especially using them as scaling factors
-- is wasteful). The vectors are longer than those representing objective
-- function parameters (e.g., @deltaCounter0@ vs @dim0@), because variables
-- represent not only parameters, but also the bindings that prevent blowup
-- via delta-expression duplication.
initializeFinMaps :: forall s r. Numeric r
                  => DeltaState r
                  -> ST s ( Data.Vector.Storable.Mutable.MVector s r
                          , Data.Vector.Mutable.MVector s (Vector r)
                          , Data.Vector.Mutable.MVector s (MO.MatrixOuter r)
                          , Data.Vector.Mutable.MVector s (OT.Array r) )
initializeFinMaps st = do
  let DeltaId counter0 = deltaCounter0 st
      DeltaId counter1 = deltaCounter1 st
      DeltaId counter2 = deltaCounter2 st
      DeltaId counterX = deltaCounterX st
  finMap0 <- VM.replicate counter0 0  -- correct value
  finMap1 <- VM.replicate counter1 (V.empty :: Vector r)  -- dummy value
  finMap2 <- VM.replicate counter2 MO.emptyMatrixOuter  -- dummy value
  finMapX <- VM.replicate counterX dummyTensor
  return (finMap0, finMap1, finMap2, finMapX)

buildFinMaps :: forall s r. (Eq r, Numeric r, Num (Vector r))
             => (CodeOut -> [r] -> [Delta0 r] -> Delta0 r)
             -> (CodeOut -> [Vector r] -> [Delta1 r] -> Delta1 r)
             -> (CodeOut -> [Matrix r] -> [Delta2 r] -> Delta2 r)
             -> (CodeOut -> [OT.Array r] -> [DeltaX r] -> DeltaX r)
             -> (forall sh. OS.Shape sh
                 => CodeOut -> [OS.Array sh r] -> [DeltaS sh r]
                 -> DeltaS sh r)
             -> DeltaState r -> Delta0 r -> r
             -> ST s ( Data.Vector.Storable.Mutable.MVector s r
                     , Data.Vector.Mutable.MVector s (Vector r)
                     , Data.Vector.Mutable.MVector s (MO.MatrixOuter r)
                     , Data.Vector.Mutable.MVector s (OT.Array r) )
buildFinMaps inlineDerivative0 inlineDerivative1 inlineDerivative2
             inlineDerivativeX inlineDerivativeS
             st deltaTopLevel dt = do
  (finMap0, finMap1, finMap2, finMapX) <- initializeFinMaps st
  let addToVector :: Vector r -> Vector r -> Vector r
      addToVector r = \v -> if V.null v then r else v + r
      addToMatrix :: MO.MatrixOuter r -> MO.MatrixOuter r -> MO.MatrixOuter r
      addToMatrix r = \v -> if MO.nullMatrixOuter v then r else MO.plus v r
      addToArray :: OT.Array r -> OT.Array r -> OT.Array r
      addToArray r = \v -> if isTensorDummy v then r else liftVT2 (+) v r
      addToArrayS :: OS.Shape sh => OS.Array sh r -> OT.Array r -> OT.Array r
      addToArrayS r = \v -> let rs = Data.Array.Convert.convert r
                            in if isTensorDummy v
                               then rs
                               else liftVT2 (+) v rs
      eval0 :: r -> Delta0 r -> ST s ()
      eval0 !r = \case
        Zero0 -> return ()
        Scale0 k d -> eval0 (k * r) d
        Add0 d e -> eval0 r d >> eval0 r e
        Var0 (DeltaId i) -> VM.modify finMap0 (+ r) i

        SumElements0 vd n -> eval1 (HM.konst r n) vd
        Index0 (Var1 (DeltaId i)) ix k -> do
          let f v = if V.null v
                    then HM.konst 0 k V.// [(ix, r)]
                    else v V.// [(ix, v V.! ix + r)]
          VM.modify finMap1 f i
            -- this would be an asymptotic optimization compared to
            -- the general case below, if not for the non-mutable update,
            -- which involves copying the whole vector, so it's just
            -- several times faster (same allocation, but not adding vectors)
        Index0 d ix k -> eval1 (HM.konst 0 k V.// [(ix, r)]) d

        Dot0 v vd -> eval1 (HM.scale r v) vd

        FromX0 d -> evalX (OT.scalar r) d
        FromS0 d -> evalS (OS.scalar r) d

        Outline0 codeOut primalArgs dualArgs ->
          eval0 r $ inlineDerivative0 codeOut primalArgs dualArgs
        Delay0 d -> eval0 r d
      eval1 :: Vector r -> Delta1 r -> ST s ()
      eval1 !r = \case
        Zero1 -> return ()
        Scale1 k d -> eval1 (k * r) d
        Add1 d e -> eval1 r d >> eval1 r e
        Var1 (DeltaId i) -> VM.modify finMap1 (addToVector r) i

        Seq1 lsd -> V.imapM_ (\i d -> eval0 (r V.! i) d) lsd
        Konst1 d _n -> V.mapM_ (`eval0` d) r
        Append1 d k e -> eval1 (V.take k r) d >> eval1 (V.drop k r) e
        Slice1 i n d len ->
          eval1 (HM.konst 0 i V.++ r V.++ HM.konst 0 (len - i - n)) d
        SumRows1 dm cols -> eval2 (MO.asColumn r cols) dm
        SumColumns1 dm rows -> eval2 (MO.asRow r rows) dm

        M_VD1 m dRow ->
          mapM_ (`eval1` dRow)
                (MO.toRows (MO.MatrixOuter (Just m) (Just r) Nothing))
        MD_V1 md row -> eval2 (MO.MatrixOuter Nothing (Just r) (Just row)) md

        FromX1 d -> evalX (OT.fromVector [V.length r] r) d
        FromS1 d -> evalS (OS.fromVector r) d

        Reverse1 d -> eval1 (V.reverse r) d
        Flatten1 rows cols d ->
          eval2 (MO.MatrixOuter (Just $ rows HM.>< cols $ V.toList r)
                                Nothing Nothing)
                d
        FlattenX1 sh d -> evalX (OT.fromVector sh r) d
        FlattenS1 d -> evalS (OS.fromVector r) d

        Outline1 codeOut primalArgs dualArgs ->
          eval1 r $ inlineDerivative1 codeOut primalArgs dualArgs
        Delay1 d -> eval1 r d
      eval2 :: MO.MatrixOuter r -> Delta2 r -> ST s ()
      eval2 !r = \case
        Zero2 -> return ()
        Scale2 k d -> eval2 (MO.multiplyWithOuter k r) d
        Add2 d e -> eval2 r d >> eval2 r e
        Var2 (DeltaId i) -> VM.modify finMap2 (addToMatrix r) i

        FromRows2 lvd -> zipWithM_ eval1 (MO.toRows r) (V.toList lvd)
        FromColumns2 lvd -> zipWithM_ eval1 (MO.toColumns r) (V.toList lvd)
        Konst2 d _sz -> mapM_ (V.mapM_ (`eval0` d)) $ MO.toRows r
        Transpose2 md -> eval2 (MO.transpose r) md  -- TODO: test!
        M_MD2 m md ->
          let mo = MO.MatrixOuter (Just $ HM.tr' m) Nothing Nothing
          in eval2 (MO.matMul mo r) md
        MD_M2 md m ->
          let mo = MO.MatrixOuter (Just $ HM.tr' m) Nothing Nothing
          in eval2 (MO.matMul r mo) md
        RowAppend2 d k e -> eval2 (MO.takeRows k r) d
                            >> eval2 (MO.dropRows k r) e
        ColumnAppend2 d k e -> eval2 (MO.takeColumns k r) d
                               >> eval2 (MO.dropColumns k r) e
        RowSlice2 i n d rows ->
          assert (MO.rows r == n) $
          let cols = MO.cols r
          in eval2 (MO.konst 0 i cols
                    `MO.rowAppend` r
                    `MO.rowAppend` MO.konst 0 (rows - i - n) cols)
                   d
        ColumnSlice2 i n d cols ->
          assert (MO.cols r == n) $
          let rows = MO.rows r
          in eval2 (MO.konst 0 rows i
                    `MO.columnAppend` r
                    `MO.columnAppend` MO.konst 0 rows (cols - i - n))
                   d

        AsRow2 dRow -> mapM_ (`eval1` dRow) (MO.toRows r)
        AsColumn2 dCol -> mapM_ (`eval1` dCol) (MO.toColumns r)

        FromX2 d -> evalX (OT.fromVector [MO.rows r, MO.cols r]
                                         (V.concat $ MO.toRows r)) d
        FromS2 d -> evalS (OS.fromVector $ V.concat $ MO.toRows r) d

        Flipud2 d -> eval2 (MO.flipud r) d
        Fliprl2 d -> eval2 (MO.fliprl r) d
        Reshape2 _cols d -> eval1 (V.concat $ MO.toRows r) d
        Conv2 m md ->
          let mor = MO.convertMatrixOuter r
              convolved = HM.corr2 m mor
              moc = MO.MatrixOuter (Just convolved) Nothing Nothing
          in eval2 moc md

        Outline2 codeOut primalArgs dualArgs ->
          eval2 r $ inlineDerivative2 codeOut primalArgs dualArgs
        Delay2 d -> eval2 r d
      evalX :: OT.Array r -> DeltaX r -> ST s ()
      evalX !r = \case
        ZeroX -> return ()
        ScaleX k d -> evalX (liftVT2 (*) k r) d
        AddX d e -> evalX r d >> evalX r e
        VarX (DeltaId i) -> VM.modify finMapX (addToArray r) i

        KonstX d _sz -> mapM_ (`eval0` d) $ OT.toList r
        AppendX d k e -> case OT.shapeL r of
          n : _ -> evalX (OT.slice [(0, k)] r) d
                   >> evalX (OT.slice [(k, n - k)] r) e
          [] -> error "evalX: appending a 0-dimensional tensor"
        SliceX i n d len -> case OT.shapeL r of
          n' : rest ->
            assert (n' == n) $
            evalX (OT.concatOuter [ OT.constant (i : rest) 0
                                  , r
                                  , OT.constant (len - i - n : rest) 0 ])
                  d
          [] -> error "evalX: slicing a 0-dimensional tensor"
        IndexX d ix len ->
          let rest = OT.shapeL r
          in evalX (OT.concatOuter [ OT.constant (ix : rest) 0
                                   , OT.reshape (1 : rest) r
                                   , OT.constant (len - ix - 1 : rest) 0 ])
                   d  -- TODO: optimize for Var case
        RavelFromListX ld -> do
          let lr = OTB.toList $ OT.unravel r
          mapM_ (uncurry evalX) (zip lr ld)
        ReshapeX sh _sh' d -> evalX (OT.reshape sh r) d

        From0X d -> eval0 (OT.unScalar r) d
        From1X d -> eval1 (OT.toVector r) d
        From2X d cols ->
          eval2 (MO.MatrixOuter (Just $ HM.reshape cols $ OT.toVector r)
                                Nothing Nothing)
                d
        FromSX d -> evalS (Data.Array.Convert.convert r) d

        OutlineX codeOut primalArgs dualArgs ->
          evalX r $ inlineDerivativeX codeOut primalArgs dualArgs
        DelayX d -> evalX r d
      evalS :: OS.Shape sh
            => OS.Array sh r -> DeltaS sh r -> ST s ()
      evalS !r = \case
        ZeroS -> return ()
        ScaleS k d -> evalS (liftVS2 (*) k r) d
        AddS d e -> evalS r d >> evalS r e
        VarS (DeltaId i) -> VM.modify finMapX (addToArrayS r) i

        KonstS d -> mapM_ (`eval0` d) $ OS.toList r
        AppendS (d :: DeltaS (k ': rest) r) (e :: DeltaS (l ': rest) r) ->
          evalS (OS.slice @'[ '(0, k) ] r) d
          >> evalS (OS.slice @'[ '(k, l) ] r) e
        SliceS (_ :: Proxy i) _ (d :: DeltaS (i_plus_n_plus_k ': rest) r) ->
          evalS (OS.constant @(i ': rest) 0
                 `OS.append` r
                 `OS.append` OS.constant 0)
                d
        IndexS (d :: DeltaS (ix_plus_1_plus_k ': rest) r) (_ :: Proxy ix) ->
          evalS (OS.constant @(ix ': rest) 0
                 `OS.append` OS.reshape r
                 `OS.append` OS.constant 0)
                d  -- TODO: optimize for Var case
        RavelFromListS ld -> do
          let lr = OSB.toList $ OS.unravel r
          mapM_ (uncurry evalS) (zip lr ld)
        ReshapeS d -> evalS (OS.reshape r) d

        From0S d -> eval0 (OS.unScalar r) d
        From1S d -> eval1 (OS.toVector r) d
        From2S proxyCols d ->
          eval2 (MO.MatrixOuter
                   (Just $ HM.reshape (fromInteger $ natVal proxyCols)
                         $ OS.toVector r)
                   Nothing Nothing)
                d
        FromXS d -> evalX (Data.Array.Convert.convert r) d

        OutlineS codeOut primalArgs dualArgs ->
          evalS r $ inlineDerivativeS codeOut primalArgs dualArgs
        DelayS d -> evalS r d

  eval0 dt deltaTopLevel

  let evalUnlessZero :: DeltaBinding r -> ST s ()
      evalUnlessZero (DeltaBinding0 (DeltaId i) d) = do
        r <- finMap0 `VM.read` i
        unless (r == 0) $  -- we init with exactly 0.0 so the comparison works
          eval0 r d
      evalUnlessZero (DeltaBinding1 (DeltaId i) d) = do
        r <- finMap1 `VM.read` i
        unless (V.null r) $
          eval1 r d
      evalUnlessZero (DeltaBinding2 (DeltaId i) d) = do
        r <- finMap2 `VM.read` i
        unless (MO.nullMatrixOuter r) $
          eval2 r d
      evalUnlessZero (DeltaBindingX (DeltaId i) d) = do
        r <- finMapX `VM.read` i
        unless (isTensorDummy r) $
          evalX r d
  mapM_ evalUnlessZero (deltaBindings st)
  return (finMap0, finMap1, finMap2, finMapX)

-- | Forward derivative computation via forward-evaluation of delta-expressions
-- (which is surprisingly competitive to the direct forward method,
-- until the allocation of deltas gets large enough to affect cache hits).
-- This is the directional derivative, calculated for the point,
-- at which the delta expression was computed (which is the point
-- represented by the parameters of the objective function and used
-- to compute it's dual number result) and along the direction vector(s)
-- given in the last parameter called @ds@.
derivativeFromDelta
  :: forall r. (Numeric r, Num (Vector r))
  => (CodeOut -> [r] -> [Delta0 r] -> Delta0 r)
  -> (CodeOut -> [Vector r] -> [Delta1 r] -> Delta1 r)
  -> (CodeOut -> [Matrix r] -> [Delta2 r] -> Delta2 r)
  -> (CodeOut -> [OT.Array r] -> [DeltaX r] -> DeltaX r)
  -> (forall sh. OS.Shape sh
      => CodeOut -> [OS.Array sh r] -> [DeltaS sh r]
      -> DeltaS sh r)
  -> DeltaState r -> Delta0 r -> Domains r -> r
derivativeFromDelta inlineDerivative0 inlineDerivative1 inlineDerivative2
                    inlineDerivativeX inlineDerivativeS
                    st deltaTopLevel
                    _ds@(params0Init, params1Init, params2Init, paramsXInit) =
  let eval0 :: Domains r -> Delta0 r -> r
      eval0 parameters@(params0, _, _, _) = \case
        Zero0 -> 0
        Scale0 k d -> k * eval0 parameters d
        Add0 d e -> eval0 parameters d + eval0 parameters e
        Var0 (DeltaId i) -> params0 V.! i

        SumElements0 vd _n -> HM.sumElements $ eval1 parameters vd
        Index0 d ix _k -> eval1 parameters d V.! ix

        Dot0 vr vd -> vr <.> eval1 parameters vd

        FromX0 d -> OT.unScalar $ evalX parameters d
        FromS0 d -> OS.unScalar $ evalS parameters d

        Outline0 codeOut primalArgs dualArgs ->
          eval0 parameters $ inlineDerivative0 codeOut primalArgs dualArgs
        Delay0 d -> eval0 parameters d
      eval1 :: Domains r -> Delta1 r -> Vector r
      eval1 parameters@(_, params1, _, _) = \case
        Zero1 -> 0
        Scale1 k d -> k * eval1 parameters d
        Add1 d e -> eval1 parameters d + eval1 parameters e
        Var1 (DeltaId i) -> params1 V.! i

        Seq1 lsd -> V.convert $ V.map (eval0 parameters) lsd
        Konst1 d n -> HM.konst (eval0 parameters d) n
        Append1 d _k e -> eval1 parameters d V.++ eval1 parameters e
        Slice1 i n d _len -> V.slice i n $ eval1 parameters d
        SumRows1 dm _cols ->
          V.fromList $ map HM.sumElements $ HM.toRows $ eval2 parameters dm
        SumColumns1 dm _rows ->
          V.fromList $ map HM.sumElements $ HM.toColumns $ eval2 parameters dm

        M_VD1 m dRow -> m #> eval1 parameters dRow
        MD_V1 md row -> eval2 parameters md #> row

        FromX1 d -> OT.toVector $ evalX parameters d
        FromS1 d -> OS.toVector $ evalS parameters d

        Reverse1 d -> V.reverse $ eval1 parameters d
        Flatten1 _rows _cols d -> HM.flatten $ eval2 parameters d
        FlattenX1 _sh d -> OT.toVector $ evalX parameters d
        FlattenS1 d -> OS.toVector $ evalS parameters d

        Outline1 codeOut primalArgs dualArgs ->
          eval1 parameters $ inlineDerivative1 codeOut primalArgs dualArgs
        Delay1 d -> eval1 parameters d
      eval2 :: Domains r -> Delta2 r -> Matrix r
      eval2 parameters@( _, _, params2, _) = \case
        Zero2 -> 0
        Scale2 k d -> k * eval2 parameters d
        Add2 d e -> eval2 parameters d + eval2 parameters e
        Var2 (DeltaId i) -> params2 V.! i

        FromRows2 lvd ->
          HM.fromRows $ map (eval1 parameters) $ V.toList lvd
        FromColumns2 lvd ->
          HM.fromColumns $ map (eval1 parameters) $ V.toList lvd
        Konst2 d sz -> HM.konst (eval0 parameters d) sz
        Transpose2 md -> HM.tr' $ eval2 parameters md
        M_MD2 m md -> m HM.<> eval2 parameters md
        MD_M2 md m -> eval2 parameters md HM.<> m
        RowAppend2 d _k e -> eval2 parameters d HM.=== eval2 parameters e
        ColumnAppend2 d _k e -> eval2 parameters d HM.||| eval2 parameters e
        RowSlice2 i n d _rows ->
          HM.takeRows n $ HM.dropRows i $ eval2 parameters d
        ColumnSlice2 i n d _cols ->
          HM.takeColumns n $ HM.dropColumns i $ eval2 parameters d

        AsRow2 dRow -> HM.asRow $ eval1 parameters dRow  -- TODO: risky
        AsColumn2 dCol -> HM.asColumn $ eval1 parameters dCol  -- TODO: risky

        FromX2 d ->
          let t = evalX parameters d
          in case OT.shapeL t of
            [_rows, cols] -> HM.reshape cols $ OT.toVector t
            _ -> error "eval2: wrong tensor dimensions"
        FromS2 d ->
          let t = evalS parameters d
          in case OS.shapeL t of
            [_rows, cols] -> HM.reshape cols $ OS.toVector t
            _ -> error "eval2: wrong tensor dimensions"

        Flipud2 d -> HM.flipud $ eval2 parameters d
        Fliprl2 d -> HM.fliprl $ eval2 parameters d
        Reshape2 cols d -> HM.reshape cols $ eval1 parameters d
        Conv2 m md -> HM.conv2 m $ eval2 parameters md

        Outline2 codeOut primalArgs dualArgs ->
          eval2 parameters $ inlineDerivative2 codeOut primalArgs dualArgs
        Delay2 d -> eval2 parameters d
      evalX :: Domains r -> DeltaX r -> OT.Array r
      evalX parameters@( _, _, _, paramsX) = \case
        ZeroX -> 0
        ScaleX k d -> k * evalX parameters d
        AddX d e -> evalX parameters d + evalX parameters e
        VarX (DeltaId i) -> paramsX V.! i

        KonstX d sz -> OT.constant sz $ eval0 parameters d
        AppendX d _k e -> evalX parameters d `OT.append` evalX parameters e
        SliceX i n d _len -> OT.slice [(i, n)] $ evalX parameters d
        IndexX d ix _len -> OT.index (evalX parameters d) ix
        RavelFromListX ld ->
          let la = map (evalX parameters) ld
              sh = case la of
                a : _ -> length la : OT.shapeL a
                [] -> []
          in OT.ravel $ OTB.fromList sh la
        ReshapeX _sh sh' d -> OT.reshape sh' $ evalX parameters d

        From0X d -> OT.scalar $ eval0 parameters d
        From1X d -> let v = eval1 parameters d
                    in OT.fromVector [V.length v] v
        From2X d cols -> let l = eval2 parameters d
                         in OT.fromVector [HM.rows l, cols] $ HM.flatten l
        FromSX d -> Data.Array.Convert.convert $ evalS parameters d

        OutlineX codeOut primalArgs dualArgs ->
          evalX parameters $ inlineDerivativeX codeOut primalArgs dualArgs
        DelayX d -> evalX parameters d
      evalS :: OS.Shape sh => Domains r -> DeltaS sh r -> OS.Array sh r
      evalS parameters@( _, _, _, paramsX) = \case
        ZeroS -> 0
        ScaleS k d -> k * evalS parameters d
        AddS d e -> evalS parameters d + evalS parameters e
        VarS (DeltaId i) -> Data.Array.Convert.convert $ paramsX V.! i

        KonstS d -> OS.constant $ eval0 parameters d
        AppendS d e -> evalS parameters d `OS.append` evalS parameters e
        SliceS (_ :: Proxy i) (_ :: Proxy n) d ->
          OS.slice @'[ '(i, n) ] $ evalS parameters d
        IndexS d proxyIx ->
          OS.index (evalS parameters d) (fromInteger $ natVal proxyIx)
        RavelFromListS ld ->
          let la = map (evalS parameters) ld
          in OS.ravel $ OSB.fromList la
        ReshapeS d -> OS.reshape $ evalS parameters d

        From0S d -> OS.scalar $ eval0 parameters d
        From1S d -> OS.fromVector $ eval1 parameters d
        From2S _ d -> OS.fromVector $ HM.flatten $ eval2 parameters d
        FromXS d -> Data.Array.Convert.convert $ evalX parameters d

        OutlineS codeOut primalArgs dualArgs ->
          evalS parameters $ inlineDerivativeS codeOut primalArgs dualArgs
        DelayS d -> evalS parameters d
      evalUnlessZero :: Domains r -> DeltaBinding r -> Domains r
      evalUnlessZero parameters@(!params0, !params1, !params2, !paramsX) = \case
        DeltaBinding0 (DeltaId i) d ->
          let v = eval0 parameters d
          in (params0 V.// [(i, v)], params1, params2, paramsX)
        DeltaBinding1 (DeltaId i) d ->
          let v = eval1 parameters d
          in (params0, params1 V.// [(i, v)], params2, paramsX)
        DeltaBinding2 (DeltaId i) d ->
          let v = eval2 parameters d
          in (params0, params1, params2 V.// [(i, v)], paramsX)
        DeltaBindingX (DeltaId i) d ->
          let v = evalX parameters d
          in (params0, params1, params2, paramsX V.// [(i, v)])
      parameters1 = runST $ do
        (finMap0, finMap1, outerFinMap2, finMapX) <- initializeFinMaps st
        -- We use normal hmatrix matrices rather than the sparse replacement.
        finMap2 <- VM.replicate (VM.length outerFinMap2) (HM.fromRows [])
        -- TODO: the following coredumps without the @VM.take@; it's a shame
        -- there's no copying of a smaller vector into a larger one in the API.
        -- Perhaps use https://hackage.haskell.org/package/base-4.16.0.0/docs/Foreign-Marshal-Array.html#v:copyArray?
        V.unsafeCopy (VM.take (V.length params0Init) finMap0) params0Init
        V.unsafeCopy (VM.take (V.length params1Init) finMap1) params1Init
        V.unsafeCopy (VM.take (V.length params2Init) finMap2) params2Init
        V.unsafeCopy (VM.take (V.length paramsXInit) finMapX) paramsXInit
        v0 <- V.unsafeFreeze finMap0
        v1 <- V.unsafeFreeze finMap1
        v2 <- V.unsafeFreeze finMap2
        vX <- V.unsafeFreeze finMapX
        return (v0, v1, v2, vX)
      parametersB = foldl' evalUnlessZero parameters1
                           (reverse $ deltaBindings st)
  in eval0 parametersB deltaTopLevel

-- | This is yet another semantics of delta-expressions and their
-- bindings --- by pretty-printing as texts.
ppBindings :: (Show r, Numeric r) => Bool -> DeltaState r -> Delta0 r -> String
ppBindings reversed st deltaTopLevel =
  let pp = if reversed
           then foldl' (\ !l b -> l ++ ppBinding "where" b)
                       ["COMPUTE " ++ ppShow deltaTopLevel ++ "\n"]
           else foldl' (\ !l b -> ppBinding "let" b ++ l)
                       ["in " ++ ppShow deltaTopLevel ++ "\n"]
  in concat $ pp $ deltaBindings st

ppBinding :: (Show r, Numeric r) => String -> DeltaBinding r -> [String]
ppBinding prefix = \case
  DeltaBinding0 (DeltaId i) d ->
    [prefix ++ "0 DeltaId_", show i, " = ", ppShow d, "\n"]
  DeltaBinding1 (DeltaId i) d ->
    [prefix ++ "1 DeltaId_", show i, " = ", ppShow d, "\n"]
  DeltaBinding2 (DeltaId i) d ->
    [prefix ++ "2 DeltaId_", show i, " = ", ppShow d, "\n"]
  DeltaBindingX (DeltaId i) d ->
    [prefix ++ "X DeltaId_", show i, " = ", ppShow d, "\n"]

bindInState0 :: Delta0 r -> DeltaState r -> (DeltaState r, DeltaId r)
{-# INLINE bindInState0 #-}
bindInState0 u' st =
  let dId = deltaCounter0 st
      !binding = DeltaBinding0 dId u'
  in ( st { deltaCounter0 = succDeltaId dId
          , deltaBindings = binding : deltaBindings st
          }
     , dId )

bindInState1 :: Delta1 r -> DeltaState r -> (DeltaState r, DeltaId (Vector r))
{-# INLINE bindInState1 #-}
bindInState1 u' st =
  let dId = deltaCounter1 st
      !binding = DeltaBinding1 dId u'
  in ( st { deltaCounter1 = succDeltaId dId
          , deltaBindings = binding : deltaBindings st
          }
     , dId )

bindInState2 :: Delta2 r -> DeltaState r -> (DeltaState r, DeltaId (Matrix r))
{-# INLINE bindInState2 #-}
bindInState2 u' st =
  let dId = deltaCounter2 st
      !binding = DeltaBinding2 dId u'
  in ( st { deltaCounter2 = succDeltaId dId
          , deltaBindings = binding : deltaBindings st
          }
     , dId )

bindInStateX :: DeltaX r -> DeltaState r -> (DeltaState r, DeltaId (OT.Array r))
{-# INLINE bindInStateX #-}
bindInStateX u' st =
  let dId = deltaCounterX st
      !binding = DeltaBindingX dId u'
  in ( st { deltaCounterX = succDeltaId dId
          , deltaBindings = binding : deltaBindings st
          }
     , dId )

{- Note [SumElements0]
~~~~~~~~~~~~~~~~~~~~~~

The second argument of SumElements0 is the length of the vector
to be summed. Given that we sum a delta-expression representing
a vector, we can't call Vector.length on it, so the length needs
to be recorded in the constructor. Alternatively, it could be
recorded in the Delta1 argument to SumElements0. This is what
shaped tensors do at the type level, so for DeltaS the argument
would not be needed.

Sum of vector elements can be implemented using a delta-expression
primitive SumElements0 as well as without this primitive, referring
only to the primitive Index0:

https://github.com/Mikolaj/horde-ad/blob/d069a45773ed849913b5ebd0345153072f304fd9/src/HordeAd/Core/DualNumber.hs#L125-L143

which is confirmed by tests to be equivalent in three different
implementations:

https://github.com/Mikolaj/horde-ad/blob/d069a45773ed849913b5ebd0345153072f304fd9/test/TestSingleGradient.hs#L116-L128

and proved to be prohibitively slow in the two implementations
that don't use the SumElements0 primitive in benchmarks (despite
an ingenious optimization of the common case of Index0 applied to a variable):

https://github.com/Mikolaj/horde-ad/blob/d069a45773ed849913b5ebd0345153072f304fd9/bench/BenchProdTools.hs#L178-L193
-}

dummyTensor :: Numeric r => OT.Array r
dummyTensor =  -- an inconsistent tensor array
  Data.Array.Internal.DynamicS.A
  $ Data.Array.Internal.DynamicG.A []
  $ Data.Array.Internal.T [] (-1) V.empty

isTensorDummy :: OT.Array r -> Bool
isTensorDummy (Data.Array.Internal.DynamicS.A
                 (Data.Array.Internal.DynamicG.A _
                    (Data.Array.Internal.T _ (-1) _))) = True
isTensorDummy _ = False

data CodeOut =
    PlusOut | MinusOut | TimesOut | NegateOut | AbsOut | SignumOut
  | DivideOut | RecipOut
  | ExpOut | LogOut | SqrtOut | PowerOut | LogBaseOut
  | SinOut | CosOut | TanOut | AsinOut | AcosOut | AtanOut
  | SinhOut | CoshOut | TanhOut | AsinhOut | AcoshOut | AtanhOut
  | Atan2Out
  deriving Show


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE BangPatterns #-}
module SDLib where

import qualified SDL
import qualified Graphics.Rendering.OpenGL.GL as GL
import           Graphics.Rendering.OpenGL.GL (($=))
import           Control.Monad
import qualified Numeric.LinearAlgebra as L
import           Numeric.LinearAlgebra.Data (flatten)
import qualified Numeric.LinearAlgebra.Data as LD
import           Control.Lens      hiding (indices)
import           Debug.Trace
import           Control.Arrow
import qualified Data.Vector as V
import Data.Vector.Strategies
import Control.Concurrent (forkOS, killThread)
import Control.Concurrent.MVar

import           Utils (step, projectionMatrix,normalize,clampSpeed,norm)

import           OBJReader
import           Data.Shaders (loadShaders)

import           Data.Mesh    ( loadTexture
                              , createMesh
                              , flatQuad
                              , quaternion)
import           Data.GameObject
import           Physics.Collision
import           Data.Scene
import           Data.PhysicsData (acc,speed)
import           Data.SkyBox
import           Physics.Mechanics
import           Tasks.Swarm
import           EventHandler
import           Utils.Quaternions

frameTime = round $ 1000/60

-- drawThread ::

debugPrint (GL.DebugMessage _ GL.DebugTypeError _ _ err) = putStrLn err
debugPrint _ = return ()

setup = do
  SDL.swapInterval $= SDL.SynchronizedUpdates
  GL.depthFunc $= Just GL.Less
  -- GL.cullFace $= Just GL.Front
  -- GL.debugMessageCallback $= Just debugPrint
  -- GL.debugOutput $= GL.Enabled

drawThread mv window = forkOS $ do
  context <- SDL.glCreateContext window
  setup
  let uniforms = ["camera","proj","pos","rot","texSampler","normSampler"] :: [String]
  shad <- loadShaders "shaders/vert.glsl" "shaders/frag.glsl" uniforms
  shadNoTex <- loadShaders "shaders/vert.glsl" "shaders/frag_no_text.glsl" uniforms
  ultraMesh <- readOBJ "media/sphere.obj"
  ultraVeh  <- readOBJ "media/cube.obj"
  vehPrim   <- readBoxFromObj "media/cube.obj"
  let miniVeh = ultraMesh
  meh       <- createMesh ultraVeh shadNoTex Nothing Nothing
  veh       <- createMesh ultraMesh shadNoTex Nothing Nothing
  -- crvMesh   <- createMesh curveVertices shadNoTex Nothing
  tex       <- loadTexture "media/bricks_dif.png"
  normalTex <- loadTexture "media/bricks_norm.png"
  sbox      <- loadSkyBox
  qd        <- flatQuad shad tex normalTex
  let quads =
        V.fromList .
        map (\(x,z) -> GameObject qd [x, -2.5, z] Nothing [] (Plane [0,1,0] (-2.5))) $
        [(x,z) | x <- [-30,-26..30], z <- [-30,-26..30]]
      camPos = [ 0, 10,20]
      vehOb pos = initializePhysics (1/10) $
        GameObject meh pos Nothing [] vehPrim
      sphOb pos = initializePhysics (1/10) .
                  GameObject veh pos Nothing [] $
                  Sphere 1 (LD.ident 4)
      dim = 1
      sep = 3
      vehs = -- V.fromList [vehOb [0.5,1,-5]]
        -- V.fromList []
                    -- initializePhysics (1/10) .
                    -- GameObject veh [0.5,1,-5] Nothing [] $
                    -- Sphere 1 (LD.ident 4)]
        V.fromList [sphOb [0,2,sep*z] | z <- [0..dim]]
      mehOb =
        -- initializePhysics (1/100) .
        --             GameObject veh [0, 0, 10] Nothing [] $
        --             Sphere 1 (LD.ident 4)
        initializePhysics (1/10) $
        GameObject meh [0,0,10] Nothing [] vehPrim
      -- curve = GameObject crvMesh [0,0,0] Nothing []
      scene = Scene camPos (Objects mehOb vehs quads) (flatten . L.tr $ projectionMatrix) (Just sbox)
  putMVar mv scene
  drawLoop
  where drawLoop = takeMVar mv >>= draw window >> drawLoop

mainLoop :: MVar Scene -> SDL.Window -> Scene -> IO ()
mainLoop mv window scene = do
  begTime <- SDL.ticks
  events <- SDL.pollEvents
  kState <- SDL.getKeyboardState
  let qpressed   = any handleQuit events
      -- newTime    = time + (0.1*step)
      -- newTime2   = if newTime > 1 then 0 else newTime
      handleKeys = foldr ((.).keyHandler) id $ filter kState keys
      loc        = scene^.objects.player.location
      -- updateNPCS = applyLinear
      --            . over (physics._Just.speed) (clamp 20)
      --            . addGravityForce loc
      !newScene  = -- set objects.player.location) (countPosition newTime2)
                 adjustCamera .
                 over objects resolveObjectCollisions .
                 traverseObjects (floorReaction . applyForce . friction) .
                 traverseObjects gravityF $
                 -- . over (objects.npc) (computeForces updateNPCS)
                 handleKeys scene
  -- print $ newScene^.campos
  endTime <- SDL.ticks
  when ((endTime - begTime) < frameTime)
    $ SDL.delay (frameTime - (endTime - begTime))
  print $ endTime - begTime
  print $ scene^.objects.player
  let cp = prepareCPrim $ scene^.objects.player
      p  = V.toList $ _cpoints cp
      d  = [norm (a- b) | a <- p, b <- p]
  print $ p
  print d
  putMVar mv newScene
  unless qpressed $ mainLoop mv window newScene

sdlMain :: IO ()
sdlMain = do
  SDL.initializeAll
  window <- SDL.createWindow "yolo" SDL.defaultWindow {
      SDL.windowOpenGL = Just SDL.defaultOpenGL
    , SDL.windowInitialSize = SDL.V2 1366 768 }
  mv <- newEmptyMVar
  dthr <- drawThread mv window
  scene <- takeMVar mv
  mainLoop mv window scene
  killThread dthr
  GL.finish
  -- SDL.glDeleteContext context
  SDL.quit


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE BangPatterns          #-}
-- |
-- Module     : Data.Matrix.Generic.Mutable
-- Copyright  : Copyright (c) 2012 Aleksey Khudyakov <alexey.skladnoy@gmail.com>
-- License    : BSD3
-- Maintainer : Aleksey Khudyakov <alexey.skladnoy@gmail.com>
-- Stability  : experimental
--
-- Interface for generic mutable matrices. For matrix transposition
-- and conjugate transposition newtype wrappers are used.
module Data.Matrix.Generic.Mutable (
    -- * Type class
    IsMMatrix(..)
    -- * Accesors
  , rows
  , cols
  , shape
    -- * Reading and writing
  , read
  , write
  , unsafeRead
  , unsafeWrite
  , cloneShape
  , clone
    -- * Newtype wrappers
  , TransposedM(..)
  , ConjugatedM(..)
  ) where

import Control.Monad.Primitive
import Data.Complex             (Complex,conjugate)
import Prelude hiding (read)

----------------------------------------------------------------
-- Type class
----------------------------------------------------------------

-- | Type class of mutable matrices. Since there's many way to lay
--   matrix in memory there isn't many operation which work for all of
--   them.
--
--   Indexing uses following form: @(row,column)@
class IsMMatrix mat a where
  -- | Number of rows.
  basicRows :: mat s a -> Int
  -- | Number of columns.
  basicCols :: mat s a -> Int
  -- | Whether index could be mutated. E.g. not every element of
  --   banded matrix could be written to. Function need not to check
  --   that indices are inside of the matrix and free to return anything if
  --   they are outside.
  basicIsIndexMutable :: mat s a -> (Int,Int) -> Bool
  -- | Retrive element at given index. It should alway succed if row
  --   and column numbers are in range.
  basicUnsafeRead :: PrimMonad m => mat (PrimState m) a -> (Int,Int) -> m a
  -- | Write element at given index.
  basicUnsafeWrite :: PrimMonad m => mat (PrimState m) a -> (Int,Int) -> a -> m ()
  -- | Create matrix with same dimensions. Elements' values are undefined
  --
  --   Unlike vector which have only length there are many kinds of
  --   matrices. it's not possible to write such function from simpler ones.
  basicCloneShape :: PrimMonad m => mat (PrimState m) a -> m (mat (PrimState m) a)
  -- | Create copy of matrix.
  basicClone :: PrimMonad m => mat (PrimState m) a -> m (mat (PrimState m) a)

----------------------------------------------------------------
-- Accesors
----------------------------------------------------------------

-- | Number of rows.
rows :: IsMMatrix mat a => mat s a -> Int
{-# INLINE rows #-}
rows = basicRows

-- | Number of columns.
cols :: IsMMatrix mat a => mat s a -> Int
{-# INLINE cols #-}
cols = basicCols

-- | Shape of the matrix.
shape :: IsMMatrix mat a => mat s a -> (Int,Int)
{-# INLINE shape #-}
shape m = (rows m, cols m)

-- | Whether index could be mutated. E.g. not every element of banded
--   matrix could be written to. Function need not to check that
--   indices are inside of the matrix and free to return anything if
--   they are outside
isIndexMutable :: IsMMatrix mat a => mat s a -> (Int, Int) -> Bool
{-# INLINE isIndexMutable #-}
isIndexMutable = basicIsIndexMutable

----------------------------------------------------------------
-- Read/write
----------------------------------------------------------------

-- | Read matrix element at given index. It should alway succed if row
--   and column numbers are in range.
read :: (PrimMonad m, IsMMatrix mat a) => mat (PrimState m) a -> (Int,Int) -> m a
{-# INLINE read #-}
read m i@(!r,!c)
  | r < 0 || r >= rows m = error "Numeric.BLAS.Matrix.Mutable.read: row out of bounds"
  | c < 0 || r >= cols m = error "Numeric.BLAS.Matrix.Mutable.read: column out of bounds"
  | otherwise            = unsafeRead m i

-- | Write element to matrix at given index. It could be impossible to
--   change values of some elements.
write :: (PrimMonad m, IsMMatrix mat a) => mat (PrimState m) a -> (Int,Int) -> a -> m ()
{-# INLINE write #-}
write m i@(!r,!c)
  | r < 0 || r >= rows m = error "Numeric.BLAS.Matrix.Mutable.write: row out of bounds"
  | c < 0 || r >= cols m = error "Numeric.BLAS.Matrix.Mutable.write: column out of bounds"
  | isIndexMutable m i   = basicUnsafeWrite m i
  | otherwise            = error "Numeric.BLAS.Matrix.Mutable.write: index is not mutable"

-- | Read element from matrix without range checking
unsafeRead :: (PrimMonad m, IsMMatrix mat a) => mat (PrimState m) a -> (Int, Int) -> m a
{-# INLINE unsafeRead #-}
unsafeRead = basicUnsafeRead

-- | Write element to matrix without range checking.
unsafeWrite :: (PrimMonad m, IsMMatrix mat a) => mat (PrimState m) a -> (Int, Int) -> a -> m ()
{-# INLINE unsafeWrite #-}
unsafeWrite = basicUnsafeWrite

-- | Create matrix with same dimensions. Elements' values are undefined
--
--   Unlike vector which have only length there are many kinds of
--   matrices. it's not possible to write such function from simpler ones.
cloneShape :: (PrimMonad m, IsMMatrix mat a) => mat (PrimState m) a -> m (mat (PrimState m) a)
{-# INLINE cloneShape #-}
cloneShape = basicCloneShape

-- | Create copy of matrix.
clone :: (PrimMonad m, IsMMatrix mat a) => mat (PrimState m) a -> m (mat (PrimState m) a)
{-# INLINE clone #-}
clone = basicClone

----------------------------------------------------------------
-- Newtypes
----------------------------------------------------------------

-- | Transposed matrix
newtype TransposedM mat s a = TransposedM { unTranspose :: mat s a }

instance IsMMatrix mat a => IsMMatrix (TransposedM mat) a where
  basicRows (TransposedM m) = cols m
  {-# INLINE basicRows #-}
  basicCols (TransposedM m) = rows m
  {-# INLINE basicCols #-}
  basicIsIndexMutable (TransposedM m) (i,j)   = isIndexMutable m (j,i)
  {-# INLINE basicIsIndexMutable #-}
  basicUnsafeRead     (TransposedM m) (i,j)   = unsafeRead     m (j,i)
  {-# INLINE basicUnsafeRead #-}
  basicUnsafeWrite    (TransposedM m) (i,j) x = unsafeWrite    m (j,i) x
  {-# INLINE basicUnsafeWrite #-}
  basicCloneShape     (TransposedM m) = do { r <- basicCloneShape m; return $ TransposedM r }
  {-# INLINE basicCloneShape #-}
  basicClone          (TransposedM m) = do { r <- basicCloneShape m; return $ TransposedM r }
  {-# INLINE basicClone #-}

-- | Conjugate-transposed matrix
newtype ConjugatedM mat s a = ConjugatedM { unConjugate :: mat s a }

instance (IsMMatrix mat (Complex a), RealFloat a) => IsMMatrix (ConjugatedM mat) (Complex a) where
  basicRows (ConjugatedM m) = basicCols m
  {-# INLINE basicRows #-}
  basicCols (ConjugatedM m) = rows m
  {-# INLINE basicCols #-}
  basicIsIndexMutable (ConjugatedM m) (i,j)   = isIndexMutable m (j,i)
  {-# INLINE basicIsIndexMutable #-}
  basicUnsafeRead (ConjugatedM m) (i,j)
    | j >= i    = unsafeRead m (j,i)
    | otherwise = do { x <- unsafeRead m (j,i); return $! conjugate x }
  {-# INLINE basicUnsafeRead #-}
  basicUnsafeWrite (ConjugatedM m) (i,j) x
    | j >= i    = unsafeWrite m (j,i) x
    | otherwise = unsafeWrite m (j,i) $! conjugate x
  {-# INLINE basicUnsafeWrite #-}
  basicCloneShape     (ConjugatedM m) = do { r <- basicCloneShape m; return $ ConjugatedM r }
  {-# INLINE basicCloneShape #-}
  basicClone          (ConjugatedM m) = do { r <- basicCloneShape m; return $ ConjugatedM r }
  {-# INLINE basicClone #-}


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_HADDOCK show-extensions #-}

module NQS.Internal.Types
  ( -- * Numeric types
    ℂ
  , ℝ

  -- , Estimate(..)
  -- , NormalError(..)
  -- , EnergyMeasurement(..)
  , RealOf
  -- , SRMeasurement(..)
  {-
  , withRbm
  , withRbmPure
  , withMRbm
  , withMRbmPure
  -}
    -- * Vectors and matrices
  , DenseVector
  , DenseMatrix
  , MDenseVector
  , MDenseMatrix
  , slice
  , withForeignPtrPrim
  , newVectorAligned
  , newDenseVector
  , newDenseMatrix

  , unsafeRow
  , unsafeColumn
  , unsafeWriteVector

  , Mutable(..)
  , Variant(..)
  , Orientation(..)
  , Transpose(..)
  , MatUpLo(..)
  , Sing(..)

  , CheckValid(..)
  , isValidVector
  , isValidMatrix
  , badVectorInfo
  , HasBuffer(..)
  , HasStride(..)
  , HasDim(..)

  , orientationOf
  , ToNative(..)

  , DenseWorkspace(..)
  , MCConfig(..)
  , defaultMCConfig
  , CGConfig(..)
  , SRConfig(..)

  , HasMc(..)
  , HasCg(..)
  , HasSteps(..)
  , HasRuns(..)
  , HasRate(..)
  , HasRestarts(..)
  , HasRegulariser(..)
  , HasMaxIter(..)
  , HasMagnetisation(..)

  , mapVectorM
  , mapMatrixM
  , zipWithVectorM
  , zipWithMatrixM
  -- , asTuple
  -- , fromTuple
  , FreezeThaw(..)

  , HasMean(..)
  , HasVar(..)
  -- , HasEnergy(..)
  -- , HasForce(..)
  -- , HasDerivatives(..)
  ) where

import Foreign.Storable
import Foreign.ForeignPtr

import Control.DeepSeq
import Control.Exception (assert)
import Control.Monad ((>=>), unless)
import Control.Monad.Primitive
import Control.Monad.ST.Strict
import Data.Complex
import Debug.Trace
import GHC.Generics (Generic)
import GHC.ForeignPtr (mallocPlainForeignPtrAlignedBytes)

import Data.Semigroup ((<>))
import Lens.Micro
import Lens.Micro.TH

import System.IO.Unsafe
import Data.Coerce
import Data.Bits

import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.ForeignPtr
import Foreign.ForeignPtr.Unsafe

import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as MV
import Data.Vector.Storable (Vector, MVector(..))

import qualified Data.Vector.Generic as GV
import qualified Data.Vector as Boxed
import Data.Aeson
import qualified Data.Aeson.Types as Aeson

import Data.Singletons

import qualified Numerical.HBLAS.MatrixTypes   as HBLAS
import           Numerical.HBLAS.MatrixTypes    ( Variant(..)
                                                , Orientation(..)
                                                , Transpose(..)
                                                , MatUpLo(..)
                                                )
import           Numerical.HBLAS.BLAS.FFI       ( CBLAS_UPLOT(..) )

import Data.Kind
import Data.Proxy

deriving instance Generic Variant
deriving instance NFData Variant

deriving instance Generic Orientation
deriving instance NFData Orientation

class ToNative a b where
  encode :: a -> b

data family Mutable (v :: * -> *) :: * -> * -> *

-- | After some prototyping, I've concluded (possibly incorrectly) that a
-- single-precision 'Float' is enough for our purposes.
--
-- [This answer](https://stackoverflow.com/a/40538415) also mentions that 32-bit
-- floats should be enought for most neural network applications.
type ℝ = Float
type ℂ = Complex ℝ

data instance Sing (orient :: Orientation) where
  SRow :: Sing 'Row
  SColumn :: Sing 'Column

instance SingI 'Row where sing = SRow
instance SingI 'Column where sing = SColumn

instance SingKind Orientation where
  type Demote Orientation = Orientation
  fromSing SRow = Row
  fromSing SColumn = Column
  toSing x = case x of
    Row -> SomeSing SRow
    Column -> SomeSing SColumn

data DenseVector (variant :: Variant) (a :: *) = DenseVector
  { _denseVectorDim    :: {-# UNPACK #-}!Int
  , _denseVectorStride :: {-# UNPACK #-}!Int
  , _denseVectorBuffer :: {-# UNPACK #-}!(Vector a)
  } deriving (Show, Generic, NFData)

data DenseMatrix (orientation :: Orientation) (a :: *) = DenseMatrix
  { _denseMatrixRows   :: {-# UNPACK #-}!Int
  , _denseMatrixCols   :: {-# UNPACK #-}!Int
  , _denseMatrixStride :: {-# UNPACK #-}!Int
  , _denseMatrixBuffer :: {-# UNPACK #-}!(Vector a)
  } deriving (Show, Generic, NFData)

data instance Mutable (DenseVector variant) s a = MDenseVector
  { _mDenseVectorDim :: {-# UNPACK #-}!Int
  , _mDenseVectorStride :: {-# UNPACK #-}!Int
  , _mDenseVectorBuffer :: {-# UNPACK #-}!(MVector s a)
  } deriving (Generic, NFData)

type MDenseVector variant = Mutable (DenseVector variant)

data instance Mutable (DenseMatrix orientation) s a = MDenseMatrix
  { _mDenseMatrixRows :: {-# UNPACK #-}!Int
  , _mDenseMatrixCols :: {-# UNPACK #-}!Int
  , _mDenseMatrixStride :: {-# UNPACK #-}!Int
  , _mDenseMatrixBuffer :: {-# UNPACK #-}!(MVector s a)
  } deriving (Generic, NFData)

type MDenseMatrix orientation = Mutable (DenseMatrix orientation)

makeFields ''DenseVector
makeFields ''DenseMatrix

unsafeIndexVector :: Storable a => DenseVector 'Direct a -> Int -> a
{-# INLINE unsafeIndexVector #-}
unsafeIndexVector !(DenseVector _ stride buff) !i = V.unsafeIndex buff (i * stride)

unsafeIndexMatrix :: forall a orient. (Storable a, SingI orient)
                  => DenseMatrix orient a -> (Int, Int) -> a
{-# INLINE unsafeIndexMatrix #-}
unsafeIndexMatrix !(DenseMatrix _ _ stride buff) !(r, c) = V.unsafeIndex buff i
  where i = case (sing :: Sing orient) of
              SRow -> r * stride + c
              SColumn -> r + c * stride

unsafeReadVector :: (Storable a, PrimMonad m)
                 => MDenseVector 'Direct (PrimState m) a -> Int -> m a
{-# INLINE unsafeReadVector #-}
unsafeReadVector !(MDenseVector _ stride buff) !i = MV.unsafeRead buff (i * stride)

unsafeReadMatrix :: forall a m orient. (Storable a, PrimMonad m, SingI orient)
                 => MDenseMatrix orient (PrimState m) a -> (Int, Int) -> m a
{-# INLINE unsafeReadMatrix #-}
unsafeReadMatrix !(MDenseMatrix _ _ stride buff) !(r, c) = MV.unsafeRead buff i
  where i = case (sing :: Sing orient) of
              SRow -> r * stride + c
              SColumn -> r + c * stride

unsafeWriteVector :: (Storable a, PrimMonad m)
                  => MDenseVector 'Direct (PrimState m) a -> Int -> a -> m ()
{-# INLINE unsafeWriteVector #-}
unsafeWriteVector !(MDenseVector _ stride buff) !i !x = MV.unsafeWrite buff (i * stride) x

unsafeWriteMatrix :: forall a m orient. (Storable a, PrimMonad m, SingI orient)
                  => MDenseMatrix orient (PrimState m) a -> (Int, Int) -> a -> m ()
{-# INLINE unsafeWriteMatrix #-}
unsafeWriteMatrix !(MDenseMatrix _ _ stride buff) !(r, c) !x = MV.unsafeWrite buff i x
  where i = case (sing :: Sing orient) of
              SRow -> r * stride + c
              SColumn -> r + c * stride

-- | Returns a row of a matrix.
unsafeRow
  :: forall orient s a
   . (Storable a, SingI orient)
  => Int
  -> MDenseMatrix orient s a
  -> MDenseVector 'Direct s a
{-# INLINE unsafeRow #-}
unsafeRow !i !(MDenseMatrix rows cols stride buff) = assert False $
  assert (i < rows) $ case (sing :: Sing orient) of
    SRow    -> MDenseVector cols 1 (MV.slice (i * stride) cols buff)
    SColumn -> MDenseVector cols stride (MV.slice i (cols * stride) buff)

-- | Returns a column of a matrix.
unsafeColumn
  :: forall orient s a
   . (Storable a, SingI orient)
  => Int
  -> MDenseMatrix orient s a
  -> MDenseVector 'Direct s a
{-# INLINE unsafeColumn #-}
unsafeColumn !i !(MDenseMatrix rows cols stride buff) =
  assert (i < cols) $ case (sing :: Sing orient) of
    SRow    -> if rows == 1
                 then MDenseVector 1 1 (MV.slice i 1 buff)
                 else MDenseVector rows stride (MV.slice i ((rows - 1)* stride + 1) buff)
    SColumn -> MDenseVector rows 1 (MV.slice (i * stride) rows buff)

data DenseWorkspace s a = DenseWorkspace
  { _denseWorkspaceForce       :: !(MDenseVector 'Direct s a) -- ^ Force
  , _denseWorkspaceDerivatives :: !(MDenseMatrix 'Row s a)    -- ^ Derivatives
  , _denseWorkspaceDelta       :: !(MDenseVector 'Direct s a) -- ^ Old delta
  } deriving (Generic, NFData)

-- | Configuration for Monte-Carlo sampling.
data MCConfig =
  MCConfig
    { _mCConfigSteps :: {-# UNPACK #-}!(Int, Int, Int)
        -- ^ A range of steps for a single run. It is very similar to Python's
        -- [range](). @(low, high, step)@ corresponds to [low, low + step, low +
        -- 2 * step, ..., high).
        --
        -- /Note:/ only positive @step@s are supported.
    , _mCConfigThreads :: {-# UNPACK #-}!(Int, Int, Int)
    , _mCConfigRuns :: {-# UNPACK #-}!Int
        -- ^ Number of Monte-Carlo runs to perform.
        --
        -- /Note:/ for optimal work scheduling, make that the total number of
        -- threads is divisible by the number of runs.
    , _mCConfigFlips :: {-# UNPACK #-}!Int
        -- ^ Number of spin-flips to do at each step.
    , _mCConfigRestarts :: {-# UNPACK #-}!Int
        -- ^ Allowed number of restarts.
        --
        -- A restart happens when the function computing local energy notices
        -- that a particular spin-flip results in a spin configuration with
        -- significantly higher probability. Monte-Carlo sampler is then reset
        -- and this new spin configuration is used as the initial one.
    , _mCConfigMagnetisation :: !(Maybe Int)
        -- ^ Specifies the magnetisation over which to sample.
    }

makeFields ''MCConfig

defaultMCConfig :: MCConfig
defaultMCConfig = MCConfig (1000, 11000, 1) (4, 1, 1) 4 2 5 Nothing

-- | Configuration for the Conjugate Gradient solver.
data CGConfig a =
  CGConfig { _cGConfigMaxIter :: {-# UNPACK #-}!Int
               -- ^ Maximum number of iterations.
           , _cGConfigTol :: !a
               -- ^ Tolerance.
           }

makeFields ''CGConfig

-- | Configuration for the Stochastic Reconfiguration algorithm.
data SRConfig a =
  SRConfig
    { _sRConfigMaxIter :: !Int
        -- ^ Maximum number of iterations to perform.
    , _sRConfigRegulariser :: !(Maybe (Int -> ℂ))
        -- ^ Regulariser for the S matrix. If it is @'Just' f@,
        -- then @λ = f i@ where @i@ is the currect iteration is used to
        -- regularise the matrix S according to @S <- S + λ1@.
    , _sRConfigRate :: !(Int -> ℂ)
        -- ^ Learning rate as a function of the iteration.
    , _sRConfigCg :: !(CGConfig ℝ)
        -- ^ Configuration for the Monte-Carlo sampling.
    , _sRConfigMc :: !MCConfig
        -- ^ Configuration for the Conjugate Gradient solver.
    }

-- | Returns a slice of the vector.
slice
  :: forall a s
   . Storable a
  => Int -- ^ Start index
  -> Int -- ^ Length
  -> MDenseVector 'Direct s a -- ^ Source vector
  -> MDenseVector 'Direct s a -- ^ Slice
slice !i !n !(MDenseVector size stride buff) =
  MDenseVector n stride (MV.slice (stride * i) (stride * n) buff)

makeFields ''SRConfig

type family RealOf a :: *

type instance RealOf Float = Float
type instance RealOf Double = Double
type instance RealOf (Complex a) = a

mallocVectorAligned :: forall a. Storable a => Int -> Int -> IO (ForeignPtr a)
mallocVectorAligned n alignment =
  mallocPlainForeignPtrAlignedBytes (n * sizeOf (undefined :: a)) alignment

newVectorAligned
  :: (Storable a, PrimMonad m) => Int -> Int -> m (MVector (PrimState m) a)
newVectorAligned n alignment =
  unsafePrimToPrim $! MVector n <$> mallocVectorAligned n alignment

newDenseVector
  :: (Storable a, PrimMonad m) => Int -> m (MDenseVector 'Direct (PrimState m) a)
newDenseVector n = MDenseVector n 1 <$> newVectorAligned n 64

-- | Default alignment used when allocating new vectors and matrices.
defaultAlignment :: Int
defaultAlignment = 64

roundUpTo :: Int -> Int -> Int
{-# INLINE roundUpTo #-}
roundUpTo !alignment !n = assert (isValidAlignment alignment && n >= 0) $
  (n + alignment - 1) .&. complement (alignment - 1)
  where isValidAlignment !x = x > 0 && (x .&. (x - 1) == 0)

newDenseMatrix
  :: forall orient a m. (Storable a, SingI orient, PrimMonad m)
  => Int
  -> Int
  -> m (MDenseMatrix orient (PrimState m) a)
newDenseMatrix rows cols = case (sing :: Sing orient) of
  SRow ->
    let ldim = assert (defaultAlignment `mod` sizeOf (undefined :: a) == 0) $
                roundUpTo (defaultAlignment `div` sizeOf (undefined :: a)) cols
    in  MDenseMatrix rows cols ldim
          <$> newVectorAligned (rows * ldim) defaultAlignment
  SColumn ->
    let ldim = roundUpTo defaultAlignment rows
    in  MDenseMatrix rows cols ldim
          <$> newVectorAligned (ldim * cols) defaultAlignment

class HasOrientation s a | s -> a where
  orientationOf :: s -> a

instance SingI orient
  => HasOrientation (DenseMatrix orient a) Orientation where
    orientationOf _ = fromSing (sing :: Sing orient)

instance SingI orient
  => HasOrientation (Mutable (DenseMatrix orient) s a) Orientation where
    orientationOf _ = fromSing (sing :: Sing orient)

isValidMatrix :: Orientation -- ^ Row- vs. column-major layout
              -> Int -- ^ Number of rows
              -> Int -- ^ Number of columns
              -> Int -- ^ Stride
              -> Int -- ^ Buffer size
              -> Bool
isValidMatrix orient rows cols i size =
  i >= 0 && (rows == 0 && cols == 0 || rows > 0 && cols > 0 && validAccesses orient)
  where validAccesses Row    = (rows - 1) * i < size && cols <= i
        validAccesses Column = (cols - 1) * i < size && rows <= i

-- | Constructs a nice message describing the problem in the BLAS vector.
badMatrixInfo :: String -- ^ Function name
              -> String -- ^ Argument name
              -> Int -- ^ Number of rows
              -> Int -- ^ Number of columns
              -> Int -- ^ Stride
              -> Int -- ^ Size of the underlying buffer
              -> String -- ^ Error message
badMatrixInfo funcName argName rows cols i size
  | rows < 0 = preamble <> " has a negative number of rows: " <> show rows <> "."
  | cols < 0 = preamble <> " has a negative number of columns: " <> show cols <> "."
  | i < 0 = preamble <> " has a negative stride: " <> show i <> "."
  | otherwise = preamble <> " has invalid range of accesses: #rows = " <> show rows <>
    ", #cols = " <> show cols <> ", stride = " <> show i <> ", bufferSize = " <> show size <> "."
  where preamble = funcName <> ": " <> argName

-- | Returns whether strides and dimensions are consistent.
isValidVector :: Int -- ^ Logical dimension
              -> Int -- ^ Stride
              -> Int -- ^ Buffer size
              -> Bool
isValidVector n i size = i >= 0 && (n == 0 || n > 0 && (n - 1) * i < size)

-- | Constructs a nice message describing the problem in the BLAS vector.
badVectorInfo :: String -- ^ Function name
              -> String -- ^ Argument name
              -> Int -- ^ Logical vector dimension
              -> Int -- ^ Vector stride
              -> Int -- ^ Size of the underlying buffer
              -> String -- ^ Error message
badVectorInfo funcName argName n i size
  | n < 0 = preamble <> " has a negative logical dimension: " <> show n <> "."
  | i < 0 = preamble <> " has a negative stride: " <> show i <> "."
  | otherwise = preamble <> " has invalid range of accesses: dim = " <> show n <>
                ", stride = " <> show i <> ", bufferSize = " <> show size <> "."
  where preamble = funcName <> ": " <> argName

class CheckValid a where
  assertValid :: String -> String -> a -> b -> b

instance Storable a => CheckValid (DenseVector variant a) where
  assertValid funcName argName !(DenseVector dim stride (V.length -> size))
    | isValidVector dim stride size = id
    | otherwise = error $! badVectorInfo funcName argName dim stride size

instance Storable a => CheckValid (Mutable (DenseVector variant) s a) where
  assertValid funcName argName !(MDenseVector dim stride (MV.length -> size))
    | isValidVector dim stride size = id
    | otherwise = error $! badVectorInfo funcName argName dim stride size

instance (Storable a, SingI orient) => CheckValid (DenseMatrix orient a) where
  assertValid funcName argName !m@(DenseMatrix rows cols stride (V.length -> size))
    | isValidMatrix (orientationOf m) rows cols stride size = id
    | otherwise = error $! badMatrixInfo funcName argName rows cols stride size

instance (Storable a, SingI orient) => CheckValid (Mutable (DenseMatrix orient) s a) where
  assertValid funcName argName !m@(MDenseMatrix rows cols stride (MV.length -> size))
    | isValidMatrix (orientationOf m) rows cols stride size = id
    | otherwise = error $! badMatrixInfo funcName argName rows cols stride size

instance Storable a => HasBuffer (MDenseVector variant s a) (MVector s a) where
  buffer inj (MDenseVector dim stride buf) = MDenseVector dim stride <$> inj buf

instance Storable a => HasBuffer (MDenseMatrix orientation s a) (MVector s a) where
  buffer inj (MDenseMatrix rows cols stride buf) = MDenseMatrix rows cols stride <$> inj buf

instance Storable a => HasStride (MDenseVector variant s a) Int where
  stride inj (MDenseVector dim stride buf) = (\x -> MDenseVector dim x buf) <$> inj stride

instance Storable a => HasStride (MDenseMatrix orientation s a) Int where
  stride inj (MDenseMatrix rows cols stride buf) = (\x -> MDenseMatrix rows cols x buf) <$> inj stride

instance Storable a => HasDim (MDenseVector variant s a) Int where
  dim inj (MDenseVector dim stride buf) = (\x -> MDenseVector x stride buf) <$> inj dim

instance Storable a => HasDim (MDenseMatrix orientation s a) (Int, Int) where
  dim inj (MDenseMatrix rows cols stride buf) =
    (\(x, y) -> MDenseMatrix x y stride buf) <$> inj (rows, cols)

instance Storable a => HasDim (DenseMatrix orientation a) (Int, Int) where
  dim inj (DenseMatrix rows cols stride buf) =
    (\(x, y) -> DenseMatrix x y stride buf) <$> inj (rows, cols)

class FreezeThaw (v :: * -> *) a where
  unsafeFreeze :: PrimMonad m => (Mutable v) (PrimState m) a -> m (v a)
  unsafeThaw :: PrimMonad m => v a -> m ((Mutable v) (PrimState m) a)

instance Storable a => FreezeThaw (DenseVector variant) a where
  unsafeFreeze (MDenseVector dim stride mv) = DenseVector dim stride <$> V.unsafeFreeze mv
  unsafeThaw (DenseVector dim stride mv) = MDenseVector dim stride <$> V.unsafeThaw mv

instance Storable a => FreezeThaw (DenseMatrix orientation) a where
  unsafeFreeze (MDenseMatrix rows cols stride mv) = DenseMatrix rows cols stride <$> V.unsafeFreeze mv
  unsafeThaw (DenseMatrix rows cols stride v) = MDenseMatrix rows cols stride <$> V.unsafeThaw v

touchForeignPtrPrim :: PrimMonad m => ForeignPtr a -> m ()
{-# NOINLINE touchForeignPtrPrim #-}
touchForeignPtrPrim fp = unsafeIOToPrim $! touchForeignPtr fp

withForeignPtrPrim :: PrimMonad m => ForeignPtr a -> (Ptr a -> m b) -> m b
{-# INLINE withForeignPtrPrim #-}
withForeignPtrPrim p func = do r <- func (unsafeForeignPtrToPtr p)
                               touchForeignPtrPrim p
                               return r

mapVectorM ::
     (PrimMonad m, Storable a, Storable b)
  => (a -> m b)
  -> MDenseVector 'Direct (PrimState m) a
  -> MDenseVector 'Direct (PrimState m) b
  -> m ()
mapVectorM f x y =
  assertValid "NQS.Rbm.zipWithVectorM" "x" x $
  assertValid "NQS.Rbm.zipWithVectorM" "x" y $
  assert (x ^. dim == y ^. dim) $ go 0
  where
    n = x ^. dim
    go !i
      | i == n = return ()
      | otherwise = do
        xi <- unsafeReadVector x i
        yi <- f xi
        unsafeWriteVector y i yi
        go (i + 1)

zipWithVectorM ::
     (PrimMonad m, Storable a, Storable b, Storable c)
  => (a -> b -> m c)
  -> MDenseVector 'Direct (PrimState m) a
  -> MDenseVector 'Direct (PrimState m) b
  -> MDenseVector 'Direct (PrimState m) c
  -> m ()
zipWithVectorM f x y z =
  assertValid "NQS.Rbm.zipWithVectorM" "x" x $
  assertValid "NQS.Rbm.zipWithVectorM" "x" y $
  assertValid "NQS.Rbm.zipWithVectorM" "x" z $
  assert (y ^. dim == n && z ^. dim == n) $ go 0
  where
    n = x ^. dim
    go !i
      | i == n = return ()
      | otherwise = do
        xi <- unsafeReadVector x i
        yi <- unsafeReadVector y i
        zi <- f xi yi
        unsafeWriteVector z i zi
        go (i + 1)

mapMatrixM ::
     forall m a b orientX orientY.
     (PrimMonad m, Storable a, Storable b, SingI orientX, SingI orientY)
  => (a -> m b)
  -> MDenseMatrix orientX (PrimState m) a
  -> MDenseMatrix orientY (PrimState m) b
  -> m ()
mapMatrixM f x y =
  assertValid "NQS.Rbm.zipWithMatrixM" "x" x $
  assertValid "NQS.Rbm.zipWithMatrixM" "y" y $
  assert (x ^. dim == y ^. dim) $
  case (sing :: Sing orientY) of
    SRow -> stepperRow 0 0
    SColumn -> stepperColumn 0 0
  where
    !(n, m) = x ^. dim
    stepperRow !i !j
      | j < m && i < n = go i j >> stepperRow i (j + 1)
      | i < n = stepperRow (i + 1) 0
      | otherwise = return ()
    stepperColumn !i !j
      | i < n && j < m = go i j >> stepperColumn (i + 1) j
      | j < m = stepperColumn 0 (j + 1)
      | otherwise = return ()
    go !i !j = do
      xij <- unsafeReadMatrix x (i, j)
      yij <- f xij
      unsafeWriteMatrix y (i, j) yij

zipWithMatrixM :: forall m a b c orientX orientY orientZ.
     (PrimMonad m, Storable a, Storable b, Storable c, SingI orientX, SingI orientY, SingI orientZ)
  => (a -> b -> m c)
  -> MDenseMatrix orientX (PrimState m) a
  -> MDenseMatrix orientY (PrimState m) b
  -> MDenseMatrix orientZ (PrimState m) c
  -> m ()
zipWithMatrixM f x y z =
  assertValid "NQS.Rbm.zipWithMatrixM" "x" x $
  assertValid "NQS.Rbm.zipWithMatrixM" "y" y $
  assertValid "NQS.Rbm.zipWithMatrixM" "z" z $
  assert (y ^. dim == (n, m) && z ^. dim == (n, m)) $
    case (sing :: Sing orientZ) of
      SRow -> stepperRow 0 0
      SColumn -> stepperColumn 0 0
  where
    !(n, m) = x ^. dim
    stepperRow !i !j
      | j < m && i < n = go i j >> stepperRow i (j + 1)
      | i < n = stepperRow (i + 1) 0
      | otherwise = return ()
    stepperColumn !i !j
      | i < n && j < m = go i j >> stepperColumn (i + 1) j
      | j < m = stepperColumn 0 (j + 1)
      | otherwise = return ()
    go !i !j = do
      xij <- unsafeReadMatrix x (i, j)
      yij <- unsafeReadMatrix y (i, j)
      zij <- f xij yij
      unsafeWriteMatrix z (i, j) zij

{-
withRbm :: PrimMonad m => Rbm a -> (Ptr (RbmCore a) -> m b) -> m b
withRbm (Rbm p) func = withForeignPtrPrim p func

withRbmPure :: Rbm a -> (Ptr (RbmCore a) -> b) -> b
withRbmPure (Rbm p) func = unsafePerformIO $! withForeignPtrPrim p (return . func)

withMRbm :: PrimMonad m => MRbm (PrimState m) a -> (Ptr (RbmCore a) -> m b) -> m b
withMRbm (MRbm p) func = withForeignPtrPrim p func

withMRbmPure :: MRbm s a -> (Ptr (RbmCore a) -> b) -> b
withMRbmPure (MRbm p) func = unsafePerformIO $! withForeignPtr p (return . func)
-}

data Estimate e a = Estimate
  { estPoint :: !a
  , estError :: !(e a)
  } deriving (Generic, NFData)

newtype NormalError a = NormalError a
  deriving (Eq, Ord, Read, Show, Generic, NFData)

data EnergyMeasurement a = EnergyMeasurement
  { _energyMeasurementMean :: !a
  , _energyMeasurementVar :: !a
  } deriving (Generic)

makeFields ''EnergyMeasurement

-- | Simple for loop. Counts from /start/ to /end/-1.
for :: Monad m => Int -> Int -> (Int -> m ()) -> m ()
for !start !end f = assert (start <= end) $ loop start
 where
  loop !i | i == end  = return ()
          | otherwise = f i >> loop (i + 1)
{-# INLINE for #-}

for' :: Monad m => a -> Int -> (a -> Int -> m a) -> m a
for' !x0 !n f = assert (n >= 0) $ loop x0 0
 where
  loop !x !i | i == n    = return x
             | otherwise = f x i >>= \x' -> loop x' (i + 1)
{-# INLINE for' #-}

generateVectorM
  :: forall a m
   . (Monad m, Storable a)
  => Int
  -> (Int -> m a)
  -> m (DenseVector 'Direct a)
{-# NOINLINE generateVectorM #-}
generateVectorM !n f
  | n < 0     = error $! "generateVectorM: Invalid dimension: " <> show n
  | n == 0    = return $! DenseVector 0 0 V.empty
  | otherwise = go (runST $ newDenseVector n >>= unsafeFreeze)
 where
  go :: DenseVector 'Direct a -> m (DenseVector 'Direct a)
  go !v' = for' v' n $ \v i -> f i >>= \x -> return (set v i x)
  set !v !i !x = runST $ do
    !v' <- unsafeThaw v
    unsafeWriteVector v' i x
    unsafeFreeze v'

generateMatrixM
  :: forall orient m a. (Monad m, Storable a, SingI orient)
  => Int
  -> Int
  -> (Int -> Int -> m a)
  -> m (DenseMatrix orient a)
{-# NOINLINE generateMatrixM #-}
generateMatrixM !rows !cols f
  | rows < 0 || cols < 0 = error $! "generateMatrixM: Invalid dimensions: " <> show (rows, cols)
  | rows == 0 || cols == 0 = return $! DenseMatrix 0 0 0 V.empty
  | otherwise = go (runST $ newDenseMatrix rows cols >>= unsafeFreeze)
 where
  go :: DenseMatrix orient a -> m (DenseMatrix orient a)
  go !m'' = case (sing :: Sing orient) of
    SRow -> for' m'' rows $ \m' i -> for' m' cols $ \m j ->
              f i j >>= \x -> return (set m i j x)
    SColumn -> for' m'' cols $ \m' j -> for' m' rows $ \m i ->
                 f i j >>= \x -> return (set m i j x)
  set !m !i !j !x = runST $ do
    !m' <- unsafeThaw m
    unsafeWriteMatrix m' (i, j) x
    unsafeFreeze m'

instance (Storable a, FromJSON a) => FromJSON (DenseVector 'Direct a) where
  parseJSON = withArray "vector elements (i.e. Array)" $ \v -> do
    generateVectorM (GV.length v) $ \i -> parseJSON (v `GV.unsafeIndex` i)

instance ToJSON a => ToJSON (Complex a) where
  toJSON (x :+ y) = toJSON [x, y]
  toEncoding (x :+ y) = toEncoding [x, y]

instance FromJSON a => FromJSON (Complex a) where
  parseJSON = withArray "Complex" $ \v ->
    (:+) <$> parseJSON (v GV.! 0)
         <*> parseJSON (v GV.! 1)

instance (Storable a, FromJSON a, SingI orient) => FromJSON (DenseMatrix orient a) where
  parseJSON = withArray "Matrix rows" $ \matrix -> do
    (n, m) <- checkDimensions matrix
    generateMatrixM n m $ \i j ->
      flip (withArray "a matrix row") (matrix `GV.unsafeIndex` i) $ \row ->
        parseJSON (row `GV.unsafeIndex` j)
    where
      checkDimensions x@(GV.length -> n)
        | n == 0 = return (0, 0)
        | otherwise = do
          !m <- withArray "first row (i.e. Array)" (return . GV.length) (GV.unsafeHead x)
          for 0 n $ \i ->
            flip (withArray "a matrix row (i.e. Array)") (x `GV.unsafeIndex` i) $ \row ->
              unless (GV.length row == m) $ fail $! "Matrix row #" <> show i <> " has wrong dimension."
          return (n, m)

vectorAsVector :: Storable a => DenseVector 'Direct a -> Boxed.Vector a
vectorAsVector !(DenseVector n stride buff) =
  Boxed.generate n (\i -> buff `V.unsafeIndex` (i * stride))

-- | Returns a row of a matrix.
unsafeRowAsVector
  :: forall orient a
   . (Storable a, SingI orient)
  => Int
  -> DenseMatrix orient a
  -> Boxed.Vector a
unsafeRowAsVector !i !(DenseMatrix rows cols stride buff) =
  assert (i < rows) $ case (sing :: Sing orient) of
    SRow    -> GV.generate cols (\j -> buff `V.unsafeIndex` (i * stride + j))
    SColumn -> GV.generate cols (\j -> buff `V.unsafeIndex` (i + j * stride))

instance (Storable a, ToJSON a) => ToJSON (DenseVector 'Direct a) where
  toEncoding = toEncoding . vectorAsVector

instance (Storable a, ToJSON a, SingI orient) => ToJSON (DenseMatrix orient a) where
  toEncoding x = toEncoding $ Boxed.generate (x ^. dim . _1) (\i -> unsafeRowAsVector i x)


module OidTableSVG where

import Data.Word
import qualified Data.ByteString.Lazy as LB
import Control.Monad hiding (forM_)
import Data.Foldable (forM_)
import Text.Printf
import Control.Arrow ((***))
import Data.String
import Data.Complex

import qualified Text.Blaze.Svg11 as S
import qualified Text.Blaze.Svg11.Attributes as A
import Text.Blaze.Svg11 ((!))
import Text.Blaze.Svg.Renderer.Utf8

import OidCode
import KnownCodes
import Types
import Utils

type Point = Complex Double

oidTableSvg :: Conf -> Bool -> String -> [(String, Word16)] -> LB.ByteString
oidTableSvg conf usePNG title entries
    | entriesPerPage < 1 = error "OID codes too large to fit on a single page"
    | otherwise = renderSvg $
        S.docTypeSvg  ! A.version (S.toValue "1.1")
                      ! A.width  (S.toValue (printf "%fmm" (a4w/mm) :: String))
                      ! A.height (S.toValue (printf "%fmm" (a4h/mm) :: String))
                      ! A.viewbox (S.toValue (printf "0 0 %f %f" a4w a4h :: String))
                      ! A.fontFamily (S.toValue "sans-serif")
                      $ do
    let patid d c | d == show c = d
                  | otherwise   = printf "%s-%d" d c

    -- Create patterns for the codes
    S.defs $ forM_ entries $ \(d,c) ->
        case code2RawCode c of
            Nothing -> return ()
            Just rc -> oidSVGPattern conf usePNG (patid d c) rc

    -- For SVG, we put all on one page (and exceed the page if it is too big)
    let chunks = [entries]
    let totalPages = length chunks

    forM_ (zip [1::Int ..] chunks) $ \(pageNum, thisPage) -> do
            S.text_ ! A.x (S.toValue (a4w / 2))
                    ! A.y (S.toValue (padTop + titleHeight))
                    ! A.textAnchor (S.toValue "middle")
                    ! A.stroke (S.toValue "black")
                    ! A.fontSize (S.toValue (printf "%f" (12*pt) :: String))
                    $ fromString title

            S.text_ ! A.x (S.toValue padLeft)
                    ! A.y (S.toValue (a4h - padBot))
                    ! A.textAnchor (S.toValue "left")
                    ! A.stroke (S.toValue "black")
                    ! A.fontSize (S.toValue (printf "%f" (8*pt) :: String))
                    $ fromString $ "Created by tttool-" ++ tttoolVersion

            forM_ (zip thisPage positions) $ \((d,c), x :+ y) -> do
                S.rect ! A.width (S.toValue imageWidth)
                       ! A.height (S.toValue imageHeight)
                       ! A.x (S.toValue x)
                       ! A.y (S.toValue y)
                       ! A.fill (S.toValue $ "url(#"++patid d c++")")

                S.text_ ! A.x (S.toValue x)
                        ! A.y (S.toValue (y + imageHeight + subtitleSep + subtitleHeight))
                        ! A.textAnchor (S.toValue "left")
                        ! A.stroke (S.toValue "black")
                        ! A.fontSize (S.toValue (printf "%f" (8*pt) :: String))
                        $ fromString d
  where
    -- Configure-dependent dimensions (all in pt)
    (imageWidth,imageHeight) = (*mm) *** (*mm) $ fromIntegral *** fromIntegral $cCodeDim conf

    -- Static dimensions (all in pt)

    -- Page paddings
    padTop, padLeft, padBot, padRight :: Double
    padTop   = 1*cm
    padBot   = 1*cm
    padLeft  = 2*cm
    padRight = 2*cm

    titleHeight  = 1*cm
    titleSep     = 0.5*cm
    footerHeight = 0.5*cm
    footerSep    = 0.5*cm

    imageSepH = 0.4*cm
    imageSepV = 0.2*cm

    subtitleHeight = 0.4*cm
    subtitleSep    = 0.2*cm

    -- Derived dimensions (all in pt)
    {-
    titleRect = Rectangle
        (padLeft          :+ (a4h - padTop - titleHeight))
        ((a4w - padRight) :+ (a4h - padTop))
    titleFont = Font (PDFFont Helvetica 12) black black

    footerRect = Rectangle
        (padLeft          :+ padBot)
        ((a4w - padRight) :+ (padBot + footerHeight))
    footerFont = Font (PDFFont Helvetica 8) black black

    bodyFont = Font (PDFFont Helvetica 8) black black
    -}

    bodyWidth  = a4w - padLeft - padRight
    bodyHeight = a4h - padTop - titleHeight - titleSep - footerSep - footerHeight - padBot

    positions = map (+(padLeft :+ (padTop + titleHeight + titleSep))) $
        calcPositions bodyWidth  bodyHeight
                      imageWidth (imageHeight + subtitleSep + subtitleHeight)
                      imageSepH  imageSepV
    entriesPerPage = length positions

    -- Derived dimensions (all in pixels)
    imageWidthPx = floor (imageWidth * px)
    imageHeightPx = floor (imageHeight * px)

    -- config-dependent conversion factors
    px :: Double
    px = fromIntegral (cDPI conf) / 72

    {-
    -- Makes sure the given point is at a coordinate that is a multiple
    -- of an pixel
    align :: Point -> Point
    align pos = alignToPx (realPart pos) :+ (a4h - alignToPx (a4h - imagPart pos))

    -- Makes sure the given distance is an interal mulitple of a pixel
    alignToPx :: Double -> Double
    alignToPx x = fromIntegral (floor (x * px)) / px
    -}

calcPositions
    :: Double -- ^ total width
    -> Double -- ^ total height
    -> Double -- ^ entry width
    -> Double -- ^ entry height
    -> Double -- ^ pad width
    -> Double -- ^ pad height
    -> [Point]
calcPositions tw th ew eh pw ph = [ x :+ ({-(th - -} y) | y <- ys , x <- xs]
  where
    xs = [0,ew+pw..tw-ew]
    ys = [0,eh+ph..th-eh]

pt :: Double
pt = 48/2.83465

-- Conversation factor
cm :: Double
cm = 10 * mm

mm :: Double
mm = 2.83465 * pt

-- A4 dimensions
a4w, a4h :: Double
a4w = 595 * pt
a4h = 842 * pt