{-# OPTIONS_GHC -Wall -Wno-partial-type-signatures #-}

{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PartialTypeSignatures #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Sundials.CVode.ODE
-- Copyright   :  Dominic Steinitz 2018,
--                Novadiscovery 2018
-- License     :  BSD
-- Maintainer  :  Dominic Steinitz
-- Stability   :  provisional
--
-- Solution of ordinary differential equation (ODE) initial value problems.
--
-- <https://computation.llnl.gov/projects/sundials/sundials-software>
--
-- A simple example:
--
-- <<diagrams/brusselator.png#diagram=brusselator&height=400&width=500>>
--
-- @
-- import           Numeric.Sundials.CVode.ODE
-- import           Numeric.LinearAlgebra
--
-- import           Plots as P
-- import qualified Diagrams.Prelude as D
-- import           Diagrams.Backend.Rasterific
--
-- brusselator :: Double -> [Double] -> [Double]
-- brusselator _t x = [ a - (w + 1) * u + v * u * u
--                    , w * u - v * u * u
--                    , (b - w) / eps - w * u
--                    ]
--   where
--     a = 1.0
--     b = 3.5
--     eps = 5.0e-6
--     u = x !! 0
--     v = x !! 1
--     w = x !! 2
--
-- lSaxis :: [[Double]] -> P.Axis B D.V2 Double
-- lSaxis xs = P.r2Axis &~ do
--   let ts = xs!!0
--       us = xs!!1
--       vs = xs!!2
--       ws = xs!!3
--   P.linePlot' $ zip ts us
--   P.linePlot' $ zip ts vs
--   P.linePlot' $ zip ts ws
--
-- main = do
--   let res1 = odeSolve brusselator [1.2, 3.1, 3.0] (fromList [0.0, 0.1 .. 10.0])
--   renderRasterific "diagrams/brusselator.png"
--                    (D.dims2D 500.0 500.0)
--                    (renderAxis $ lSaxis $ [0.0, 0.1 .. 10.0]:(toLists $ tr res1))
-- @
--
-----------------------------------------------------------------------------
module Numeric.Sundials.CVode.ODE ( odeSolve
                                   , odeSolveV
                                   , odeSolveVWith
                                   , odeSolveVWith'
                                   , odeSolveRootVWith'
                                   , odeSolveWithEvents
                                   , ODEMethod(..)
                                   , StepControl(..)
                                   , SolverResult(..)
                                   , SundialsSolution(..)
                                   , EventSpec(..)
                                   , EventInfo(..)
                                   , CrossingDirection(..)
                                   ) where

import qualified Language.C.Inline as C
import qualified Language.C.Inline.Unsafe as CU

import           Data.Monoid ((<>))
import           Data.Maybe (isJust, fromJust)
import           Data.List (genericLength)

import           Foreign.C.Types (CDouble, CInt, CLong)
import           Foreign.Ptr (Ptr)
import           Foreign.Storable (peek, poke)

import qualified Data.Vector.Storable as V

import           Data.Coerce (coerce)
import           System.IO.Unsafe (unsafePerformIO)

import           Numeric.LinearAlgebra.Devel (createVector)

import           Numeric.LinearAlgebra.HMatrix (Vector, Matrix, toList, rows,
                                                cols, toLists, size, reshape,
                                                subVector, subMatrix, toColumns)

import           Numeric.Sundials.Arkode (cV_ADAMS, cV_BDF,
                                          vectorToC, cV_SUCCESS,
                                          SunVector(..))
import qualified Numeric.Sundials.Arkode as T
import           Numeric.Sundials.ODEOpts

C.context (C.baseCtx <> C.vecCtx <> C.funCtx <> T.sunCtx)

C.include "<stdlib.h>"
C.include "<stdio.h>"
C.include "<math.h>"
C.include "<cvode/cvode.h>"               -- prototypes for CVODE fcts., consts.
C.include "<nvector/nvector_serial.h>"    -- serial N_Vector types, fcts., macros
C.include "<sunmatrix/sunmatrix_dense.h>" -- access to dense SUNMatrix
C.include "<sunlinsol/sunlinsol_dense.h>" -- access to dense SUNLinearSolver
C.include "<cvode/cvode_direct.h>"        -- access to CVDls interface
C.include "<cvode/cvode_diag.h>"          -- CVDiag
C.include "<sundials/sundials_types.h>"   -- definition of type realtype
C.include "<sundials/sundials_math.h>"
C.include "../../../helpers.h"
C.include "Numeric/Sundials/Arkode_hsc.h"

-- | The direction in which a function should cross the x axis
data CrossingDirection = Upwards | Downwards | AnyDirection
  deriving (Eq, Show)

-- Contrary to the documentation, it appears that CVodeGetRootInfo
-- may use both 1 and -1 to indicate a root, depending on the
-- direction of the sign change. See near the end of cvRootfind.
intToDirection :: Integral d => d -> Maybe CrossingDirection
intToDirection d =
  case d of
    1  -> Just Upwards
    -1 -> Just Downwards
    _  -> Nothing

-- | Almost inverse of 'intToDirection'. Map 'Upwards' to 1, 'Downwards' to
-- -1, and 'AnyDirection' to 0.
directionToInt :: Integral d => CrossingDirection -> d
directionToInt d =
  case d of
    Upwards -> 1
    Downwards -> -1
    AnyDirection -> 0

data SundialsSolution =
  SundialsSolution
  { actualTimeGrid :: V.Vector Double                 -- ^ actual time grid returned by the solver (with duplicated event times)
  , solutionMatrix :: Matrix Double                -- ^ matrix of solutions: each column is an unknwown (add also the time vector?)
  , eventInfo      :: [EventInfo]         -- ^ event infos, as many items as triggered events during the simulation
  , diagnostics    :: SundialsDiagnostics -- ^ usual Sundials diagnostics
  }

data EventInfo =
  EventInfo
  { eventTime     :: !Double            -- ^ time at which event was triggered
  , eventIndex    :: !Int               -- ^ which index was triggered
  , rootDirection :: !CrossingDirection -- ^ in which direction ((+)->(-) or (-)->(+)) the root is crossed
  }
  deriving Show

getMethod :: ODEMethod -> Int
getMethod (ADAMS) = cV_ADAMS
getMethod (BDF)   = cV_BDF

getJacobian :: ODEMethod -> Maybe Jacobian
getJacobian _ = Nothing

-- | A version of 'odeSolveVWith' with reasonable default step control.
odeSolveV
    :: ODEMethod
    -> Maybe Double      -- ^ initial step size - by default, CVode
                         -- estimates the initial step size to be the
                         -- solution \(h\) of the equation
                         -- \(\|\frac{h^2\ddot{y}}{2}\| = 1\), where
                         -- \(\ddot{y}\) is an estimated value of the
                         -- second derivative of the solution at \(t_0\)
    -> Double            -- ^ absolute tolerance for the state vector
    -> Double            -- ^ relative tolerance for the state vector
    -> (Double -> Vector Double -> Vector Double) -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
    -> Vector Double     -- ^ initial conditions
    -> Vector Double     -- ^ desired solution times
    -> Matrix Double     -- ^ solution
odeSolveV meth hi epsAbs epsRel f y0 ts =
  odeSolveVWith meth (X epsAbs epsRel) hi g y0 ts
  where
    g t x0 = coerce $ f t x0

-- | A version of 'odeSolveV' with reasonable default parameters and
-- system of equations defined using lists. FIXME: we should say
-- something about the fact we could use the Jacobian but don't for
-- compatibility with hmatrix-gsl.
odeSolve :: (Double -> [Double] -> [Double]) -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
         -> [Double]                         -- ^ initial conditions
         -> Vector Double                    -- ^ desired solution times
         -> Matrix Double                    -- ^ solution
odeSolve f y0 ts =
  -- FIXME: These tolerances are different from the ones in GSL
  odeSolveVWith BDF (XX' 1.0e-6 1.0e-10 1 1)  Nothing g (V.fromList y0) (V.fromList $ toList ts)
  where
    g t x0 = V.fromList $ f t (V.toList x0)

-- | A version of 'odeSolveVWith'' with reasonable default solver
-- options.
odeSolveVWith ::
  ODEMethod
  -> StepControl
  -> Maybe Double -- ^ initial step size - by default, CVode
                  -- estimates the initial step size to be the
                  -- solution \(h\) of the equation
                  -- \(\|\frac{h^2\ddot{y}}{2}\| = 1\), where
                  -- \(\ddot{y}\) is an estimated value of the second
                  -- derivative of the solution at \(t_0\)
  -> (Double -> V.Vector Double -> V.Vector Double) -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
  -> V.Vector Double                     -- ^ Initial conditions
  -> V.Vector Double                     -- ^ Desired solution times
  -> Matrix Double                       -- ^ Error code or solution
odeSolveVWith method control initStepSize f y0 tt =
  case odeSolveVWith' opts f y0 tt of
    Left  (c, _v) -> error $ show c -- FIXME
    Right (v, _d) -> v
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod   = method
                   , stepControl = control
                   , initStep    = initStepSize
                   }

odeSolveVWith' ::
  ODEOpts
  -> (Double -> V.Vector Double -> V.Vector Double) -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
  -> V.Vector Double                     -- ^ Initial conditions
  -> V.Vector Double                     -- ^ Desired solution times
  -> Either (Matrix Double, Int) (Matrix Double, SundialsDiagnostics) -- ^ Error code or solution
odeSolveVWith' opts f y0 tt =
  case solveOdeC (fromIntegral $ maxFail opts)
                  (fromIntegral $ maxNumSteps opts) (coerce $ minStep opts)
                  (fromIntegral . getMethod . odeMethod $ opts) (coerce $ initStep opts) jacH (scise $ stepControl opts)
                  (coerce f) (coerce y0)
                  0 (\_ x -> x) [] 0 (\_ _ y -> y) (coerce tt) of
    -- Remove the time column for backwards compatibility
    SolverError m c         -> Left
                               ( subMatrix (0, 1) (V.length tt, l) m
                               , fromIntegral c
                               )
    SolverSuccess _ m d     -> Right
                               ( subMatrix (0, 1) (V.length tt, l) m
                               , d
                               )
  where
    l = size y0
    scise (X aTol rTol)                          = coerce (V.replicate l aTol, rTol)
    scise (X' aTol rTol)                         = coerce (V.replicate l aTol, rTol)
    scise (XX' aTol rTol yScale _yDotScale)      = coerce (V.replicate l aTol, yScale * rTol)
    -- FIXME; Should we check that the length of ss is correct?
    scise (ScXX' aTol rTol yScale _yDotScale ss) = coerce (V.map (* aTol) ss, yScale * rTol)
    jacH = fmap (\g t v -> matrixToSunMatrix $ g (coerce t) (coerce v)) $
           getJacobian $ odeMethod opts

matrixToSunMatrix :: Matrix Double -> T.SunMatrix
matrixToSunMatrix m = T.SunMatrix { T.rows = nr, T.cols = nc, T.vals = vs }
  where
    nr = fromIntegral $ rows m
    nc = fromIntegral $ cols m
    -- FIXME: efficiency
    vs = V.fromList $ map coerce $ concat $ toLists m

solveOdeC ::
  CInt ->
  CLong ->
  CDouble ->
  CInt ->
  Maybe CDouble ->
  (Maybe (CDouble -> V.Vector CDouble -> T.SunMatrix)) ->
  (V.Vector CDouble, CDouble) ->
  (CDouble -> V.Vector CDouble -> V.Vector CDouble) -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
  -> V.Vector CDouble -- ^ Initial conditions
  -> CInt -- ^ Number of event equations
  -> (CDouble -> V.Vector CDouble -> V.Vector CDouble) -- ^ The event equations themselves
  -> [CrossingDirection] -- ^ The required crossing direction for each event
  -> CInt -- ^ Maximum number of events
  -> (Int -> CDouble -> V.Vector CDouble -> V.Vector CDouble)
      -- ^ Function to reset/update the state when an event occurs. The
      -- 'Int' argument is the 0-based number of the event that has
      -- occurred. If multiple events have occurred at the same time, they
      -- are handled in the increasing order of the event index. The other
      -- arguments are the time and the point in the state space. Return
      -- the updated point in the state space.
  -> V.Vector CDouble -- ^ Desired solution times
  -> SolverResult
solveOdeC maxErrTestFails maxNumSteps_ minStep_ method initStepSize
          jacH (aTols, rTol) fun f0 nr event_fn directions max_events apply_event ts =
  unsafePerformIO $ do

  let isInitStepSize :: CInt
      isInitStepSize = fromIntegral $ fromEnum $ isJust initStepSize
      ss :: CDouble
      ss = case initStepSize of
             -- It would be better to put an error message here but
             -- inline-c seems to evaluate this even if it is never
             -- used :(
             Nothing -> 0.0
             Just x  -> x

  let dim = V.length f0
      nEq :: CLong
      nEq = fromIntegral dim
      nTs :: CInt
      nTs = fromIntegral $ V.length ts
  output_mat_mut :: V.MVector _ CDouble <- V.thaw =<< createVector ((1 + fromIntegral dim) * (fromIntegral (2 * max_events) + fromIntegral nTs))
  diagMut :: V.MVector _ CLong <- V.thaw =<< createVector 10 -- FIXME
  -- We need the types that sundials expects.
  -- FIXME: The Haskell type is currently empty!
  let funIO :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr () -> IO CInt
      funIO t y f _ptr = do
        sv <- peek y
        poke f $ SunVector { sunVecN = sunVecN sv
                           , sunVecVals = fun t (sunVecVals sv)
                           }
        return 0

  let nrPre = fromIntegral nr
  gResults :: V.Vector CInt <- createVector nrPre
  -- FIXME: Do we need to do this here? Maybe as it will get GC'd and
  -- we'd have to do a malloc in C otherwise :(
  gResMut <- V.thaw gResults
  event_index_mut :: V.MVector _ CInt <- V.thaw =<< createVector (fromIntegral max_events)
  event_time_mut :: V.MVector _ CDouble <- V.thaw =<< createVector (fromIntegral max_events)
  -- Total number of events. This is *not* directly re
  n_events_mut :: V.MVector _ CInt <- V.thaw =<< createVector 1
  -- Total number of rows in the output_mat_mut matrix. It *cannot* be
  -- inferred from n_events_mut because when an event occurs k times, it
  -- contributes k to n_events_mut but only 2 to n_rows_mut.
  n_rows_mut :: V.MVector _ CInt <- V.thaw =<< createVector 1
  actual_event_direction_mut :: V.MVector _ CInt <- V.thaw =<< createVector (fromIntegral max_events)

  let event_fn_c :: CDouble -> Ptr T.SunVector -> Ptr CDouble -> Ptr () -> IO CInt
      event_fn_c x y f _ptr = do
        vals <- event_fn x <$> (sunVecVals <$> peek y)
        -- FIXME: We should be able to use poke somehow
        vectorToC vals nrPre f
        return 0

      apply_event_c :: CInt -> CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> IO CInt
      apply_event_c event_index t y y' = do
        sv <- peek y
        poke y' $ SunVector
          { sunVecN = sunVecN sv
          , sunVecVals = apply_event (fromIntegral event_index) t (sunVecVals sv)
          }
        return 0

      requested_event_directions :: V.Vector CInt
      requested_event_directions = V.fromList $ map directionToInt directions

  let isJac :: CInt
      isJac = fromIntegral $ fromEnum $ isJust jacH
      jacIO :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr T.SunMatrix ->
               Ptr () -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr T.SunVector ->
               IO CInt
      jacIO t y _fy jacS _ptr _tmp1 _tmp2 _tmp3 = do
        case jacH of
          Nothing   -> error "Numeric.Sundials.CVode.ODE: Jacobian not defined"
          Just jacI -> do j <- jacI t <$> (sunVecVals <$> peek y)
                          poke jacS j
                          -- FIXME: I don't understand what this comment means
                          -- Unsafe since the function will be called many times.
                          [CU.exp| int{ 0 } |]

  res :: Int <- fromIntegral <$> [C.block| int {
                         /* general problem variables */

                         int flag;                  /* reusable error-checking flag                 */

                         int i, j;                  /* reusable loop indices                        */
                         N_Vector y = NULL;         /* empty vector for storing solution            */
                         N_Vector tv = NULL;        /* empty vector for storing absolute tolerances */

                         SUNMatrix A = NULL;        /* empty matrix for linear solver               */
                         SUNLinearSolver LS = NULL; /* empty linear solver object                   */
                         void *cvode_mem = NULL;    /* empty CVODE memory structure                 */
                         realtype t;
                         long int nst, nfe, nsetups, nje, nfeLS, nni, ncfn, netf, nge;

                         realtype tout;

                         /* input_ind tracks the current index into the ts array */
                         int input_ind = 1;
                         /* output_ind tracks the current row into the output_mat_mut matrix.
                            If differs from input_ind because of the extra rows corresponding to events. */
                         int output_ind = 1;
                         /* We need to update n_rows_mut every time we update output_ind because
                            of the possibility of early return (in which case we still need to assemble
                            the partial results matrix). We could even work with n_rows_mut only and ditch
                            output_ind, but the inline-c expression is quite verbose, and output_ind is
                            more convenient to use in index calculations.
                         */
			 ($vec-ptr:(int *n_rows_mut))[0] = output_ind;
                         /* event_ind tracks the current event number */
                         int event_ind = 0;

                         /* general problem parameters */

                         realtype T0 = RCONST(($vec-ptr:(double *ts))[0]); /* initial time              */
                         sunindextype NEQ = $(sunindextype nEq);           /* number of dependent vars. */

                         /* Initialize data structures */

                         y = N_VNew_Serial(NEQ); /* Create serial vector for solution */
                         if (check_flag((void *)y, "N_VNew_Serial", 0)) return 1;
                         /* Specify initial condition */
                         for (i = 0; i < NEQ; i++) {
                           NV_Ith_S(y,i) = ($vec-ptr:(double *f0))[i];
                         };

                         cvode_mem = CVodeCreate($(int method), CV_NEWTON);
                         if (check_flag((void *)cvode_mem, "CVodeCreate", 0)) return(1);

                         /* Call CVodeInit to initialize the integrator memory and specify the
                          * user's right hand side function in y'=f(t,y), the inital time T0, and
                          * the initial dependent variable vector y. */
                         flag = CVodeInit(cvode_mem, $fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), T0, y);
                         if (check_flag(&flag, "CVodeInit", 1)) return(1);

                         tv = N_VNew_Serial(NEQ); /* Create serial vector for absolute tolerances */
                         if (check_flag((void *)tv, "N_VNew_Serial", 0)) return 1;
                         /* Specify tolerances */
                         for (i = 0; i < NEQ; i++) {
                           NV_Ith_S(tv,i) = ($vec-ptr:(double *aTols))[i];
                         };

                         flag = CVodeSetMinStep(cvode_mem, $(double minStep_));
                         if (check_flag(&flag, "CVodeSetMinStep", 1)) return 1;
                         flag = CVodeSetMaxNumSteps(cvode_mem, $(long int maxNumSteps_));
                         if (check_flag(&flag, "CVodeSetMaxNumSteps", 1)) return 1;
                         flag = CVodeSetMaxErrTestFails(cvode_mem, $(int maxErrTestFails));
                         if (check_flag(&flag, "CVodeSetMaxErrTestFails", 1)) return 1;

                         /* Call CVodeSVtolerances to specify the scalar relative tolerance
                          * and vector absolute tolerances */
                         flag = CVodeSVtolerances(cvode_mem, $(double rTol), tv);
                         if (check_flag(&flag, "CVodeSVtolerances", 1)) return(1);

                         /* Call CVodeRootInit to specify the root function event_fn_c with nr components */
                         flag = CVodeRootInit(cvode_mem, $(int nr), $fun:(int (* event_fn_c) (double t, SunVector y[], double gout[], void * params)));

                         if (check_flag(&flag, "CVodeRootInit", 1)) return(1);

                         /* Set the initial step size if there is one */
                         if ($(int isInitStepSize)) {
                           /* FIXME: We could check if the initial step size is 0 */
                           /* or even NaN and then throw an error                 */
                           flag = CVodeSetInitStep(cvode_mem, $(double ss));
                           if (check_flag(&flag, "CVodeSetInitStep", 1)) return 1;
                         }

                         /* Set the Jacobian if there is one */
                         if ($(int isJac)) {
                           /* Initialize dense matrix data structure and solver */
                           A = SUNDenseMatrix(NEQ, NEQ);
                           if (check_flag((void *)A, "SUNDenseMatrix", 0)) return 1;
                           LS = SUNDenseLinearSolver(y, A);
                           if (check_flag((void *)LS, "SUNDenseLinearSolver", 0)) return 1;

                           /* Attach matrix and linear solver */
                           flag = CVDlsSetLinearSolver(cvode_mem, LS, A);
                           if (check_flag(&flag, "CVDlsSetLinearSolver", 1)) return 1;

                           /* Set the Jacobian */
                           flag = CVDlsSetJacFn(cvode_mem, $fun:(int (* jacIO) (double t, SunVector y[], SunVector fy[], SunMatrix Jac[], void * params, SunVector tmp1[], SunVector tmp2[], SunVector tmp3[])));
                           if (check_flag(&flag, "CVDlsSetJacFn", 1)) return 1;
                         } else /* Else we use CVDiag directly */ {
                           flag = CVDiag(cvode_mem);
                           if (check_flag(&flag, "CVDiag", 1)) return 1;
                         }

                         /* Store initial conditions */
			 ($vec-ptr:(double *output_mat_mut))[0 * (NEQ + 1) + 0] = ($vec-ptr:(double *ts))[0];
                         for (j = 0; j < NEQ; j++) {
                           ($vec-ptr:(double *output_mat_mut))[0 * (NEQ + 1) + (j + 1)] = NV_Ith_S(y,j);
                         }

                         while (1) {
                           flag = CVode(cvode_mem, ($vec-ptr:(double *ts))[input_ind], y, &t, CV_NORMAL); /* call integrator */
                           if (check_flag(&flag, "CVode solver failure, stopping integration", 1)) return 1;

                           /* Store the results for Haskell */
			   ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + 0] = t;
                           for (j = 0; j < NEQ; j++) {
                             ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + (j + 1)] = NV_Ith_S(y,j);
                           }
                           output_ind++;
                           ($vec-ptr:(int *n_rows_mut))[0] = output_ind;

                           if (flag == CV_ROOT_RETURN) {
                             if (event_ind >= $(int max_events)) {
                               /* We don't have any more space for events. Return an error. */
                               return 1;
                             }

                             /* Are we interested in this event?
                                If not, continue without any observable side-effects.
                             */
                             int good_event = 0;
			     flag = CVodeGetRootInfo(cvode_mem, ($vec-ptr:(int *gResMut)));
			     if (check_flag(&flag, "CVodeGetRootInfo", 1)) return 1;
                             for (i = 0; i < $(int nr); i++) {
                               int ev = ($vec-ptr:(int *gResMut))[i];
                               int req_dir = ($vec-ptr:(const int *requested_event_directions))[i];
                               if (ev != 0 && ev * req_dir >= 0) {
                                 good_event = 1;

                                 ($vec-ptr:(int *actual_event_direction_mut))[event_ind] = ev;
                                 ($vec-ptr:(int *event_index_mut))[event_ind] = i;
                                 ($vec-ptr:(double *event_time_mut))[event_ind] = t;
                                 event_ind++;

                                 /* Update the state with the supplied function */
                                 $fun:(int (* apply_event_c) (int, double, SunVector y[], SunVector z[]))(i, t, y, y);
                               }
                             }

                             if (good_event) {
                               ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + 0] = t;
                               for (j = 0; j < NEQ; j++) {
                                 ($vec-ptr:(double *output_mat_mut))[output_ind * (NEQ + 1) + (j + 1)] = NV_Ith_S(y,j);
                               }
                               output_ind++;
                               ($vec-ptr:(int *n_rows_mut))[0] = output_ind;

                               flag = CVodeReInit(cvode_mem, t, y);
                               if (check_flag(&flag, "CVodeReInit", 1)) return(1);
                             } else {
                               /* Since this is not a wanted event, it shouldn't get a row */
                               output_ind--;
                               ($vec-ptr:(int *n_rows_mut))[0] = output_ind;
                             }
                           }
			   else {
			     if (++input_ind >= $(int nTs))
                               break;
			   }
                         }

			 /* The number of actual roots we found */
			 ($vec-ptr:(int *n_events_mut))[0] = event_ind;

                         /* Get some final statistics on how the solve progressed */
                         flag = CVodeGetNumSteps(cvode_mem, &nst);
                         check_flag(&flag, "CVodeGetNumSteps", 1);
                         ($vec-ptr:(long int *diagMut))[0] = nst;

                         /* FIXME */
                         ($vec-ptr:(long int *diagMut))[1] = 0;

                         flag = CVodeGetNumRhsEvals(cvode_mem, &nfe);
                         check_flag(&flag, "CVodeGetNumRhsEvals", 1);
                         ($vec-ptr:(long int *diagMut))[2] = nfe;
                         /* FIXME */
                         ($vec-ptr:(long int *diagMut))[3] = 0;

                         flag = CVodeGetNumLinSolvSetups(cvode_mem, &nsetups);
                         check_flag(&flag, "CVodeGetNumLinSolvSetups", 1);
                         ($vec-ptr:(long int *diagMut))[4] = nsetups;

                         flag = CVodeGetNumErrTestFails(cvode_mem, &netf);
                         check_flag(&flag, "CVodeGetNumErrTestFails", 1);
                         ($vec-ptr:(long int *diagMut))[5] = netf;

                         flag = CVodeGetNumNonlinSolvIters(cvode_mem, &nni);
                         check_flag(&flag, "CVodeGetNumNonlinSolvIters", 1);
                         ($vec-ptr:(long int *diagMut))[6] = nni;

                         flag = CVodeGetNumNonlinSolvConvFails(cvode_mem, &ncfn);
                         check_flag(&flag, "CVodeGetNumNonlinSolvConvFails", 1);
                         ($vec-ptr:(long int *diagMut))[7] = ncfn;

                         flag = CVDlsGetNumJacEvals(cvode_mem, &nje);
                         check_flag(&flag, "CVDlsGetNumJacEvals", 1);
                         ($vec-ptr:(long int *diagMut))[8] = ncfn;

                         flag = CVDlsGetNumRhsEvals(cvode_mem, &nfeLS);
                         check_flag(&flag, "CVDlsGetNumRhsEvals", 1);
                         ($vec-ptr:(long int *diagMut))[9] = ncfn;

                         /* Clean up and return */

                         N_VDestroy(y);          /* Free y vector          */
                         N_VDestroy(tv);         /* Free tv vector         */
                         CVodeFree(&cvode_mem);  /* Free integrator memory */
                         SUNLinSolFree(LS);      /* Free linear solver     */
                         SUNMatDestroy(A);       /* Free A matrix          */

                         return CV_SUCCESS;
                       } |]
  preD <- V.freeze diagMut
  let d = SundialsDiagnostics (fromIntegral $ preD V.!0)
                              (fromIntegral $ preD V.!1)
                              (fromIntegral $ preD V.!2)
                              (fromIntegral $ preD V.!3)
                              (fromIntegral $ preD V.!4)
                              (fromIntegral $ preD V.!5)
                              (fromIntegral $ preD V.!6)
                              (fromIntegral $ preD V.!7)
                              (fromIntegral $ preD V.!8)
                              (fromIntegral $ preD V.!9)
  n_rows <- fromIntegral . V.head <$> V.freeze n_rows_mut
  output_mat <- coerce . reshape (dim + 1) . subVector 0 ((dim + 1) * n_rows) <$>
    V.freeze output_mat_mut
  n_events <- fromIntegral . V.head <$> V.freeze n_events_mut
  event_time             :: V.Vector Double
    <- coerce . V.take n_events <$> V.freeze event_time_mut
  event_index            :: V.Vector Int
    <- V.map fromIntegral . V.take n_events <$> V.freeze event_index_mut
  actual_event_direction :: V.Vector CInt
    <- V.take n_events <$> V.freeze actual_event_direction_mut
  let
    events :: [EventInfo]
    events = zipWith3 EventInfo
      (V.toList event_time)
      (V.toList event_index)
      (map (fromJust . intToDirection) $ V.toList actual_event_direction)
  return $
    if res == cV_SUCCESS
      then
        SolverSuccess events output_mat d
      else
        SolverError output_mat res

data SolverResult
  = SolverError !(Matrix Double) !Int
      -- ^ Partial results and error code
  | SolverSuccess
      [EventInfo]
      !(Matrix Double)
      !SundialsDiagnostics
      -- ^ Times at which the event was triggered, information about which root and the
                                                   -- results and diagnostics.
    deriving Show

odeSolveRootVWith' ::
  ODEOpts
  -> (Double -> V.Vector Double -> V.Vector Double)
      -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
  -> Maybe (Double -> Vector Double -> Matrix Double)
      -- ^ The Jacobian (optional)
  -> V.Vector Double                      -- ^ Initial conditions
  -> [EventSpec]                          -- ^ Event specifications
  -> Int                                  -- ^ Maximum number of events
  -> V.Vector Double                      -- ^ Desired solution times
  -> SolverResult
odeSolveRootVWith' opts f mb_jacobian y0 event_specs nRootEvs tt =
  solveOdeC (fromIntegral $ maxFail opts)
                 (fromIntegral $ maxNumSteps opts) (coerce $ minStep opts)
                 (fromIntegral . getMethod . odeMethod $ opts) (coerce $ initStep opts) jacH (scise $ stepControl opts)
                 (coerce f) (coerce y0)
                 (genericLength event_specs) event_equations event_directions
                 (fromIntegral nRootEvs) reset_state
                 (coerce tt)
  where
    l = size y0
    scise (X aTol rTol)                          = coerce (V.replicate l aTol, rTol)
    scise (X' aTol rTol)                         = coerce (V.replicate l aTol, rTol)
    scise (XX' aTol rTol yScale _yDotScale)      = coerce (V.replicate l aTol, yScale * rTol)
    -- FIXME; Should we check that the length of ss is correct?
    scise (ScXX' aTol rTol yScale _yDotScale ss) = coerce (V.map (* aTol) ss, yScale * rTol)
    jacH = fmap (\g t v -> matrixToSunMatrix $ g (coerce t) (coerce v)) $ mb_jacobian
    event_equations :: CDouble -> Vector CDouble -> Vector CDouble
    event_equations t y = V.fromList $
      map (\ev -> coerce (eventCondition ev) t y) event_specs
    event_directions :: [CrossingDirection]
    event_directions = map eventDirection event_specs
    reset_state :: Int -> CDouble -> Vector CDouble -> Vector CDouble
    reset_state n_event = coerce $ eventUpdate (event_specs !! n_event)

data EventSpec = EventSpec
  { eventCondition :: Double -> V.Vector Double -> Double
  , eventDirection :: CrossingDirection
  , eventUpdate :: Double -> V.Vector Double -> V.Vector Double
  }

odeSolveWithEvents
  :: ODEOpts
  -> [EventSpec]
    -- ^ Event specifications
  -> Int
    -- ^ Maximum number of events
  -> (Double -> V.Vector Double -> V.Vector Double)
    -- ^ The RHS of the system \(\dot{y} = f(t,y)\)
  -> Maybe (Double -> Vector Double -> Matrix Double)
    -- ^ The Jacobian (optional)
  -> V.Vector Double
    -- ^ Initial conditions
  -> V.Vector Double
    -- ^ Desired solution times
  -> Either Int SundialsSolution
    -- ^ Either an error code or a solution
odeSolveWithEvents opts event_specs max_events rhs mb_jacobian initial sol_times =
  let
    result :: SolverResult
    result =
      odeSolveRootVWith' opts rhs mb_jacobian initial event_specs
        max_events sol_times
  in
    case result of
      SolverError _ code -> Left code
      SolverSuccess events mx diagn ->
        Right $ SundialsSolution
            { actualTimeGrid = extractTimeGrid mx
            , solutionMatrix = mx
            , eventInfo = events
            , diagnostics = diagn
            }
  where
    -- The time grid is the first column of the result matrix
    extractTimeGrid :: Matrix Double -> Vector Double
    extractTimeGrid = head . toColumns


{-# LANGUAGE BangPatterns, DeriveDataTypeable, DeriveGeneric #-}

-- |
-- Module    : Statistics.Resampling
-- Copyright : (c) 2009, 2010 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Resampling statistics.

module Statistics.Resampling
    (
      Resample(..)
    , jackknife
    , jackknifeMean
    , jackknifeVariance
    , jackknifeVarianceUnb
    , jackknifeStdDev
    , resample
    , estimate
    , splitGen
    ) where

import Data.Aeson (FromJSON, ToJSON)
import Control.Concurrent (forkIO, newChan, readChan, writeChan)
import Control.Monad (forM_, liftM, replicateM, replicateM_)
import Data.Binary (Binary(..))
import Data.Data (Data, Typeable)
import Data.Vector.Algorithms.Intro (sort)
import Data.Vector.Binary ()
import Data.Vector.Generic (unsafeFreeze)
import Data.Word (Word32)
import GHC.Conc (numCapabilities)
import GHC.Generics (Generic)
import Numeric.Sum (Summation(..), kbn)
import Statistics.Function (indices)
import Statistics.Sample (mean, stdDev, variance, varianceUnbiased)
import Statistics.Types (Estimator(..), Sample)
import System.Random.MWC (GenIO, initialize, uniform, uniformVector)
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU

-- | A resample drawn randomly, with replacement, from a set of data
-- points.  Distinct from a normal array to make it harder for your
-- humble author's brain to go wrong.
newtype Resample = Resample {
      fromResample :: U.Vector Double
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON Resample
instance ToJSON Resample

instance Binary Resample where
    put = put . fromResample
    get = fmap Resample get

-- | /O(e*r*s)/ Resample a data set repeatedly, with replacement,
-- computing each estimate over the resampled data.
--
-- This function is expensive; it has to do work proportional to
-- /e*r*s/, where /e/ is the number of estimation functions, /r/ is
-- the number of resamples to compute, and /s/ is the number of
-- original samples.
--
-- To improve performance, this function will make use of all
-- available CPUs.  At least with GHC 7.0, parallel performance seems
-- best if the parallel garbage collector is disabled (RTS option
-- @-qg@).
resample :: GenIO
         -> [Estimator]         -- ^ Estimation functions.
         -> Int                 -- ^ Number of resamples to compute.
         -> Sample              -- ^ Original sample.
         -> IO [Resample]
resample gen ests numResamples samples = do
  let !numSamples = U.length samples
      ixs = scanl (+) 0 $
            zipWith (+) (replicate numCapabilities q)
                        (replicate r 1 ++ repeat 0)
          where (q,r) = numResamples `quotRem` numCapabilities
  results <- mapM (const (MU.new numResamples)) ests
  done <- newChan
  gens <- splitGen numCapabilities gen
  forM_ (zip3 ixs (tail ixs) gens) $ \ (start,!end,gen') -> do
    forkIO $ do
      let loop k ers | k >= end = writeChan done ()
                     | otherwise = do
            re <- U.replicateM numSamples $ do
                    r <- uniform gen'
                    return (U.unsafeIndex samples (r `mod` numSamples))
            forM_ ers $ \(est,arr) ->
                MU.write arr k . est $ re
            loop (k+1) ers
      loop start (zip ests' results)
  replicateM_ numCapabilities $ readChan done
  mapM_ sort results
  mapM (liftM Resample . unsafeFreeze) results
 where
  ests' = map estimate ests

-- | Run an 'Estimator' over a sample.
estimate :: Estimator -> Sample -> Double
estimate Mean             = mean
estimate Variance         = variance
estimate VarianceUnbiased = varianceUnbiased
estimate StdDev           = stdDev
estimate (Function est) = est

-- | /O(n) or O(n^2)/ Compute a statistical estimate repeatedly over a
-- sample, each time omitting a successive element.
jackknife :: Estimator -> Sample -> U.Vector Double
jackknife Mean sample             = jackknifeMean sample
jackknife Variance sample         = jackknifeVariance sample
jackknife VarianceUnbiased sample = jackknifeVarianceUnb sample
jackknife StdDev sample = jackknifeStdDev sample
jackknife (Function est) sample
  | G.length sample == 1 = singletonErr "jackknife"
  | otherwise            = U.map f . indices $ sample
  where f i = est (dropAt i sample)

-- | /O(n)/ Compute the jackknife mean of a sample.
jackknifeMean :: Sample -> U.Vector Double
jackknifeMean samp
  | len == 1  = singletonErr "jackknifeMean"
  | otherwise = G.map (/l) $ G.zipWith (+) (pfxSumL samp) (pfxSumR samp)
  where
    l   = fromIntegral (len - 1)
    len = G.length samp

-- | /O(n)/ Compute the jackknife variance of a sample with a
-- correction factor @c@, so we can get either the regular or
-- \"unbiased\" variance.
jackknifeVariance_ :: Double -> Sample -> U.Vector Double
jackknifeVariance_ c samp
  | len == 1  = singletonErr "jackknifeVariance"
  | otherwise = G.zipWith4 go als ars bls brs
  where
    als = pfxSumL . G.map goa $ samp
    ars = pfxSumR . G.map goa $ samp
    goa x = v * v where v = x - m
    bls = pfxSumL . G.map (subtract m) $ samp
    brs = pfxSumR . G.map (subtract m) $ samp
    m = mean samp
    n = fromIntegral len
    go al ar bl br = (al + ar - (b * b) / q) / (q - c)
      where b = bl + br
            q = n - 1
    len = G.length samp

-- | /O(n)/ Compute the unbiased jackknife variance of a sample.
jackknifeVarianceUnb :: Sample -> U.Vector Double
jackknifeVarianceUnb = jackknifeVariance_ 1

-- | /O(n)/ Compute the jackknife variance of a sample.
jackknifeVariance :: Sample -> U.Vector Double
jackknifeVariance = jackknifeVariance_ 0

-- | /O(n)/ Compute the jackknife standard deviation of a sample.
jackknifeStdDev :: Sample -> U.Vector Double
jackknifeStdDev = G.map sqrt . jackknifeVarianceUnb

pfxSumL :: U.Vector Double -> U.Vector Double
pfxSumL = G.map kbn . G.scanl add zero

pfxSumR :: U.Vector Double -> U.Vector Double
pfxSumR = G.tail . G.map kbn . G.scanr (flip add) zero

-- | Drop the /k/th element of a vector.
dropAt :: U.Unbox e => Int -> U.Vector e -> U.Vector e
dropAt n v = U.slice 0 n v U.++ U.slice (n+1) (U.length v - n - 1) v

singletonErr :: String -> a
singletonErr func = error $
                    "Statistics.Resampling." ++ func ++ ": singleton input"

-- | Split a generator into several that can run independently.
splitGen :: Int -> GenIO -> IO [GenIO]
splitGen n gen
  | n <= 0    = return []
  | otherwise =
  fmap (gen:) . replicateM (n-1) $
  initialize =<< (uniformVector gen 256 :: IO (U.Vector Word32))


{-# language BangPatterns #-}
{-# language ConstraintKinds #-}
{-# language DefaultSignatures #-}
{-# language FlexibleContexts #-}
{-# language FlexibleInstances #-}
{-# language GeneralizedNewtypeDeriving #-}
{-# language RankNTypes #-}
{-# language RebindableSyntax #-}
{-# language TypeFamilies #-}
module Prob where

import Prelude hiding (Functor(..), Applicative(..), Monad(..), sequence, (=<<), mapM, mapM_)
import qualified Prelude

import GHC.Exts (Constraint)

import Control.DeepSeq (NFData(..), force)
import Control.Parallel.Strategies (withStrategy, rpar, rseq, parList, parListChunk)
import Control.Monad.State (State)

import Data.List (groupBy, sort, sortBy)
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.Random                          as R
import qualified Data.Random.Distribution.Categorical as R

import Numeric.SpecFunctions (choose)

type QQ = Rational

data Event a = Event !a !QQ deriving (Eq, Ord)

eventToPair :: Event a -> (a, QQ)
eventToPair (Event a p) = (a, p)
{-# inline eventToPair #-}

mapEvent :: (a -> b) -> Event a -> Event b
mapEvent f (Event a p) = Event (f a) p
{-# inline mapEvent #-}

instance NFData a => NFData (Event a) where
    rnf (Event a _) = rnf a

instance Show a => Show (Event a) where
    show (Event a p) = show a ++ ": " ++ show p -- (realToFrac p :: Float)

newtype Prob a = Prob { density :: [Event a] } deriving (Eq, Ord, Show, NFData)

events :: Prob a -> [Event a]
events (Prob es) = es

eventValues :: Prob a -> [a]
eventValues prob = [a | Event a _ <- events prob]

eventProbs :: Prob a -> [QQ]
eventProbs prob = [p | Event _ p <- events prob]

class NoConstr a
instance NoConstr a

class ConstrMonad m where
    type Constr m :: * -> Constraint
    type Constr m = NoConstr

    fmap :: Constr m b => (a -> b) -> m a -> m b
    fmap f ma = do
        a <- ma
        return (f a)

    return :: Constr m a => a -> m a
    (>>=)  :: Constr m b => m a -> (a -> m b) -> m b

    (|>>=|) :: (Constr m a, Constr m b) => m a -> (a -> m b) -> m b
    (|>>=|) = (>>=)

    default return :: Prelude.Monad m => a -> m a
    return = Prelude.return
    default (>>=) :: Prelude.Monad m => m a -> (a -> m b) -> m b
    ma >>= f = ma Prelude.>>= f

infixl 1 >>=

liftA2 :: (ConstrMonad m, Constr m a, Constr m b, Constr m d) => (a -> b -> d) -> m a -> m b -> m d
liftA2 f ma mb = do
    a <- ma
    b <- mb
    return (f a b)
{-# inline liftA2 #-}

liftA3 :: (ConstrMonad m, Constr m a, Constr m b, Constr m d, Constr m e) => (a -> b -> d -> e) -> m a -> m b -> m d -> m e
liftA3 f ma mb md = do
    a <- ma
    b <- mb
    d <- md
    return (f a b d)
{-# inlinable liftA3 #-}

(=<<)  :: (ConstrMonad m, Constr m a, Constr m b) => (a -> m b) -> m a -> m b
f =<< ma = ma >>= f
infixr 1 =<<
{-# inline (=<<) #-}

instance ConstrMonad Maybe
instance ConstrMonad IO
instance ConstrMonad (State s)

instance Prelude.Monad m => ConstrMonad (R.RVarT m)
instance Prelude.Monad m => ProbMonad (R.RVarT m) where
    fromEventsSorted evts = R.categoricalT [(realToFrac p :: Double, a) | Event a p <- evts]

weightedEvents :: Real b => [(a, b)] -> [Event a]
weightedEvents es = [Event a (realToFrac w / s) | (a,w) <- es]
  where
    s = realToFrac $ sum [w | (_,w) <- es]

rnfList :: [a] -> ()
rnfList []     = ()
rnfList (a:as) = a `seq` rnfList as

forceList :: [a] -> [a]
forceList as = rnfList as `seq` as

fmapProb :: Ord b => (a -> b) -> Prob a -> Prob b
fmapProb f (Prob [Event a p]) = Prob [Event (f a) p]
fmapProb f (Prob evts)        =
    Prob $ forceList $ group1 $ sortBy (\(Event b _) (Event b' _) -> compare b b') [Event (f a) p | Event a p <- evts]
  where
    group1 :: Ord a => [Event a] -> [Event a]
    group1 []     = []
    group1 (e:es) = group e es

    group :: Ord a => Event a -> [Event a] -> [Event a]
    group e@(Event a p) ees =
      case ees of
        [] -> [e]
        e'@(Event a' p') : es
          | a == a'   -> group (Event a (p + p')) es
          | otherwise -> e : group e' es

{-# inlinable fmapProb #-}

bindProbWithStrat :: Ord b => (forall c. [c] -> [c]) -> Prob a -> (a -> Prob b) -> Prob b
bindProbWithStrat evalList (Prob [Event a _]) f = Prob $ evalList $ events (f a)
bindProbWithStrat evalList (Prob evts) f =
    Prob $ forceList $ group1 $ merge1 $ evalList $ totalProb [Event (f a) p | Event a p <- evts]
  where
    totalProb :: [Event (Prob a)] -> [[Event a]]
    totalProb ess = [[Event b (p*p') | Event b p' <- es] | Event (Prob es) p <- ess]

    merge1 :: Ord a => [[Event a]] -> [Event a]
    merge1 []       = []
    merge1 (es:ess) = merge es ess

    merge :: Ord a => [Event a] -> [[Event a]] -> [Event a]
    merge es []             = es
    merge es (es':[])       = merge2 es es'
    merge es (es':es'':ess) = merge2 (merge2 es es') (merge es'' ess)

    merge2 :: Ord a => [Event a] -> [Event a] -> [Event a]
    merge2 as         []         = as
    merge2 []         bs         = bs
    merge2 aas@(a:as) bbs@(b:bs)
      | a <= b    = a : merge2 as bbs
      | otherwise = b : merge2 aas bs

    group1 :: Ord a => [Event a] -> [Event a]
    group1 []     = []
    group1 (e:es) = group e es

    group :: Ord a => Event a -> [Event a] -> [Event a]
    group e@(Event a p) ees =
      case ees of
        [] -> [e]
        e'@(Event a' p') : es
          | a == a'   -> group (Event a (p + p')) es
          | otherwise -> e : group e' es

{-# inlinable bindProbWithStrat #-}

instance ConstrMonad Prob where
    type Constr Prob = Ord

    fmap = fmapProb
    {-# inline fmap #-}

    return a = Prob [Event a 1]
    {-# inline return #-}

    (>>=) = bindProbWithStrat forceList
    {-# inline (>>=) #-}

    -- (|>>=|) :: (Ord a, Ord b) => Prob a -> (a -> Prob b) -> Prob b
    (|>>=|) = bindProbWithStrat (withStrategy $ parList rseq)
    {-# inline (|>>=|) #-}

(|=<<|) :: (Ord a, Ord b) => (a -> Prob b) -> Prob a -> Prob b
f |=<<| ma = ma |>>=| f
{-# inline (|=<<|) #-}

ifThenElse :: Bool -> a -> a -> a
ifThenElse True  t _ = t
ifThenElse False _ f = f
{-# inline ifThenElse #-}

fail :: a
fail = error "fail called"

(>>) :: (ConstrMonad m, Constr m (), Constr m b) => m a -> m b -> m b
ma >> mb = fmap (\_ -> ()) ma >>= \() -> mb
infixl 1 >>
{-# inline (>>) #-}

(<<) :: (ConstrMonad m, Constr m (), Constr m b) => m b -> m a -> m b
mb << ma = ma >> mb
infixr 1 <<
{-# inline (<<) #-}

sequence :: (ConstrMonad m, Constr m a, Constr m [a]) => [m a] -> m [a]
sequence []       = return []
sequence (ma:mas) = liftA2 (:) ma (sequence mas)
{-# inlinable sequence #-}

mapM :: (ConstrMonad m, Constr m b, Constr m [b]) => (a -> m b) -> [a] -> m [b]
mapM f []     = return []
mapM f (a:as) = liftA2 (:) (f a) (mapM f as)
{-# inlinable mapM #-}

mapM_ :: (ConstrMonad m, Constr m ()) => (a -> m ()) -> [a] -> m ()
mapM_ f []     = return ()
mapM_ f (a:as) = f a >> mapM_ f as
{-# inlinable mapM_ #-}

forM :: (ConstrMonad m, Constr m b, Constr m [b]) => [a] -> (a -> m b) -> m [b]
forM as f = mapM f as
{-# inline forM #-}

forM_ :: (ConstrMonad m, Constr m ()) => [a] -> (a -> m ()) -> m ()
forM_ as f = mapM_ f as
{-# inline forM_ #-}

sequence1 :: (ConstrMonad m, Constr m a, Constr m [a], Constr m (NonEmpty a)) => NonEmpty (m a) -> m (NonEmpty a)
sequence1 (ma:|mas) = liftA2 (:|) ma (sequence mas)
{-# inline sequence1 #-}

mapM1 :: (ConstrMonad m, Constr m b, Constr m [b], Constr m (NonEmpty b)) => (a -> m b) -> NonEmpty a -> m (NonEmpty b)
mapM1 f (a:|as) = liftA2 (:|) (f a) (mapM f as)
{-# inline mapM1 #-}

forM1 :: (ConstrMonad m, Constr m b, Constr m [b], Constr m (NonEmpty b)) => NonEmpty a -> (a -> m b) -> m (NonEmpty b)
forM1 as f = mapM1 f as
{-# inline forM1 #-}

liftFoldl' :: (ConstrMonad m, Constr m a, Constr m b) => (b -> a -> b) -> m b -> [m a] -> m b
liftFoldl' _ !z []     = z
liftFoldl' f !z (p:ps) = liftFoldl' f (liftA2 f z p) ps

liftFoldr :: (ConstrMonad m, Constr m a, Constr m b) => (a -> b -> b) -> m b -> [m a] -> m b
liftFoldr f z = foldr (liftA2 f) z

liftFold :: (ConstrMonad m, Constr m a, Monoid a) => [m a] -> m a
liftFold = liftFoldl' mappend (return mempty)

liftSum :: (ConstrMonad m, Constr m a, Num a) => [m a] -> m a
liftSum = liftFoldl' (+) (return 0)

class ConstrMonad m => ProbMonad m where
    fromEventsSorted :: Constr m a => [Event a] -> m a

instance ProbMonad Prob where
    fromEventsSorted = Prob

uniformly :: (ProbMonad m, Constr m a, Ord a) => [a] -> m a
uniformly = uniformlySorted . sort

uniformlySorted :: (ProbMonad m, Constr m a) => [a] -> m a
uniformlySorted as = fromEventsSorted [Event a (1/n) | a <- as]
  where
    n = fromIntegral (length as)

probTrue :: Prob Bool -> QQ
probTrue prob =
    case events prob of
        [Event False _, Event True p] -> p
        [Event True p]                -> p
        [Event False q]               -> 1 - q
        []                            -> error "empty Prob!"
        _                             -> error "unnormalized Prob!"

probFalse :: Prob Bool -> QQ
probFalse prob = 1 - probTrue prob

eventProb :: Eq a => a -> Prob a -> QQ
eventProb e prob =
    case [p | Event a p <- events prob, a == e] of
      []  -> 0
      [p] -> p
      _   -> error "eventProb: prob is not normalized"

probOf :: (a -> Bool) -> Prob a -> QQ
probOf test prob = sum [p | Event a p <- events prob, test a]

given :: (a -> Bool) -> Prob a -> Prob a
given hyp prob =
    Prob [Event a (p / probHyp) | Event a p <- events prob, hyp a]
  where
    probHyp = probOf hyp prob

bernoulli :: QQ -> Prob Bool
bernoulli p = Prob [Event False (1-p), Event True p]

binomial :: Int -> QQ -> Prob Int
binomial 0 _ = return 0
binomial n p =
    case p of
      0 -> force $ return 0
      1 -> force $ return n
      _ -> force $ Prob [Event k (binomialProb n p k) | k <- [0..n]]

binomialProb :: Int -> QQ -> Int -> QQ
binomialProb n p k
  | k < 0 || k > n = 0
  | n == 0         = 1
  | otherwise      = realToFrac (choose n k) * p^k * (1-p)^fromIntegral (n-k)

distribution :: Ord a => Prob a -> [Event a]
distribution (Prob evts) = scanl1 sumEvents evts
  where
    sumEvents (Event _ p) (Event a' p') = Event a' (p + p')

revDistribution :: Ord a => Prob a -> [Event a]
revDistribution (Prob evts) = scanr1 sumEvents evts
  where
    sumEvents (Event a p) (Event _ p') = Event a (p + p')

summary :: Prob QQ -> IO ()
summary p =
  let mu    = "µ=" ++ show (realToFrac (mean p) :: Double)
      sigma = "σ=" ++ show (sqrt $ realToFrac (variance p) :: Double)
  in
      putStrLn $ mu ++ ", " ++ sigma

summaryInt :: Prob Int -> IO ()
summaryInt = summary . fmap fromIntegral

mean :: Prob QQ -> QQ
mean prob = sum [k * p | Event k p <- events prob]

variance :: Prob QQ -> QQ
variance prob = sum [k^2 * p | Event k p <- events prob] - mean prob^2

maxEvent :: Ord a => Prob a -> a
maxEvent (Prob es) = maximum [a | Event a _ <- es]


{-# LANGUAGE PackageImports,BangPatterns,StandaloneDeriving,TypeFamilies,FlexibleContexts,FlexibleInstances,TypeOperators,ExistentialQuantification,NoMonomorphismRestriction,DeriveTraversable,DeriveFunctor,DeriveFoldable #-}
module Kalman where
import Space.Class hiding(dim)
import Solver.RungeKutta
import qualified Space.Class as M
import Data.Typeable
import Data.Distributive
import Data.Functor.Compose
import Linear.Matrix hiding(trace)
import qualified Linear.Metric as Metric
import Vectorization
import Control.Applicative
import Data.Maybe
import Data.Foldable (Foldable,foldl')
import qualified  Data.Foldable as F
import Data.Traversable (Traversable,fmapDefault,mapAccumL)
import qualified Data.Traversable as T
import Linear.V0
import Linear.V1
import Foreign.Storable
import Foreign.FStorable
import Foreign.Cholesky
import qualified Numeric.LinearAlgebra as M
import qualified Data.Packed.Matrix as M
import Local
import Linear.Vector
import CartesianProduct
import Control.Monad.Trans.State

-- Covariance Version

fullM
  :: (RealFloat a, FStorable (Local b),
      Vectorize (Local b), Vectorize (Local b1), Space b,
      Space b1, Cholesky a, M.Field a) =>
     (b a -> b a)
     -> Local b (Local b a)
     -> (b a -> b1 a)
     -> Local b1 (Local b1 a)
     -> b1 a
     -> (b a, Local b (Local b a))
     -> (b a, Local b (Local b a))
fullM f r h q  m s= execState  ( do
      modify (prediction f r)
      modify (measure h q m)
      --modify (smoothing f r s)
      --modify (prediction f r)
      )  s

full
  :: (RealFloat a, FStorable (Local b2), FStorable (Local b),
      Vectorize (Local b), Vectorize (Local b1), Space b2, Space b,
      Space b1, Cholesky a, M.Field a) =>
     (b2 a -> b a)
     -> Local b (Local b a)
     -> (b a -> b1 a)
     -> Local b1 (Local b1 a)
     -> b1 a
     -> (b2 a, Local b2 (Local b2 a))
     -> (b a, Local b (Local b a))
full f r h q  m s= sm
  where sp = prediction f r s
        sm = measure h q m sp
        ss = smoothing f r s sm
        ssp = prediction f r ss
        ssm = measure h q m ssp

predictionM f = modify . prediction f
measureM h q  =  modify . measure  h q
smoothingM f r  = modify . smoothing f r

prediction
  :: (M.Field a , Vectorize (Local b) ,RealFloat a,
      FStorable (Local b1), Space b1, Space b, Cholesky a) =>
     (b1 a -> b a)
     -> Local b (Local b a)
     -> (b1 a, Local b1 (Local b1 a))
     -> (b a, Local b (Local b a))
prediction f r (x ,p)= (mean ,pxx )
    where
        (mean,pxx,_,_) = sigmaSamplingCovGain f r (x,p)

measure
  :: (RealFloat a, FStorable (Local b),
      Vectorize (Local b1), Space b, Space b1, Cholesky a, M.Field a) =>
     (b a -> b1 a)
     -> Local b1 (Local b1 a)
     -> b1 a
     -> (b a, Local b (Local b a))
     -> (b a, Local b (Local b a))
measure h q m (x ,p) = ( xb  , pb )
    where
        (mmean,pxx,pxy,k) =sigmaSamplingCovGain h q (x,p)
        xb = x |+| ( k !*  ( m |-| mmean ) )
        pb = p !-! ( k !*! pxx !*! distribute k )

smoothing
  :: (RealFloat a, FStorable (Local b), Vectorize (Local b1),
      Space b, Space b1, Cholesky a,
      M.Field a) =>
     (b a -> b1 a)
     -> Local b1 (Local b1 a)
     -> (b a, Local b (Local b a))
     -> (b1 a, Local b1 (Local b1 a))
     -> (b a, Local b (Local b a))
smoothing f r (x,p) (x1,p1) = (xb ,pb)
    where
        (xmean ,pxx,pxy,k) = sigmaSamplingCovGain f r (x,p)
        xb = x |+| (k !* ( x1 |-| xmean ) )
        pb = p !+! (k !*! ( p1 !-! pxx ) !*! distribute k )

sigmaSamplingCovGain
  :: (RealFloat a,M.Field a ,Vectorize (Local b),
      FStorable (Local b1), Space b1, Space b, Cholesky a) =>
     (b1 a -> b a)
     ->  (Local b (Local b a))
     -> (b1 a, Local b1 (Local b1 a))
     -> (b a, Local b (Local b a), Local b (Local b1 a),Local b1 (Local b a ))
sigmaSamplingCovGain f r (x,p) = (mean , pxx,pxy,k)
    where
        spoints = (sigmaPoints1 x p)
        points = fmap f spoints
        mean = sigmaMean1 points
        pxx = sigmaCov mean points !+! r
        pxy = sigmaCovXY x mean spoints  points
        k =  distribute pxy  !*! inverseM pxx

sigmaSamplingCov
  :: (RealFloat a, Show (Local b1 (Local b1 a)),
      FStorable (Local b1), Space b1, Space b, Cholesky a) =>
     (b1 a -> b a)
     ->  (Local b (Local b a))
     -> (b1 a, Local b1 (Local b1 a))
     -> (b a, Local b (Local b a), Local b (Local b1 a))

sigmaSamplingCov f r (x,p) = (mean , pxx,pxy)
    where
        spoints = (sigmaPoints1 x p)
        points = fmap f spoints
        mean = sigmaMean1 points
        pxx = sigmaCov mean points !+! r
        pxy = sigmaCovXY x mean spoints  points

inverseM = fromLists .  M.toLists . M.inv  . M.fromLists . toLists

{-
measureDelayed h m (x,p) (xlag,plag) =  (xnew,pxk)
    where
        -- Xlag
        sxlag = sigmaPoints1 xlag plag
        -- X
        sx = sigmaPoints1 x p
        -- XXlag
        pxxlag = sigmaCovXY x xlag sx sxlag
        -- Xk
        xk = x :|: xlag
        pxk = prodSym p plag pxxlag
        sxk = sigmaPoints1 xk pxk
        -- Y
        (sy,ym,py) = sigmaTransform h xlag plag
        pxlagy = sigmaCovXY xk ym sxk sy
        k = (distribute pxlagy) !*! (inverseM py)
        xnew = xk |+| (k !* (m |-| ym) )
-}

prodSym x y xy = prod x y xy (distribute xy)

prod x y xy yx = (liftA2 (:|:) x xy ) :|: (liftA2 (:|:) yx  y)

-- Helpers

square :: (Num a, R f) => f a -> f (f a)
square v = mult v v

squareT v = (distribute v) !*! v

mult x m = fmap (x ^*) m

sigmaMean
  :: (RealFloat a, Traversable t, FStorable t, Space b) =>
       t (b a) -> b a -> b a
sigmaMean p p0= case sigmaMean' p p0 of
                  Just x -> x
                  Nothing -> error $ "sigmaMean nonConvergence"

sigmaMean1 (V1 p0 :|: p )= sigmaMean p p0

sigmaMean'
  :: (RealFloat a, Traversable t, FStorable t, Space b) =>
       t (b a) -> b a -> Maybe (b a)
sigmaMean' p p0 = go 1000 p0
    where
        go ! i ! mik
            | i == 0 = Nothing
            | Metric.quadrance m > 1.0e-14 = go (i-1) mik1
            | otherwise = Just mik1
            where
                mik1 = mik |+| m
                s = F.foldr1 (^+^) $  fmap (|-|mik) p
                m =  s ^/ n
                n = 2 * (fromIntegral $ fsize p)

sigmaCovXY xmean ymean samplex sampley  =  covxy
    where
        xstd = fmap (|-| xmean) samplex
        ystd = fmap (|-| ymean) sampley
        covxy =  fmap (^/ 2) $ F.foldr1 (!+!) (liftA2 mult xstd ystd)

sigmaCov mean sample = cov
    where
        std  = fmap (|-| mean) sample
        cov  =  fmap (^/2) $ F.foldr1 (!+!) (fmap square std)

sigmaPoints1 x p = V1 x :|: points
  where
    m = case potrf p of
            Right x -> x
            Left i -> error $ "cholesky factorization error: sigmapoints1 : "
    points = fmap (x|+|) $  m  :|: fmap ((-1) *^) m

sqrtPrediction
  :: (Applicative g1, Applicative g, Traversable g, Traversable g1,
      FStorable g, FStorable g1, FStorable (Local b1), Space b,
      Space b1) =>
     (b Double -> b1 Double)
     -> g1 (Local b1 Double)
     -> (b Double, g (Local b Double))
     -> (b1 Double, Local b1 (Local b1 Double))
sqrtPrediction f r (x,s)= (xmean,sk'')
    where
        (y,w0m,w0c,w1) = weigthsGamma (fromIntegral $ fsize s)
        sp = fmap (x|+|) $ fmap (fmap (*y)) s :|: fmap (fmap (*(-y))) s
        sp' = fmap f sp
        x'  = f x
        xmean = sigmaMean1 $ V1 x' :|: sp'
        std  =fmap (|-| xmean) sp'
        sk' = dgeqrf $ distribute $ (fmap (fmap (*sqrt w1)) std):|: r
        sk'' = distribute $  ch1upJ sk' $ fmap (*sqrt w0c) $ x' |-| xmean

weigthsGamma l =  weigths 1 2 0 l

weigths a b k l= (y,w0m,w0c,w1) where
    lam =  a^2 * (l + k ) - l
    y =  sqrt ( l + lam)
    w0m = lam /(l + lam)
    w0c = w0m  +  (1 - a^2 + b)
    w1 = 1 / (2*(l+ lam))

ch1dnR x = unRight . ch1dn  x
sqrtMeasure
  :: (Show (Local b Double),Show (b Double),Applicative g, Traversable g, FStorable g, FStorable (Local b),
      FStorable (Local b1), Space b1, Space b) =>
     (b1 Double -> b Double)
     -> g (Local b Double)
     -> b Double
     -> (b1 Double, Local b1 (Local b1 Double))
     -> (b1 Double, Local b1 (Local b1 Double))
sqrtMeasure h r m (x,s) =  (x'',s''')
    where
{-
        spx = V1 x :|: (fmap (x|+|) $ fmap (fmap (*y)) s :|: fmap (fmap (*(-y))) s )
        wc = V1 w0c :|: pure w1
        --wm = V1 w0m :|: pure w1
        V1 y' :|: _ =  spy
        xmean = sigmaMean1  spx
        stdx  =  liftA2 (*^) wc $ fmap (|-| xmean ) spx
        spy = fmap h spx
        ymean = sigmaMean1  spy
        stdy  =  fmap (|-| ymean) spy
        pxy =   distribute stdx !*!  stdy
        sk' =  dgeqrf $ distribute $ fmap (fmap (*sqrt w1))  stdy :|: r
        sk''= distribute .fromJust  $   ch1up sk' (fmap (*w0c) $ ymean |-| y')
-}
        (ymean,pxx,pxy) = sigmaSampling h r (x,s)
        k  =  pseudoInverse pxy pxx
        -- x' = xmean |+| (k !* (m |-| y')  )
        x' = x |+| (k !* (m |-| ymean)  )
        u  = pxx !*! distribute k
        s' = distribute $ F.foldl ch1dnR  (distribute s)  u
        (x'',s''',_) = sigmaSampling id V0 (x',s')
{-      spx' = fmap (x' |+| ) $ fmap (fmap (*y))  s' :|: fmap (fmap (*(-y))) s'
        x'' =  sigmaMean1 $ V1 x' :|: spx'
        stdx' =  fmap (|-| x'') spx'
        s'' = dgeqrf $ distribute $ fmap (fmap (*sqrt w1))  stdx'
        s''' = distribute $ fromJust $ ch1up s'' $ fmap (*sqrt w0c) $ x'' |-| x'
-}

pseudoInverse
  :: (Applicative nhrs, Applicative m, Applicative m1,
      Traversable nhrs, Traversable m, Traversable m1, Distributive m1,
      FStorable nhrs, FStorable m, FStorable m1) =>
     nhrs (m1 Double) -> m (m1 Double) -> nhrs (m1 Double)
pseudoInverse y x  = dgelssR (distribute x )  (dgelssR x  y )

sqrtBackward f p (x,s) (x1,s1) = (xb ,sb)
    where
        (xmean ,pxx,pxy) = sigmaSampling f p (x,s)
        k =  pseudoInverse pxy pxx
        xb = x |+| (k !* (x1 |-| xmean ))
        s1pxx = distribute $ F.foldl (\y-> ch1dnR y)  (distribute s1 )  (distribute pxx )
        u  = s1pxx !*! distribute k
        sb = distribute $ F.foldl  ch1upJ  (distribute s)  u

sigmaSampling h r (x,s) = (ymean,sxx,sxy)
    where
        (y,w0m,w0c,w1) = weigthsGamma (fromIntegral $ fsize s)
        spx = V1 x :|: (fmap (x|+|) $ fmap (fmap (*y)) s :|: fmap (fmap (*(-y))) s )
        wc = V1 w0c :|: pure w1
        --wm = V1 w0m :|: pure w1
        V1 y' :|: _ =  spy
        xmean = sigmaMean1  spx
        stdx  =  liftA2 (*^) wc $ fmap (|-| xmean ) spx
        spy = fmap h spx
        ymean = sigmaMean1  spy
        stdy  =  fmap (|-| ymean) spy
        sxy =   distribute stdx !*!  stdy
        sk' =  dgeqrf $ distribute $ fmap (fmap (*sqrt w1))  stdy :|: r
        sxx = distribute $ ch1upJ sk' (fmap (*w0c) $ ymean |-| y')

ch1upJ x = fromJust . ch1up x
dgelssR x = unRight . dgelss x

unRight (Right x) = x
unRight (Left x) = error $ "no right: "


module Data.SQP
       ( Problem (..)
       , optimize
       ) where

import Numeric.LinearAlgebra.HMatrix
import Numeric.Minimization.QuadProgPP

import Debug.Trace

data Problem = Problem
               { _cost :: Vector Double -> Double
               , _approxCost :: Vector Double
                             -> (Matrix Double, Vector Double, Double)
               , _trueIneqs :: Vector Double -> Vector Double
               , _approxAffineIneqs :: Vector Double
                                    -> (Matrix Double, Vector Double)
               , _eqs :: Vector Double -> Vector Double
               , _approxEqs :: Vector Double -> (Matrix Double, Vector Double)
               , _numVariables :: Int
               , _numIneqs :: Int
               , _numEqs :: Int
               }

stepAcceptanceThreshold :: Double
stepAcceptanceThreshold = 0.25

trustShrinkFactor :: Double
trustShrinkFactor = 0.1

trustExpandFactor :: Double
trustExpandFactor = 1.5

constraintPenaltyScalingFactor :: Double
constraintPenaltyScalingFactor = 10

minTrustSize :: Double
minTrustSize = 1e-4

minModelImprove :: Double
minModelImprove = 1e-4

minModelImproveRatio :: Double
minModelImproveRatio = negate $ read "Infinity"

constraintSatisfactionThreshold :: Double
constraintSatisfactionThreshold = 1e-4

initTrustSize :: Double
initTrustSize = 0.1

initPenalty :: Double
initPenalty = 10

evalMerit :: Problem -> Double -> Vector Double -> Double
evalMerit problem penaltyParam x =
  let cost = _cost problem x
      penalty =
        penaltyParam * sum (map (max 0.0) $ toList $ _trueIneqs problem x) +
        penaltyParam * sum (map abs $ toList $ _eqs problem x)
  in  cost + penalty

optimize :: Problem
         -> Vector Double -- xInitial
         -> (Vector Double, Double)
optimize problem xInit =
  findSuitableConstraintPenalty problem xInit initTrustSize initPenalty

findSuitableConstraintPenalty :: Problem
                              -> Vector Double -- x
                              -> Double -- trust region size
                              -> Double -- constraint penalty
                              -- xNew, trueMerit
                              -> (Vector Double, Double)
findSuitableConstraintPenalty problem x trustSize penaltyParam =
  let (xNew, trueMerit, newTrustSize) =
        reconvexify problem x trustSize penaltyParam
      ineqsSatisfied = all (<= constraintSatisfactionThreshold) $
                         toList $ _trueIneqs problem xNew
      eqsSatisfied = all ((<= constraintSatisfactionThreshold) . abs) $
                       toList $ _eqs problem xNew
      constraintsSatisfied = ineqsSatisfied && eqsSatisfied
  in  if constraintsSatisfied
      then (xNew, trueMerit)
      else let penalty' = penaltyParam * constraintPenaltyScalingFactor
           in  findSuitableConstraintPenalty
                 problem xNew newTrustSize penalty'

evalQuadratic :: (Matrix Double, Vector Double, Double)
              -> Vector Double
              -> Double
evalQuadratic (m, b, c) x =
  (x `dot` ((0.5 * (m #> x)) + b)) + c

evalLinear :: (Matrix Double, Vector Double)
           -> Vector Double
           -> Vector Double
evalLinear (m, b) x = (m #> x) + b

reconvexify :: Problem
            -> Vector Double -- x
            -> Double -- trust region size
            -> Double -- constraint penalty
            -- xNew, newTrueMerit, newTrustSize
            -> (Vector Double, Double, Double)
reconvexify problem x trustSize penaltyParam =
  let convexCost = _approxCost problem x
      convexIneqs = _approxAffineIneqs problem x
      convexEqs = _approxEqs problem x
      trueMerit = evalMerit problem penaltyParam x
      trustResult = findSuitableTrustStep problem x convexCost
                      convexIneqs convexEqs trueMerit trustSize
                      penaltyParam
  in  case trustResult of
        Reconvexify xNew newMerit newTrustSize ->
          reconvexify problem xNew newTrustSize penaltyParam
        Finished xNew newMerit newTrustSize ->
          (xNew, newMerit, newTrustSize)

data TrustStepResult = -- x merit trustSize
                       Reconvexify (Vector Double) Double Double
                     | Finished (Vector Double) Double Double
  deriving (Show)

findSuitableTrustStep :: Problem
                      -> Vector Double -- x
                      -- convexified cost
                      -> (Matrix Double, Vector Double, Double)
                      -- convexified inequality constraints
                      -> (Matrix Double, Vector Double)
                      -- convexified equality constraints
                      -> (Matrix Double, Vector Double)
                      -> Double -- old true merit
                      -> Double -- trust region size
                      -> Double -- constraint penalty parameter
                      -- xNew, newTrueMerit, newTrustSize
                      -> TrustStepResult
findSuitableTrustStep
  problem x convexCost convexIneq convexEq oldTrueMerit trustSize penaltyParam =
  let trustStep =
        trustRegionStep problem x convexCost convexIneq convexEq oldTrueMerit
          trustSize penaltyParam
  in  case trace (show trustStep) trustStep of
        Reject -> let newTrustSize = trustSize * trustShrinkFactor
                  in  findSuitableTrustStep problem x convexCost convexIneq
                        convexEq oldTrueMerit newTrustSize penaltyParam
        Accept xNew newTrueMerit ->
          let newTrustSize = trustSize * trustExpandFactor
          in  Reconvexify xNew newTrueMerit newTrustSize
        Converged xNew newTrueMerit ->
          let newTrustSize =
                max trustSize $ (minTrustSize / trustShrinkFactor) * 1.5
          in  Finished xNew newTrueMerit newTrustSize

data IterationResult = Reject
                     -- xNew newMerit
                     | Accept (Vector Double) Double
                     | Converged (Vector Double) Double
  deriving (Show)

trustRegionStep :: Problem
                -> Vector Double -- x
                -- convexified cost
                -> (Matrix Double, Vector Double, Double)
                -- convexified inequality constraints
                -> (Matrix Double, Vector Double)
                -- convexified equality constraints
                -> (Matrix Double, Vector Double)
                -> Double -- old true merit
                -> Double -- trust region size
                -> Double -- constraint penalty parameter
                -> IterationResult
trustRegionStep
  problem x convexCost convexIneq convexEq oldTrueMerit trustSize penaltyParam =
    let (xNew, modelMerit) = solveQuadraticSubproblem problem x convexCost
                               convexIneq convexEq trustSize penaltyParam
        trueMerit =  evalMerit problem penaltyParam xNew
        trueImprove = oldTrueMerit - trueMerit
        modelImprove = oldTrueMerit - modelMerit
    in  if modelImprove < -1e-2 -- why so high? our convexification?
        then error $ "Model improvement got worse: " ++
               show (modelImprove, modelMerit, oldTrueMerit, trueMerit)
        else
          if trustSize < minTrustSize ||
             modelImprove < minModelImprove ||
             modelImprove / oldTrueMerit < minModelImproveRatio
          then Converged x oldTrueMerit
          else
            if trueImprove > 0.0 &&
               trueImprove / modelImprove > stepAcceptanceThreshold
            then Accept xNew trueMerit
            else Reject

solveQuadraticSubproblem :: Problem
                         -> Vector Double
                         -- convex cost
                         -> (Matrix Double, Vector Double, Double)
                         -- convexified inequality constraints
                         -> (Matrix Double, Vector Double)
                         -- convexified equality constraints
                         -> (Matrix Double, Vector Double)
                         -> Double -- trust region size
                         -> Double -- constraint penalty parameter
                         -> (Vector Double, Double) -- new x, model merit
solveQuadraticSubproblem
  problem x (costMatrix, costVector, costConstant)
  (approxIneqMatrix, approxIneqVector) (approxEqMatrix, approxEqVector)
  trustSize penaltyParam =
  -- We approximate each nonlinear inequality as |ax + b|^+. For each
  -- of these, we introduce a new optimization variable t (i.e., a
  -- slack variable) that comes with two inequalities:
  --
  -- 0 <= t
  --
  -- ax + b <= t
  --
  -- Each nonlinear equality as |ax + b|. For each of these, two new
  -- optimization variables s and t with following properties:
  --
  -- -s <= 0
  -- -t <= 0
  -- s - t = ax + b
  let numIneqs = _numIneqs problem
      numEqs = _numEqs problem
      numVariables = _numVariables problem

      costMatrixWithSlacks =
        diagBlock [ costMatrix
                  , konst 0.0 (numIneqs, numIneqs)
                  , konst 0.0 (2 * numEqs, 2 * numEqs)]

      -- LET'S MAKE COST SPD
      (l, v) = eigSH costMatrixWithSlacks
      positiveEigVals = cmap (max 1e-12) l
      augmentedCostMatrix = v <> diag positiveEigVals <> tr v

      augmentedCostVector = vjoin [costVector
                                  , konst 1.0 numIneqs
                                  , konst 1.0 $ 2 * numEqs]

      augmentedIneqMatrix =
        ((scalar penaltyParam * approxIneqMatrix) |||
         negate (ident numIneqs) |||
         konst 0.0 (numIneqs, 2 * numEqs)) ===
         (konst 0.0 (numIneqs + 2 * numEqs, numVariables) |||
          negate (ident $ numIneqs + 2 * numEqs))
      augmentedIneqVector =
        vjoin [ scalar penaltyParam * approxIneqVector
              , konst 0.0 $ numIneqs + 2 * numEqs ]

      ineqMatrixWithTrustConstraints =
        augmentedIneqMatrix ===
        (ident numVariables ||| konst 0.0 (numVariables, numIneqs + 2 * numEqs)) ===
        (negate (ident numVariables) ||| konst 0.0 (numVariables, numIneqs + 2 * numEqs))
      ineqVectorWithTrustConstraints =
        vjoin [ augmentedIneqVector
              , negate (x + scalar trustSize)
              , x - scalar trustSize ]

      augmentedEqMatrix =
        scalar penaltyParam * approxEqMatrix |||
        konst 0.0 (numEqs, numIneqs) |||
        negate (ident numEqs) |||
        ident numEqs
      augmentedEqVector = scalar penaltyParam * approxEqVector

      -- inequalities are given as ax + b <= 0, but quadprog++ wants
      -- a'x + b' >= 0
      result = solveQuadProg
        (augmentedCostMatrix, augmentedCostVector)
        (Just (augmentedEqMatrix, augmentedEqVector)) $
        Just ((-ineqMatrixWithTrustConstraints),
              (-ineqVectorWithTrustConstraints))

  in  case result of
        Left e -> error $ show e
        Right (xNewAugmented, newMeritHomogenous) ->
          let xNew = subVector 0 numVariables xNewAugmented
              newMerit = newMeritHomogenous + costConstant

              -- Compute the merit of the approximation without regularization
              cost = evalQuadratic (costMatrix, costVector, costConstant) xNew
              ineqViolations =
                evalLinear (approxIneqMatrix, approxIneqVector) xNew
              ineqPenalty =
                penaltyParam * (sum $ map (max 0.0) $ toList ineqViolations)
              eqViolations = evalLinear (approxEqMatrix, approxEqVector) xNew
              eqPenalty =
                penaltyParam * (sum $ map abs $ toList eqViolations)
              meritNoReg = cost + ineqPenalty + eqPenalty

              -- DEBUG
              meritError = abs $ newMerit - meritNoReg
          in  if meritError / (abs meritNoReg) > 0.5
              then error $ "merits don't match!" ++ show (meritNoReg, newMerit)
              else (xNew, meritNoReg)


{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# OPTIONS_GHC -Wno-deprecations #-}
module BlueRipple.Model.PreferenceBayes
  (
    -- * MLE
    cgOptimize
  , cgOptimizeAD
  , invFisher
  , mleCovEigens
  , variances
  , covar
  , correlFromCov

    -- * MCMC
  , Sample
  , Chain
  , runMany
    -- * Re-exports
  , dispf
  , disps
  )
where

import qualified Statistics.Types              as S
import qualified Control.Foldl                 as FL
import           Control.Lens.At                ( IxValue
                                                , Index
                                                , Ixed
                                                )
import           Control.Lens.Indexed           ( FunctorWithIndex )
import           Control.Monad                  ( sequence )

import qualified Numeric.AD                    as AD
import qualified Numeric.LinearAlgebra         as LA
import           Numeric.LinearAlgebra          ( dispf
                                                , disps
                                                )
import           Numeric.MathFunctions.Constants
                                                ( m_ln_sqrt_2_pi )
import qualified Numeric.MCMC                  as MC
import Numeric (log)

import           Math.Gamma                     ( gamma
                                                , Gamma
                                                )
import           System.Random                  ( randomRIO )
import qualified Control.Concurrent            as CC
import qualified Control.Concurrent.MVar       as CC
import qualified Data.Vector.Unboxed           as VU
import qualified Data.Vector.Storable          as VS
import qualified Data.Vector.Generic           as VG
import qualified Data.Vector                   as VB

import qualified Numeric.Optimization.Algorithms.HagerZhang05
                                               as CG

import qualified Numeric.Optimization.Algorithms.HagerZhang05.AD
                                               as CGAD

newtype ObservedVote = ObservedVote { dem :: Int}

data Pair a b = Pair !a !b

binomialNormalParams :: RealFrac a => VB.Vector Int -> VB.Vector a -> (a, a)
binomialNormalParams turnoutCounts demProbs =
  let np = VB.zip turnoutCounts demProbs
      meanVarUpdate :: RealFrac a => Pair a a -> (Int, a) -> Pair a a
      meanVarUpdate (Pair m v) (n, p) =
        let m' = realToFrac n * p in Pair (m + m') (v + m' * (1 - p))
      foldMeanVar = FL.Fold meanVarUpdate (Pair 0 0) id
      Pair m v    = FL.fold foldMeanVar np
  in  (m, v)

logBinomialObservedVote
  :: (RealFrac a, Floating a) => VB.Vector a -> (Int, VB.Vector Int) -> a
logBinomialObservedVote demProbs (demVote, turnoutCounts) =
  let (m, v) = binomialNormalParams turnoutCounts demProbs
  in  negate $ log v + ((realToFrac demVote - m) ^ 2 / (2 * v))

gradLogBinomialObservedVote
  :: (Floating a, RealFrac a)
  => VB.Vector a
  -> (Int, VB.Vector Int)
  -> VB.Vector a
gradLogBinomialObservedVote demProbs (demVote, turnoutCounts) =
  let (m, v) = binomialNormalParams turnoutCounts demProbs
      np     = VB.zip turnoutCounts demProbs
      dv     = fmap (\(n, p) -> let n' = realToFrac n in n' * (1 - 2 * p)) np
      dm     = turnoutCounts
      dmv    = VB.zip dm dv
      a1     = negate (1 / v) -- d (log v)
      a2     = (realToFrac demVote - m)
      a3     = (a2 * a2) / (2 * v * v)
      a4     = (a2 / v)
  in  fmap (\(dm, dv) -> (a1 + a3) * dv + a4 * realToFrac dm) dmv

logBinomialObservedVotes
  :: (Functor f, Foldable f, Floating a, RealFrac a)
  => f (Int, VB.Vector Int)
  -> VB.Vector a
  -> a
logBinomialObservedVotes votesAndTurnout demProbs =
  FL.fold FL.sum $ fmap (logBinomialObservedVote demProbs) votesAndTurnout

gradLogBinomialObservedVotes
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> VB.Vector Double
gradLogBinomialObservedVotes votesAndTurnout demProbs =
  let
    n           = VG.length demProbs
    indexVector = VG.generate n id
    sumEach =
      sequenceA $ (\n -> FL.premap (VG.! n) (FL.sum @Double)) <$> indexVector
  in
    FL.fold sumEach
      $ fmap (gradLogBinomialObservedVote demProbs) votesAndTurnout

cgOptimize
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> IO (VS.Vector Double, CG.Result, CG.Statistics)
cgOptimize votesAndVoters guess = do
  let params = CG.defaultParameters { CG.printFinal  = False
                                    , CG.printParams = False
                                    , CG.verbose     = CG.Quiet
                                    }
      grad_tol = 0.0000001
  CG.optimize
    params
    grad_tol
    guess
    (CG.VFunction (negate . logBinomialObservedVotes votesAndVoters))
    (CG.VGradient (fmap negate . gradLogBinomialObservedVotes votesAndVoters))
    Nothing

cgOptimizeAD
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> IO (VB.Vector Double, CG.Result, CG.Statistics)
cgOptimizeAD votesAndVoters guess = do
  let params = CG.defaultParameters { CG.printFinal  = False
                                    , CG.printParams = False
                                    , CG.verbose     = CG.Quiet
                                    }
      grad_tol = 0.0000001
  CGAD.optimize params
                grad_tol
                guess
                (negate . logBinomialObservedVotes votesAndVoters)

-- I'm not sure about row/column order her. But it doesn't
-- matter since the Hessian is symmetric.
hessianLL
  :: (Foldable f, Functor f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> LA.Matrix Double
hessianLL votesAndVoters x =
  LA.fromRows $ fmap VS.convert $ VB.toList $ AD.hessian
    (negate . logBinomialObservedVotes votesAndVoters)
    x

invFisher
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> LA.Matrix Double
invFisher votesAndVoters x = LA.inv $ hessianLL votesAndVoters x

variances
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> LA.Vector Double
variances votesAndVoters x = LA.takeDiag $ invFisher votesAndVoters x

covar
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> LA.Matrix Double
covar = invFisher

correlFromCov :: LA.Matrix Double -> LA.Matrix Double
correlFromCov cov =
  let vars = LA.takeDiag cov
      mISD = LA.diag $ LA.cmap (\x -> 1 / sqrt x) vars
  in  mISD LA.<> cov LA.<> mISD

correl
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> LA.Matrix Double
correl votesAndVoters x =
  let cov          = invFisher votesAndVoters x
      sds          = LA.cmap sqrt $ LA.takeDiag cov
      (rows, cols) = LA.size cov
      rho n m = cov `LA.atIndex` (n, m) / ((sds LA.! n) * (sds LA.! m))
  in  LA.assoc
        (rows, cols)
        0
        [ ((n, m), rho n m) | n <- [0 .. (rows - 1)], m <- [0 .. (cols - 1)] ]

mleCovEigens
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> (LA.Vector Double, LA.Matrix Double)
mleCovEigens votesAndVoters x =
  LA.eigSH $ LA.trustSym $ invFisher votesAndVoters x

betaDist :: (Floating a, RealFrac a, Gamma a) => a -> a -> a -> a
betaDist alpha beta x =
  let b = gamma (alpha + beta) / (gamma alpha + gamma beta)
  in  if (x >= 0) && (x <= 1)
        then b * (x ** (alpha - 1)) * ((1 - x) ** (beta - 1))
        else 0

logBetaDist :: (Floating a, RealFrac a, Gamma a) => a -> a -> a -> a
logBetaDist alpha beta x =
  let lb = log $ gamma (alpha + beta) / (gamma alpha + gamma beta)
  in  lb + (alpha - 1) * log x + (beta - 1) * log (1 - x)

betaPrior :: (Floating a, RealFrac a, Gamma a) => a -> a -> VB.Vector a -> a
betaPrior a b xs = FL.fold FL.product $ fmap (betaDist a b) xs

{-
--logBetaPrior :: Double -> Double -> [Double] -> Double
--logBetaPrior = FL.fold FL.product $ fmap (logBetabDist a b) xs

f :: [(Int, [Int])] -> [Double] -> Double
f votesAndTurnout demProbs =
  let v = demProbs
  in  (exp $ logProbObservedVotes votesAndTurnout v) * (betaPrior 2 2 v)
-}

logBinomialWithPrior
  :: (Functor f, Foldable f, Floating a, RealFrac a, Gamma a)
  => f (Int, VB.Vector Int)
  -> VB.Vector a
  -> a
logBinomialWithPrior votesAndTurnout demProbs =
  let x = demProbs
  in  logBinomialObservedVotes votesAndTurnout x + log (betaPrior 2 2 x)

type Sample = VB.Vector Double
type Chain = [Sample]

-- TODO: withSystemRandom is deprecated in mwc-random but comes from "declarative" so would be a pain to fix.
runMCMC
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> Int
  -> Sample
  -> IO Chain
runMCMC votesAndTurnout numIters start =
  fmap (fmap MC.chainPosition) . MC.withSystemRandom . MC.asGenIO $ MC.chain
    numIters
    start
    (MC.frequency [(1, MC.hamiltonian 0.0001 10), (1, MC.hamiltonian 0.00001 5)]
    )
    (MC.Target (logBinomialWithPrior votesAndTurnout)
               (Just $ gradLogBinomialObservedVotes votesAndTurnout)
    )

-- launch each action on its own thread, writing the result to an MVar.
-- as each MVar is written to,
-- place each result in the appropriate place in the structure.
-- return the new structure when all results are available.
sequenceConcurrently :: Traversable t => t (IO a) -> IO (t a)
sequenceConcurrently actions = do
  let f :: IO a -> IO (CC.MVar a, CC.ThreadId)
      f action = do
        mvar     <- newEmptyMVar
        threadId <- CC.forkIO (action >>= putMVar mvar)
        return (mvar, threadId)
      g :: (CC.MVar a, CC.ThreadId) -> IO a
      g (mvar, _) = takeMVar mvar
  forked <- traverse f actions -- IO (t (MVar a, ThreadId))
  traverse g forked

runMany
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> Int
  -> Int
  -> Int
  -> Int
  -> IO [Chain]
runMany votesAndTurnout nParams nChains nSamplesPerChain nBurnPerChain = do
  let
    randomStart :: Int -> IO Sample
    randomStart n = VG.fromList <$> replicateM n (randomRIO (0, 1))
    randomStarts :: Int -> Int -> IO [Sample]
    randomStarts n m = replicateM m (randomStart n)
    doEach =
      fmap (drop nBurnPerChain) . runMCMC votesAndTurnout nSamplesPerChain
  starts <- randomStarts nParams nChains
  sequenceConcurrently $ fmap doEach starts

{-
data Normal = Normal { m :: Double, v :: Double } deriving (Show)

instance Semigroup Normal where
  (Normal m1 v1) <> (Normal m2 v2) = Normal (m1 + m2) (v1 + v2)

logNormalObservedVote :: [Normal] -> (Int, [Int]) -> Double
logNormalObservedVote demProbs (demVote, turnoutCounts) =
  let Normal m' v' = mconcat $ fmap (\(t,Normal m v) -> Normal (t*m) (v*m))$ zip turnoutCounts demProbs
  in  negate $ log v'  + ((realToFrac demVote - m')) ^ 2 / (2 * v'))

gradLogNormalObservedVote :: [Normal] -> (Int, [Int]) -> [Normal]
gradLogNormalObservedVote demProbs (demVote, turnoutCounts) =
  let Normal m' v' = mconcat $ fmap (\(t,Normal m v) -> Normal (t*m) (v*m))$ zip turnoutCounts demProbs
      np     = zip turnoutCounts demProbs
      dv     = fmap (\(n, p) -> let n' = realToFrac n in n' * (1 - 2 * p)) np
      dm     = turnoutCounts
      dmv    = zip dm dv
      a1     = negate (1 / v) -- d (log v)
      a2     = (realToFrac demVote - m)
      a3     = (a2 * a2) / (2 * v * v)
      a4     = (a2 / v)
  in  fmap (\(dm, dv) -> (a1 + a3) * dv + a4 * realToFrac dm) dmv

logBinomialObservedVotes :: [(Int, [Int])] -> [Double] -> Double
logBinomialObservedVotes votesAndTurnout demProbs =
  FL.fold FL.sum $ fmap (logBinomialObservedVote demProbs) votesAndTurnout

gradLogBinomialObservedVotes :: [(Int, [Int])] -> [Double] -> [Double]
gradLogBinomialObservedVotes votesAndTurnout demProbs =
  let n = length demProbs
      sumEach =
        sequenceA
          $ fmap (\n -> FL.premap (!! n) (FL.sum @Double))
          $ [0 .. (n - 1)]
  in  FL.fold sumEach $ fmap (gradLogBinomialObservedVote demProbs) votesAndTurnout

-}
{-
--gradLogProbObservedVotes2 :: [(Int, [Int])] -> [Double] -> [Double]
gradLogProbObservedVotes2 votesAndTurnout =
  grad (logProbObservedVotes votesAndTurnout)
-}

{-
data VPSummary = VPSummary { mean :: Double, stDev :: Double, rHat :: Double } deriving (Show)

summarize :: (Sample -> Double) -> [Chain] -> VPSummary -- FL.Fold Chain VPSummary
summarize integrandF chains =
  -- build up the fold
  let n = (length $ head chains) `div` 2
      splitChains :: [Chain]
      splitChains =
        concat $ fmap (\c -> let (a, b) = splitAt n c in [a, b]) chains
      meanVarF = (,) <$> FL.mean <*> FL.variance
      meanVars = fmap (FL.fold (FL.premap integrandF meanVarF)) splitChains
      bVar     = FL.fold (FL.premap fst FL.variance) meanVars -- variance of the means
      wVar     = FL.fold (FL.premap snd FL.mean) meanVars -- mean of the variances
      v        = (realToFrac n / realToFrac (n + 1)) * wVar + bVar
      rhat     = sqrt (v / wVar)
  in  VPSummary (FL.fold (FL.premap fst FL.mean) meanVars)
                (FL.fold (FL.premap snd FL.mean) meanVars)
                rhat

-- . foldTo (meanOfVariances, varianceOfMeans) . fmap (compute meanVar . fmap integrandF) . splitChains
-}


{-# language FlexibleInstances      #-}
{-# language KindSignatures         #-}
{-# language FunctionalDependencies #-}
{-# language LambdaCase             #-}
{-# language MultiParamTypeClasses  #-}
{-# language Rank2Types             #-}
{-# language TupleSections          #-}
module Main where

import Control.Monad.Reader
import Data.Complex
import Data.Functor.Const
import Data.Functor.Contravariant
import Data.Functor.Identity
import Data.Monoid                 (First(..))
import Data.Profunctor
import Data.Profunctor.Unsafe      ((#.), (.#))
import Data.Tagged
import Data.Void                   (absurd, vacuous)

type Equality s t a b
  = forall (p :: * -> * -> *) f.
    p a (f b) -> p s (f t)

type Lens      s t a b
  = forall   f. Functor f
  => (a -> f b) -> s -> f t

type Traversal s t a b
  = forall   f. Applicative f
  => (a -> f b) -> s -> f t

type Iso       s t a b
  = forall p f. (Functor f, Profunctor p)
  => p a (f b)  -> p s (f t)

type Prism     s t a b
  = forall p f. (Applicative f, Choice p)
  => p a (f b)  -> p s (f t)

type Getter    s   a
  = forall   f. (Functor f, Contravariant f)
  => (a -> f a) -> s -> f s

type Fold      s   a
  = forall   f. (Applicative f, Contravariant f)
  => (a -> f a) -> s -> f s

type Optic p f s t a b
  = p a (f b)  -> p s (f t)
type Optic' p f t b
  = p b (f b)  -> p t (f t)

type Review t b
  = Tagged b (Identity b) -> Tagged t (Identity t)
-- in other words,
type Review' t b
  = Optic' Tagged Identity t b

-- Traversal:
--
-- We strengthen `Functor` to `Applicative`. Why?
--
-- Clue: note that
--
-- traverse :: Applicative f => (a -> f b) -> t a -> f (t b)
--           ~ Traversable f => Traversal (f a) (f b) a b
--
-- So, instead we can ask why does `traverse` require `Applicative`?
--
-- The instance for lists is informative:
--
-- instance Traversable [] where
--   traverse _ []     = pure []
--   traverse f (x:xs) = (:) <$> f x <*> traverse f xs
--
-- Every lens is a valid Traversal. Witness:

lensToTraversal :: Lens s t a b -> Traversal s t a b
lensToTraversal = id

--
-- Iso:
--
-- We generalize `->` to `(Profunctor p =>) p`. Why?
--
-- One answer is that using `Exchange` (a `Profunctor`) allows us to reverse
-- the `Iso` and extract its parts.
--
-- Prism:
--
-- Here `f` is an `Applicative` and `p` is a `Choice`. Why?
--
-- First of all:
-- * Every `Prism`: Applicative f, Choice p
--   is a valid          v            v
--     `Traversal`: Applicative f
--
--   witness:

prismToTraversal :: Prism s t a b -> Traversal s t a b
prismToTraversal = id

--
--   Here `Prism` adds `Choice p`, meaning TODO
--
-- * Every `Iso`:   Functor f,     Profunctor p
--   is a valid           v             v
--         `Prism`: Applicative f, Choice p
--
--   witness:

isoToPrism :: Iso s t a b -> Prism s t a b
isoToPrism = id

-- Iso's are everything actually:

isoToLens :: Iso s t a b -> Lens s t a b
isoToLens = id

isoToTraversal :: Iso s t a b -> Traversal s t a b
isoToTraversal = id

isoToGetter :: Iso s s a a -> Getter s a
isoToGetter = id

isoToFold :: Iso s s a a -> Fold s a
isoToFold = id

isoToReview :: Iso s s a a -> Review s a
isoToReview = id

-- Except an `Equality`, that is. An `Equality` is an `Iso`:

equalityToIso :: Equality s t a b -> Iso s t a b
equalityToIso = id

-- An Equality is a witness that a ~ s and b ~ t.
--
--
--   Here `Prism` strenthens the constraints:
--   * `Functor` to `Applicative` because it doesn't touch exactly one position
--   * `Profunctor` to `Choice` because TODO
--
-- Getter:
--
-- A getter describes how to retrieve a single value
--
-- > for f to be both Functor and Contravariant implies `f a` doesn't contain
-- > any `a`s at all!
--
-- citation: https://www.reddit.com/r/haskell/comments/5vb6x1/how_do_i_learn_lensinternals/de0uz1v
-- Witness:

coerceGetterF :: (Functor f, Contravariant f) => f a -> f b
coerceGetterF = vacuous . contramap absurd

-- Fold:
--
-- A `Fold` describes how to retrieve multiple values
--
-- Note how conspicuously `Fold` and `Getter` are:
--
-- TODO: understand
-- "A `Getter` is a legal `Fold` that just ignores the supplied `Monoid`
--
-- * Every `Getter`: Functor f,   Contravariant f
--   is a valid          v              v
--           `Fold`: Applicative f, Contravariant f
--
-- Witness:

getterToFold :: Getter s a -> Fold s a
getterToFold = id

-- Like `Getter`, note that the functor used in the `Fold` can't contain any
-- values:

coerceFoldF :: (Applicative f, Contravariant f) => f a -> f b
coerceFoldF = vacuous . contramap absurd

-- A "is a limited form of a `Prism` that can only be used for `re` operations.
-- Witness:

prismToReview :: Prism' t b -> Review t b
prismToReview = id

type Simple f s a  = f s s a a
type Equality' s a = Simple Equality s a
type Lens' s a     = Simple Lens s a
type Iso'  s a     = Simple Iso  s a
type Prism' s a    = Simple Prism s a
type Getting r s a = (a -> Const r a) -> s -> Const r s

iso :: (s -> a) -> (b -> t) -> Iso s t a b
iso sa bt = dimap sa (fmap bt)

-- used to provide access to the two parts of an iso
data Exchange a b s t = Exchange (s -> a) (b -> t)

instance Functor (Exchange a b s) where
  fmap g' (Exchange f g) = Exchange f (g' . g)

instance Profunctor (Exchange a b) where
  dimap f' g' (Exchange f g) = Exchange (f . f') (g' . g)

withIso :: Iso s t a b -> ((s -> a) -> (b -> t) -> r) -> r
withIso ai k = case ai (Exchange id Identity) of
  Exchange sa bt -> k sa (runIdentity #. bt)

-- used to provide access to the two parts of a prism
data Market a b s t = Market (b -> t) (s -> Either t a)

instance Functor (Market a b s) where
  fmap h (Market f g) = Market (h . f) (either (Left . h) Right . g)

instance Profunctor (Market a b) where
  dimap f' g' (Market f g) = Market (g' . f) (left' g' . g . f')

instance Choice (Market a b) where
  left' (Market f g) = Market (Left . f) $ \case
    Left x -> case g x of
      Left y -> Left (Left y)
      Right y -> Right y
    Right c -> Left (Right c)
  -- TODO: implement right'

withPrism :: Prism s t a b -> ((b -> t) -> (s -> Either t a) -> r) -> r
withPrism p k = case p (Market Identity Right) of
  Market bt sa -> k (runIdentity #. bt) (left' runIdentity . sa)

prism :: (b -> t) -> (s -> Either t a) -> Prism s t a b
prism bt seta = dimap seta (either pure (fmap bt)) . right'

matching :: Prism s t a b -> s -> Either t a
matching k = withPrism k $ \_ seta -> seta

_Left :: Prism (Either a c) (Either b c) a b
_Left = prism Left $ either Right (Left . Right)

_Right :: Prism (Either c a) (Either c b) a b
_Right = prism Right $ either (Left . Left) Right

-- TODO: rewrite  / become more comfortable with this section

re :: Review t b -> Getter b t
re p = to (runIdentity #. unTagged #. p .# Tagged .# Identity)

to :: (Profunctor p, Contravariant f) => (s -> a) -> Optic' p f s a
to k = dimap k (contramap k)

reviewSimple :: Prism' s a -> a -> s
reviewSimple r = runIdentity . unTagged . r . Tagged . Identity

review :: MonadReader b m => Optic' Tagged Identity t b -> m t
review r = asks $ runIdentity . unTagged . r . Tagged . Identity

foldMapOf :: Getting r s a -> (a -> r) -> s -> r
foldMapOf l f = getConst #. l (Const #. f)

preview :: Prism' s a -> s -> Maybe a
preview p s = getFirst #. foldMapOf p (First #. Just) $ s

(^?) :: s -> Prism' s a -> Maybe a
(^?) s p = preview p s

--

from :: Iso s t a b -> Iso b a t s
from l = withIso l $ \ sa bt -> iso bt sa

view :: Lens s t a b -> s -> a
view l s = getConst (l Const s)

over :: Lens s t a b -> (a -> b) -> s -> t
over l f = runIdentity . l (Identity . f)

set :: Lens s t a b -> b -> s -> t
set l = over l . const

lens :: (s -> a) -> (s -> b -> t) -> Lens s t a b
lens get' set' f s = set' s <$> f (get' s)

realLens, realLens' :: RealFloat a => Lens' (Complex a) a
realLens f (r :+ i) = fmap (:+ i) (f r)

imagLens, imagLens' :: RealFloat a => Lens' (Complex a) a
imagLens f (r :+ i) = fmap (r :+) (f i)

realLens' = lens (\(r :+ _) -> r) (\(_ :+ i) r -> r :+ i)
imagLens' = lens (\(_ :+ i) -> i) (\(r :+ _) i -> r :+ i)

(<&>) :: Functor f => f a -> (a -> b) -> f b
(<&>) = flip fmap

class Field1 s t a b | s -> a, t -> b, s b -> t, t a -> s where
  _1 :: Lens s t a b

instance Field1 (a, b) (a', b) a a' where
  _1 f ~(a, b) = f a <&> \a' -> (a', b)

instance Field1 (a, b, c) (a', b, c) a a' where
  _1 k ~(a, b, c) = k a <&> \a' -> (a', b, c)

class Field2 s t a b | s -> a, t -> b, s b -> t, t a -> s where
  _2 :: Lens s t a b

instance Field2 (a, b) (a, b') b b' where
  _2 f ~(a, b) = f b <&> \b' -> (a, b')

instance Field2 (a, b, c) (a, b', c) b b' where
  _2 k ~(a, b, c) = k b <&> \b' -> (a, b', c)

main :: IO ()
main = do
  print $ set realLens (1 :: Double) (0 :+ 1)
  print $ view _1 ('x', 'y', 'z')
  print $ view _2 $ set _2 False ('x', 'y', 'z')