{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS_GHC -Wall #-}

-- |
-- Module      :  Numeric.NLOPT
-- Copyright   :  (c) Matthew Peddie 2017
-- License     :  BSD3
-- Maintainer  :  Matthew Peddie <mpeddie@gmail.com>
-- Stability   :  provisional
-- Portability :  GHC
--
-- This module provides a high-level, @hmatrix@-compatible interface to
-- the <http://ab-initio.mit.edu/wiki/index.php/NLopt NLOPT> library by
-- Steven G. Johnson.
--
-- = Documentation
--
-- Most non-numerical details are documented, but for specific
-- information on what the optimization methods do, how constraints are
-- handled, etc., you should consult:
--
--   * The <http://ab-initio.mit.edu/wiki/index.php/NLopt_Introduction NLOPT introduction>
--
--   * The <http://ab-initio.mit.edu/wiki/index.php/NLopt_Reference NLOPT reference manual>
--
--   * The <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms NLOPT algorithm manual>
--
-- = Example program
--
-- The following interactive session example uses the Nelder-Mead simplex
-- algorithm, a derivative-free local optimizer, to minimize a trivial
-- function with a minimum of 22.0 at @(0, 0)@.
--
-- >>> import Numeric.LinearAlgebra ( dot, fromList )
-- >>> let objf x = x `dot` x + 22                         -- define objective
-- >>> let stop = ObjectiveRelativeTolerance 1e-6 :| []    -- define stopping criterion
-- >>> let algorithm = NELDERMEAD objf [] Nothing          -- specify algorithm
-- >>> let problem = LocalProblem 2 stop algorithm         -- specify problem
-- >>> let x0 = fromList [5, 10]                           -- specify initial guess
-- >>> minimizeLocal problem x0
-- Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED})
module Numeric.NLOPT
  ( -- * Specifying the objective function
    Objective,
    ObjectiveD,
    Preconditioner,

    -- * Specifying the constraints

    -- ** Bound constraints
    Bounds (..),

    -- ** Nonlinear constraints

    --
    -- $nonlinearconstraints

    -- *** Constraint functions
    ScalarConstraint,
    ScalarConstraintD,
    VectorConstraint,
    VectorConstraintD,

    -- *** Constraint types
    Constraint (..),
    EqualityConstraint (..),
    InequalityConstraint (..),

    -- *** Collections of constraints
    EqualityConstraints,
    EqualityConstraintsD,
    InequalityConstraints,
    InequalityConstraintsD,

    -- * Stopping conditions

    --
    -- $nonempty
    StoppingCondition (..),
    NonEmpty (..),

    -- * Additional configuration
    RandomSeed (..),
    Population (..),
    VectorStorage (..),
    InitialStep (..),

    -- * Minimization problems

    -- ** Local minimization
    LocalAlgorithm (..),
    LocalProblem (..),
    minimizeLocal,

    -- ** Global minimization
    GlobalAlgorithm (..),
    GlobalProblem (..),
    minimizeGlobal,

    -- ** Minimization by augmented Lagrangian
    AugLagAlgorithm (..),
    AugLagProblem (..),
    minimizeAugLag,

    -- ** Results
    Solution (..),
    N.Result (..),
  )
where

import Control.Exception (Exception)
import qualified Control.Exception as Ex
import Data.Foldable (traverse_)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Typeable (Typeable)
import qualified Data.Vector.Storable as V
import Numeric.LinearAlgebra as HM
import qualified Numeric.Optimization.NLOPT.Bindings as N
import System.IO.Unsafe (unsafePerformIO)

{- Function wrapping for the immutable HMatrix interface -}
wrapScalarFunction :: (Vector Double -> Double) -> N.ScalarFunction ()
wrapScalarFunction f params _ _ = return $ f params

wrapScalarFunctionD ::
  (Vector Double -> (Double, Vector Double)) ->
  N.ScalarFunction ()
wrapScalarFunctionD f params grad _ = do
  case grad of
    Nothing -> return ()
    Just g -> V.copy g usergrad
  return result
  where
    (result, usergrad) = f params

wrapVectorFunction ::
  (Vector Double -> Word -> Vector Double) ->
  Word ->
  N.VectorFunction ()
wrapVectorFunction f n params vout _ _ = V.copy vout $ f params n

wrapVectorFunctionD ::
  (Vector Double -> Word -> (Vector Double, Matrix Double)) ->
  Word ->
  N.VectorFunction ()
wrapVectorFunctionD f n params vout jac _ = do
  V.copy vout result
  case jac of
    Nothing -> return ()
    Just j -> V.copy j (HM.flatten userjac)
  where
    (result, userjac) = f params n

wrapPreconditionerFunction ::
  (Vector Double -> Vector Double -> Vector Double) ->
  N.PreconditionerFunction ()
wrapPreconditionerFunction f params v vpre _ = V.copy vpre (f params v)

{- Objective functions -}

-- | An objective function that calculates the objective value at the
-- given parameter vector.
type Objective =
  -- | Parameter vector
  Vector Double ->
  -- | Objective function value
  Double

-- | An objective function that calculates both the objective value
-- and the gradient of the objective with respect to the input
-- parameter vector, at the given parameter vector.
type ObjectiveD =
  -- | Parameter vector
  Vector Double ->
  -- | (Objective function value, gradient)
  (Double, Vector Double)

-- | A preconditioner function, which computes @vpre = H(x) v@, where
-- @H@ is the Hessian matrix: the positive semi-definite second
-- derivative at the given parameter vector @x@, or an approximation
-- thereof.
type Preconditioner =
  -- | Parameter vector @x@
  Vector Double ->
  -- | Vector @v@ to precondition at @x@
  Vector Double ->
  -- | Preconditioned vector @vpre@
  Vector Double

data ObjectiveFunction f
  = MinimumObjective f
  | PreconditionedMinimumObjective Preconditioner f

applyObjective :: N.Opt -> ObjectiveFunction Objective -> IO N.Result
applyObjective opt (MinimumObjective f) =
  N.set_min_objective opt (wrapScalarFunction f) ()
applyObjective opt (PreconditionedMinimumObjective p f) =
  N.set_precond_min_objective
    opt
    (wrapScalarFunction f)
    (wrapPreconditionerFunction p)
    ()

applyObjectiveD :: N.Opt -> ObjectiveFunction ObjectiveD -> IO N.Result
applyObjectiveD opt (MinimumObjective f) =
  N.set_min_objective opt (wrapScalarFunctionD f) ()
applyObjectiveD opt (PreconditionedMinimumObjective p f) =
  N.set_precond_min_objective
    opt
    (wrapScalarFunctionD f)
    (wrapPreconditionerFunction p)
    ()

{- Constraint functions -}

-- | A constraint function which returns @c(x)@ given the parameter
-- vector @x@.  The constraint will enforce that @c(x) == 0@ (equality
-- constraint) or @c(x) <= 0@ (inequality constraint).
type ScalarConstraint =
  -- | Parameter vector @x@
  Vector Double ->
  -- | Constraint violation (deviation from 0)
  Double

-- | A constraint function which returns @c(x)@ given the parameter
-- vector @x@ along with the gradient of @c(x)@ with respect to @x@ at
-- that point.  The constraint will enforce that @c(x) == 0@ (equality
-- constraint) or @c(x) <= 0@ (inequality constraint).
type ScalarConstraintD =
  -- | Parameter vector
  Vector Double ->
  -- | (Constraint violation, constraint gradient)
  (Double, Vector Double)

-- | A constraint function which returns a vector @c(x)@ given the
-- parameter vector @x@.  The constraint will enforce that @c(x) == 0@
-- (equality constraint) or @c(x) <= 0@ (inequality constraint).
type VectorConstraint =
  -- | Parameter vector
  Vector Double ->
  -- | Constraint vector size
  Word ->
  -- | Constraint violation vector
  Vector Double

-- | A constraint function which returns @c(x)@ given the parameter
-- vector @x@ along with the Jacobian (first derivative) matrix of
-- @c(x)@ with respect to @x@ at that point.  The constraint will
-- enforce that @c(x) == 0@ (equality constraint) or @c(x) <= 0@
-- (inequality constraint).
type VectorConstraintD =
  -- | Parameter vector
  Vector Double ->
  -- | Constraint vector size
  Word ->
  -- | (Constraint violation vector,
  -- constraint Jacobian)
  (Vector Double, Matrix Double)

-- $nonlinearconstraints
--
-- Note that most NLOPT algorithms do not support nonlinear
-- constraints natively; if you need to enforce nonlinear constraints,
-- you may want to use the 'AugLagAlgorithm' family of solvers, which
-- can add nonlinear constraints to some algorithm that does not
-- support them by a principled modification of the objective
-- function.
--
-- == Example program
--
-- The following interactive session example enforces a scalar
-- constraint on the problem given in the beginning of the module: the
-- parameters must always sum to 1.  The minimizer finds a constrained
-- minimum of 22.5 at @(0.5, 0.5)@.
--
-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList )
-- >>> let objf x = x `dot` x + 22
-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| []
-- >>>          -- define constraint function:
-- >>> let constraintf x = sum (toList x) - 1.0
-- >>>          -- define constraint object to pass to the algorithm:
-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6
-- >>> let algorithm = COBYLA objf [] [] [constraint] Nothing
-- >>> let problem = LocalProblem 2 stop algorithm
-- >>> let x0 = fromList [5, 10]
-- >>> minimizeLocal problem x0
-- Right (Solution {solutionCost = 22.500000000013028, solutionParams = [0.5000025521533521,0.49999744784664796], solutionResult = FTOL_REACHED})

data Constraint s v
  = -- | A scalar constraint.
    Scalar s
  | -- | A vector constraint.
    Vector Word v
  | -- | A scalar constraint with an attached preconditioning function.
    Preconditioned Preconditioner s

-- | An equality constraint, comprised of both the constraint function
-- (or functions, if a preconditioner is used) along with the desired
-- tolerance.
data EqualityConstraint s v = EqualityConstraint
  { eqConstraintFunctions :: Constraint s v,
    eqConstraintTolerance :: Double
  }

-- | An inequality constraint, comprised of both the constraint
-- function (or functions, if a preconditioner is used) along with the
-- desired tolerance.
data InequalityConstraint s v = InequalityConstraint
  { ineqConstraintFunctions :: Constraint s v,
    ineqConstraintTolerance :: Double
  }

-- | A collection of equality constraints that do not supply
-- constraint derivatives.
type EqualityConstraints =
  [EqualityConstraint ScalarConstraint VectorConstraint]

-- | A collection of inequality constraints that do not supply
-- constraint derivatives.
type InequalityConstraints =
  [InequalityConstraint ScalarConstraint VectorConstraint]

-- | A collection of equality constraints that supply constraint
-- derivatives.
type EqualityConstraintsD = [EqualityConstraint ScalarConstraintD VectorConstraintD]

-- | A collection of inequality constraints that supply constraint
-- derivatives.
type InequalityConstraintsD = [InequalityConstraint ScalarConstraintD VectorConstraintD]

class ApplyConstraint constraint where
  applyConstraint :: N.Opt -> constraint -> IO N.Result

instance ApplyConstraint (EqualityConstraint ScalarConstraint VectorConstraint) where
  applyConstraint opt (EqualityConstraint ty tol) = case ty of
    Scalar s ->
      N.add_equality_constraint opt (wrapScalarFunction s) () tol
    Vector n v ->
      N.add_equality_mconstraint opt n (wrapVectorFunction v n) () tol
    Preconditioned p s ->
      N.add_precond_equality_constraint
        opt
        (wrapScalarFunction s)
        (wrapPreconditionerFunction p)
        ()
        tol

instance ApplyConstraint (InequalityConstraint ScalarConstraint VectorConstraint) where
  applyConstraint opt (InequalityConstraint ty tol) = case ty of
    Scalar s ->
      N.add_inequality_constraint opt (wrapScalarFunction s) () tol
    Vector n v ->
      N.add_inequality_mconstraint opt n (wrapVectorFunction v n) () tol
    Preconditioned p s ->
      N.add_precond_inequality_constraint
        opt
        (wrapScalarFunction s)
        (wrapPreconditionerFunction p)
        ()
        tol

instance ApplyConstraint (EqualityConstraint ScalarConstraintD VectorConstraintD) where
  applyConstraint opt (EqualityConstraint ty tol) = case ty of
    Scalar s ->
      N.add_equality_constraint opt (wrapScalarFunctionD s) () tol
    Vector n v ->
      N.add_equality_mconstraint opt n (wrapVectorFunctionD v n) () tol
    Preconditioned p s ->
      N.add_precond_equality_constraint
        opt
        (wrapScalarFunctionD s)
        (wrapPreconditionerFunction p)
        ()
        tol

instance ApplyConstraint (InequalityConstraint ScalarConstraintD VectorConstraintD) where
  applyConstraint opt (InequalityConstraint ty tol) = case ty of
    Scalar s ->
      N.add_inequality_constraint opt (wrapScalarFunctionD s) () tol
    Vector n v ->
      N.add_inequality_mconstraint opt n (wrapVectorFunctionD v n) () tol
    Preconditioned p s ->
      N.add_precond_inequality_constraint
        opt
        (wrapScalarFunctionD s)
        (wrapPreconditionerFunction p)
        ()
        tol

{- Bounds -}

-- | Bound constraints are specified by vectors of the same dimension
-- as the parameter space.
--
-- == Example program
--
-- The following interactive session example enforces lower bounds on
-- the example from the beginning of the module.  This prevents the
-- optimizer from locating the true minimum at @(0, 0)@; a slightly
-- higher constrained minimum at @(1, 1)@ is found.  Note that the
-- optimizer returns 'N.XTOL_REACHED' rather than 'N.FTOL_REACHED',
-- because the bound constraint is active at the final minimum.
--
-- >>> import Numeric.LinearAlgebra ( dot, fromList )
-- >>> let objf x = x `dot` x + 22                           -- define objective
-- >>> let stop = ObjectiveRelativeTolerance 1e-6 :| []      -- define stopping criterion
-- >>> let lowerbound = LowerBounds $ fromList [1, 1]        -- specify bounds
-- >>> let algorithm = NELDERMEAD objf [lowerbound] Nothing  -- specify algorithm
-- >>> let problem = LocalProblem 2 stop algorithm           -- specify problem
-- >>> let x0 = fromList [5, 10]                             -- specify initial guess
-- >>> minimizeLocal problem x0
-- Right (Solution {solutionCost = 24.0, solutionParams = [1.0,1.0], solutionResult = XTOL_REACHED})
data Bounds
  = -- | Lower bound vector @v@ means we want @x >= v@.
    LowerBounds (Vector Double)
  | -- | Upper bound vector @u@ means we want @x <= u@.
    UpperBounds (Vector Double)
  deriving (Eq, Show, Read)

applyBounds :: N.Opt -> Bounds -> IO N.Result
applyBounds opt (LowerBounds lbvec) = N.set_lower_bounds opt lbvec
applyBounds opt (UpperBounds ubvec) = N.set_upper_bounds opt ubvec

{- Stopping conditions -}

-- | A 'StoppingCondition' tells NLOPT when to stop working on a
-- minimization problem.  When multiple 'StoppingCondition's are
-- provided, the problem will stop when any one condition is met.
data StoppingCondition
  = -- | Stop minimizing when an objective value @J@ less than or equal
    -- to the provided value is found.
    MinimumValue Double
  | -- | Stop minimizing when an optimization step changes the objective
    -- value @J@ by less than the provided tolerance multiplied by @|J|@.
    ObjectiveRelativeTolerance Double
  | -- | Stop minimizing when an optimization step changes the objective
    -- value by less than the provided tolerance.
    ObjectiveAbsoluteTolerance Double
  | -- | Stop when an optimization step changes /every element/ of the
    -- parameter vector @x@ by less than @x@ scaled by the provided
    -- tolerance.
    ParameterRelativeTolerance Double
  | -- | Stop when an optimization step changes /every element/ of the
    -- parameter vector @x@ by less than the corresponding element in
    -- the provided vector of tolerances values.
    ParameterAbsoluteTolerance (Vector Double)
  | -- | Stop when the number of evaluations of the objective function
    -- exceeds the provided count.
    MaximumEvaluations Word
  | -- | Stop when the optimization time exceeds the provided time (in
    -- seconds).  This is not a precise limit.
    MaximumTime Double
  deriving (Eq, Show, Read)

-- $nonempty
--
-- The 'NonEmpty' data type from 'Data.List.NonEmpty' is re-exported
-- here, because it is used to ensure that you always specify at least
-- one stopping condition.

applyStoppingCondition :: N.Opt -> StoppingCondition -> IO N.Result
applyStoppingCondition opt (MinimumValue x) = N.set_stopval opt x
applyStoppingCondition opt (ObjectiveRelativeTolerance x) = N.set_ftol_rel opt x
applyStoppingCondition opt (ObjectiveAbsoluteTolerance x) = N.set_ftol_abs opt x
applyStoppingCondition opt (ParameterRelativeTolerance x) = N.set_xtol_rel opt x
applyStoppingCondition opt (ParameterAbsoluteTolerance v) = N.set_xtol_abs opt v
applyStoppingCondition opt (MaximumEvaluations n) = N.set_maxeval opt n
applyStoppingCondition opt (MaximumTime deltat) = N.set_maxtime opt deltat

{- Random seed control -}

-- | This specifies how to initialize the random number generator for
-- stochastic algorithms.
data RandomSeed
  = -- | Seed the RNG with the provided value.
    SeedValue Word
  | -- | Seed the RNG using the system clock.
    SeedFromTime
  | -- | Don't perform any explicit initialization of the RNG.
    Don'tSeed
  deriving (Eq, Show, Read)

applyRandomSeed :: RandomSeed -> IO ()
applyRandomSeed Don'tSeed = return ()
applyRandomSeed (SeedValue n) = N.srand n
applyRandomSeed SeedFromTime = N.srand_time

{- Random stuff -}

-- | This specifies the population size for algorithms that use a pool
-- of solutions.
newtype Population = Population Word deriving (Eq, Show, Read)

applyPopulation :: N.Opt -> Population -> IO N.Result
applyPopulation opt (Population n) = N.set_population opt n

-- | This specifies the memory size to be used by algorithms like
-- 'LBFGS' which store approximate Hessian or Jacobian matrices.
newtype VectorStorage = VectorStorage Word deriving (Eq, Show, Read)

applyVectorStorage :: N.Opt -> VectorStorage -> IO N.Result
applyVectorStorage opt (VectorStorage n) = N.set_vector_storage opt n

-- | This vector with the same dimension as the parameter vector @x@
-- specifies the initial step for the optimizer to take.  (This
-- applies to local gradient-free algorithms, which cannot use
-- gradients to estimate how big a step to take.)
newtype InitialStep = InitialStep (Vector Double) deriving (Eq, Show, Read)

applyInitialStep :: N.Opt -> InitialStep -> IO N.Result
applyInitialStep opt (InitialStep v) = N.set_initial_step opt v

{- Algorithms -}

data GlobalProblem = GlobalProblem
  { -- | Lower bounds for @x@
    lowerBounds :: Vector Double,
    -- | Upper bounds for @x@
    upperBounds :: Vector Double,
    -- | At least one stopping
    -- condition
    gstop :: NonEmpty StoppingCondition,
    -- | Algorithm specification
    galgorithm :: GlobalAlgorithm
  }

-- | These are the global minimization algorithms provided by NLOPT.  Please see
-- <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms the NLOPT algorithm manual>
-- for more details on how the methods work and how they relate to one another.
--
-- Optional parameters are wrapped in a 'Maybe'; for example, if you
-- see 'Maybe' 'Population', you can simply specify 'Nothing' to use
-- the default behavior.
--
-- N.B.: There may be at most n equality constraints for problems with n dimensions due to restrictions
-- in NLopt. This may be circumvented by combination of equality constraints according to:
-- h1(x)==0 && h2(x)==0 <=> h1(x)^2+h2(x)^2==0.
data GlobalAlgorithm
  = -- | DIviding RECTangles
    DIRECT Objective
  | -- | DIviding RECTangles, locally-biased variant
    DIRECT_L Objective
  | -- | DIviding RECTangles, "slightly randomized"
    DIRECT_L_RAND Objective RandomSeed
  | -- | DIviding RECTangles, unscaled version
    DIRECT_NOSCAL Objective
  | -- | DIviding RECTangles, locally-biased and unscaled
    DIRECT_L_NOSCAL Objective
  | -- | DIviding RECTangles, locally-biased, unscaled and "slightly
    -- randomized"
    DIRECT_L_RAND_NOSCAL Objective RandomSeed
  | -- | DIviding RECTangles, original FORTRAN implementation
    ORIG_DIRECT Objective InequalityConstraints
  | -- | DIviding RECTangles, locally-biased, original FORTRAN
    -- implementation
    ORIG_DIRECT_L Objective InequalityConstraints
  | -- | Stochastic Global Optimization.
    -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__
    STOGO ObjectiveD
  | -- | Stochastic Global Optimization, randomized variant.
    -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__
    STOGO_RAND ObjectiveD RandomSeed
  | -- | Controlled Random Search with Local Mutation
    CRS2_LM Objective RandomSeed (Maybe Population)
  | -- | Improved Stochastic Ranking Evolution Strategy, the population defaults to 20×(n+1) in n dimensions
    ISRES Objective InequalityConstraints EqualityConstraints RandomSeed (Maybe Population)
  | -- | Evolutionary Algorithm
    ESCH Objective
  | -- | Original Multi-Level Single-Linkage
    MLSL Objective LocalProblem (Maybe Population)
  | -- | Multi-Level Single-Linkage with Sobol Low-Discrepancy
    -- Sequence for starting points
    MLSL_LDS Objective LocalProblem (Maybe Population)

algorithmEnumOfGlobal :: GlobalAlgorithm -> N.Algorithm
algorithmEnumOfGlobal (DIRECT _) = N.GN_DIRECT
algorithmEnumOfGlobal (DIRECT_L _) = N.GN_DIRECT_L
algorithmEnumOfGlobal (DIRECT_L_RAND _ _) = N.GN_DIRECT_L_RAND
algorithmEnumOfGlobal (DIRECT_NOSCAL _) = N.GN_DIRECT_NOSCAL
algorithmEnumOfGlobal (DIRECT_L_NOSCAL _) = N.GN_DIRECT_L_NOSCAL
algorithmEnumOfGlobal (DIRECT_L_RAND_NOSCAL _ _) = N.GN_DIRECT_L_RAND_NOSCAL
algorithmEnumOfGlobal (ORIG_DIRECT _ _) = N.GN_ORIG_DIRECT
algorithmEnumOfGlobal (ORIG_DIRECT_L _ _) = N.GN_ORIG_DIRECT_L
algorithmEnumOfGlobal (STOGO _) = N.GD_STOGO
algorithmEnumOfGlobal (STOGO_RAND _ _) = N.GD_STOGO_RAND
algorithmEnumOfGlobal (CRS2_LM _ _ _) = N.GN_CRS2_LM
algorithmEnumOfGlobal (ISRES _ _ _ _ _) = N.GN_ISRES
algorithmEnumOfGlobal (ESCH _) = N.GN_ESCH
algorithmEnumOfGlobal (MLSL _ _ _) = N.G_MLSL
algorithmEnumOfGlobal (MLSL_LDS _ _ _) = N.G_MLSL_LDS

applyGlobalObjective :: N.Opt -> GlobalAlgorithm -> IO ()
applyGlobalObjective opt alg = go alg
  where
    obj = tryTo . applyObjective opt . MinimumObjective
    objD = tryTo . applyObjectiveD opt . MinimumObjective

    go (DIRECT o) = obj o
    go (DIRECT_L o) = obj o
    go (DIRECT_NOSCAL o) = obj o
    go (DIRECT_L_NOSCAL o) = obj o
    go (ESCH o) = obj o
    go (STOGO o) = objD o
    go (DIRECT_L_RAND o _) = obj o
    go (DIRECT_L_RAND_NOSCAL o _) = obj o
    go (ORIG_DIRECT o _) = obj o
    go (ORIG_DIRECT_L o _) = obj o
    go (STOGO_RAND o _) = objD o
    go (CRS2_LM o _ _) = obj o
    go (ISRES o _ _ _ _) = obj o
    go (MLSL o _ _) = obj o
    go (MLSL_LDS o _ _) = obj o

applyGlobalAlgorithm :: N.Opt -> GlobalAlgorithm -> IO ()
applyGlobalAlgorithm opt alg = do
  applyGlobalObjective opt alg
  go alg
  where
    seed = applyRandomSeed
    pop = maybe (return ()) (tryTo . applyPopulation opt)
    ic = traverse_ (tryTo . applyConstraint opt)
    ec = traverse_ (tryTo . applyConstraint opt)

    local lp = setupLocalProblem lp >>= N.set_local_optimizer opt

    go (DIRECT_L_RAND _ s) = seed s
    go (DIRECT_L_RAND_NOSCAL _ s) = seed s
    go (ORIG_DIRECT _ ineq) = ic ineq
    go (ORIG_DIRECT_L _ ineq) = ic ineq
    go (STOGO_RAND _ s) = seed s
    go (CRS2_LM _ s p) = seed s *> pop p
    go (ISRES _ ineq eq s p) = ic ineq *> ec eq *> seed s *> pop p
    go (MLSL _ lp p) = local lp *> pop p
    go (MLSL_LDS _ lp p) = local lp *> pop p
    go _ = return ()

tryTo :: IO N.Result -> IO ()
tryTo act = do
  result <- act
  if (N.isSuccess result)
    then return ()
    else Ex.throw $ NloptException result

data NloptException = NloptException N.Result deriving (Show, Typeable)

instance Exception NloptException

-- | Solve the specified global optimization problem.
--
-- = Example program
--
-- The following interactive session example uses the 'ISRES'
-- algorithm, a stochastic, derivative-free global optimizer, to
-- minimize a trivial function with a minimum of 22.0 at @(0, 0)@.
-- The search is conducted within a box from -10 to 10 in each
-- dimension.
--
-- >>> import Numeric.LinearAlgebra ( dot, fromList )
-- >>> let objf x = x `dot` x + 22                         -- define objective
-- >>> let stop = ObjectiveRelativeTolerance 1e-12 :| []   -- define stopping criterion
-- >>> let algorithm = ISRES objf [] [] (SeedValue 22)     -- specify algorithm
-- >>> let lowerbounds = fromList [-10, -10]               -- specify bounds
-- >>> let upperbounds = fromList [10, 10]                 -- specify bounds
-- >>> let problem = GlobalProblem lowerbounds upperbounds stop algorithm
-- >>> let x0 = fromList [5, 8]                            -- specify initial guess
-- >>> minimizeGlobal problem x0
-- Right (Solution {solutionCost = 22.000000000002807, solutionParams = [-1.660591102367038e-6,2.2407062393213684e-7], solutionResult = FTOL_REACHED})
minimizeGlobal ::
  -- | Problem specification
  GlobalProblem ->
  -- | Initial parameter guess
  Vector Double ->
  -- | Optimization results
  Either N.Result Solution
minimizeGlobal prob x0 =
  unsafePerformIO $ (Right <$> minimizeGlobal' prob x0) `Ex.catch` handler
  where
    handler :: NloptException -> IO (Either N.Result a)
    handler (NloptException retcode) = return $ Left retcode

applyGlobalProblem :: N.Opt -> GlobalProblem -> IO ()
applyGlobalProblem opt (GlobalProblem lb ub stop alg) = do
  tryTo $ applyBounds opt (LowerBounds lb)
  tryTo $ applyBounds opt (UpperBounds ub)
  traverse_ (tryTo . applyStoppingCondition opt) stop
  applyGlobalAlgorithm opt alg

newOpt :: N.Algorithm -> Word -> IO N.Opt
newOpt alg sz = do
  opt' <- N.create alg sz
  case opt' of
    Nothing -> Ex.throw $ NloptException N.FAILURE
    Just opt -> return opt

setupGlobalProblem :: GlobalProblem -> IO N.Opt
setupGlobalProblem gp@(GlobalProblem _ _ _ alg) = do
  opt <- newOpt (algorithmEnumOfGlobal alg) (problemSize gp)
  applyGlobalProblem opt gp
  return opt

solveProblem :: N.Opt -> Vector Double -> IO Solution
solveProblem opt x0 = do
  (N.Output outret outcost outx) <- N.optimize opt x0
  if (N.isSuccess outret)
    then return $ Solution outcost outx outret
    else Ex.throw $ NloptException outret

minimizeGlobal' :: GlobalProblem -> Vector Double -> IO Solution
minimizeGlobal' gp x0 = do
  opt <- setupGlobalProblem gp
  solveProblem opt x0

data LocalProblem = LocalProblem
  { -- | The dimension of the
    -- parameter vector.
    lsize :: Word,
    -- | At least one stopping
    -- condition
    lstop :: NonEmpty StoppingCondition,
    -- | Algorithm specification
    lalgorithm :: LocalAlgorithm
  }

-- | These are the local minimization algorithms provided by NLOPT.  Please see
-- <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms the NLOPT algorithm manual>
-- for more details on how the methods work and how they relate to one
-- another.  Note that some local methods require you provide
-- derivatives (gradients or Jacobians) for your objective function
-- and constraint functions.
--
-- Optional parameters are wrapped in a 'Maybe'; for example, if you
-- see 'Maybe' 'VectorStorage', you can simply specify 'Nothing' to
-- use the default behavior.
data LocalAlgorithm
  = -- | Limited-memory BFGS
    LBFGS_NOCEDAL ObjectiveD (Maybe VectorStorage)
  | -- | Limited-memory BFGS
    LBFGS ObjectiveD (Maybe VectorStorage)
  | -- | Shifted limited-memory variable-metric, rank-2
    VAR2 ObjectiveD (Maybe VectorStorage)
  | -- | Shifted limited-memory variable-metric, rank-1
    VAR1 ObjectiveD (Maybe VectorStorage)
  | -- | Truncated Newton's method
    TNEWTON ObjectiveD (Maybe VectorStorage)
  | -- | Truncated Newton's method with automatic restarting
    TNEWTON_RESTART ObjectiveD (Maybe VectorStorage)
  | -- | Preconditioned truncated Newton's method
    TNEWTON_PRECOND ObjectiveD (Maybe VectorStorage)
  | -- | Preconditioned truncated Newton's method with automatic
    -- restarting
    TNEWTON_PRECOND_RESTART ObjectiveD (Maybe VectorStorage)
  | -- | Method of moving averages
    MMA ObjectiveD InequalityConstraintsD
  | -- | Sequential Least-Squares Quadratic Programming
    SLSQP ObjectiveD [Bounds] InequalityConstraintsD EqualityConstraintsD
  | -- | Conservative Convex Separable Approximation
    CCSAQ ObjectiveD Preconditioner
  | -- | PRincipal AXIS gradient-free local optimization
    PRAXIS Objective [Bounds] (Maybe InitialStep)
  | -- | Constrained Optimization BY Linear Approximations
    COBYLA
      Objective
      [Bounds]
      InequalityConstraints
      EqualityConstraints
      (Maybe InitialStep)
  | -- | Powell's NEWUOA algorithm
    NEWUOA Objective (Maybe InitialStep)
  | -- | Powell's NEWUOA algorithm with bounds by SGJ
    NEWUOA_BOUND Objective [Bounds] (Maybe InitialStep)
  | -- | Nelder-Mead Simplex gradient-free method
    NELDERMEAD Objective [Bounds] (Maybe InitialStep)
  | -- | NLOPT implementation of Rowan's Subplex algorithm
    SBPLX Objective [Bounds] (Maybe InitialStep)
  | -- | Bounded Optimization BY Quadratic Approximations
    BOBYQA Objective [Bounds] (Maybe InitialStep)

algorithmEnumOfLocal :: LocalAlgorithm -> N.Algorithm
algorithmEnumOfLocal (LBFGS_NOCEDAL _ _) = N.LD_LBFGS_NOCEDAL
algorithmEnumOfLocal (LBFGS _ _) = N.LD_LBFGS
algorithmEnumOfLocal (VAR2 _ _) = N.LD_VAR2
algorithmEnumOfLocal (VAR1 _ _) = N.LD_VAR1
algorithmEnumOfLocal (TNEWTON _ _) = N.LD_TNEWTON
algorithmEnumOfLocal (TNEWTON_RESTART _ _) = N.LD_TNEWTON_RESTART
algorithmEnumOfLocal (TNEWTON_PRECOND _ _) = N.LD_TNEWTON_PRECOND
algorithmEnumOfLocal (TNEWTON_PRECOND_RESTART _ _) = N.LD_TNEWTON_PRECOND_RESTART
algorithmEnumOfLocal (MMA _ _) = N.LD_MMA
algorithmEnumOfLocal (SLSQP _ _ _ _) = N.LD_SLSQP
algorithmEnumOfLocal (CCSAQ _ _) = N.LD_CCSAQ
algorithmEnumOfLocal (PRAXIS _ _ _) = N.LN_PRAXIS
algorithmEnumOfLocal (COBYLA _ _ _ _ _) = N.LN_COBYLA
algorithmEnumOfLocal (NEWUOA _ _) = N.LN_NEWUOA
algorithmEnumOfLocal (NEWUOA_BOUND _ _ _) = N.LN_NEWUOA
algorithmEnumOfLocal (NELDERMEAD _ _ _) = N.LN_NELDERMEAD
algorithmEnumOfLocal (SBPLX _ _ _) = N.LN_SBPLX
algorithmEnumOfLocal (BOBYQA _ _ _) = N.LN_BOBYQA

applyLocalObjective :: N.Opt -> LocalAlgorithm -> IO ()
applyLocalObjective opt alg = go alg
  where
    obj = tryTo . applyObjective opt . MinimumObjective
    objD = tryTo . applyObjectiveD opt . MinimumObjective
    precond p = tryTo . applyObjectiveD opt . PreconditionedMinimumObjective p

    go (LBFGS_NOCEDAL o _) = objD o
    go (LBFGS o _) = objD o
    go (VAR2 o _) = objD o
    go (VAR1 o _) = objD o
    go (TNEWTON o _) = objD o
    go (TNEWTON_RESTART o _) = objD o
    go (TNEWTON_PRECOND o _) = objD o
    go (TNEWTON_PRECOND_RESTART o _) = objD o
    go (MMA o _) = objD o
    go (SLSQP o _ _ _) = objD o
    go (CCSAQ o prec) = precond prec o
    go (PRAXIS o _ _) = obj o
    go (COBYLA o _ _ _ _) = obj o
    go (NEWUOA o _) = obj o
    go (NEWUOA_BOUND o _ _) = obj o
    go (NELDERMEAD o _ _) = obj o
    go (SBPLX o _ _) = obj o
    go (BOBYQA o _ _) = obj o

applyLocalAlgorithm :: N.Opt -> LocalAlgorithm -> IO ()
applyLocalAlgorithm opt alg = do
  applyLocalObjective opt alg
  go alg
  where
    ic = traverse_ (tryTo . applyConstraint opt)
    icd = traverse_ (tryTo . applyConstraint opt)
    ec = traverse_ (tryTo . applyConstraint opt)
    ecd = traverse_ (tryTo . applyConstraint opt)
    store = maybe (return ()) (tryTo . applyVectorStorage opt)
    bound = traverse_ (tryTo . applyBounds opt)
    step0 = maybe (return ()) (tryTo . applyInitialStep opt)

    go (LBFGS_NOCEDAL _ vs) = store vs
    go (LBFGS _ vs) = store vs
    go (VAR2 _ vs) = store vs
    go (VAR1 _ vs) = store vs
    go (TNEWTON _ vs) = store vs
    go (TNEWTON_RESTART _ vs) = store vs
    go (TNEWTON_PRECOND _ vs) = store vs
    go (TNEWTON_PRECOND_RESTART _ vs) = store vs
    go (MMA _ ineqd) = icd ineqd
    go (SLSQP _ b ineqd eqd) =
      bound b *> icd ineqd *> ecd eqd
    go (CCSAQ _ _) = return ()
    go (PRAXIS _ b s) = bound b *> step0 s
    go (COBYLA _ b ineq eq s) =
      bound b *> ic ineq *> ec eq *> step0 s
    go (NEWUOA _ s) = step0 s
    go (NEWUOA_BOUND _ b s) = bound b *> step0 s
    go (NELDERMEAD _ b s) = bound b *> step0 s
    go (SBPLX _ b s) = bound b *> step0 s
    go (BOBYQA _ b s) = bound b *> step0 s

applyLocalProblem :: N.Opt -> LocalProblem -> IO ()
applyLocalProblem opt (LocalProblem _ stop alg) = do
  traverse_ (tryTo . applyStoppingCondition opt) stop
  applyLocalAlgorithm opt alg

setupLocalProblem :: LocalProblem -> IO N.Opt
setupLocalProblem lp@(LocalProblem sz _ alg) = do
  opt <- newOpt (algorithmEnumOfLocal alg) sz
  applyLocalProblem opt lp
  return opt

minimizeLocal' :: LocalProblem -> Vector Double -> IO Solution
minimizeLocal' lp x0 = do
  opt <- setupLocalProblem lp
  solveProblem opt x0

-- |
-- == Example program
--
-- The following interactive session example enforces the same scalar
-- constraint as the nonlinear constraint example, but this time it
-- uses the SLSQP solver to find the minimum.
--
-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList, scale )
-- >>> let objf x = (x `dot` x + 22, 2 `scale` x)
-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| []
-- >>> let constraintf x = (sum (toList x) - 1.0, fromList [1, 1])
-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6
-- >>> let algorithm = SLSQP objf [] [] [constraint]
-- >>> let problem = LocalProblem 2 stop algorithm
-- >>> let x0 = fromList [5, 10]
-- >>> minimizeLocal problem x0
-- Right (Solution {solutionCost = 22.5, solutionParams = [0.4999999999999998,0.5000000000000002], solutionResult = FTOL_REACHED})
minimizeLocal :: LocalProblem -> Vector Double -> Either N.Result Solution
minimizeLocal prob x0 =
  unsafePerformIO $ (Right <$> minimizeLocal' prob x0) `Ex.catch` handler
  where
    handler :: NloptException -> IO (Either N.Result a)
    handler (NloptException retcode) = return $ Left retcode

class ProblemSize c where
  problemSize :: c -> Word

instance ProblemSize LocalProblem where
  problemSize = lsize

instance ProblemSize GlobalProblem where
  problemSize = fromIntegral . HM.size . lowerBounds

instance ProblemSize AugLagProblem where
  problemSize (AugLagProblem _ _ alg) = case alg of
    AUGLAG_LOCAL lp _ _ -> problemSize lp
    AUGLAG_EQ_LOCAL lp -> problemSize lp
    AUGLAG_GLOBAL gp _ _ -> problemSize gp
    AUGLAG_EQ_GLOBAL gp -> problemSize gp

-- | __IMPORTANT NOTE__
--
-- For augmented lagrangian problems, you, the user, are responsible
-- for providing the appropriate type of constraint.  If the
-- subsidiary problem requires an `ObjectiveD`, then you should
-- provide constraint functions with derivatives.  If the subsidiary
-- problem requires an `Objective`, you should provide constraint
-- functions without derivatives.  If you don't do this, you may get a
-- runtime error.
data AugLagProblem = AugLagProblem
  { -- | Possibly empty set of
    -- equality constraints
    alEquality :: EqualityConstraints,
    -- | Possibly empty set of
    -- equality constraints with
    -- derivatives
    alEqualityD :: EqualityConstraintsD,
    -- | Algorithm specification.
    alalgorithm :: AugLagAlgorithm
  }

-- | The Augmented Lagrangian solvers allow you to enforce nonlinear
-- constraints while using local or global algorithms that don't
-- natively support them.  The subsidiary problem is used to do the
-- minimization, but the @AUGLAG@ methods modify the objective to
-- enforce the constraints.  Please see
-- <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms the NLOPT algorithm manual>
-- for more details on how the methods work and how they relate to one another.
--
-- See the documentation for 'AugLagProblem' for an important note
-- about the constraint functions.
data AugLagAlgorithm
  = -- | AUGmented LAGrangian with a local subsidiary method
    AUGLAG_LOCAL LocalProblem InequalityConstraints InequalityConstraintsD
  | -- | AUGmented LAGrangian with a local subsidiary method and with
    -- penalty functions only for equality constraints
    AUGLAG_EQ_LOCAL LocalProblem
  | -- | AUGmented LAGrangian with a global subsidiary method
    AUGLAG_GLOBAL GlobalProblem InequalityConstraints InequalityConstraintsD
  | -- | AUGmented LAGrangian with a global subsidiary method and with
    -- penalty functions only for equality constraints.
    AUGLAG_EQ_GLOBAL GlobalProblem

algorithmEnumOfAugLag :: AugLagAlgorithm -> N.Algorithm
algorithmEnumOfAugLag (AUGLAG_LOCAL _ _ _) = N.AUGLAG
algorithmEnumOfAugLag (AUGLAG_EQ_LOCAL _) = N.AUGLAG_EQ
algorithmEnumOfAugLag (AUGLAG_GLOBAL _ _ _) = N.AUGLAG
algorithmEnumOfAugLag (AUGLAG_EQ_GLOBAL _) = N.AUGLAG_EQ

-- | This structure is returned in the event of a successful
-- optimization.
data Solution = Solution
  { -- | The objective function value
    -- at the minimum
    solutionCost :: Double,
    -- | The parameter vector which
    -- minimizes the objective
    solutionParams :: Vector Double,
    -- | Why the optimizer stopped
    solutionResult :: N.Result
  }
  deriving (Eq, Show, Read)

applyAugLagAlgorithm :: N.Opt -> AugLagAlgorithm -> IO ()
applyAugLagAlgorithm opt alg = go alg
  where
    ic = traverse_ (tryTo . applyConstraint opt)
    icd = traverse_ (tryTo . applyConstraint opt)
    -- AUGLAG won't work at all if you don't pass it the same
    -- objective as the subproblem -- here we pull out the subproblem
    -- objectives from the algorithm spec and set the same objective
    -- function so the user can't mess it up.
    local lp = tryTo $ do
      localopt <- setupLocalProblem lp
      applyLocalObjective opt (lalgorithm lp)
      N.set_local_optimizer opt localopt
    global gp = do
      tryTo $ setupGlobalProblem gp >>= N.set_local_optimizer opt
      applyGlobalObjective opt (galgorithm gp)

    go (AUGLAG_LOCAL lp ineq ineqd) = local lp *> ic ineq *> icd ineqd
    go (AUGLAG_EQ_LOCAL lp) = local lp
    go (AUGLAG_GLOBAL gp ineq ineqd) = global gp *> ic ineq *> icd ineqd
    go (AUGLAG_EQ_GLOBAL gp) = global gp

applyAugLagProblem :: N.Opt -> AugLagProblem -> IO ()
applyAugLagProblem opt (AugLagProblem eq eqd alg) = do
  traverse_ (tryTo . applyConstraint opt) eq
  traverse_ (tryTo . applyConstraint opt) eqd
  applyAugLagAlgorithm opt alg

minimizeAugLag' :: AugLagProblem -> Vector Double -> IO Solution
minimizeAugLag' ap@(AugLagProblem _ _ alg) x0 = do
  opt <- newOpt (algorithmEnumOfAugLag alg) (problemSize ap)
  applyAugLagProblem opt ap
  solveProblem opt x0

-- |
-- == Example program
--
-- The following interactive session example enforces the same scalar
-- constraint as the nonlinear constraint example, but this time it
-- uses the augmented Lagrangian method to enforce the constraint and
-- the 'SBPLX' algorithm, which does not support nonlinear constraints
-- itself, to perform the minimization.  As before, the parameters
-- must always sum to 1, and the minimizer finds the same constrained
-- minimum of 22.5 at @(0.5, 0.5)@.
-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList )
-- >>> let objf x = x `dot` x + 22
-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| []
-- >>> let algorithm = SBPLX objf [] Nothing
-- >>> let subproblem = LocalProblem 2 stop algorithm
-- >>> let x0 = fromList [5, 10]
-- >>> minimizeLocal subproblem x0
-- Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED})
-- >>>          -- define constraint function:
-- >>> let constraintf x = sum (toList x) - 1.0
-- >>>          -- define constraint object to pass to the algorithm:
-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6
-- >>> let problem = AugLagProblem [constraint] [] (AUGLAG_EQ_LOCAL subproblem)
-- >>> minimizeAugLag problem x0
-- Right (Solution {solutionCost = 22.500000015505844, solutionParams = [0.5000880506776678,0.4999119493223323], solutionResult = FTOL_REACHED})
minimizeAugLag :: AugLagProblem -> Vector Double -> Either N.Result Solution
minimizeAugLag prob x0 =
  unsafePerformIO $ (Right <$> minimizeAugLag' prob x0) `Ex.catch` handler
  where
    handler :: NloptException -> IO (Either N.Result a)
    handler (NloptException retcode) = return $ Left retcode


{-# 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]}

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

run :: Computer -> Computer -- (Integer, Integer) -> Instructions -> [Integer] -> [Integer]
run c
  | state c == Ready = run $ c {state = Running}
  | state c == Done = c
  | state c == Blocked = if input c == [] then c else run $ c {state = Running}
  | otherwise = -- running
    --traceShow (iptr c, base c, memory c) $
    case instr `mod` 100 of
      1 -> -- add
        run $ c {iptr = i+4, memory = insert (addr 3) (arg 1 + arg 2) instructions}
      2 -> -- multiply
        run $ 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 run $ c {iptr = i+2, memory = insert (addr 1) (head $ input c) instructions, input = tail $ input c}
      4 -> -- output
        --traceShow (arg 1) $
        run $ c {iptr = i+2, output = output c ++ [arg 1]}
      5 -> -- jump-if-true
        run $ c {iptr = if arg 1 == 0 then i+3 else arg 2}
      6 -> -- jump-if-false
        run $ c {iptr = if arg 1 == 0 then arg 2 else i+3}
      7 -> -- less than
        run $ c {iptr = i+4, memory = insert (addr 3) (if arg 1 < arg 2 then 1 else 0) instructions}
      8 -> -- equals
        run $ c {iptr = i+4, memory = insert (addr 3) (if arg 1 == arg 2 then 1 else 0) instructions}
      9 -> -- set relative base
        run $ c {iptr = i+2, base = base c + arg 1}
      99 -> -- halt
        run $ 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"

instance (Show a) => Show (Complex a) where
  show (C x y) = "C " ++ show x ++ " " ++ show y

instance (Eq a) => Eq (Complex a) where
  (C a ai) == (C b bi) = a == b && ai == bi

instance (Eq a, Ord a) => Ord (Complex a) where
  compare (C a ai) (C b bi) | a < b || a == b && ai < bi = LT
                            | a > b || a == b && ai > bi = GT
                            | otherwise = EQ

-- Note that we are using a plane where imaginary (y-axis) numbers
-- increase downward. I think this means that right and left are
-- reversed.

data Complex a = C a a

-- instance Num (Complex Int) where
--   abs c = error "Can't take abs of Complex Int"

instance Num a => Num (Complex a) where
  (C x y) + (C u v) = C (x+u) (y+v)
  (C x y) * (C u v) = C (x*u-y*v) (x*v+y*u)
  fromInteger n = C (fromInteger n) 0
  abs (C x y) = C (abs x) y -- sqrt $ x*x + y*y
  signum (C x y) = C (signum x) 0
  negate (C x y) = C (negate x) (negate y)

type Point = Complex Int
type Direction = Complex Int
type Scaffold = Set Point
data Move = R | L | F Int | Err
  deriving (Show, Eq)

show' :: [Move] -> String
show' ms = intercalate "," $ map showMove ms
  where showMove R = "R"
        showMove L = "L"
        showMove (F n) = show n
        showMove Err = error "showing err"

charToDirection '>' = C 1 0
charToDirection '<' = C (-1) 0
charToDirection '^' = C 0 (-1)
charToDirection 'v' = C 0 1
charToDirection c = error $ "bad direction" ++ [c]

fromJust' :: String -> Maybe a -> a
fromJust' s Nothing = error s
fromJust' s (Just x) = x

path :: Scaffold -> (Point, Direction) -> [Move]
path scaffold (pt, dir)
  | S.size scaffold < 2 = []
  -- | otherwise = (traceShowId $ rotation ++ [F steps]) ++ path scaffold' (pt', dir')
  | otherwise = rotation ++ [F steps] ++ path scaffold' (pt', dir')
  where left :: Complex Int
        left = C 0 (-1)
        dir' = fromJust' "dir'" $ find (\d -> (pt + d) `S.member` scaffold) [left^i | i <- [1 .. 4]]
        directions = [(left, [L]), (left^2, [L,L]), (left^3, [R]), (left^4, [Err])]
        rotation = snd $ fromJust' ("rotation: "++ show (dir', directions))  $
          find (\(d, m) -> dir * d == dir') directions
        (pt', steps, scaffold') = walk (pt, 0, scaffold)
        walk (p, n, s) | not $ (p+dir') `S.member` s = (p, n, s)
                       | (p+dir'*left) `S.member` s = walk (p+dir', n+1, s) -- intersection, don't delete
                       | otherwise = walk (p+dir', n+1, S.delete p s)

testIntersection scaffold pt = all (`S.member` scaffold) [ pt + d | d <- [l, l*l, l*l*l, l*l*l*l]]
  where l = C 0 1

data Function = A | B | C' deriving (Show)

substitute :: [Move] -> [Function]
-- substitute (L: F n: ms) = A: (substitute $ F (n-8): ms)
-- substitute (F n: ms) = B n: substitute ms
-- substitute (R: F n: ms) = C': (substitute $ F (n-8): ms)
-- substitute [] = []
-- substitute ms = error $ "substitute: can't do " ++ show ms
substitute (L: F 12: L: F 12: R: F 12: ms) = A: substitute ms
substitute (L: F 8: L: F 8: R: F 12: L: F 8: L: F 8: ms) = B: substitute ms
substitute (L: F 10: R: F 8: R: F 12: ms) = C': substitute ms
substitute [] = []
substitute ms = error $ "substitute: can't do " ++ show ms

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

  putStrLn "Part 1"
  let c = run $ computer0{memory = instructions}
  let view = output c
  let viewStr = map (chr . fromIntegral) view
  putStrLn $ viewStr
  let listView = [(C i j, c) | (j, line) <- zip [0 ..] (lines viewStr)
                              , (i, c) <- zip [0 ..] line
                              , c `elem` "#<>^vV"]

  let scaffold = S.fromList $ map fst listView

  let intersections = S.filter (testIntersection scaffold) scaffold
  print $ S.foldl' (\s (C x y) -> s+x*y) 0 intersections

  putStrLn "Part 2"
  let (r, d) = fromJust' "find robot" $ find (\(pt, v) -> v `elem` "<>^vV") listView
  let p = path scaffold (r, charToDirection d)
  -- print p
  -- let cs = substitute p
  -- print cs
  -- print $ length cs
  let ls = map (fromIntegral . ord) $ unlines
        [ "A,A,B,C,C,A,B,C,A,B"
        , "L,12,L,12,R,12"
        , "L,8,L,8,R,12,L,8,L,8"
        , "L,10,R,8,R,12"
        , "n"
        ]
  print ls
  -- let c = run $ computer0 { memory = insert 0 2 instructions
  --                         , input = ls }

  -- print $ output c

  print "Part 2 - alt"
  print p
  print $ length p
  let substrings xs = concatMap inits $ tails xs
  let functions = Prelude.filter (\l -> length l > 1 && (length $ show' l) < 21) $ substrings p
  print $ length $ functions
  -- let functions3 = [(a,b,c) | a <- functions, b <- functions, show' a < show' b, c <- functions, show' b < show' c]
  let f xs ys = show' xs < show' ys
  let functions3 = [[a,b,c] | a <- inits $ take 10 p, length a > 0, c <- functions, b <- Prelude.filter (f c) functions]
  print $ length $ functions3

  -- print $ Prelude.filter (\[a,b,c] -> (length $ show' a) == 14 && (length $ show' b) == 13 && (length $ show' c) == 20) functions3
  let test3 [] fs = True
      test3 pth fs = any id [test3 (drop (length f) pth) fs | f <- fs, isPrefixOf f pth]
  let goodFunctions = Prelude.filter (test3 p) $ functions3
  let goodFunction = head $ goodFunctions
  print $ goodFunction
  let programs [] fs = [[]]
      programs pth fs = [ (l: ls) | (l, f) <- zip ['A' .. 'C'] fs
                                  , f `isPrefixOf` pth
                                  , ls <- programs (drop (length f) pth) fs
                                  ]
  let program = intersperse ',' $ head $ programs p goodFunction
  let functions = map show' goodFunction
  print program
  print functions
  let ls = unlines ((program: functions) ++ ["n"])
  print ls
  print $ map (fromIntegral . ord) ls
  let c = run $ computer0 { memory = insert 0 2 instructions
                          , input = map (fromIntegral . ord) ls }

  print $ last $ output c
  print $ map (chr . fromIntegral) $ init $ output c

-- After first run:
-- [A,B 4,
--  A,B 4,
--  C',B 4,

--  A,B 4,
--  A,B 4,
--  C',B 4,

--  A,B 0,
--  A,B 0,
--  C',B 4,
--  A,B 0,
--  A,B 0,

--  A,B 2,
--  C',B 0,
--  C',B 4,

--  A,B 2,
--  C',B 0,
--  C',B 4,

--  A,B 4,
--  A,B 4,
--  C',B 4,

--  A,B 0,
--  A,B 0,
--  C',B 4,
--  A,B 0,
--  A,B 0,

--  A,B 2,
--  C',B 0,
--  C',B 4,

--  A,B 4,
--  A,B 4,
--  C',B 4,

--  A,B 0,
--  A,B 0,
--  C',B 4,
--  A,B 0,
--  A,B 0]


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

module CPVO.IO.Reader.Ecalj.Common where
import CPVO.Numeric
import CPVO.IO

import qualified Data.Text as T
import qualified Data.Text.IO as T
import Data.List.Split
import Data.List
import Data.Maybe
import Numeric.LinearAlgebra
import Data.Char (ord)
import Control.Monad.IO.Class (MonadIO)

readCtrlAtoms :: String -> String -> IO [T.Text]
readCtrlAtoms tailer foldernya = do
    fCtrl <- T.readFile $ foldernya ++ "/ctrl." ++ tailer
    return $ catMaybes
      $ map ( T.stripPrefix "ATOM=" .  head)
      $ filter (/=[])
      $ map ( T.words . T.takeWhile (/='#') )
      $ head
      $ splitWhen (T.isPrefixOf "SPEC")
      $ last $ splitWhen (T.isPrefixOf "SITE")
      $ T.lines fCtrl

    -- uniqAtoms : [(count,noFirstAtom,atomicSymbol)]
type CntN1Atsym = (Int,Int,T.Text)
readUniqAtoms :: [T.Text] -> [CntN1Atsym]
readUniqAtoms allAtoms =
          map (\a -> (length a, snd $ head a, fst $ head a)) $
          groupBy (\a b -> fst a == fst b) $
          zip  allAtoms ([1..] :: [Int])

-- daftarCetak :: [((notCtk,((count,noFrstAtom,atomSymbol),label  ,[    AOs ])),spin)]
-- daftarCetak :: [(( 1    ,((  2  ,    13    ,  "Ni"    ),"Ni#2p",["3","4","5"])),1)]
genDaftarCetak :: [CntN1Atsym] -> String -> String -> [String]
               -> IO [((Integer,(CntN1Atsym,String,[String])),Integer)]
genDaftarCetak listAtoms _ _ aos = do
  return $ [ (i,j) | i <- daftarCetak' listAtoms , j <- [1,2] ]
  where
    daftarCetak' lAtoms = zip [1..]
                      $ map (\(a,label,b) -> (head $ filter (\(_,_,aa) -> aa == (T.pack a)) lAtoms , label, b) )
                      $ map ( (\(a:label:as) -> (a,label,as)) . splitOn ":") aos

-- ctrlAtomicAOs :: [(atomNumber,(atomSym,(label ,[intAOs])))]
-- ctrlAtomicAOs :: [(    1     ,(  "O"  ,("O#2p",[ 2,3,4])))]
-- ctrlAtomicAOs :: 14 atoms
genCtrlAtomicAOs :: [(String, String, [Int])]
                 -> [T.Text]
                 -> [(Int, (String, (String, [Int])))]
genCtrlAtomicAOs aoSet ctrlAtoms =  map (\x -> (head $ takeAOs x aoSet))
          $ concat
          $ groupBy (\(_,a:_) (_,b:_) -> (ord a) == (ord b))
          $ zip ([1..]::[Int]) $ map T.unpack ctrlAtoms

-- totalDOS :: Matrix Double [ energy, DOSspinUp, DOSspinDown ]
readTotalDOSText :: String -> String -> IO (Matrix Double)
readTotalDOSText tailer foldernya = loadMatrix $ foldernya ++ "/dos.tot." ++ tailer

-----------------------------------------------------------
getLastLLMF :: MonadIO io =>
  String -> io [T.Text]
getLastLLMF foldernya = inshell2text $ concat ["ls -laht ", foldernya,"/llmf{,_gwscend.*} | head -1|awk '{print $NF}'" ]
-----------------------------------------------------------

--readHeaderData (texFile:jd:jdHead:colAlign:xr:ymax':wTot::invS:tailer:foldernya:aos) = do
readHeaderData :: [String]
               -> IO ( Either String ( Double, Double, Double, Double, [T.Text]
                      , [CntN1Atsym]
                      , [(Int, (String, (String,[Int])))]
                      , String, [String], String, String, String, String))
readHeaderData (texFile:jd:jdHead:colAlign:xr:ymax':_:_:invS:tailer:foldernya:aos) = do
  -------------------------------reading data------------------------
    let invStat = if (invS == "flipSpin") then (-1) else 1
    let ymax = read ymax' :: Double
    let [xmin,xmax] = map (read :: String -> Double) $ splitOn ":" xr
    ctrlAtoms <- readCtrlAtoms tailer foldernya
    let jdHeads = splitOn "|" jdHead
    let uniqAtoms = readUniqAtoms ctrlAtoms
    putStrLn $ show ctrlAtoms
    putStrLn $ show uniqAtoms
    -- daftarCetak : [(nourut,,jumlah,nourut,symbol)]
    daftarCetak <- genDaftarCetak uniqAtoms tailer foldernya aos
    putStrLn $ show aos
    putStrLn $ show $ last daftarCetak
      -------------------------------generating DOS data------------------------
    totalDOS <- readTotalDOSText tailer foldernya
      -------------------------------integrating DOS data------------------------
    let intgTot = map (\i -> integrateToZero $ totalDOS ¿ [0,i]) [1,2] -- run it on spin [1,2]
    putStrLn $ show intgTot
    let aoSet = map ( (\(n:l:as) -> (n,l,map ( ((+) (-1)) . read :: String -> Int) as) ) . splitOn ":") aos
    {-
      -------------------------------generating PDOS data------------------------
              -- map ditambah -1 karena input mengikuti gnuplot
              -- input : d kolom 6-10
              -- gnuplot : d kolom 6-10
              -- hmatrix : d kolom 5-9

              -- ((namaAtom,jdAtom,[intAOs]),[(nourutAtom,namaAtom)])
              -- ((String , String,[ Int  ]),[(Int       , String )])
              -- (("O"    ,"O#2p" ,[2,3,4 ]),[(1         ,"O"     )])
      -}
    let ctrlAtomicAOs = genCtrlAtomicAOs aoSet ctrlAtoms
    let jdTable = "Table: " ++ jd
    putStrLn $ show $ head ctrlAtomicAOs
--    pdosAtomicPilihan <- readPDOS invStat tailer foldernya $ take 2 ctrlAtomicAOs
--    let integratedAtomicPDOS = integrateAtomicPDOS pdosAtomicPilihan
--    putStrLn $ show $ integratedAtomicPDOS
    putStrLn "===done:readHeaderData@CPVO/IO/Reader/Common ====================="
    return $ Right
      (invStat, ymax, xmin, xmax, ctrlAtoms, uniqAtoms, ctrlAtomicAOs,jdTable, jdHeads, foldernya, tailer, colAlign, texFile)

readHeaderData _ = return $ Left
  "===Error:readHeaderData@CPVO/IO/Reader/Common wrong args ========"


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}

import           Data.Bits
import           Data.Convertible
import           Data.List                            (maximumBy)
import           Data.Maybe                           (fromJust)
import           Data.Ord
import qualified Data.Vector.Storable                 as V
import           Data.Word
import           Foreign.ForeignPtr

import           Numeric.LinearAlgebra
import qualified Numeric.LinearAlgebra.Devel          as U
import qualified Numeric.LinearAlgebra.Static         as SA

import           Graphics.Gloss
import           Graphics.Gloss.Interface.IO.Game
import qualified Graphics.Gloss.Interface.IO.Interact as GI

import           Grenade
import           Grenade.Demos.MNIST                  hiding (runNet')
import           Grenade.Layers.Internal.Shrink

data MouseState = MouseDown | MouseUp

-- | all of the information needed to create the canvas to draw on
data World      = Canvas (S ('D2 224 224)) MouseState MNIST

-- | turns the matrix in the world structure into a picture. To do this, we use some pointer
--   magic to turn our matrix (which is represented internally as an array in memory) into
--   a bitmap, where each pixel is represented by four word8s for the R, G, B, A channels
--   (which we construct as a Word32)
renderCanvas :: World -> IO Picture
renderCanvas (Canvas (S2D !arr) _ _)
    = return $ bitmapOfForeignPtr 224 224 (BitmapFormat BottomToTop PxABGR) bitmapPtr False
  where
    -- | given a RealNum p, creates an RGBA word32 of the form (p, p, p, 255) - so it is essentially
    --   creating a greyscale pixel with no transparency.
    convColor :: RealNum -> Word32
    convColor p = let p'  =  convert p :: Word32
                      !w  =  unsafeShiftL p' 24
                         .|. unsafeShiftL p' 16
                         .|. unsafeShiftL p' 8
                         .|. 255
                  in w

    vec         = flatten $ SA.extract arr
    pixels      = V.map convColor vec
    (ptr, _, _) = U.unsafeToForeignPtr pixels
    -- | this last step is to turn the array of word32s into an array of word8s by
    --   interpreting each element as 4 word8s
    bitmapPtr   = castForeignPtr ptr

-- | updates the world with some event
handleInput :: GI.Event -> World -> IO World
-- if the mouse button is pressed down, update the mouse state
handleInput (GI.EventKey (GI.MouseButton GI.LeftButton) (GI.Down) _ _) (Canvas arr _ net)
  = return $ Canvas arr MouseDown net

-- if the mouse button is unpressed, update the mouse state and run a neural network inference
-- on the current canvas image by shrinking the image to the correct input size. Prints
-- the result to the console.
handleInput (GI.EventKey (GI.MouseButton GI.LeftButton) (GI.Up) _ _) (Canvas (S2D arr') _ net) =
  let extractedMatrix = SA.extract arr'
      shrunkenMatrix  = shrink_2d 224 224 28 28 extractedMatrix
      staticMatrix    = fromJust . fromStorableMatrix $ shrunkenMatrix
  in do
    putStrLn $ runNet' net staticMatrix
    return $ Canvas (S2D arr') MouseUp net

-- If the 'c' key is pressed and then let go of, wipe the canvas
handleInput (GI.EventKey (GI.Char 'c') (GI.Up) _ _) (Canvas (S2D !arr) mb net)
  = return $ Canvas cleanCanvas mb net
  where
    cleanCanvas = fromJust . fromStorableMatrix $ cleanMatrix
    cleanMatrix = U.mapMatrixWithIndex (const (const 255)) (SA.extract arr)

-- if moving the mouse while the left mouse button is held down, draw to the
-- canvas (note that we add 112, since the origin of a canvas is at the middle)
handleInput (GI.EventMotion (y, x)) (Canvas arr MouseDown net)
  = return $ Canvas (draw arr (xAsInt + 112) (yAsInt + 112)) MouseDown net
  where
    xAsInt = convert x
    yAsInt = convert y

handleInput _ c = return c

-- | draws a circle with squared radius 50 at the given coordinates
draw :: S ('D2 224 224) -> Int -> Int -> S ('D2 224 224)
draw (S2D arr) x y = fromJust . fromStorableMatrix $ m
  where
    m    = U.mapMatrixWithIndex f arr'
    arr' = SA.extract arr :: Matrix RealNum

    f :: (Int, Int) -> RealNum -> RealNum
    f (x', y') p = if (x - x') ^ (2 :: Int) + (y - y') ^ (2 :: Int) <= 50 then 0 else p

-- | runs a neural network inference and pretty prints the most probable label
--   with its probability
runNet' :: MNIST -> S ('D2 28 28) -> String
runNet' net m
  = let S1D ps = runNet net m
        (p, i) = (getProb . V.toList) (SA.extract ps)
    in  "This number is " ++ show i ++ " with probability " ++ show (p * 100) ++ "%"
  where
    getProb :: [RealNum] -> (RealNum, Int)
    getProb xs = maximumBy (comparing fst) (zip xs [0..])

main :: IO ()
main = do
    mnistPath <- getPathForNetwork MNIST
    net <- (loadSerializedNetwork mnistPath :: IO MNIST)
    putStrLn "Successfully loaded model"

    let initialCanvas = Canvas initialMat MouseUp net

    playIO window backgroundColor 30 initialCanvas renderCanvas handleInput emptyStep
  where
    window            = InWindow "Draw here!" (224, 224) (100, 100)
    backgroundColor   = makeColor 255 255 255 0
    initialMat        = fromJust . fromStorable . V.fromList $ replicate (224 * 224) 255
    emptyStep _ world = return world


{-# LANGUAGE TemplateHaskell, RecordWildCards, ScopedTypeVariables #-}
module TestEvolution where

import qualified Data.ByteString.Lazy          as B
import           Data.MonoTraversable           ( MonoFoldable(olength) )
import           Data.Ord                       ( Down(Down) )
import qualified Data.Set                      as S
import           Data.Typeable                  ( Typeable )
import           Data.Word                      ( Word8 )
import           Instances                      ( )
import           Internal.Evolution
import           Statistics.ConfidenceInt       ( binomialCI )
import           Statistics.Types               ( ConfInt(ConfInt)
                                                , Estimate(Estimate)
                                                , cl95
                                                )
import           Test.QuickCheck                ( Arbitrary(arbitrary)
                                                , CoArbitrary
                                                , Gen
                                                , NonNegative(getNonNegative)
                                                , NonPositive(getNonPositive)
                                                , Positive(getPositive)
                                                , Property
                                                , choose
                                                , cover
                                                , elements
                                                , listOf1
                                                , quickCheckAll
                                                , suchThat
                                                , vectorOf
                                                )
import           Util

newtype ReproduceArgs a = ReproduceArgs (EvolutionParams, [a] -> [Float], [Agent a]) deriving (Show)

instance (Show a, Typeable a, Eq a, Arbitrary a, CoArbitrary a) => Arbitrary (ReproduceArgs a) where
    arbitrary = do
        params@EvolutionParams {..} <- arbitrary
        agents                      <- genPop evolvePopSize =<< arbitrary
        -- generate a fitness fn that strictly increases fitness
        -- fitness starts at 0 so this also guarantees fitness will be > 0
        -- necessary because fitnesses are passed as weights to `frequency`
        f                           <- (arbitrary :: Gen (a -> Positive Float))
        return $ ReproduceArgs (params, map (getPositive . f), agents)

newtype MutateArgs = MutateArgs (Double, [Word8], B.ByteString) deriving (Show, Eq)

instance Arbitrary MutateArgs where
    arbitrary = do
        p     <- choose (0, 1)
        genes <- listOf1 arbitrary
            `suchThat` combineWith (&&) [unique, (> 1) . length]
        genome <- B.pack <$> vectorOf 1000 (elements genes)
        return $ MutateArgs (p, genes, genome)

matchup :: [Agent a] -> [[Agent a]]
--matchup = matchWithBest 2 10 (Down . agentFitness) . S.fromDistinctAscList
matchup = matchups 2 . S.fromDistinctAscList

genPop :: (Arbitrary a) => Int -> Agent a -> Gen [Agent a]
genPop n Agent {..} =
    genPopulation n agentGenes agentEncoder agentDecoder arbitrary

prop_mergeAgentsUnique :: [Agent a] -> Bool
prop_mergeAgentsUnique = unique . mergeAgents

prop_mergeAgentsKeepAll :: [Agent a] -> Bool
prop_mergeAgentsKeepAll xs = all (`elem` mergeAgents xs) xs

prop_mergeAgentsNoNew :: [Agent a] -> Bool
prop_mergeAgentsNoNew xs = all (`elem` xs) (mergeAgents xs)

prop_populationLength :: Arbitrary a => NonNegative Int -> Agent a -> Gen Bool
prop_populationLength n agent = do
    pop <- genPop (getNonNegative n) agent
    return $ length pop == getNonNegative n

prop_populationIds :: Arbitrary a => NonNegative Int -> Agent a -> Gen Bool
prop_populationIds n agent = do
    pop <- genPop (getNonNegative n) agent
    return . and $ zipWith (==) (map agentId pop) [0 ..]

prop_populationUniform
    :: (Eq a, Arbitrary a) => NonNegative Int -> Agent a -> a -> Gen Bool
prop_populationUniform n agent c = do
    pop <- genPop (getNonNegative n) agent
    return $ uniform pop True
  where
    uniform []       acc = acc
    uniform (x : xs) acc = uniform xs (acc && all (agentCommon x) xs)
    agentCommon x y = fEq enc && fEq dec
      where
        fEq f = f x == f y
        enc z = agentEncoder z c
        dec z = agentDecoder z $ enc z

prop_mutate :: MutateArgs -> Gen Property
prop_mutate (MutateArgs (p, genes, genome)) = do
    mutated <- mutate p genes genome
    let Estimate p' (ConfInt lower upper _) = binomialCI
            cl95
            (olength genome)
            (olength $ filter id (B.zipWith (/=) genome mutated))
    -- proportion of mutations == p with 95% confidence
    return $ cover 95 ((p > p' - lower) && (p < p' + upper)) "near p" True

prop_reproduceLength :: ReproduceArgs a -> Gen Bool
prop_reproduceLength (ReproduceArgs (params, f, pop)) = do
    pop' <- reproduce params f matchup pop
    return $ length pop == length pop'

prop_reproduceIds :: ReproduceArgs a -> Gen Bool
prop_reproduceIds (ReproduceArgs (params, f, pop)) =
    fIncreasing True agentId <$> reproduce params f matchup pop
  where
    fIncreasing acc _ []           = acc
    fIncreasing acc _ [_         ] = acc
    fIncreasing acc g (x : y : xs) = fIncreasing (acc && g x < g y) g (y : xs)

prop_evolveId :: ReproduceArgs a -> NonPositive Int -> Gen Bool
prop_evolveId (ReproduceArgs (params, f, pop)) n =
    (pop ==)
        <$> evolve params { evolveGenerations = getNonPositive n } f matchup pop

prop_evolvePreserve :: ReproduceArgs a -> Gen Bool
prop_evolvePreserve (ReproduceArgs (params, f, pop)) = do
    pop' <- evolve params f matchup pop
    let args' = ReproduceArgs (params, f, pop')
    len <- prop_reproduceLength args'
    ids <- prop_reproduceIds args'
    return $ len && ids

return []
runTests :: IO Bool
runTests = $quickCheckAll


{-# LANGUAGE ExtendedDefaultRules #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE BangPatterns         #-}

module TestChapter4DP (
       testChapter4
      )where

import           Control.Concurrent
import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.State

import           Data.Configurator
import           Data.Configurator.Types
import           Data.List.Split (chunksOf)
import           Data.Foldable (toList)

import           Data.Random
import           Data.Random.Distribution
import           Data.Random.Distribution.Poisson

import           Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import           Data.Text (Text)
import qualified Data.Vector.Unboxed as VU

import           Graphics.Matplotlib hiding (def)

import           Numeric.LinearAlgebra (Vector, Matrix)
import qualified Numeric.LinearAlgebra as LA
import           System.Console.AsciiProgress(Options(..), Stats(..),
                                              displayConsoleRegions, complete,
                                              getProgressStats, def, newProgressBar, tickN)
import           Text.Printf
-- project
import           Utils
import           Chapter4DP

testChapter4 :: FilePath -> IO ()
testChapter4 configPath = do
  print "Chapter 4 Experiment Starting "
  (config, _) <- autoReload autoConfig [Required configPath]
  (bCarRental :: Bool) <- require config "enable.bCarRental"
  when bCarRental (doCarRentalTest config)
  (bGambler :: Bool) <- require config "enable.bGamblerProblem"
  when bGambler (doGamblerTest config)

----------------------------------------------------------------------------------------------------
-- Car Rental Test
doCarRentalTest :: Config -> IO ()
doCarRentalTest config = do
  (experimentName::String) <- require config "carRental.experiment"
  (theTheta::Double) <- require config "carRental.theta"
  (discountGamma::Double) <- require config "carRental.discount"
  (maxCars::[Int]) <- require config "carRental.maxCars"
  (rentalCredit::[Double]) <- require config "carRental.rentalCredit"
  (transCost::[Double]) <- require config "carRental.transferCost"
  (maxTransferCars::[Int]) <- require config "carRental.maxTransferCars"
  (freeParkingLimit::[Int]) <- require config "carRental.freeParkingLimit"
  (additionalParkingCost::[Double]) <- require config "carRental.additionalParkingCost"
  (additionalTransferSaving::[Double]) <- require config "carRental.additionalTransferSaving"
  (rentalCars::[Double]) <- require config "carRental.rentalCars"
  (returnCars::[Double]) <- require config "carRental.returnCars"

  let carRental = mkCarRental (length maxCars) theTheta discountGamma maxCars rentalCredit transCost
                              maxTransferCars freeParkingLimit additionalParkingCost
                              additionalTransferSaving rentalCars returnCars
  -- do experiments
  carRental' <- goLoop carRental
  when (_locationNum carRental == 2) (drawTwoLocationCarRental carRental' experimentName)
  threadDelay 100000
  where
  goLoop carRental = displayConsoleRegions $ do
    putStrLn "Will do car rental experiment " >> putStrLn (show carRental)
    pg <- newProgressBar def { pgWidth = 80
                             , pgTotal = 100
                             , pgOnCompletion = Just "Done :percent after :elapsed seconds"
                             }
    carRental' <- loop pg carRental
    putStrLn "car rental experiment finish. Final Results: "
    putStrLn $ showCarRentalResult carRental'
    pure carRental'
    where
    loop pg carRental = do
      let (!(bFinish, percent), !carRental') = runState carRentalStep carRental
      case bFinish of
         False -> do
           stat <- getProgressStats pg
           let ticked = fromInteger $ stCompleted stat
               willTick = percent - ticked
           if (willTick > 0) then(tickN pg willTick)
              else print ("Finish One Improvement " ++ show percent)
           loop pg carRental'
         True -> complete pg >> pure carRental'

drawTwoLocationCarRental :: CarRental -> String -> IO ()
drawTwoLocationCarRental carRental experimentName = do
  putStrLn "Draw Two Location Car Rental Result Graph"
  -- dataX is the second location, whereas dataY is the first location
  let dataX = toList $ fmap (\ (f:s:[]) -> s) (_states carRental)
      dataY = toList $ fmap (\ (f:s:[]) -> f) (_states carRental)
      (dataZ::[Int]) = toList (Seq.zipWith (\ i saPair -> calcZ . foldl (zipWith (+)) [0, 0] $ Seq.index saPair i)
                                           (_actions carRental) (_possibleActions carRental))
      dataValue = toList $ _stateValues carRental
      maxZ = maximum dataZ
      figure = readData ([dataX], [dataY], [dataZ])
                 % mp # "ax = plot.gca(projection='3d')"
                 % mp # "ax.scatter(np.array(data[" # 0 # "]), np.array(data[" # 1 # "]), np.array(data[" # 2 # "]))"
                 % xlabel "#Cars at second location"
                 % ylabel "#Cars at first location"
                 % zlabel "#Cars to move (action)"
                 % xticks [0, 1..(_maxCars carRental !! 1)]
                 % yticks [0, 1..(_maxCars carRental !! 0)]
                 % zticks [negate maxZ, (negate $ maxZ - 1)..maxZ]
                 % title experimentName
      valueFigure = readData ([dataX], [dataY], [dataValue])
                      % mp # "ax = plot.gca(projection='3d')"
                      % mp # "ax.scatter(np.array(data[" # 0 # "]), np.array(data[" # 1 # "]), np.array(data[" # 2 # "]))"
                      % xlabel "#Cars at second location"
                      % ylabel "#Cars at first location"
                      % zlabel "State Values"
                      % xticks [0, 1..(_maxCars carRental !! 1)]
                      % yticks [0, 1..(_maxCars carRental !! 0)]
                      % title (experimentName ++ " - State Values")
  -- avoid Matplotlib's bug
  code figure >> code valueFigure >> onscreen figure >> onscreen valueFigure
  threadDelay 100000
  -- also output table format
  let moveResult = showAsTable [0..(_maxCars carRental)!!0]
                               [0..(_maxCars carRental)!!1] (map fromIntegral dataZ) False
      stateResult = showAsTable [0..(_maxCars carRental)!!0]
                                [0..(_maxCars carRental)!!1] dataValue True
  putStrLn moveResult
  putStrLn stateResult
  where
  calcZ :: [Int] -> Int
  calcZ (firstLocation:secondLocation:[]) =
            (firstLocation > 0) ? (negate firstLocation, secondLocation)
  calcZ _ = error "Not correct action move"

-- output in github markdown format
showAsTable :: [Int] -> [Int] -> [Double] -> Bool -> String
showAsTable col row vals bFloat =
  let colLen = length col
      header = "|First/second| " ++ (concat $ map ((++ "|") . show) col) ++ "\n"
      alignHeader = (concat . take (colLen + 1) $ repeat "|:-----:") ++ "|\n"
      showRows = concat $ map (go colLen vals) row
  in  header ++ alignHeader ++ showRows
  where
  go len vals rowIdx =
    let first = show rowIdx ++ "|"
        rows = take len $ drop (len * rowIdx) vals
        format = bFloat ? ("%7.2f", "%7.0f")
    in  first ++ (concat $ map (\ x -> (printf format x :: String) ++ "|") rows) ++ "|\n"

----------------------------------------------------------------------------------------------------
-- Gambler Problem Test

doGamblerTest :: Config -> IO ()
doGamblerTest config = do
  (prob::Double) <- require config "gambler.headProb"
  (gamblerGoal::Int) <- require config "gambler.goal"
  (theTheta::Double) <- require config "gambler.theta"
  let gambler = mkGambler prob theTheta gamblerGoal
      (_, !gambler') = runState gamblerStep gambler
  drawGamblerGraph gambler'
  threadDelay 100000
  pure ()

drawGamblerGraph :: Gambler -> IO ()
drawGamblerGraph gambler = do
  putStrLn "Draw Gambler Problem Result Graph"
  -- dataX is the second location, whereas dataY is the first location
  let policyFigure = plot [0..(goal gambler)] (stateVals gambler)
                       % xlabel ("Capital (head probability = " ++ (show $ headProb gambler) ++ ")")
                       % ylabel "Value Estimates"
                       % xticks [0, 10..(goal gambler)]
                       % title ("Figure 4.3 - State Values")
      valueFigure = scatter [0..(goal gambler - 1)] (stateActs gambler)
                      % xlabel ("Capital (head probability = " ++ (show $ headProb gambler) ++ ")")
                      % ylabel "Final Policy (Stack)"
                      % xticks [0, 10..(goal gambler)]
                      % title ("Figure 4.3 - Optimal Actions")
  -- avoid Matplotlib's bug
  code policyFigure >> code valueFigure >> onscreen policyFigure >> onscreen valueFigure
  threadDelay 100000


{-# LANGUAGE AllowAmbiguousTypes, CPP, ConstraintKinds, DataKinds,
             FlexibleInstances, GADTs, MultiParamTypeClasses, PolyKinds,
             QuantifiedConstraints, TypeFamilyDependencies, TypeOperators,
             UndecidableInstances #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | The class defining dual components of dual numbers and related classes,
-- type families, constraints and instances. This is a low-level API
-- used to define types and operations in "HordeAd.Core.DualNumber"
-- that is the high-level API.
module HordeAd.Core.DualClass
  ( IsPrimalWithScalar, IsPrimalAndHasFeatures, IsScalar, HasDelta
  , DMode(..), Dual, IsPrimal(..), HasRanks(..)
  , HasVariables(..)  -- use sparringly
  ) where

import Prelude

import qualified Data.Array.Convert
import qualified Data.Array.Dynamic as OTB
import qualified Data.Array.DynamicS as OT
import qualified Data.Array.Shaped as OSB
import qualified Data.Array.ShapedS as OS
import           Data.MonoTraversable (Element, MonoFunctor)
import           Data.Proxy (Proxy)
import qualified Data.Strict.Vector as Data.Vector
import qualified Data.Vector.Generic as V
import           GHC.TypeLits (KnownNat, natVal, type (+))
import           Numeric.LinearAlgebra (Matrix, Numeric, Vector)
import qualified Numeric.LinearAlgebra as HM

import HordeAd.Internal.Delta

-- * Abbreviations to export (not used anywhere below)

-- | The intended semantics (not fully enforced by the constraint in isolation)
-- is that the second type is the primal component of a dual number type
-- at an unknown rank, with the given differentiation mode
-- and underlying scalar.
type IsPrimalWithScalar (d :: DMode) a r =
  (ScalarOf a ~ r, IsPrimal d a, HasVariables a)

-- | A shorthand for a useful set of constraints.
type IsPrimalAndHasFeatures (d :: DMode) a r =
  (IsPrimalWithScalar d a r, RealFloat a, MonoFunctor a, Element a ~ r)

-- | A mega-shorthand for a bundle of connected type constraints.
-- The @Scalar@ in the name means that the second argument is the underlying
-- scalar type of a well behaved (wrt the differentiation mode in the first
-- argument) collection of primal and dual components of dual numbers.
type IsScalar (d :: DMode) r =
  ( HasRanks d r, Ord r, Numeric r, Show r
  , IsPrimalAndHasFeatures d r r
  , IsPrimalAndHasFeatures d (Vector r) r
  , IsPrimalAndHasFeatures d (Matrix r) r
  , IsPrimalAndHasFeatures d (OT.Array r) r
  -- This fragment is for @OS.Array@ and it's irregular, because we can't
  -- mention @sh@ and so fully apply the type constructor.
  , IsPrimalS d r  -- TODO: Floating (OS.Array sh r), MonoFunctor
  )

-- | Is a scalar and will be used to compute gradients via delta-expressions.
type HasDelta r = ( IsScalar 'DModeGradient r
                  , Dual 'DModeGradient r ~ Delta0 r )

-- * Class definitions

-- | The enumeration of all possible differentiation (and more generally,
-- computation with dual numbers) schemes.
data DMode =
    DModeGradient
  | DModeDerivative

-- | The type family that enumerates all possible "ranks"
-- for each differentiation mode.
-- The second type argument is meant to be the primal component
-- of dual numbers. The result is the dual component.
--
-- Rank 0 is special because, in derivatives mode, the dual component
-- is not the primal component wrapped in a datatype or newtype constructor.
-- This makes impossible a representation of primal and dual components as
-- the primal plus the type constructor for creating the dual.
--
-- Rank S is special, because of the extra type parameter @sh@ representing
-- a shape. This is another obstacle to a dual number representation via
-- a single-argument type constructor.
type family Dual (d :: DMode) a = result | result -> d a where
  Dual 'DModeGradient Double = Delta0 Double
  Dual 'DModeGradient Float = Delta0 Float
  Dual 'DModeGradient (Vector r) = Delta1 r
  Dual 'DModeGradient (Matrix r) = Delta2 r
  Dual 'DModeGradient (OT.Array r) = DeltaX r
  Dual 'DModeGradient (OS.Array sh r) = DeltaS sh r
-- not injective:  Dual 'DModeDerivative r = r
  Dual 'DModeDerivative Double = Double
  Dual 'DModeDerivative Float = Float
  Dual 'DModeDerivative (Vector r) = Vector r
  Dual 'DModeDerivative (Matrix r) = Matrix r
  Dual 'DModeDerivative (OT.Array r) = OT.Array r
  Dual 'DModeDerivative (OS.Array sh r) = OS.Array sh r

-- | The underlying scalar of a given primal component of a dual number.
-- A long name to remember not to use, unless necessary, and not to export.
type family ScalarOf a where
  ScalarOf Double = Double
  ScalarOf Float = Float
  ScalarOf (Vector r) = r
  ScalarOf (Matrix r) = r
  ScalarOf (OT.Array r) = r
  ScalarOf (OS.Array sh r) = r

-- | Second argument is a primal component of dual numbers at some rank
-- wrt the differentiation mode given in the first argument.
class IsPrimal d a where
  dZero :: Dual d a
  dScale :: a -> Dual d a -> Dual d a
  dAdd :: Dual d a -> Dual d a -> Dual d a
  dDelay :: Dual d a -> Dual d a

-- | Part 1/2 of a hack to squeeze the shaped tensors rank,
-- with its extra @sh@ parameter, into the 'IsPrimal' class.
class IsPrimalS d r where
  dZeroS :: forall sh. OS.Shape sh => Dual d (OS.Array sh r)
  dScaleS :: forall sh. OS.Shape sh
          => OS.Array sh r -> Dual d (OS.Array sh r) -> Dual d (OS.Array sh r)
  dAddS :: forall sh. OS.Shape sh
        => Dual d (OS.Array sh r) -> Dual d (OS.Array sh r)
        -> Dual d (OS.Array sh r)
  dDelayS :: forall sh. OS.Shape sh
             => Dual d (OS.Array sh r) -> Dual d (OS.Array sh r)

-- | Part 2/2 of a hack to squeeze the shaped tensors rank,
-- with its extra @sh@ parameter, into the 'IsPrimal' class.
instance (IsPrimalS d r, OS.Shape sh) => IsPrimal d (OS.Array sh r) where
  dZero = dZeroS
  dScale = dScaleS
  dAdd = dAddS
  dDelay = dDelayS

-- | Assuming that the first argument is the primal component of dual numbers
-- with the underyling scalar in the second argument and with differentiation
-- mode `DModeGradient`, it additionally admits delta-variable
-- introduction and binding as defined by the methods of the class.
class HasVariables a where
  dVar :: DeltaId a -> Dual 'DModeGradient a
  bindInState :: Dual 'DModeGradient a
              -> DeltaState (ScalarOf a)
              -> (DeltaState (ScalarOf a), DeltaId a )
  dOutline :: CodeOut -> [a] -> [Dual 'DModeGradient a]
           -> Dual 'DModeGradient a

-- | The class provides methods required for the second type parameter
-- to be the underlying scalar of a well behaved collection of dual numbers
-- of various ranks wrt the differentation mode given in the first argument.
class HasRanks (d :: DMode) r where
  dSumElements0 :: Dual d (Vector r) -> Int -> Dual d r
  dIndex0 :: Dual d (Vector r) -> Int -> Int -> Dual d r
  dDot0 :: Vector r -> Dual d (Vector r) -> Dual d r
  dFromX0 :: Dual d (OT.Array r) -> Dual d r
  dFromS0 :: Dual d (OS.Array '[] r) -> Dual d r

  dSeq1 :: Data.Vector.Vector (Dual d r) -> Dual d (Vector r)
  dKonst1 :: Dual d r -> Int -> Dual d (Vector r)
  dAppend1 :: Dual d (Vector r) -> Int -> Dual d (Vector r) -> Dual d (Vector r)
  dSlice1 :: Int -> Int -> Dual d (Vector r) -> Int -> Dual d (Vector r)
  dSumRows1 :: Dual d (Matrix r) -> Int -> Dual d (Vector r)
  dSumColumns1 :: Dual d (Matrix r) -> Int -> Dual d (Vector r)
  dM_VD1 :: Matrix r -> Dual d (Vector r) -> Dual d (Vector r)
  dMD_V1 :: Dual d (Matrix r) -> Vector r -> Dual d (Vector r)
  dFromX1 :: Dual d (OT.Array r) -> Dual d (Vector r)
  dFromS1 :: KnownNat len
          => Dual d (OS.Array '[len] r) -> Dual d (Vector r)
  dReverse1 :: Dual d (Vector r) -> Dual d (Vector r)
  dFlatten1 :: Int -> Int -> Dual d (Matrix r) -> Dual d (Vector r)
  dFlattenX1 :: OT.ShapeL -> Dual d (OT.Array r) -> Dual d (Vector r)
  dFlattenS1 :: OS.Shape sh
             => Dual d (OS.Array sh r) -> Dual d (Vector r)

  dFromRows2 :: Data.Vector.Vector (Dual d (Vector r)) -> Dual d (Matrix r)
  dFromColumns2 :: Data.Vector.Vector (Dual d (Vector r)) -> Dual d (Matrix r)
  dKonst2 :: Dual d r -> (Int, Int) -> Dual d (Matrix r)
  dTranspose2 :: Dual d (Matrix r) -> Dual d (Matrix r)
  dM_MD2 :: Matrix r -> Dual d (Matrix r) -> Dual d (Matrix r)
  dMD_M2 :: Dual d (Matrix r) -> Matrix r -> Dual d (Matrix r)
  dRowAppend2 :: Dual d (Matrix r) -> Int -> Dual d (Matrix r) -> Dual d (Matrix r)
  dColumnAppend2 :: Dual d (Matrix r) -> Int -> Dual d (Matrix r) -> Dual d (Matrix r)
  dRowSlice2 :: Int -> Int -> Dual d (Matrix r) -> Int -> Dual d (Matrix r)
  dColumnSlice2 :: Int -> Int -> Dual d (Matrix r) -> Int -> Dual d (Matrix r)
  dAsRow2 :: Dual d (Vector r) -> Dual d (Matrix r)
  dAsColumn2 :: Dual d (Vector r) -> Dual d (Matrix r)
  dFromX2 :: Dual d (OT.Array r) -> Dual d (Matrix r)
  dFromS2 :: (KnownNat rows, KnownNat cols)
          => Dual d (OS.Array '[rows, cols] r) -> Dual d (Matrix r)

  dFlipud2 :: Dual d (Matrix r) -> Dual d (Matrix r)
  dFliprl2 :: Dual d (Matrix r) -> Dual d (Matrix r)
  dReshape2 :: Int -> Dual d (Vector r) -> Dual d (Matrix r)
  dConv2 :: Matrix r -> Dual d (Matrix r) -> Dual d (Matrix r)

  dKonstX :: Dual d r -> OT.ShapeL -> Dual d (OT.Array r)
  dAppendX :: Dual d (OT.Array r) -> Int -> Dual d (OT.Array r) -> Dual d (OT.Array r)
  dSliceX :: Int -> Int -> Dual d (OT.Array r) -> Int -> Dual d (OT.Array r)
  dIndexX :: Dual d (OT.Array r) -> Int -> Int -> Dual d (OT.Array r)
  dRavelFromListX :: [Dual d (OT.Array r)] -> Dual d (OT.Array r)
  dReshapeX :: OT.ShapeL -> OT.ShapeL -> Dual d (OT.Array r) -> Dual d (OT.Array r)
  dFrom0X :: Dual d r -> Dual d (OT.Array r)
  dFrom1X :: Dual d (Vector r) -> Dual d (OT.Array r)
  dFrom2X :: Dual d (Matrix r) -> Int -> Dual d (OT.Array r)
  dFromSX :: OS.Shape sh
          => Dual d (OS.Array sh r) -> Dual d (OT.Array r)

  dKonstS :: OS.Shape sh
          => Dual d r -> Dual d (OS.Array sh r)
  dAppendS :: (OS.Shape sh, KnownNat m, KnownNat n)
           => Dual d (OS.Array (m ': sh) r) -> Dual d (OS.Array (n ': sh) r)
           -> Dual d (OS.Array ((m + n) ': sh) r)
  dSliceS :: (KnownNat i, KnownNat n, KnownNat k, OS.Shape rest)
          => Proxy i -> Proxy n -> Dual d (OS.Array (i + n + k ': rest) r)
          -> Dual d (OS.Array (n ': rest) r)
  dIndexS :: (KnownNat ix, KnownNat k, OS.Shape rest)
          => Dual d (OS.Array (ix + 1 + k ': rest) r)
          -> Proxy ix -> Dual d (OS.Array rest r)
  dRavelFromListS :: (KnownNat k, OS.Shape rest)
                  => [Dual d (OS.Array rest r)]
                  -> Dual d (OS.Array (k : rest) r)
  dReshapeS :: (OS.Shape sh, OS.Shape sh', OS.Size sh ~ OS.Size sh')
            => Dual d (OS.Array sh r) -> Dual d (OS.Array sh' r)
  dFrom0S :: Dual d r -> Dual d (OS.Array '[] r)
  dFrom1S :: KnownNat n => Dual d (Vector r) -> Dual d (OS.Array '[n] r)
  dFrom2S :: (KnownNat rows, KnownNat cols)
          => Proxy cols
          -> Dual d (Matrix r) -> Dual d (OS.Array '[rows, cols] r)
  dFromXS :: OS.Shape sh => Dual d (OT.Array r) -> Dual d (OS.Array sh r)

-- * Backprop gradient method instances

instance IsPrimal 'DModeGradient Double where
  dZero = Zero0
  dScale = Scale0
  dAdd = Add0
  dDelay = Delay0

instance IsPrimal 'DModeGradient Float where
  -- Identical as above:
  dZero = Zero0
  dScale = Scale0
  dAdd = Add0
  dDelay = Delay0

instance IsPrimal 'DModeGradient (Vector r) where
  dZero = Zero1
  dScale = Scale1
  dAdd = Add1
  dDelay = Delay1

instance IsPrimal 'DModeGradient (Matrix r) where
  dZero = Zero2
  dScale = Scale2
  dAdd = Add2
  dDelay = Delay2

instance IsPrimal 'DModeGradient (OT.Array r) where
  dZero = ZeroX
  dScale = ScaleX
  dAdd = AddX
  dDelay = DelayX

instance IsPrimalS 'DModeGradient r where
  dZeroS = ZeroS
  dScaleS = ScaleS
  dAddS = AddS
  dDelayS = DelayS

instance HasVariables Double where
  dVar = Var0
  {-# INLINE bindInState #-}
  bindInState = bindInState0
  dOutline = Outline0

instance HasVariables Float where
  dVar = Var0
  {-# INLINE bindInState #-}
  bindInState = bindInState0
  dOutline = Outline0

instance HasVariables (Vector r) where
  dVar = Var1
  {-# INLINE bindInState #-}
  bindInState = bindInState1
  dOutline = Outline1

instance HasVariables (Matrix r) where
  dVar = Var2
  {-# INLINE bindInState #-}
  bindInState = bindInState2
  dOutline = Outline2

instance HasVariables (OT.Array r) where
  dVar = VarX
  {-# INLINE bindInState #-}
  bindInState = bindInStateX
  dOutline = OutlineX

instance OS.Shape sh => HasVariables (OS.Array sh r) where
  dVar = VarS
  {-# INLINE bindInState #-}
  bindInState u' st = let (st2, did) = bindInStateX (FromSX u') st
                      in (st2, convertDeltaId did)
  dOutline = OutlineS

instance Dual 'DModeGradient r ~ Delta0 r
         => HasRanks 'DModeGradient r where
  dSumElements0 = SumElements0
  dIndex0 = Index0
  dDot0 = Dot0
  dFromX0 = FromX0
  dFromS0 = FromS0
  dSeq1 = Seq1
  dKonst1 = Konst1
  dAppend1 = Append1
  dSlice1 = Slice1
  dSumRows1 = SumRows1
  dSumColumns1 = SumColumns1
  dM_VD1 = M_VD1
  dMD_V1 = MD_V1
  dFromX1 = FromX1
  dFromS1 = FromS1
  dReverse1 = Reverse1
  dFlatten1 = Flatten1
  dFlattenX1 = FlattenX1
  dFlattenS1 = FlattenS1
  dFromRows2 = FromRows2
  dFromColumns2 = FromColumns2
  dKonst2 = Konst2
  dTranspose2 = Transpose2
  dM_MD2 = M_MD2
  dMD_M2 = MD_M2
  dRowAppend2 = RowAppend2
  dColumnAppend2 = ColumnAppend2
  dRowSlice2 = RowSlice2
  dColumnSlice2 = ColumnSlice2
  dAsRow2 = AsRow2
  dAsColumn2 = AsColumn2
  dFromX2 = FromX2
  dFromS2 = FromS2
  dFlipud2 = Flipud2
  dFliprl2 = Fliprl2
  dReshape2 = Reshape2
  dConv2 = Conv2
  dKonstX = KonstX
  dAppendX = AppendX
  dSliceX = SliceX
  dIndexX = IndexX
  dRavelFromListX = RavelFromListX
  dReshapeX = ReshapeX
  dFrom0X = From0X
  dFrom1X = From1X
  dFrom2X = From2X
  dFromSX = FromSX
  dKonstS = KonstS
  dAppendS = AppendS
  dSliceS = SliceS
  dIndexS = IndexS
  dRavelFromListS = RavelFromListS
  dReshapeS = ReshapeS
  dFrom0S = From0S
  dFrom1S = From1S
  dFrom2S = From2S
  dFromXS = FromXS

-- * Alternative instances: forward derivatives computed on the spot

instance IsPrimal 'DModeDerivative Double where
  dZero = 0
  dScale k d = k * d
  dAdd d e = d + e
  dDelay = id  -- no delaying

instance IsPrimal 'DModeDerivative Float where
  dZero = 0
  dScale k d = k * d
  dAdd d e = d + e
  dDelay = id

-- These constraints force @UndecidableInstances@.
instance Num (Vector r)
         => IsPrimal 'DModeDerivative (Vector r) where
  dZero = 0
  dScale k d = k * d
  dAdd d e = d + e
  dDelay = id

instance Num (Matrix r)
         => IsPrimal 'DModeDerivative (Matrix r) where
  dZero = 0
  dScale k d = k * d
  dAdd d e = d + e
  dDelay = id

instance Num (OT.Array r)
         => IsPrimal 'DModeDerivative (OT.Array r) where
  dZero = 0
  dScale k d = k * d
  dAdd d e = d + e
  dDelay = id

instance (Numeric r, Num (Vector r))
         => IsPrimalS 'DModeDerivative r where
  dZeroS = 0
  dScaleS k d = k * d
  dAddS d e = d + e
  dDelayS = id

instance ( Numeric r, Num (Vector r)
         , Dual 'DModeDerivative r ~ r )
         => HasRanks 'DModeDerivative r where
  dSumElements0 vd _ = HM.sumElements vd
  dIndex0 d ix _ = d V.! ix
  dDot0 = (HM.<.>)
  dFromX0 = OT.unScalar
  dFromS0 = OS.unScalar
  dSeq1 = V.convert
  dKonst1 = HM.konst
  dAppend1 d _k e = d V.++ e
  dSlice1 i n d _len = V.slice i n d
  dM_VD1 = (HM.#>)
  dMD_V1 = (HM.#>)
  dSumRows1 dm _cols = V.fromList $ map HM.sumElements $ HM.toRows dm
  dSumColumns1 dm _rows = V.fromList $ map HM.sumElements $ HM.toColumns dm
  dFromX1 = OT.toVector
  dFromS1 = OS.toVector
  dReverse1 = V.reverse
  dFlatten1 _rows _cols = HM.flatten
  dFlattenX1 _sh = OT.toVector
  dFlattenS1 = OS.toVector
  dFromRows2 = HM.fromRows . V.toList
  dFromColumns2 = HM.fromColumns . V.toList
  dKonst2 = HM.konst
  dTranspose2 = HM.tr'
  dM_MD2 = (HM.<>)
  dMD_M2 = (HM.<>)
  dAsRow2 = HM.asRow
  dAsColumn2 = HM.asColumn
  dRowAppend2 d _k e = d HM.=== e
  dColumnAppend2 d _k e = d HM.||| e
  dRowSlice2 i n d _rows = HM.takeRows n $ HM.dropRows i d
  dColumnSlice2 i n d _cols = HM.takeColumns n $ HM.dropColumns i d
  dFromX2 d = case OT.shapeL d of
    [_rows, cols] -> HM.reshape cols $ OT.toVector d
    _ -> error "dFromX2: wrong tensor dimensions"
  dFromS2 d = case OS.shapeL d of
    [_rows, cols] -> HM.reshape cols $ OS.toVector d
    _ -> error "dFromS2: wrong tensor dimensions"
  dFlipud2 = HM.flipud
  dFliprl2 = HM.fliprl
  dReshape2 = HM.reshape
  dConv2 = HM.conv2
  dKonstX d sz = OT.constant sz d
  dAppendX d _k e = d `OT.append` e
  dSliceX i n d _len = OT.slice [(i, n)] d
  dIndexX d ix _len = OT.index d ix
  dRavelFromListX ld =
    let sh = case ld of
          d : _ -> length ld : OT.shapeL d
          [] -> []
    in OT.ravel $ OTB.fromList sh ld
  dReshapeX _sh = OT.reshape
  dFrom0X = OT.scalar
  dFrom1X d = OT.fromVector [V.length d] d
  dFrom2X d cols = OT.fromVector [HM.rows d, cols] $ HM.flatten d
  dFromSX = Data.Array.Convert.convert
  dKonstS = OS.constant
  dAppendS = OS.append
  dSliceS (_ :: Proxy i) (_ :: Proxy n) = OS.slice @'[ '(i, n) ]
  dIndexS d proxyIx = OS.index d (fromInteger $ natVal proxyIx)
  dRavelFromListS = OS.ravel . OSB.fromList
  dReshapeS = OS.reshape
  dFrom0S = OS.scalar
  dFrom1S = OS.fromVector
  dFrom2S _ = OS.fromVector . HM.flatten
  dFromXS = Data.Array.Convert.convert


{-|
  Mapping of basic Matlab types to Haskell types.
-}
module Foreign.Matlab.Types (
    -- * Representations of Matlab types
    -- |Most types are simple aliases of Haskell types, given simply to identify them and provide a consistent naming scheme.

    MIO,

    MXClass(..),
    MChar,
    MLogical,
    MScalar,

    MNumeric,
    -- * Numeric types
    MDouble,
    MSingle,
    MInt8,
    MInt16,
    MInt32,
    MInt64,
    MUint8,
    MUint16,
    MUint32,
    MUint64,
    MComplex,

    -- * Array indexing
    MSize, MSubs,
    MIndex(..),
    mStart, mOffset,
    mSizeRange, mRangeSize,
    normMSize, realMSize,

    -- * Opaque types
    MXArray,
    MAny, MAnyArray,
    MCell(MCell), mCell,
    MStruct, mStruct,
    MFun
  ) where

import Data.Complex
import Data.Ix
import Foreign.Matlab.Internal

-- |The class of types which are simple Matlab scalars and can be array elements
class MScalar a
-- |The class of types which Matlab consisters \"numeric\"
class (MScalar a, Num a) => MNumeric a

instance MScalar MChar
instance MScalar MLogical
instance MScalar MDouble
instance MScalar MSingle
instance MScalar MInt8
instance MScalar MInt16
instance MScalar MInt32
instance MScalar MInt64
instance MScalar MUint8
instance MScalar MUint16
instance MScalar MUint32
instance MScalar MUint64
instance MNumeric a => MScalar (MComplex a)

instance MNumeric MDouble
instance MNumeric MSingle
instance MNumeric MInt8
instance MNumeric MInt16
instance MNumeric MInt32
instance MNumeric MInt64
instance MNumeric MUint8
instance MNumeric MUint16
instance MNumeric MUint32
instance MNumeric MUint64

-- |Complex numeric types.  Unfortunately, 'Complex' only applies to 'RealFloat' types, whereas Matlab allows any numeric type, so some types are (currently) unaccessable.
type MComplex = Complex

-- |The type of array sizes, which are the column-major lengths of each dimensions
type MSize = [Int]
-- |The type of array index subscripts, which are the column-major, 0-based indices in each dimension with normal Matlab semantics (flatten along last dimension)
type MSubs = [Int]
-- |Ways to index an array.  Using Matlab semantics, a singleton MSubs [n] is equivalent to a raw 0-based offset (MOffset n).
newtype MIndex = MSubs MSubs
    -- MOffset Int -- ^0-based offset from the beginning of the array (equivalent to singleton MSubs)

-- |First index in array
mStart :: MIndex
mStart = MSubs []

-- |Raw, 0-based array offset index
mOffset :: Int -> MIndex
mOffset i = MSubs [i]

instance Show MIndex where
  showsPrec _ (MSubs []) = id
  showsPrec _ (MSubs l) = showChar '(' . foldr1 (\s r -> s . showChar ',' . r) (map shows l) . showChar ')'
instance Eq MIndex where
  MSubs a == MSubs b = eq a b where
    eq [] [] = True
    eq (x:a) (y:b) = (x == y) && eq a b
    eq [] b = eq [0] b
    eq a [] = eq a [0]
instance Ord MIndex where
  compare (MSubs a) (MSubs b) = cmp a b where
    cmp [] [] = EQ
    cmp (x:a) (y:b) = case cmp a b of { EQ -> compare x y ; r -> r }
    cmp [] b = cmp [0] b
    cmp a [] = cmp a [0]
instance Ix MIndex where
  range (MSubs a, MSubs b) = map MSubs $ rng a b where
    rng [] [] = [[]]
    rng (x:a) (y:b) = concatMap (\l -> map (:l) [x..y]) $ rng a b
    rng [] b = rng [0] b
    rng _ _ = error "MIndex.range: length mismatch"
  index (MSubs a, MSubs b) (MSubs l) = idx a b l where
    idx _ _ [] = 0
    idx (x:a) (y:b) (i:l) = i-x+((y-x+1)*idx a b l)
    idx [] [] (0:l) = idx [] [] l
    idx [] b l = idx [0] b l
    idx _ _ _ = error "MIndex.index: length mismatch"
  inRange (MSubs a, MSubs b) (MSubs l) = inr a b l where
    inr _ _ [] = True
    inr a@(x:_) b@(_:_) [i] = i >= x && i-x < rangeSize (MSubs a, MSubs b)
    inr (x:a) (y:b) (i:l) = i >= x && i <= y && inr a b l
    inr [] [] (0:l) = inr [] [] l
    inr [] b l = inr [0] b l
    inr _ _ _ = error "MIndex.inRange: length mismatch"
  rangeSize (MSubs a, MSubs b) = rsz a b where
    rsz [] b = product $ map succ b
    rsz (x:a) (y:b) = (y-x+1)*rsz a b
    rsz _ _ = error "MIndex.rangeSize: length mismatch"

-- |Convert an array size to an index range, which will be of the form ((0,0,0...),(i-1,j-1,k-1,...))
mSizeRange :: MSize -> (MIndex,MIndex)
mSizeRange l = (MSubs [], MSubs $ map pred $ normMSize l)

-- |Convert an index range to an array size
mRangeSize :: (MIndex,MIndex) -> MSize
mRangeSize (MSubs [], MSubs l) = map succ l
mRangeSize _ = error "mRangeSize: invalid lower bound"

-- |Get the size form that Matlab likes (@length (realMSize s) >= 2@)
realMSize :: MSize -> MSize
realMSize [] = [1,1]
realMSize [n] = [n,1]
realMSize s = s

-- |Get a more useful size form (no trailing singletons)
normMSize :: MSize -> MSize
normMSize [] = []
normMSize (1:l) = case normMSize l of { [] -> [] ; l -> (1:l) }
normMSize (x:l) = x : normMSize l