{-# LANGUAGE GeneralisedNewtypeDeriving #-}
-- |
-- Module Orbital
--
-- A probability distribution for driver laptimes, initialised from
-- Elo-like ratings.
--
-- In the PDF, t = 0 is taken to be the ideal laptime in a track.
--
-- The name "orbital" alludes to this PDF originally having been, modulo
-- constant factors, the radial probability density of the 1s hydrogen atom
-- orbital. The corresponding distribution can be obtained by setting the
-- gamma shape parameter to 3.
module Orbital
    ( deltaWP
    , eloWP
    , initialRating
    , kFromRating
    , ratingFromK
    , OrbitalDistribution
    , orbitalDistr
    , orbitalK
    ) where

import Numeric.SpecFunctions (choose)
import Statistics.Distribution
import Statistics.Distribution.Gamma
import Data.MemoTrie
import Data.List (foldl')

-- | The base implementation of victory probabilities for performances
-- modelled by gamma distributions with the same (integer) shape.
ratioWP :: Int -> Double -> Double
ratioWP gsh = \w -> w^gsh * evalPoly (memoWpCoefficient gsh) (gsh-1) w

-- | A strict pair type.
data SP a b = SP !a !b

-- | A polynomial evaluator. For our immediate purposes, it performs slightly
-- better than the one available in 'Numeric.Polynomial'.
evalPoly :: (Int -> Double) -> Int -> Double -> Double
evalPoly coef n x = y
    where
    SP xn y = foldl' alg (SP 1 0) [0..n]
    alg (SP xx s) m = SP (x * xx) (coef m * xx + s)

-- | Coefficients of the polynomial used in the gamma winning probability
-- calculation.
wpCoefficient :: Int -> Int -> Double
wpCoefficient gsh m =
    (-1)^m * choose gsh m * choose (2*gsh-1) (gsh-1)
        * fromIntegral (gsh-m) / fromIntegral (gsh+m)

-- | Memoised version of 'wpCoefficient'.
memoWpCoefficient :: Int -> Int -> Double
memoWpCoefficient = memo2 wpCoefficient

-- | Initial rating for new racers. Defined as a constant here for
-- expediteness, taking into account that making it configurable isn't as
-- essential as it is making other engine parameters configurable.
initialRating :: Double
initialRating = 1500

-- | The base Elo exponent. In essence, a logistic growth factor for expected
-- scores/winning probabilities. The denominator used here, 400, is standard
-- for chess rating computations. Using it, a 0.6 expected score corresponds
-- to a ~70.4 rating gap, and a 0.9 expected score, to a ~382 gap.
eloAlpha :: Double
eloAlpha = log 10 / 400

-- | Adjusts the k values obtained from ratings depending on the gamma shape so
-- that the same rating gaps correspond to approximately the same winning
-- probabilities.
--
-- The formula used here is made of two factors:
--
-- - A base correction, which on its own would adjust the derivative of
--   @deltaWP gsh@ to match that of @deltaWP 1@ at zero. This factor, which is
--   the best adjustment for rating gaps close to zero, slightly increases
--   the absolute excess winning probabilities over the entire range of gaps.
--
-- - An additional correction, which reduces the aforementioned bias. Its value
--   is chosen to (approximately) minimise the sum of the squares of the
--   differences between @ratioWP gsh@ and @ratioWP 1@ over the entire range
--   of ratios.
eloGammaCorrection :: Int -> Double
eloGammaCorrection gsh = extraAdjustment * baseCorrection
    where
    gsh' = fromIntegral gsh
    baseCorrection = 4^(gsh-1) / (gsh' * choose (2*gsh-1) (gsh-1))
    fittedConstant = (2*sqrt 29 - 3)*pi
    extraAdjustment = 1 - (gsh'-1) / (gsh' * fittedConstant)

-- | Elo conversion factor. Amounts to eloAlpha for shape 1, in which case the
-- gamma model coincides with the conventional Elo system.
eloFactor :: Int -> Double
eloFactor gsh = eloGammaCorrection gsh * eloAlpha

-- | Memoised Elo conversion factor.
memoEloFactor :: Int -> Double
memoEloFactor = memo eloFactor

-- | Winning probabilities for (scaled) rating differences. Ratings are
-- essentially logarithms of the gamma distribution rate parameters.
deltaWP :: Int -> Double -> Double
deltaWP gsh = \d -> ratioWP gsh (1 / (1 + exp(- memoEloFactor gsh * d)))
-- If ru = log u, rv = log v, and d = ru - rv, then
-- u / (u + v) = 1 / (1 + exp (-d))
-- Another way of expressing the ratio is given by
-- 1 / (1 + exp(-d)) = (1 + tanh (d/2)) / 2

-- | Elo-based winning probability. This version takes the ratings of the two
-- racers. Only here for reference, at least for now.
eloWP :: Double -> Double -> Double
eloWP ru rv = deltaWP 1 (ru-rv)

-- | Rating that corresponds to a gamma rate parameter of 1.
--
-- As far as the results go, the effect of 'referenceRating' is superficial.
-- Since it adjusts all ratings in the same manner, changing it does not affect
-- the overall results.
--
-- There used to be a counterpart @referenceK@ definition, which was removed
-- for being redundant. It is worth noting that scaling the conversions to make
-- the gamma rate ("k") values larger slows down the calculations, probably by
-- making it harder for the gamma distribution algorithms.
referenceRating :: Double
referenceRating = initialRating

-- | k-to-rating conversion. Only for reference, at least for now.
ratingFromK :: Int -> Double -> Double
ratingFromK gsh u = referenceRating + log u / eloFactor gsh

-- | rating-to-k conversion.
kFromRating :: Int -> Double -> Double
kFromRating gsh r = exp ((r - referenceRating) * eloFactor gsh)
-- Using memoEloFactor instead of eloFactor here doesn't seem to improve
-- performance.

-- | The model distribution is a special case of the gamma distribution, with
-- theta = 1/k.
newtype OrbitalDistribution = OrbitalDistribution
    { orbitalGamma :: GammaDistribution }
    deriving (Eq, Show, Distribution, ContDistr, ContGen, MaybeEntropy
        , Variance, MaybeVariance, Mean, MaybeMean)

-- | Sets up the model distribution given a rating.
orbitalDistr
    :: Int    -- ^ Shape parameter to be used.
    -> Double -- ^ Elo-like rating.
    -> OrbitalDistribution
orbitalDistr gsh r = OrbitalDistribution
    (gammaDistr (fromIntegral gsh) (1 / kFromRating gsh r))

-- | Recovers the k factor from the distribution.
orbitalK :: OrbitalDistribution -> Double
orbitalK distr = 1 / gdScale (orbitalGamma distr)

-- | Recovers the shape parameter from the distribution.
orbitalShape :: OrbitalDistribution -> Double
orbitalShape distr = gdShape (orbitalGamma distr)


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

{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Sundials.ARKode.ODE
-- Copyright   :  Dominic Steinitz 2018,
--                Novadiscovery 2018
-- License     :  BSD
-- Maintainer  :  Dominic Steinitz
-- Stability   :  provisional
--
-- Solution of ordinary differential equation (ODE) initial value problems.
-- See <https://computation.llnl.gov/projects/sundials/sundials-software> for more detail.
--
-- A simple example:
--
-- <<diagrams/brusselator.png#diagram=brusselator&height=400&width=500>>
--
-- @
-- import           Numeric.Sundials.ARKode.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))
-- @
--
-- With Sundials ARKode, it is possible to retrieve the Butcher
-- tableau for the solver. FIXME: Not available just now and hopefully
-- normal service will be resumed soon.
--
-- @
-- import           Numeric.Sundials.ARKode.ODE
-- import           Numeric.LinearAlgebra
--
-- import           Data.List (intercalate)
--
-- import           Text.PrettyPrint.HughesPJClass
--
--
-- butcherTableauTex :: ButcherTable -> String
-- butcherTableauTex (ButcherTable m c b b2) =
--   render $
--   vcat [ text ("\n\\begin{array}{c|" ++ (concat $ replicate n "c") ++ "}")
--        , us
--        , text "\\hline"
--        , text bs <+> text "\\\\"
--        , text b2s <+> text "\\\\"
--        , text "\\end{array}"
--        ]
--   where
--     n = rows m
--     rs = toLists m
--     ss = map (\r -> intercalate " & " $ map show r) rs
--     ts = zipWith (\i r -> show i ++ " & " ++ r) (toList c) ss
--     us = vcat $ map (\r -> text r <+> text "\\\\") ts
--     bs  = " & " ++ (intercalate " & " $ map show $ toList b)
--     b2s = " & " ++ (intercalate " & " $ map show $ toList b2)
--
-- main :: IO ()
-- main = do
--
--   let res = butcherTable (SDIRK_2_1_2 undefined)
--   putStrLn $ show res
--   putStrLn $ butcherTableauTex res
--
--   let resA = butcherTable (KVAERNO_4_2_3 undefined)
--   putStrLn $ show resA
--   putStrLn $ butcherTableauTex resA
--
--   let resB = butcherTable (SDIRK_5_3_4 undefined)
--   putStrLn $ show resB
--   putStrLn $ butcherTableauTex resB
-- @
--
--  Using the code above from the examples gives
--
-- KVAERNO_4_2_3
--
-- \[
-- \begin{array}{c|cccc}
-- 0.0 & 0.0 & 0.0 & 0.0 & 0.0 \\
-- 0.871733043 & 0.4358665215 & 0.4358665215 & 0.0 & 0.0 \\
-- 1.0 & 0.490563388419108 & 7.3570090080892e-2 & 0.4358665215 & 0.0 \\
-- 1.0 & 0.308809969973036 & 1.490563388254106 & -1.235239879727145 & 0.4358665215 \\
-- \hline
--  & 0.308809969973036 & 1.490563388254106 & -1.235239879727145 & 0.4358665215 \\
--  & 0.490563388419108 & 7.3570090080892e-2 & 0.4358665215 & 0.0 \\
-- \end{array}
-- \]
--
-- SDIRK_2_1_2
--
-- \[
-- \begin{array}{c|cc}
-- 1.0 & 1.0 & 0.0 \\
-- 0.0 & -1.0 & 1.0 \\
-- \hline
--  & 0.5 & 0.5 \\
--  & 1.0 & 0.0 \\
-- \end{array}
-- \]
--
-- SDIRK_5_3_4
--
-- \[
-- \begin{array}{c|ccccc}
-- 0.25 & 0.25 & 0.0 & 0.0 & 0.0 & 0.0 \\
-- 0.75 & 0.5 & 0.25 & 0.0 & 0.0 & 0.0 \\
-- 0.55 & 0.34 & -4.0e-2 & 0.25 & 0.0 & 0.0 \\
-- 0.5 & 0.2727941176470588 & -5.036764705882353e-2 & 2.7573529411764705e-2 & 0.25 & 0.0 \\
-- 1.0 & 1.0416666666666667 & -1.0208333333333333 & 7.8125 & -7.083333333333333 & 0.25 \\
-- \hline
--  & 1.0416666666666667 & -1.0208333333333333 & 7.8125 & -7.083333333333333 & 0.25 \\
--  & 1.2291666666666667 & -0.17708333333333334 & 7.03125 & -7.083333333333333 & 0.0 \\
-- \end{array}
-- \]
-----------------------------------------------------------------------------
module Numeric.Sundials.ARKode.ODE ( odeSolve
                                   , odeSolveV
                                   , odeSolveVWith
                                   , odeSolveVWith'
                                   , odeSolveWithEvents
                                   , ODEMethod(..)
                                   , StepControl(..)
                                   ) where

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

import           Data.Monoid ((<>))
import           Data.Maybe (isJust)

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

import qualified Data.Vector.Storable as V

import           Data.Coerce (coerce)
import           System.IO.Unsafe (unsafePerformIO)
import           GHC.Generics (C1, Constructor, (:+:)(..), D1, Rep, Generic, M1(..),
                               from, conName)

import           Numeric.LinearAlgebra.Devel (createVector)

import           Numeric.LinearAlgebra.HMatrix (Vector, Matrix, toList, rows,
                                                cols, toLists, size, reshape,
                                                (><))

import           Numeric.Sundials.Types
import qualified Numeric.Sundials.Arkode as T
import           Numeric.Sundials.Arkode (SunIndexType)
import           Numeric.Sundials.Arkode (sDIRK_2_1_2,
                                          bILLINGTON_3_3_2,
                                          tRBDF2_3_3_2,
                                          kVAERNO_4_2_3,
                                          aRK324L2SA_DIRK_4_2_3,
                                          cASH_5_2_4,
                                          cASH_5_3_4,
                                          sDIRK_5_3_4,
                                          kVAERNO_5_3_4,
                                          aRK436L2SA_DIRK_6_3_4,
                                          kVAERNO_7_4_5,
                                          aRK548L2SA_DIRK_8_4_5,
                                          hEUN_EULER_2_1_2,
                                          bOGACKI_SHAMPINE_4_2_3,
                                          aRK324L2SA_ERK_4_2_3,
                                          zONNEVELD_5_3_4,
                                          aRK436L2SA_ERK_6_3_4,
                                          sAYFY_ABURUB_6_3_4,
                                          cASH_KARP_6_4_5,
                                          fEHLBERG_6_4_5,
                                          dORMAND_PRINCE_7_4_5,
                                          aRK548L2SA_ERK_8_4_5,
                                          vERNER_8_5_6,
                                          fEHLBERG_13_7_8)

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

C.include "<stdlib.h>"
C.include "<stdio.h>"
C.include "<math.h>"
C.include "<arkode/arkode_arkstep.h>"                 -- prototypes for ARKODE fcts., consts.
C.include "<arkode/arkode_erkstep.h>"
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 "<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"

-- | Stepping functions
data ODEMethod = SDIRK_2_1_2            Jacobian
               | SDIRK_2_1_2'
               | BILLINGTON_3_3_2       Jacobian
               | BILLINGTON_3_3_2'
               | TRBDF2_3_3_2           Jacobian
               | TRBDF2_3_3_2'
               | KVAERNO_4_2_3          Jacobian
               | KVAERNO_4_2_3'
               | ARK324L2SA_DIRK_4_2_3  Jacobian
               | ARK324L2SA_DIRK_4_2_3'
               | CASH_5_2_4             Jacobian
               | CASH_5_2_4'
               | CASH_5_3_4             Jacobian
               | CASH_5_3_4'
               | SDIRK_5_3_4            Jacobian
               | SDIRK_5_3_4'
               | KVAERNO_5_3_4          Jacobian
               | KVAERNO_5_3_4'
               | ARK436L2SA_DIRK_6_3_4  Jacobian
               | ARK436L2SA_DIRK_6_3_4'
               | KVAERNO_7_4_5          Jacobian
               | KVAERNO_7_4_5'
               | ARK548L2SA_DIRK_8_4_5  Jacobian
               | ARK548L2SA_DIRK_8_4_5'
               | HEUN_EULER_2_1_2         Jacobian
               | HEUN_EULER_2_1_2'
               | BOGACKI_SHAMPINE_4_2_3   Jacobian
               | BOGACKI_SHAMPINE_4_2_3'
               | ARK324L2SA_ERK_4_2_3     Jacobian
               | ARK324L2SA_ERK_4_2_3'
               | ZONNEVELD_5_3_4          Jacobian
               | ZONNEVELD_5_3_4'
               | ARK436L2SA_ERK_6_3_4     Jacobian
               | ARK436L2SA_ERK_6_3_4'
               | SAYFY_ABURUB_6_3_4       Jacobian
               | SAYFY_ABURUB_6_3_4'
               | CASH_KARP_6_4_5          Jacobian
               | CASH_KARP_6_4_5'
               | FEHLBERG_6_4_5         Jacobian
               | FEHLBERG_6_4_5'
               | DORMAND_PRINCE_7_4_5     Jacobian
               | DORMAND_PRINCE_7_4_5'
               | ARK548L2SA_ERK_8_4_5     Jacobian
               | ARK548L2SA_ERK_8_4_5'
               | VERNER_8_5_6            Jacobian
               | VERNER_8_5_6'
               | FEHLBERG_13_7_8         Jacobian
               | FEHLBERG_13_7_8'
  deriving Generic

constrName :: (HasConstructor (Rep a), Generic a)=> a -> String
constrName = genericConstrName . from

class HasConstructor (f :: * -> *) where
  genericConstrName :: f x -> String

instance HasConstructor f => HasConstructor (D1 c f) where
  genericConstrName (M1 x) = genericConstrName x

instance (HasConstructor x, HasConstructor y) => HasConstructor (x :+: y) where
  genericConstrName (L1 l) = genericConstrName l
  genericConstrName (R1 r) = genericConstrName r

instance Constructor c => HasConstructor (C1 c f) where
  genericConstrName x = conName x

instance Show ODEMethod where
  show x = constrName x

-- FIXME: We can probably do better here with generics
getMethod :: ODEMethod -> Int
getMethod (SDIRK_2_1_2 _)            = sDIRK_2_1_2
getMethod (SDIRK_2_1_2')             = sDIRK_2_1_2
getMethod (BILLINGTON_3_3_2 _)       = bILLINGTON_3_3_2
getMethod (BILLINGTON_3_3_2')        = bILLINGTON_3_3_2
getMethod (TRBDF2_3_3_2 _)           = tRBDF2_3_3_2
getMethod (TRBDF2_3_3_2')            = tRBDF2_3_3_2
getMethod (KVAERNO_4_2_3  _)         = kVAERNO_4_2_3
getMethod (KVAERNO_4_2_3')           = kVAERNO_4_2_3
getMethod (ARK324L2SA_DIRK_4_2_3 _)  = aRK324L2SA_DIRK_4_2_3
getMethod (ARK324L2SA_DIRK_4_2_3')   = aRK324L2SA_DIRK_4_2_3
getMethod (CASH_5_2_4 _)             = cASH_5_2_4
getMethod (CASH_5_2_4')              = cASH_5_2_4
getMethod (CASH_5_3_4 _)             = cASH_5_3_4
getMethod (CASH_5_3_4')              = cASH_5_3_4
getMethod (SDIRK_5_3_4 _)            = sDIRK_5_3_4
getMethod (SDIRK_5_3_4')             = sDIRK_5_3_4
getMethod (KVAERNO_5_3_4 _)          = kVAERNO_5_3_4
getMethod (KVAERNO_5_3_4')           = kVAERNO_5_3_4
getMethod (ARK436L2SA_DIRK_6_3_4 _)  = aRK436L2SA_DIRK_6_3_4
getMethod (ARK436L2SA_DIRK_6_3_4')   = aRK436L2SA_DIRK_6_3_4
getMethod (KVAERNO_7_4_5 _)          = kVAERNO_7_4_5
getMethod (KVAERNO_7_4_5')           = kVAERNO_7_4_5
getMethod (ARK548L2SA_DIRK_8_4_5 _)  = aRK548L2SA_DIRK_8_4_5
getMethod (ARK548L2SA_DIRK_8_4_5')   = aRK548L2SA_DIRK_8_4_5
getMethod (HEUN_EULER_2_1_2 _)       = hEUN_EULER_2_1_2
getMethod (HEUN_EULER_2_1_2')        = hEUN_EULER_2_1_2
getMethod (BOGACKI_SHAMPINE_4_2_3 _) = bOGACKI_SHAMPINE_4_2_3
getMethod (BOGACKI_SHAMPINE_4_2_3')  = bOGACKI_SHAMPINE_4_2_3
getMethod (ARK324L2SA_ERK_4_2_3 _)   = aRK324L2SA_ERK_4_2_3
getMethod (ARK324L2SA_ERK_4_2_3')    = aRK324L2SA_ERK_4_2_3
getMethod (ZONNEVELD_5_3_4 _)        = zONNEVELD_5_3_4
getMethod (ZONNEVELD_5_3_4')         = zONNEVELD_5_3_4
getMethod (ARK436L2SA_ERK_6_3_4 _)   = aRK436L2SA_ERK_6_3_4
getMethod (ARK436L2SA_ERK_6_3_4')    = aRK436L2SA_ERK_6_3_4
getMethod (SAYFY_ABURUB_6_3_4 _)     = sAYFY_ABURUB_6_3_4
getMethod (SAYFY_ABURUB_6_3_4')      = sAYFY_ABURUB_6_3_4
getMethod (CASH_KARP_6_4_5 _)        = cASH_KARP_6_4_5
getMethod (CASH_KARP_6_4_5')         = cASH_KARP_6_4_5
getMethod (FEHLBERG_6_4_5 _)         = fEHLBERG_6_4_5
getMethod (FEHLBERG_6_4_5' )         = fEHLBERG_6_4_5
getMethod (DORMAND_PRINCE_7_4_5 _)   = dORMAND_PRINCE_7_4_5
getMethod (DORMAND_PRINCE_7_4_5')    = dORMAND_PRINCE_7_4_5
getMethod (ARK548L2SA_ERK_8_4_5 _)   = aRK548L2SA_ERK_8_4_5
getMethod (ARK548L2SA_ERK_8_4_5')    = aRK548L2SA_ERK_8_4_5
getMethod (VERNER_8_5_6 _)           = vERNER_8_5_6
getMethod (VERNER_8_5_6')            = vERNER_8_5_6
getMethod (FEHLBERG_13_7_8 _)        = fEHLBERG_13_7_8
getMethod (FEHLBERG_13_7_8')         = fEHLBERG_13_7_8

getJacobian :: ODEMethod -> Maybe Jacobian
getJacobian (SDIRK_2_1_2 j)            = Just j
getJacobian (BILLINGTON_3_3_2 j)       = Just j
getJacobian (TRBDF2_3_3_2 j)           = Just j
getJacobian (KVAERNO_4_2_3  j)         = Just j
getJacobian (ARK324L2SA_DIRK_4_2_3 j)  = Just j
getJacobian (CASH_5_2_4 j)             = Just j
getJacobian (CASH_5_3_4 j)             = Just j
getJacobian (SDIRK_5_3_4 j)            = Just j
getJacobian (KVAERNO_5_3_4 j)          = Just j
getJacobian (ARK436L2SA_DIRK_6_3_4 j)  = Just j
getJacobian (KVAERNO_7_4_5 j)          = Just j
getJacobian (ARK548L2SA_DIRK_8_4_5 j)  = Just j
getJacobian (HEUN_EULER_2_1_2 j)       = Just j
getJacobian (BOGACKI_SHAMPINE_4_2_3 j) = Just j
getJacobian (ARK324L2SA_ERK_4_2_3 j)   = Just j
getJacobian (ZONNEVELD_5_3_4 j)        = Just j
getJacobian (ARK436L2SA_ERK_6_3_4 j)   = Just j
getJacobian (SAYFY_ABURUB_6_3_4 j)     = Just j
getJacobian (CASH_KARP_6_4_5 j)        = Just j
getJacobian (FEHLBERG_6_4_5 j)         = Just j
getJacobian (DORMAND_PRINCE_7_4_5 j)   = Just j
getJacobian (ARK548L2SA_ERK_8_4_5 j)   = Just j
getJacobian (VERNER_8_5_6 j)           = Just j
getJacobian (FEHLBERG_13_7_8 j)        = Just j
getJacobian _                          = Nothing

-- | A version of 'odeSolveVWith' with reasonable default step control.
odeSolveV
    :: ODEMethod
    -> Maybe Double      -- ^ initial step size - by default, ARKode
                         -- 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 SDIRK_5_3_4' (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)

odeSolveVWith ::
  ODEMethod
  -> StepControl
  -> Maybe Double -- ^ initial step size - by default, ARKode
                  -- 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 method control initStepSize 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   = error "ARKode: unexpected use of ODEOpts.odeMethod"
                   , stepControl = error "ARKode: unexpected use of ODEOpts.stepControl"
                   , initStep    = error "ARKode: unexpected use of ODEOpts.initStep"
                   }

odeSolveVWith' ::
  ODEOpts ODEMethod
  -> ODEMethod
  -> StepControl
  -> Maybe Double -- ^ initial step size - by default, ARKode
                  -- 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
  -> Either (Matrix Double, Int) (Matrix Double, SundialsDiagnostics) -- ^ Error code or solution
odeSolveVWith' opts method control initStepSize f y0 tt =
  case solveOdeC (fromIntegral $ maxFail opts)
                 (fromIntegral $ maxNumSteps opts) (coerce $ minStep opts)
                 (fromIntegral $ getMethod method) (coerce initStepSize) jacH (scise control)
                 (coerce f) (coerce y0) (coerce tt) of
    Left  (v, c) -> Left  (reshape l (coerce v), fromIntegral c)
    Right (v, d)
      | V.null y0 -> Right ((V.length tt >< 0) [], emptyDiagnostics)
      | otherwise -> Right (reshape l (coerce v), 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 method
    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

-- | This function implements the same interface as
-- 'Numeric.Sundials.CVode.ODE.odeSolveWithEvents', although it does not
-- currently support events.
odeSolveWithEvents
  :: ODEOpts ODEMethod
  -> [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 events _ rhs _mb_jac y0 times
  | (not . null) events =
      -- Call error rather than return a Left because this is a programming
      -- error, not just a runtime issue.
      error $ "ARKode called with a non-empty list of events (" ++ show (length events) ++
      " in total).\
      \ ARKode does not support events at this point and should not be passed any."
  | otherwise =
      let
        result :: Either (Matrix Double, Int)
                         (Matrix Double, SundialsDiagnostics)
        result =
          odeSolveVWith' opts
            (odeMethod opts)
            (stepControl opts)
            (initStep opts)
            rhs y0 times
      in
        case result of
          Left (_, code) -> Left code
          Right (mx, diagn) ->
            Right $ SundialsSolution
                { actualTimeGrid = times
                , solutionMatrix =
                    -- Note: at this time, ARKode's output matrix does not
                    -- include the time column, so we're not dropping it
                    -- here unlike in CVode. If/when we add event support
                    -- to ARKode, this is going to change.
                    mx
                , eventInfo = []
                , diagnostics = diagn
                }

solveOdeC ::
  CInt ->
  SunIndexType ->
  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
  -> V.Vector CDouble -- ^ Desired solution times
  -> Either (V.Vector CDouble, CInt) (V.Vector CDouble, SundialsDiagnostics) -- ^ Partial solution and error code or
                                                                             -- solution and diagnostics
solveOdeC maxErrTestFails maxNumSteps_ minStep_ method initStepSize
          jacH (aTols, rTol) fun f0 ts
  | V.null f0 = -- 0-dimensional (empty) system
    Right (V.empty, emptyDiagnostics)
  | otherwise =
  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 :: SunIndexType
      nEq = fromIntegral dim
      nTs :: CInt
      nTs = fromIntegral $ V.length ts
  quasiMatrixRes <- createVector ((fromIntegral dim) * (fromIntegral nTs))
  qMatMut <- V.thaw quasiMatrixRes
  diagMut :: V.MVector _ SunIndexType <- V.thaw =<< createVector 10 -- FIXME
  -- We need the types that sundials expects. These are tied together
  -- in 'CLangToHaskellTypes'. 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 $ T.SunVector { T.sunVecN = T.sunVecN sv
                             , T.sunVecVals = fun t (T.sunVecVals sv)
                             }
        [CU.exp| int{ 0 } |]
  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.ARKode.ODE: Jacobian not defined"
          Just jacI -> do j <- jacI t <$> (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 <- [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 *arkode_mem = NULL;   /* empty ARKode memory structure                */
                         realtype t;
                         long nst, nst_a, nfe, nfi, nsetups, nje, nfeLS, nni, ncfn, netf;

                         /* 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];
                         };

                         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];
                         };

                         /* Call ARKStepCreate to initialize the ARK timestepper module and */
                         /* specify the right-hand side function in y'=f(t,y), the inital time */
                         /* T0, and the initial dependent variable vector y.  Note: since this */
                         /* problem is fully implicit, we set f_E to NULL and f_I to f. */

                         /* Here we use the C types defined in helpers.h which tie up with */
                         /* the Haskell types defined in CLangToHaskellTypes               */
                         if ($(int method) < MIN_DIRK_NUM) {
                           arkode_mem = ARKStepCreate($fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), NULL, T0, y);
                           if (check_flag((void *)arkode_mem, "ARKStepCreate", 0)) return 1;
                             } else {
                           arkode_mem = ARKStepCreate(NULL, $fun:(int (* funIO) (double t, SunVector y[], SunVector dydt[], void * params)), T0, y);
                           if (check_flag(&flag, "ARKStepCreate", 0)) return 1;
                         }

                         flag = ARKStepSetMinStep(arkode_mem, $(double minStep_));
                         if (check_flag(&flag, "ARKStepSetMinStep", 1)) return 1;
                         flag = ARKStepSetMaxNumSteps(arkode_mem, $(sunindextype maxNumSteps_));
                         if (check_flag(&flag, "ARKStepSetMaxNumSteps", 1)) return 1;
                         flag = ARKStepSetMaxErrTestFails(arkode_mem, $(int maxErrTestFails));
                         if (check_flag(&flag, "ARKStepSetMaxErrTestFails", 1)) return 1;

                         /* Set routines */
                         flag = ARKStepSVtolerances(arkode_mem, $(double rTol), tv);
                         if (check_flag(&flag, "ARKStepSVtolerances", 1)) return 1;

                         /* 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 = ARKStepSetLinearSolver(arkode_mem, LS, A);
                         if (check_flag(&flag, "ARKStepSetLinearSolver", 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 = ARKStepSetInitStep(arkode_mem, $(double ss));
                           if (check_flag(&flag, "ARKStepSetInitStep", 1)) return 1;
                         }

                         /* Set the Jacobian if there is one */
                         if ($(int isJac)) {
                           flag = ARKStepSetJacFn(arkode_mem, $fun:(int (* jacIO) (double t, SunVector y[], SunVector fy[], SunMatrix Jac[], void * params, SunVector tmp1[], SunVector tmp2[], SunVector tmp3[])));
                           if (check_flag(&flag, "ARKStepSetJacFn", 1)) return 1;
                         }

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

                         /* Explicitly set the method */
                         if ($(int method) >= MIN_DIRK_NUM) {
                           flag = ARKStepSetTableNum(arkode_mem, $(int method), -1);
                           if (check_flag(&flag, "ARKStepSetTableNum", 1)) return 1;
                         } else {
                           flag = ARKStepSetTableNum(arkode_mem, -1, $(int method));
                           if (check_flag(&flag, "ERKStepSetTableNum", 1)) return 1;
                         }

                         /* Main time-stepping loop: calls ARKStep to perform the integration */
                         /* Stops when the final time has been reached                       */
                         for (i = 1; i < $(int nTs); i++) {

                           flag = ARKStepEvolve(arkode_mem, ($vec-ptr:(double *ts))[i], y, &t, ARK_NORMAL); /* call integrator */
                           if (check_flag(&flag, "ARKStep solver failure, stopping integration", 1)) return 1;

                           /* Store the results for Haskell */
                           for (j = 0; j < NEQ; j++) {
                             ($vec-ptr:(double *qMatMut))[i * NEQ + j] = NV_Ith_S(y,j);
                           }
                         }

                         /* Get some final statistics on how the solve progressed */

                         flag = ARKStepGetNumSteps(arkode_mem, &nst);
                         check_flag(&flag, "ARKStepGetNumSteps", 1);
                         ($vec-ptr:(sunindextype *diagMut))[0] = nst;

                         flag = ARKStepGetNumStepAttempts(arkode_mem, &nst_a);
                         check_flag(&flag, "ARKStepGetNumStepAttempts", 1);
                         ($vec-ptr:(sunindextype *diagMut))[1] = nst_a;

                         flag = ARKStepGetNumRhsEvals(arkode_mem, &nfe, &nfi);
                         check_flag(&flag, "ARKStepGetNumRhsEvals", 1);
                         ($vec-ptr:(sunindextype *diagMut))[2] = nfe;
                         ($vec-ptr:(sunindextype *diagMut))[3] = nfi;

                         flag = ARKStepGetNumLinSolvSetups(arkode_mem, &nsetups);
                         check_flag(&flag, "ARKStepGetNumLinSolvSetups", 1);
                         ($vec-ptr:(sunindextype *diagMut))[4] = nsetups;

                         flag = ARKStepGetNumErrTestFails(arkode_mem, &netf);
                         check_flag(&flag, "ARKStepGetNumErrTestFails", 1);
                         ($vec-ptr:(sunindextype *diagMut))[5] = netf;

                         flag = ARKStepGetNumNonlinSolvIters(arkode_mem, &nni);
                         check_flag(&flag, "ARKStepGetNumNonlinSolvIters", 1);
                         ($vec-ptr:(sunindextype *diagMut))[6] = nni;

                         flag = ARKStepGetNumNonlinSolvConvFails(arkode_mem, &ncfn);
                         check_flag(&flag, "ARKStepGetNumNonlinSolvConvFails", 1);
                         ($vec-ptr:(sunindextype *diagMut))[7] = ncfn;

                         flag = ARKStepGetNumJacEvals(arkode_mem, &nje);
                         check_flag(&flag, "ARKStepGetNumJacEvals", 1);
                         ($vec-ptr:(sunindextype *diagMut))[8] = nje;

                         flag = ARKStepGetNumLinRhsEvals(arkode_mem, &nfeLS);
                         check_flag(&flag, "ARKStepGetNumLinRhsEvals", 1);
                         ($vec-ptr:(sunindextype *diagMut))[9] = nfeLS;

                         /* Clean up and return */
                         N_VDestroy(y);            /* Free y vector          */
                         N_VDestroy(tv);           /* Free tv vector         */
                         ARKStepFree(&arkode_mem);  /* Free integrator memory */
                         SUNLinSolFree(LS);        /* Free linear solver     */
                         SUNMatDestroy(A);         /* Free A matrix          */

                         return flag;
                       } |]
  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)
  m <- V.freeze qMatMut
  if res == 0
    then do
      return $ Right (m, d)
    else do
      return $ Left  (m, res)


import Data.Complex

#include "../Prelude.hs"

f :: Int -> Complex Double
f n = mkPolar 1 ((2 Prelude.* Prelude.pi) / Prelude.fromIntegral n) Prelude.^ n

realPart'Double = realPart :: Complex Double -> Double
add'ComplexDouble = (Prelude.+) :: Complex Double -> Complex Double -> Complex Double
lit0'ComplexDouble = 0 :: Complex Double


module Solve (
	verboseSolve
	) where

import qualified Data.DList as D
import qualified Data.IntMap.Lazy as M
import Data.Complex (realPart, Complex((:+)))
import Control.Monad.Trans.Writer
import Control.Monad.Trans.State

import Data.MyPolynomial
import Data.MyPolynomial.Print

type ComplexF = Complex Float
type ShowD = D.DList Char
data Solvability b = IsBroad b | Degree | Clear

sToShowD :: String -> D.DList ShowD
sToShowD s = (D.singleton . D.fromList) s

toShowD :: [String] -> D.DList ShowD
toShowD ls = ( D.fromList . ( fmap D.fromList ) ) ls

discriminantSpeech :: Float -> Writer (D.DList ShowD) ()
discriminantSpeech d = let s = case (compare d 0) of	LT -> "strictly negative, the two complex solutions are:\n"
							GT -> "strictly positive, the two real solutions are:\n"
							EQ -> "null, the unique solution is:\n"
				in tell $ toShowD ["The discriminant is ", s]

unrootableSpeech :: Solvability BroadSol -> Writer (D.DList ShowD) ()
unrootableSpeech (IsBroad Absurd) = tell $ sToShowD "Cannot solve: absurd.\n"
unrootableSpeech (IsBroad Real) = tell $ sToShowD "The solution is |R itself.\n"
unrootableSpeech Degree = tell $ sToShowD "The polynomial degree is strictly greater than 2, I can't solve."
unrootableSpeech _ = return ()

degSolvability :: Int -> Solvability b
degSolvability d | d >= 0 && d <= 2	= Clear
		 | otherwise		= Degree

 -- Tells solution/root(s) with mention of exclusion.
tellSolutions :: Maybe Float -> Roots -> [ComplexF] -> Writer (D.DList ShowD) ()
tellSolutions delta (c1, c2) badRoots
	| delta == Just 0 = tell $ toShowD $ tellBad [c1] badRoots "root" []
	| Nothing <- delta = tell $ toShowD $ tellBad [c1] badRoots "solution" []
	| otherwise = tell $ toShowD $ tellBad [c1, c2] badRoots "root" []
  where
    tellBad (rt:rts) bad w =
    	case rt `elem` bad of	True -> ([prettyComplex rt, " (excluded " ++ w ++ ")"] ++) . next
    				_ -> ([prettyComplex rt] ++) . next
    	  where
    	    next = case rts of	(r:_) -> (["\n"] ++) . tellBad rts bad  w
    	    			[] -> ([] ++)

tellReduced :: Int -> M.IntMap Float -> Writer (D.DList ShowD) ()
tellReduced deg map = tell $ toShowD ["Reduced form: ", porcelainPolynomialSM dMap "", " = 0\n"]
	where dMap = M.filterWithKey (\k _ -> k <= deg) map

tellNatural :: M.IntMap Float -> Writer (D.DList ShowD) ()
tellNatural map = tell $ toShowD ["Natural reduced form: ", prettyPolynomialM map, " = 0\n"]

tellForbidden :: [ComplexF] -> Writer (D.DList ShowD) ()
tellForbidden [] = return ()
tellForbidden frts = tell $ toShowD $ ["Excluded roots/solutions:\n"] ++ (showForbidden frts [])
  where
    showForbidden [] = ([] ++)
    showForbidden (fr:fs) = ([prettyComplex fr, "\n"]++) . showForbidden fs

verboseSolve :: Equation -> Writer (D.DList ShowD) (Equation, Maybe Solution)
verboseSolve eq = do
	tellReduced deg mapPol
	tellNatural mapPol
	tell $ toShowD ["Polynomial degree: ", show deg, "\n"]
	tellForbidden badRoots
	case degSolvability deg of	Clear -> doSolve
					Degree -> quit Degree
  where
    Eq (cl, cr) = canonify eq
    (straightP, badRoots) = runState ( handleNegativePowers cl ) []
    ceq = Eq (straightP, cr)	-- canonified equation with negative powers lifted
    mapPol = toMap straightP	-- left polynomial as an IntMap
    deg = degree mapPol		-- left polynomial degree

    -- doSolve :: Writer (D.DList ShowD) (Equation, Maybe Solution)
    doSolve = do
    	let sol = solveEquation ceq
	  in case sol of	Quadratic _	-> result sol
				Simple _	-> result sol
				Broad b		-> quit (IsBroad b)

    -- quit :: Solvability b -> Writer (D.DList ShowD) (Equation, Maybe Solution)
    quit sb = do
    	unrootableSpeech sb
	let ret = case sb of	Degree -> Nothing
				IsBroad b -> Just (Broad b)
	return (ceq, ret)

    -- result :: Solution s -> Writer (D.DList ShowD) (Equation, Maybe Solution)
    result sol@(Quadratic rts) = do
	let delta = discriminantQuadratic mapPol
	discriminantSpeech delta
	tellSolutions (Just delta) rts badRoots
	return ( ceq, Just sol )
    result sol@(Simple fl) = do
	tell $ toShowD [ "The solution is:\n" ]
	tellSolutions Nothing (fl :+ 0, 0 :+ 0) badRoots
	return ( ceq, Just sol )


module Math.IRT.Model.TwoPLM
  ( TwoPLM (..)
  ) where

import Statistics.Distribution

import Math.IRT.Internal.Distribution
import Math.IRT.Internal.LogLikelihood
import Math.IRT.Model.FourPLM ( FourPLM(..) )
import Math.IRT.Model.Generic

data TwoPLM = TwoPLM { discrimination :: !Double
                     , difficulty     :: !Double
                     } deriving (Show)

instance Distribution TwoPLM where
    cumulative = cumulative . toFourPLM

instance ContDistr TwoPLM where
    density    = density . toFourPLM
    quantile _ = error "This shouldn't be needed"

instance DensityDeriv TwoPLM where
    densityDeriv = densityDeriv . toFourPLM

instance GenericModel TwoPLM where
    fromRasch         b = TwoPLM 1.0 b
    fromOnePLM        b = TwoPLM 1.7 b
    fromTwoPLM          = TwoPLM
    fromThreePLM  a b _ = TwoPLM   a b
    fromFourPLM a b _ _ = TwoPLM   a b

instance LogLikelihood TwoPLM where
    logLikelihood b = logLikelihood b . toFourPLM

toFourPLM :: TwoPLM -> FourPLM
toFourPLM (TwoPLM sa sb) = FourPLM sa sb 0.0 1.0


module HLearn.Models.Distributions.Multivariate.MultiNormalFast
    ( MultiNormal
    , covar
    )
    where

import Control.DeepSeq
import qualified Data.Foldable as F
import qualified Data.Semigroup as SG

import Foreign.Storable
import Numeric.LinearAlgebra hiding ((<>))
import qualified Numeric.LinearAlgebra as LA

import HLearn.Algebra
import HLearn.Models.Distributions.Common

-------------------------------------------------------------------------------
-- data types

type MultiNormal dp = AddUnit1 MultiNormal' dp

data MultiNormal' dp = MultiNormal'
    { raw       :: !(RawMoments dp)
    , meanV     :: Vector (Ring dp)
    , covarM    :: Matrix (Ring dp)
    , invcovarM :: Matrix (Ring dp)
    , pdfdenom  :: Ring dp
    }

deriving instance (Element (Ring dp), Read (Ring dp)) => Read (MultiNormal' dp)
deriving instance (Element (Ring dp), Show (Ring dp)) => Show (MultiNormal' dp)

instance (NFData (Ring dp), Storable (Ring dp)) => NFData (MultiNormal' dp) where
    rnf m = deepseq (raw m)
          $ deepseq (meanV m)
          $ deepseq (covarM m)
          $ deepseq (invcovarM m)
          $ rnf (pdfdenom m)

---------------------------------------

data RawMoments dp = RawMoments
    { m0 :: !(Ring dp)
    , m1 :: !(Vector (Ring dp))
    , m2 :: !(Matrix (Ring dp))
    }

deriving instance (Element (Ring dp), Read (Ring dp)) => Read (RawMoments dp)
deriving instance (Element (Ring dp), Show (Ring dp)) => Show (RawMoments dp)

instance NFData (RawMoments dp) where
    rnf rm = seq (m0 rm)
           $ seq (m1 rm)
           $ seq (m2 rm) ()

---------------------------------------

covar :: MultiNormal dp -> Matrix (Ring dp)
covar (UnitLift1 m) = covarM m

mkMultiNormal' ::
    ( Storable (Ring dp)
    , Element (Ring dp)
    , Field (Ring dp)
    , Floating (Ring dp)
    ) => RawMoments dp -> MultiNormal' dp
mkMultiNormal' raw = MultiNormal'
    { raw = raw
    , meanV = meanV'
    , covarM = covarM'
    , invcovarM = inv covarM'
    , pdfdenom = 1/(sqrt $ det covarM' * (2*pi)^d)
    }
    where
        n = m0 raw
        d = dim $ m1 raw
        meanV' = mapVector (/n) (m1 raw)
        covarM' = mapMatrixWithIndex (\(i,j) mij -> (mij/n-(meanV' @> j)*(meanV' @> i))) (m2 raw)

-------------------------------------------------------------------------------
-- algebra

instance
    ( Container Matrix (Ring dp)
    , Container Vector (Ring dp)
    , Field (Ring dp)
    , Floating (Ring dp)
    ) => SG.Semigroup (MultiNormal' dp)
        where
    d1 <> d2 = mkMultiNormal' $ RawMoments
        { m0 = m0 (raw d1) + m0 (raw d2)
        , m1 = m1 (raw d1) `LA.add` m1 (raw d2)
        , m2 = m2 (raw d1) `LA.add` m2 (raw d2)
        }

-------------------------------------------------------------------------------
-- training

instance Num a => HasRing [a] where
    type Ring [a] = a

instance
    ( Element r
    , Container Vector r
    , Field r
    , Floating r
    , Ring (dp r) ~ r
    , F.Foldable dp
    ) => HomTrainer (MultiNormal (dp r))
        where
    type Datapoint (MultiNormal (dp r)) = dp r

    train1dp dp = UnitLift1 . mkMultiNormal' $ RawMoments
        { m0 = 1
        , m1 = fromList $ F.toList dp
        , m2 = n><n $ [i*j | i <- F.toList dp, j <- F.toList dp]
        }
        where
            n = length $ F.toList dp

-------------------------------------------------------------------------------
-- distribution

instance Probabilistic (MultiNormal dp) where
    type Probability (MultiNormal dp) = Ring dp

instance
    ( Element r
    , Container Vector r
    , Floating r
    , LA.Product r
    , Ring (dp r) ~ r
    , F.Foldable dp
    ) => PDF (MultiNormal (dp r))
        where
    pdf (UnitLift1 dist) dp = pdfdenom dist * (exp $ (-1/2) * top)
        where
            x = fromList $ F.toList dp
            top=minElement $ ((asRow  x) `sub` (asRow $ meanV dist))
                       LA.<> (invcovarM dist)
                       LA.<> ((asColumn x) `sub` (asColumn $ meanV dist))


-- Copyright 2021 Google LLC
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--      http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-# OPTIONS_GHC -Wall #-}
module LeastSquaresVoting
  ( vote
  ) where

import           Control.Arrow
import           Data.HashMap.Strict            ( HashMap )
import qualified Data.HashMap.Strict           as M
import           Data.Hashable                  ( Hashable(..) )
import           Data.List                      ( sortBy )
import           Data.Maybe                     ( fromMaybe )
import           Data.Monoid                    ( Sum(..) )
import           Data.Ord                       ( comparing )
import           Numeric.LinearAlgebra          ( (><)
                                                , Matrix
                                                , flatten
                                                , linearSolveSVD
                                                , toList
                                                )

import           Types

-- | Let _(u, v)_ be a pair of candidates in a voter's ordered list. This means
-- that the voter's intent is that _u_ is ideally ranked a unit _D_ (1V) higher
-- than _v_. When their potentials are _U_ and _V_, the difference from the intent
-- will be _(U-V)-1_.
--
-- The **heat** (corresponding to the heat dissipated by an electrical
-- resistor) is the square of the difference multiplied by a unit conductivity
-- _G_ (which represents the voter's ability to influence the outcome).
--
-- ((U-V)-D)^2 G.
--
-- Partial derivatives
--   d/dU ... =   2 G (U-V-D)
--   d/dV ... = - 2 G (U-V-D)
edge :: (Integral c, Num n) => c -> a -> a -> [(a, [(a, n)], n)]
edge count u v =
  let c = fromIntegral count
  in  [(u, [(u, c), (v, -c)], -c), (v, [(u, -c), (v, c)], c)]

newtype IndexedMonoid i m = IndexedMonoid (HashMap i m)
  deriving (Eq, Show, Read)

instance (Eq i, Hashable i, Semigroup m) => Semigroup (IndexedMonoid i m) where
  IndexedMonoid m1 <> IndexedMonoid m2 = IndexedMonoid $ M.unionWith (<>) m1 m2

instance (Eq i, Hashable i, Monoid m) => Monoid (IndexedMonoid i m) where
  mempty = IndexedMonoid M.empty

insert :: (Eq i, Hashable i, Semigroup m) => i -> m -> IndexedMonoid i m
insert i m = IndexedMonoid $ M.singleton i m

ilookup :: (Eq i, Hashable i, Monoid m) => i -> IndexedMonoid i m -> m
ilookup i (IndexedMonoid m) = fromMaybe mempty $ M.lookup i m

size :: (Hashable i) => IndexedMonoid i m -> Int
size (IndexedMonoid m) = M.size m

indices :: (Hashable i) => IndexedMonoid i m -> [i]
indices (IndexedMonoid m) = M.keys m

type LinearSystem a n = IndexedMonoid a (IndexedMonoid a (Sum n), (Sum n))

sumEdges :: (Hashable a, Eq a, Num n) => [(a, [(a, n)], n)] -> LinearSystem a n
sumEdges = foldMap
  (\(c, ls, r) -> insert c (foldMap (uncurry insert . second Sum) ls, Sum r))

equations
  :: (Hashable a, Eq a)
  => [a]
  -> LinearSystem a Double
  -> (Matrix Double, Matrix Double)
equations is m =
  ( (n >< n)
    [ getSum $ ilookup col (fst (ilookup row m)) | row <- is, col <- is ]
  , (n >< 1) [ getSum $ snd (ilookup row m) | row <- is ]
  )
  where n = size m

solve :: (Hashable a, Eq a) => LinearSystem a Double -> [(a, Double)]
solve m =
  sortBy (comparing snd)
    . zip is
    . toList
    . flatten
    . uncurry linearSolveSVD
    . equations is
    $ m
  where is = indices m

vote :: (Eq a, Hashable a, Integral n) => Prefs a n -> [(a, Double)]
vote (Prefs []) = []
vote (Prefs prefs@((_, cs) : _)) | null edges = zip cs (repeat 0)
                                 | otherwise  = solve . sumEdges $ edges
 where
  edges =
    [ e
    | (count, order) <- prefs
    , (u    , v    ) <- zip order (tail order)
    , e              <- edge count u v
    ]


{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE StrictData #-}

------------------------------------------------
-- |
-- Module    :  Numeric.Matrix.Integral.NormalForms
-- Copyright :  (c) Jun Yoshida 2019
-- License   :  BSD3
--
-- Hermite Normal form and Smith normal form.
-- The functions are implemented with C; see the source files in src/Internal/C.
--
------------------------------------------------

module Numeric.Matrix.Integral.NormalForms (
  hermiteNFST,
  hermiteNF,
  smithNFST,
  smithNF,
  smithRepST,
  smithRep
  ) where

import Control.Monad
import Control.Monad.ST
import Control.Monad.ST.Unsafe

import Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra.Devel

import System.IO.Unsafe (unsafePerformIO)
import Foreign.C.Types (CLong(..), CInt(..))
import Foreign.Ptr (Ptr)

-------------------------------------------------------------
-- Copied from the source of hmatrix package.
infixr 5 :>, ::>
type (:>)  t r = CInt -> Ptr t -> r
type (::>) t r =  CInt -> CInt -> CInt -> CInt -> Ptr t -> r
type Ok = IO CInt

--------------------------------------------------------------

-- | Hermite normal form
foreign import ccall unsafe "c_hermiteNF_LLL" cHermiteNF :: Z ::> Z ::> Z ::> Ok

hermiteNFST :: STMatrix s Z -> STMatrix s Z -> STMatrix s Z -> ST s ()
hermiteNFST stU stUi stH = do
  u <- unsafeFreezeMatrix stU
  ui <- unsafeFreezeMatrix stUi
  h <- unsafeFreezeMatrix stH
  unsafeIOToST (apply u (apply ui (apply h id)) cHermiteNF #| "cHermiteNF")

hermiteNFDecomp :: Matrix Z -> (Matrix Z, Matrix Z, Matrix Z)
hermiteNFDecomp mat = runST $ do
  stMat <- thawMatrix mat
  stU <- thawMatrix $ ident (rows mat)
  stUi <- thawMatrix $ ident (rows mat)
  hermiteNFST stU stUi stMat
  (\x y z -> (x,y,z)) <$!> unsafeFreezeMatrix stU <*> unsafeFreezeMatrix stUi <*> unsafeFreezeMatrix stMat

hermiteNF :: Matrix Z -> (Matrix Z, Matrix Z)
hermiteNF = p . hermiteNFDecomp
  where
    p (x,y,z) = (y,z)

hermiteDecomp :: Matrix Z -> (Matrix Z, Matrix Z)
hermiteDecomp = p . hermiteNFDecomp
  where
    p (x,y,z) = (x,z)

-- | Smith normal form
foreign import ccall unsafe "c_smithNF" cSmithNF :: Z ::> Z ::> Z ::> Ok

smithNFST :: STMatrix s LA.Z -> STMatrix s LA.Z ->  STMatrix s LA.Z -> ST s ()
smithNFST stU stM stV = do
  u <- unsafeFreezeMatrix stU
  m <- unsafeFreezeMatrix stM
  v <- unsafeFreezeMatrix stV
  unsafeIOToST (apply u (apply m (apply v id)) cSmithNF #| "cSmithNF")

smithNF :: Matrix Z -> (Matrix Z, Matrix Z, Matrix Z)
smithNF mat = runST $ do
  stMat <- thawMatrix mat
  stU <- thawMatrix $ ident (rows mat)
  stV <- thawMatrix $ ident (cols mat)
  smithNFST stU stMat stV
  (\x y z -> (x,y,z))
    <$> unsafeFreezeMatrix stU
    <*> unsafeFreezeMatrix stMat
    <*> unsafeFreezeMatrix stV

-- | Smith representation
foreign import ccall unsafe "c_smithRep" cSmithRep :: Z ::> Z ::> Z::> Ok

smithRepST :: STMatrix s LA.Z -> STMatrix s LA.Z -> STMatrix s LA.Z -> ST s ()
smithRepST stA stM stB = do
  a <- unsafeFreezeMatrix stA
  m <- unsafeFreezeMatrix stM
  b <- unsafeFreezeMatrix stB
  unsafeIOToST (apply a (apply m (apply b id)) cSmithRep #| "cSmithRep")

smithRep :: Matrix Z -> (Matrix Z, Matrix Z, Matrix Z)
smithRep mat = runST $ do
  stMat <- thawMatrix mat
  stU <- thawMatrix $ ident (rows mat)
  stV <- thawMatrix $ ident (cols mat)
  smithRepST stU stMat stV
  (\x y z -> (x,y,z))
    <$> unsafeFreezeMatrix stU
    <*> unsafeFreezeMatrix stMat
    <*> unsafeFreezeMatrix stV


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE InstanceSigs #-}

module BraKet where

import Data.Complex
import qualified Numeric.LinearAlgebra as L

import LinAlg

-- Ket is a ket vector
-- is the coefficient type
-- in which the Ket is written
newtype Ket c = Ket (L.Vector c)
      deriving (Show, Eq)

-- Function used to act on a Ket with an operator
act :: MatOp -> Ket (Complex Double) -> Ket (Complex Double)
act op ket = Ket $ repr op L.#> getCoeffs ket

getCoeffs :: Ket c -> L.Vector c
getCoeffs (Ket cs) = cs

mapKet :: (L.Vector a -> L.Vector a) -> Ket a -> Ket a
mapKet f (Ket cs) = Ket $ f cs

instance Vector (Ket (Complex Double)) where
    -- Negation: just negate all coefficients
    vneg = mapKet negate
    -- Addition: add coefficients element-wise
    (Ket exp1) <+> (Ket exp2) = Ket $ exp1 + exp2
    -- Scalar multiplication: multiply all coefficients
    (<**>) a = mapKet (L.cmap (*a))

instance Hilbert (Ket (Complex Double)) where
    -- Inner product: Like l2 inner product; multiply coefficients
    -- and sum
    (Ket coeffs1) <.> (Ket coeffs2) = coeffs1 L.<.> coeffs2

outerProduct :: Ket (Complex Double) -> Ket (Complex Double)
                                     -> MatOp
outerProduct (Ket cs1) (Ket cs2) = Operator $ cs1 `L.outer` cs2


module Mnist where

import Data.List (nub, sort, intercalate)
import Codec.Compression.GZip (decompress)
import qualified Data.ByteString.Lazy as BS
import Numeric.LinearAlgebra (vector, Vector, R, Z, reshape, Matrix, toList, toRows, fromRows)

-- | typeclass that represents a dataset
class DataSet a where
  groupByLabel :: a -> [Matrix R]
  distinctLabels :: a -> [R]
  dataLabel :: a -> [(Vector R, R)]  -- pair the data and label for each sample
  -- oneHotLabel :: a -> Matrix Z

img_header_size = 16
label_header_size = 8
width = 28
img_size = 784

-- labeled data consists of data and their label
data LabeledData = LabeledData
    { dat :: !(Matrix R)  -- this has num_data >< data_dimension
    , label :: !(Vector R)
    }

-- define LabeledData as instance of DataSet typeclass
instance DataSet LabeledData where
    distinctLabels = sort . nub . toList . label
    groupByLabel ld = [
        -- extract index that matches the label, and cherrypicks them from data rows
        fromRows [dataRows !! idx | idx <- extractIdx lab]
        | lab <- distinctLabels ld ]
      where
        dataRows = toRows $ dat ld  -- represent matrix as list of rows
        labels = toList $ label ld  -- represent vector to list of values
        labelIdx = zip labels [0..]  -- zip it with index
        extractIdx targetLabel = map snd $ filter ((== targetLabel) . fst) labelIdx
    dataLabel ld = zip (toRows $ dat ld) (toList $ label ld)

-- mnist dataset consists of train set and a test set
data Mnist = Mnist
    { trainData :: LabeledData
    , testData :: LabeledData
    }

-- render a single number
render :: Vector R -> String
render s = intercalate "\n" $ splitted $ map (intensity . floor) (toList s)  -- insert \n every 28 chars
  where
    chars = " ..:oO0@"
    intensity n = chars !! (n * length chars `div` 256)
    splitted :: [a] -> [[a]]
    splitted [] = []
    splitted v = (take width v) : (splitted $ drop width v)  -- split by 28

-- convert byteString to vector -- in this case, bytestring is read as numbers byte by byte
byteStringToVector :: BS.ByteString -> Vector R
byteStringToVector = vector . map fromIntegral . BS.unpack

readDataFile :: String -> IO BS.ByteString
readDataFile filepath = decompress <$> BS.readFile filepath

-- make label data represented as Vector
makeLabelData :: BS.ByteString -> Vector R
makeLabelData = byteStringToVector . BS.drop 8  -- 8-byte header

-- make image data represented as Matrix
makeImgData :: BS.ByteString -> Matrix R
makeImgData = reshape img_size . byteStringToVector . BS.drop 16  -- 16-byte header

-- reads the file and creates a labeled data
readData :: String -> String -> IO LabeledData
readData imgfile labelfile = do
    imgData <- makeImgData <$> decompress <$> BS.readFile imgfile
    labelData <- makeLabelData <$> decompress <$> BS.readFile labelfile
    return $ LabeledData imgData labelData

readMnist :: IO Mnist
readMnist = do
    trainDataset <- readData "train-images-idx3-ubyte.gz" "train-labels-idx1-ubyte.gz"
    testDataset <- readData "t10k-images-idx3-ubyte.gz" "t10k-labels-idx1-ubyte.gz"
    return $ Mnist trainDataset testDataset

-- ByteString is an optimized representation of Word8.
-- BS.readFile :: FilePath -> IO ByteString
-- decompress :: ByteString -> ByteString
readMnistAndShow :: IO ()
readMnistAndShow = do
  imgData <- makeImgData <$> decompress <$> BS.readFile "train-images-idx3-ubyte.gz"
  labelData <- makeLabelData <$> decompress <$> BS.readFile "train-labels-idx1-ubyte.gz"
  putStrLn $ show $ labelData


-- | The Metropolis transition.

module Strategy.Metropolis (metropolis) where

import Control.Monad.State.Strict
import Data.Maybe
import Numeric.LinearAlgebra
import Math.Probably.Sampler
import Math.Probably.Types
import Math.Probably.Utils

-- | The Metropolis transition for a symmetric proposal.  Uses a simple
--   spherical Gaussian with standard deviation specified by the first
--   argument.
metropolis :: Maybe Double -> Transition Double
metropolis e = do
  Chain current@(ds, cs) target _ t <- get
  let sd = fromMaybe t e
  pcs <- lift $ perturb cs sd
  zc  <- lift unit
  let next = nextState target current (ds, pcs) sd zc
  put $ Chain next target (logObjective target next) sd
  return next

perturb :: ContinuousParams -> Double -> Prob (ContinuousParams)
perturb q sd = fmap fromList $ mapM (`normal` sd) (toList q)

acceptRatio :: Target -> Parameters -> Parameters -> Double -> Double
acceptRatio target c@(_, current) p@(_, proposed) sd = exp . min 0 $
    logObjective target p + log (sphereGauss current proposed sd)
  - logObjective target c - log (sphereGauss proposed current sd)

nextState
  :: Target
  -> Parameters
  -> Parameters
  -> Double
  -> Double
  -> Parameters
nextState target current proposed sd z
    | z < acceptProb = proposed
    | otherwise      = current
  where
    ratio = acceptRatio target current proposed sd
    acceptProb | isNaN ratio = 0
               | otherwise   = ratio


module Malliavin where

import Data.Random.Distribution.Normal
import Numeric.LinearAlgebra.HMatrix as H
import Data.Ratio

data L2Map a = L2Map {eval :: a -> Double,
                      norm_l2 :: a -> Double}
l2_1 = L2Map {eval = (\t -> 1), norm_l2 = id}
l2id = L2Map {eval = id, norm_l2 = (\x -> x ^ 3 / 3)}

data L1Map a = L1Map {norm_l1 :: a -> Double}
l1_1 = L1Map {eval = (\t -> 1), norm_l1 = id}
l1id = L1Map {eval = id, norm_l1 = (\x -> x ^ 2 / 2)}

-- | 'hermite' is an infinite list of Hermite polynomials for given x
hermite :: (Enum a, Num a) => a -> [a]
hermite x = s
  where s@(_:ts) = 1 : x : zipWith3 (\hn2 hn1 n1 -> x * hn1 - n1 * hn2) s ts [1..]

-- | 'gaussianProcess' samples from Gaussian process
gaussianProcess :: Seed                   -- random state
                -> Int                    -- number of samples m
                -> Int                    -- number of dimensions n
                -> ((Int, Int) -> Double) -- function that maps indices of the matrix into dot products of its elements
                -> [Vector Double]        -- m n-th dimensional samples of Gaussian process
gaussianProcess seed 0 _ _ = []
gaussianProcess seed m 0 _ = replicate m $ vector []
gaussianProcess seed m n dotProducts = toRows $ gaussianSample seed m mean cov_matrix
  where mean = n |> repeat 0
        cov_matrix = H.sym $ (n><n) $ map (\i -> dotProducts (quot i n, rem i n)) [0..]

-- | 'second' function takes a list and gives each second element of it
second (x:y:xs) = y : second xs
second _ = []

-- | 'coinTossExpansion' is a Wiener chaos expansion for coin-toss r.v. to n-th element
coinTossExpansion :: Int    -- number of elements in the sum
                  -> Double -- gaussian random variable
                  -> Double -- the sum
coinTossExpansion n xi = sum (take n $ 0.5 : zipWith (*) fn (second $ hermite xi))
  where fn = 1.0 / (sqrt $ 2 * pi) : zipWith (\fn1 k -> -fn1 * k / (k ^ 2 + 3 * k + 2)) fn [1, 3..]

-- | 'coinTossSequence' is a coin-toss sequence of given size
coinTossSequence :: Seed  -- random state ω
                 -> Int   -- size of resulting sequence m
                 -> [Int] -- coin-toss sequence of size m
coinTossSequence seed m = map (round.coinTossExpansion 100) (toList nvec)
  where nvec = gaussianProcess seed m 1 (\(i,j) -> 1) !! 0

-- | 'mapOverInterval' map function f over the interval [0, T]
mapOverInterval :: (Fractional a) => Int      -- n, size of the output list
                                  -> a        -- T, end of the time interval
                                  -> (a -> b) -- f, function that maps from fractional numbers to some abstract space
                                  -> [b]      -- list of values f(t), t \in [0, T]
mapOverInterval n endT fn = [fn $ (endT * fromIntegral i) / fromIntegral (n - 1) | i <- [0..(n-1)]]

type ItoIntegral = Seed         -- ω, random state
                -> Int          -- n, sample size
                -> Double       -- T, end of the time interval
                -> L2Map Double -- h, L2-function
                -> [Double]     -- list of values sampled from Ito integral

-- | 'itoIntegral'' trajectory of Ito integral on the time interval [0, T]
itoIntegral' :: ItoIntegral
itoIntegral' seed n endT h = 0 : (toList $ gp !! 0)
  where gp = gaussianProcess seed 1 (n-1) (\(i, j) -> norm_l2 h $ fromIntegral (min i j + 1) * t)
        t = endT / fromIntegral n

-- | 'itoIntegral' faster implementation of itoIntegral' function
itoIntegral :: ItoIntegral
itoIntegral seed 0 _    _ = []
itoIntegral seed n endT h = scanl (+) 0 increments
  where increments = toList $ sigmas * gaussianVector
        gaussianVector = flatten $ gaussianSample seed (n-1) (vector [0]) (H.sym $ matrix 1 [1])
        sigmas = fromList $ zipWith (\x y -> sqrt(x-y)) (tail hnorms) hnorms
        hnorms = mapOverInterval n endT $ norm_l2 h

-- | 'brownianMotion' trajectory of Brownian motion a.k.a. Wiener process on the time interval [0, T]
brownianMotion :: Seed     -- ω, random state
               -> Int      -- n, sample size
               -> Double   -- T, end of the time interval
               -> [Double] -- list of values sampled from Brownian motion
brownianMotion seed n endT = itoIntegral seed n endT l2_1

type DiffusionProcess = Seed         -- ω, random state
                     -> Int          -- n, sample size
                     -> Double       -- T, end of the time interval
                     -> L1Map Double -- function \mu with L1 norm
                     -> L2Map Double -- function \sigma with L2 norm
                     -> [Double]     -- list of values sampled from stochastic diffusion process X(t)

-- | 'diffusionProcess' stochastic process, defined by SDE: dX(t) = \mu(t)dt + \sigma(t)dB(t)
diffusionProcess :: DiffusionProcess
diffusionProcess seed n endT mu sigma = zipWith (+) trend noise
  where trend = mapOverInterval n endT $ norm_l1 mu
        noise = itoIntegral seed n endT sigma

-- | 'multipleItoIntegral' multiple m-th Ito integral of m-th tensor power of h
multipleItoIntegral :: Seed         -- ω, random state
                    -> Int          -- m, tensor power
                    -> Int          -- n, sample size
                    -> Double       -- T, end of the time interval
                    -> L2Map Double -- function h with L2 norm
                    -> [Double]     -- list of values sampled from multiple m-th Wiener-Ito integral of h^(⊗m)
multipleItoIntegral seed m n endT h = zipWith fn itoTrajectory hnorm
  where itoTrajectory = itoIntegral seed n endT h
        hnorm = mapOverInterval n endT $ norm_l2 h
        fn :: Double -> Double -> Double
        fn _ 0 = 0
        fn x y = (y ^ m) * (hermite (x / y) !! m)

multipleBrownianMotion :: Seed     -- ω, random state
               -> Int      -- m, tensor power
               -> Int      -- n, sample size
               -> Double   -- T, end of the time interval
               -> [Double] -- list of values sampled from Brownian motion
multipleBrownianMotion seed m n endT = multipleItoIntegral seed m n endT l2_1

-- | 'totalVariation' calculates total variation of the sample process X(t)
totalVariation :: [Double] -- process trajectory X(t)
               -> Double   -- V(X(t))
totalVariation x@(_:tx) = sum $ zipWith (\x xm1 -> abs $ x - xm1) tx x

-- | 'quadraticVariation' calculates quadratic variation of the sample process X(t)
quadraticVariation :: [Double] -- process trajectory X(t)
                   -> Double   -- [X(t)]
quadraticVariation x@(_:tx) = sum $ zipWith (\x xm1 -> (x - xm1) ^ 2) tx x


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards       #-}

module Q.Options.BlackScholes (
    BlackScholes(..)
  , atmf
  , euOption
  , eucall
  , euput
  , module Q.Options
) where

import           Control.Monad.State
import           Data.Random                    hiding (Gamma)
import           Data.Time
import           Numeric.RootFinding
import           Q.ContingentClaim.Options
import           Q.MonteCarlo
import           Q.Options
import           Q.Stochastic.Discretize
import           Q.Stochastic.Process
import           Q.Time
import           Q.Types
import           Statistics.Distribution        (cumulative, density)
import           Statistics.Distribution.Normal (standard)
import qualified Q.Options.Black76 as B76

dcf = dcYearFraction ThirtyUSA

-- | Parameters for a simplified black scholes equation.
data BlackScholes = BlackScholes {
    bsSpot :: Spot -- ^ The asset's spot on the valuation date.
  , bsRate :: Rate   -- ^ Risk free rate.
  , bsVol  :: Vol    -- ^ Volatility.
} deriving Show

instance Model BlackScholes Double where
  discountFactor BlackScholes{..} t1 t2 = return $ exp (scale dt bsRate)
    where dt = t2 - t1

  evolve (BlackScholes spot (Rate r) (Vol sigma)) (YearFrac t) = do
    (YearFrac t0, s0) <- get
    let dt = t - t0
    dw <- (lift stdNormal)::StateT (YearFrac, Double) RVar Double
    let st = s0 * exp ((r - 0.5 * sigma * sigma) * dt + sigma * dw * sqrt dt)
    put (YearFrac t, st)
    return st

atmf :: BlackScholes -> YearFrac -> Strike
atmf BlackScholes{..} t = Strike $ s / d where
  (Rate d) = exp (scale t (-bsRate))
  (Spot s) = bsSpot

-- | European option valuation with black scholes.
euOption ::  BlackScholes ->  YearFrac -> OptionType -> Strike ->Valuation
euOption bs@BlackScholes{..} t cp k =
  let b76 = B76.Black76 {
          b76F  = forward bs t
        , b76DF = Q.Types.discountFactor t bsRate
        , b76T  = t
        , b76Vol = bsVol
        }
  in B76.euOption b76 cp k

-- | see 'euOption'
euput bs t = euOption bs t Put

-- | see 'euOption'
eucall bs t = euOption bs t Call

forward BlackScholes{..} (YearFrac t) = Forward $ s * exp (r * t)
  where (Spot s) = bsSpot
        (Rate r) = bsRate

corradoMillerIniitalGuess bs@BlackScholes{..} cp (Strike k) (YearFrac t) (Premium premium) =
  (recip $ sqrt t) * ((sqrt (2 * pi)/ (s + discountedStrike)) + (premium - (s - discountedStrike)/2) + sqrt ((premium - (s - discountedStrike)/2)**2 - ((s - discountedStrike)**2/pi))) where
    discountedStrike = k * (exp $ (-r) * t)
    (Rate r) = bsRate
    (Spot s) = bsSpot


-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Pass.Util.Beta
-- Copyright   :  (C) 2012-2013 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable (GADTs, Rank2Types)
--
----------------------------------------------------------------------------
module Data.Pass.Util.Beta
  ( Beta(..)
  , cumulative
  , density
  , quantile
  ) where

import Numeric.SpecFunctions
import Numeric.MathFunctions.Constants (m_NaN)

data Beta = Beta {-# UNPACK #-} !Double {-# UNPACK #-} !Double
  deriving (Eq,Read,Show)

cumulative :: Beta -> Double -> Double
cumulative (Beta a b) x
  | x <= 0 = 0
  | otherwise = incompleteBeta a b x
{-# INLINE cumulative #-}

-- invert a monotone function
invertMono :: (Double -> Double) -> Double -> Double -> Double -> Double
invertMono f l0 h0 b = go l0 h0 where
  go l h
    | h - l < epsilon = m
    | otherwise = case compare (f m) b of
      LT -> go m h
      EQ -> m
      GT -> go l m
    where m = l + (h-l)/2
          epsilon = 1e-12
{-# INLINE invertMono #-}

density :: Beta -> Double -> Double
density (Beta a b) x
  | a <= 0 || b <= 0 = m_NaN
  | x <= 0 = 0
  | x >= 1 = 0
  | otherwise = exp $ (a-1)*log x + (b-1)*log (1-x) - logBeta a b
{-# INLINE density #-}

quantile :: Beta -> Double -> Double
quantile d p
  | p < 0 = error $ "probability must be positive. Got " ++ show p
  | p > 1 = error $ "probability must be less than 1. Got " ++ show p
  | otherwise = invertMono (cumulative d) 0 1 p
{-# INLINE quantile #-}


{-# LANGUAGE DataKinds                                #-}

module MNIST where
import           Control.DeepSeq
import           Control.Exception.Base
import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Maybe
import           Control.Monad.Trans.State
import           Data.Bitraversable
import           Data.IDX
import           Data.Traversable
import           Data.Tuple
import qualified Data.Vector                         as V
import qualified Data.Vector.Generic                 as VG
import qualified Data.Vector.Unboxed                 as VU
import qualified Numeric.LinearAlgebra               as HM  -- non-backprop hmatrix
import qualified Numeric.LinearAlgebra.Data          as HMD
import qualified Numeric.LinearAlgebra.Static        as HMS -- hmatrix with type-checked operations :D

loadMNIST
   :: FilePath
   -> FilePath
   -> IO (Maybe [(HMS.L 28 28, HMS.R 10)])
loadMNIST fpI fpL = runMaybeT $ do
   i <- MaybeT          $ decodeIDXFile       fpI
   l <- MaybeT          $ decodeIDXLabelsFile fpL
   d <- MaybeT . return $ labeledIntData l i
   r <- MaybeT . return $ for d (bitraverse mkImage mkLabel . swap)
   liftIO . evaluate $ force r
      where
         mkImage :: VU.Vector Int -> Maybe (HMS.L 28 28)
         mkImage = HMS.create . HMD.reshape 28 . VG.convert . VG.map (\i -> fromIntegral i / 255)
         mkLabel :: Int -> Maybe (HMS.R 10)
         mkLabel n = HMS.create $ HM.build 10 (\i -> if round i == n then 1 else 0)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE AllowAmbiguousTypes #-}

{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}

{-|
Module      : Numeric.LinearAlgebra.Static.RList
Description : List of static vectors (`R`s)
Copyright   : (c) Ryan Orendorff, 2020
License     : BSD3
Stability   : experimental
-}
module Numeric.LinearAlgebra.Static.RList
    ( RList(..)
    , LList(..)
    , Sum
    , exRList
    , concat
    , blockDiag
    , mapRListSquare
    , applyLList
    )
where

import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static
import Prelude hiding (concat)

-- TODO: Make a fold function, then this can be abstracted.
-- | Sum a type level list of naturals
type family Sum (xs :: [Nat]) where
    Sum '[] = 0
    Sum (n ': ns) = n + Sum ns

-- | A non-empty list of `R` vectors of different sizes.
-- The list is non-empty since we cannot create a vector of size 0 in a
-- reasonable way when using Data.Vector as the backend (as HMatrix).
data RList (ns :: [Nat]) where
    ROne :: (KnownNat n) => R n -> RList '[n]
    (:::) :: (KnownNat n) => R n -> RList ns -> RList (n ': ns)

infixr 6 :::

-- | An example of how to construct an RList.
exRList :: RList '[2, 3]
exRList = konst 1 ::: ROne (konst 2)

-- There is no null vector. Also the HMatrix library does the wrong thing
-- with this example.
--
-- λ> (konst 4 :: R 2) # (konst 0 :: R 0)
-- (vector [4.0,4.0,0.0] :: R 2)

-- | Concatenate an RList of vectors
--
-- Thank you Christiaan Baaij for figuring out how to do concat! This is
-- basically his solution to the problem, with a modification to work with
-- the non-empty version of RList. See the following gist for more details:
-- https://gist.github.com/ryanorendorff/d05c378b71829e3c0c33de462cb9a973
concat :: RList ns -> R (Sum ns)
concat rs = case concatWorker rs of
  RWorker rsW -> rsW

data RWorker n where
  RWorker :: KnownNat n => R n -> RWorker n

concatWorker :: RList ns -> RWorker (Sum ns)
concatWorker (ROne r) = RWorker r
concatWorker (r ::: rs) = case concatWorker rs of
  RWorker rsW -> RWorker (r # rsW)

-- | A non-empty list of `L` matrices of different sizes.
-- The list is non-empty since we cannot create a matrix of size 0 in a
-- reasonable way when using Data.Vector as the backend (as HMatrix).
data LList (ms :: [Nat]) (ns :: [Nat]) where
    LOne :: (KnownNat m, KnownNat n) => L m n -> LList '[m] '[n]
    (:-:) :: (KnownNat m, KnownNat n) => L m n -> LList ms ns
                                      -> LList (m ': ms) (n ': ns)

-- | Concatenate a LList of matrices as a block diagonal.
blockDiag :: LList ms ns -> L (Sum ms) (Sum ns)
blockDiag ls = case blockDiagWorker ls of
  LWorker lsW -> lsW

data LWorker m n where
  LWorker :: (KnownNat m, KnownNat n) => L m n -> LWorker m n

blockDiagWorker :: LList ms ns -> LWorker (Sum ms) (Sum ns)
blockDiagWorker (LOne l) = LWorker l
blockDiagWorker (l :-: ls) = case blockDiagWorker ls of
  LWorker lsW -> LWorker (l ||| konst 0 === konst 0 ||| lsW)

-- | Map some function that preserves the dimension of the vector (a square
-- matrix) over a list of vectors.
mapRListSquare :: (forall n. KnownNat n => R n -> R n)
               -> RList ns -> RList ns
mapRListSquare f (ROne r) = ROne (f r)
mapRListSquare f (r ::: rs) = f r ::: mapRListSquare f rs

-- I have not proved that the LList and RList are the same length, so GHC
-- complains that I have non-exhaustive patterns that cannot be constructed.
-- TODO: Add comment on how this would work in Agda to show how this type
-- of error would be avoided.
-- | Apply a list of static matrices over a list of static vectors
applyLList :: LList ms ns -> RList ns -> RList ms
applyLList (LOne _) (_ ::: _) = error "This case is impossible"
applyLList (_ :-: _) (ROne _) = error "This case is impossible"
applyLList (LOne l) (ROne r) = ROne (l #> r)
applyLList (l :-: ls) (r ::: rs) = (l #> r) ::: applyLList ls rs

-- This function is more complicated, may need to reflect out the position
-- we are in the list in order to write this function.
-- splitRList :: R (Sum ns) -> RList ns
-- splitRList vec = let (r, rs) = split vec in r ::: (splitRList rs)


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedLists #-}

module Main where

import qualified Data.Vector.Unboxed         as U
import           Data.Vector.Unboxed         (Vector)
import qualified Data.Vector as V

-- DLA
import qualified Statistics.Matrix as M
import qualified Statistics.Matrix.Fast as MF
import qualified Statistics.Matrix.Fast.Algorithms as A

-- hmatrix
import qualified Numeric.LinearAlgebra as H

-- numhask
import qualified NumHask.Array as NH
import qualified NumHask.Prelude as NP

-- data.matrix from matrix
import qualified Data.Matrix as DMX

-- massiv
import qualified Data.Massiv.Array as MA

import qualified System.Random.MWC as Mwc

import qualified Weigh as W

n :: Int
n = 100

vectorGen :: IO (Vector Double)
vectorGen =  do
    gen <- Mwc.create
    Mwc.uniformVector gen (n*n)

matrixDLA :: IO M.Matrix
matrixDLA = do
    vec <- vectorGen
    return $ M.Matrix n n vec

matrixH :: IO (H.Matrix H.R)
matrixH = do
    vec <- vectorGen
    return $ (n H.>< n) $ U.toList $ vec

identH :: Int -> H.Matrix Double
identH = H.ident

identDMX :: Int -> DMX.Matrix Double
identDMX = DMX.identity

elemZero :: Double -> Double
elemZero = const 0

elemSqr :: Double -> Double
elemSqr x = x*x

mapH :: (Double -> Double) -> H.Matrix Double -> H.Matrix Double
mapH = H.cmap

main :: IO ()
main = do

    let

    vDLA <- vectorGen
    uDLA <- vectorGen

    let

    --
      subDLA = U.take n vDLA
      aDLA = M.Matrix n n vDLA
      bDLA = M.Matrix n n uDLA

    --
      vList = U.toList vDLA
      uList = U.toList uDLA

    --
      aH = (n H.>< n) vList
      bH = (n H.>< n) uList

      subH = H.fromList . take n $ vList
      vH = H.fromList vList

    --
      aNH = NP.fromList vList :: NH.Array V.Vector '[50, 50] Double
      bNH = NP.fromList uList :: NH.Array V.Vector '[50, 50] Double

      subNH = NP.fromList . take n $ vList :: NH.Array V.Vector '[50] Double
      vNH = NP.fromList vList :: NH.Array V.Vector '[2500] Double

    --
      aDMX = DMX.fromList n n vList
      bDMX = DMX.fromList n n uList

    --
      vMA = MA.fromList MA.Seq vList :: MA.Array MA.P MA.Ix1 Double
      aMA = MA.resize' (MA.Sz (n MA.:. n)) vMA :: MA.Array MA.P MA.Ix2 Double
      bMA = MA.resize' (MA.Sz (n MA.:. n)) $ MA.fromList MA.Seq uList :: MA.Array MA.P MA.Ix2 Double

    W.mainWith (do
               W.func "DLA - multiplication" (MF.multiply aDLA) bDLA
               W.func "DLA - qr factorization" A.qr aDLA
               W.func "DLA - transpose" M.transpose aDLA
               W.func "DLA - norm" MF.norm vDLA
               W.func "DLA - row" (M.row  aDLA) 0
               W.func "DLA - column" (M.column  aDLA) 0
               W.func "DLA - identity" M.ident n

               W.func "Hmatrix - multiplication" ((<>) aH) bH
               W.func "Hmatrix - qr factorization" H.qr aH
               W.func "Hmatrix - transpose" H.tr aH
               W.func "Hmatrix - norm" H.norm_2 vH
               W.func "Hmatrix - row" ((H.?) aH) [0]
               W.func "Hmatrix - column" ((H.¿) aH) [0]
               W.func "Hmatrix - identity" identH n

               W.func "NumHask Array - multiplication" (NH.mmult aNH) bNH
               W.func "NumHask Array - transpose" NH.transpose aNH
               W.func "NumHask Array - norm" (sqrt . (NP.<.> vNH)) vNH
               W.func "NumHask Array - row" (NH.row (NP.Proxy :: NP.Proxy 0)) aNH
               W.func "NumHask Array - column" (NH.col (NP.Proxy :: NP.Proxy 0)) aNH

               W.func "Massiv - norm" (sqrt . MA.foldlS (+) 0 . MA.zipWith (*) vMA) vMA
               W.func "Massiv - multiplication" ((MA.|*|) aMA) bMA
               W.func "Massiv - transpose" (MA.computeAs MA.P . MA.transpose) aMA
               W.func "Massiv - row" (MA.computeAs MA.P . (MA.!>) aMA) 0

               W.func "matrix - multiplication" (DMX.multStrassenMixed aDMX) bDMX
               W.func "matrix - transpose" DMX.transpose aDMX
               W.func "matrix - row" (DMX.getRow 1) aDMX
               W.func "matrix - column" (DMX.getCol 1) aDMX
               W.func "matrix - identity" identDMX n
               )


import Data.Complex

-- Polynomial type
type Poly a = [Complex a]

__fftAux :: (RealFloat a) => a -> Poly a -> Poly a
__fftAux _ [] = []
__fftAux _ [x] = [x]
__fftAux c poly = uncurry (++) res
  where
    n = length poly
    ls = __fftAux c [ x | (x, i) <- zip poly [1..],  odd i ]
    rs = __fftAux c [ x | (x, i) <- zip poly [1..], even i ]
    res = combine ls rs 1
      where
        ang = c * 2 * pi / fromIntegral n
        wn = cos ang :+ sin ang
        combine [] [] _ = ([], [])
        combine (x:xs) (y:ys) w = (x + w * y : rxs, x - w * y : rys)
          where
            (rxs, rys) = combine xs ys (w * wn)

__closestPowerOf2 :: Int -> Int
__closestPowerOf2 n = compute 1
  where
    compute x
      | x >= n = x
      | otherwise = compute $ 2 * x

__fftWrapper :: (RealFloat a) => a -> Poly a -> Poly a
__fftWrapper c poly
  | oddf (length poly) == 1 = __fftAux c poly
  | otherwise = error "FFT: size is not a power of 2!"
    where
      oddf x
        | odd x = x
        | otherwise = oddf $ x `div` 2

fft :: (RealFloat a) => Poly a -> Poly a
fft = __fftWrapper 1

ifft :: (RealFloat a) => Poly a -> Poly a
ifft poly = map (/n) $ __fftWrapper (-1) poly
  where
    n = fromIntegral $ length poly

multiply :: (RealFloat a) => Poly a -> Poly a -> Poly a
multiply ps qs = ifft $ zipWith (*) (fft $ pad ps) (fft $ pad qs)
  where
    size = __closestPowerOf2 $ length ps + length qs - 1
    pad xs = xs ++ replicate (size - length xs) 0

multiplyReal :: (RealFloat a) => [a] -> [a] -> [a]
multiplyReal ps qs = map realPart $ multiply (conv ps) (conv qs)
  where
    conv = map (:+ 0)

multiplyIntegral :: (Integral a) => [a] -> [a] -> [a]
multiplyIntegral ps qs = map round $ multiplyReal (conv ps) (conv qs)
  where
    conv = map fromIntegral


module Reanimate.Morph.Triagulate where

-- import Graphics.SvgTree
import Linear.V2
import Linear.Vector
-- import Reanimate
import Numeric.LinearAlgebra
-- import Numeric.LinearAlgebra.Sparse
-- import Data.Sparse.SpMatrix
-- import Data.Sparse.SpVector

type Poly = [V2 Double]

-- renderPolygon :: Poly -> SVG
-- renderPolygon = undefined

triangulate :: Poly -> Poly -> ([Poly],[Poly])
triangulate = undefined

type Point = V2 Double

p1,p2,p3,p4 :: Point
p1 = V2 0 0
p2 = V2 2 0
p3 = V2 1 2
p4 = V2 1 1

-- let (t1,t2,t3) = barycentricCoordinates p1 p2 p3 p4
-- p4 = t1*p1 + t2*p2 + t3*p3
--
--
--                       [ t1
--                       , t2
--                       , t3 ]
-- [ [p1_x, p2_x, p3_x ] [ p4_x
-- , [p1_y, p2_y, p3_y ] , p4_y
-- , [1, 1, 1] ]         , 1 ]
-- barycentricCoordinates :: Point -> Point -> Point -> Point -> (Double, Double, Double)
barycentricCoordinates (V2 p1x p1y) (V2 p2x p2y) (V2 p3x p3y) (V2 p4x p4y) =
    linearSolve m v
    -- (m, v)
  where
    m = (3><3) [p1x, p2x, p3x, p1y, p2y, p3y, 1, 1, 1 :: Double]
    v = (3><1) [ p4x, p4y, 1 :: Double ]

bCoords (V2 x1 y1) (V2 x2 y2) (V2 x3 y3) (V2 x y) =
    (lam1, lam2, lam3)
  where
    lam1 = ((y2-y3)*(x-x3) + (x3 - x2)*(y-y3)) /
           ((y2-y3)*(x1-x3) + (x3-x2)*(y1-y3))
    lam2 = ((y3-y1)*(x-x3) + (x1-x3)*(y-y3)) /
           ((y2-y3)*(x1-x3) + (x3-x2)*(y1-y3))
    lam3 = 1 - lam1 - lam2


{-# LANGUAGE TypeOperators #-}

--------------------------------------------------------------------------------
{- |
Module      :  Numeric.LinearAlgebra.Data
Copyright   :  (c) Alberto Ruiz 2015
License     :  BSD3
Maintainer  :  Alberto Ruiz
Stability   :  provisional

This module provides functions for creation and manipulation of vectors and matrices, IO, and other utilities.

-}
--------------------------------------------------------------------------------

module Numeric.LinearAlgebra.Data(

    -- * Elements
    R,C,I,Z,type(./.),

    -- * Vector
    {- | 1D arrays are storable vectors directly reexported from the vector package.
    -}

    fromList, toList, (|>), vector, range, idxs,

    -- * Matrix

    {- | The main data type of hmatrix is a 2D dense array defined on top of
         a storable vector. The internal representation is suitable for direct
         interface with standard numeric libraries.
    -}

    (><), matrix, tr, tr',

    -- * Dimensions

    size, rows, cols,

    -- * Conversion from\/to lists

    fromLists, toLists,
    row, col,

    -- * Conversions vector\/matrix

    flatten, reshape, asRow, asColumn,
    fromRows, toRows, fromColumns, toColumns,

    -- * Indexing

    atIndex,
    Indexable(..),

    -- * Construction
    scalar, Konst(..), Build(..), assoc, accum, linspace,  -- ones, zeros,

    -- * Diagonal
    ident, diag, diagl, diagRect, takeDiag,

    -- * Vector extraction
    subVector, takesV, vjoin,

    -- * Matrix extraction
    Extractor(..), (??),

    (?), (¿), fliprl, flipud,

    subMatrix, takeRows, dropRows, takeColumns, dropColumns,

    remap,

    -- * Block matrix
    fromBlocks, (|||), (===), diagBlock, repmat, toBlocks, toBlocksEvery,

    -- * Mapping functions
    conj, cmap, cmod,

    step, cond,

    -- * Find elements
    find, maxIndex, minIndex, maxElement, minElement,
    sortVector, sortIndex,

    -- * Sparse
    AssocMatrix, toDense,
    mkSparse, mkDiagR, mkDense,

    -- * IO
    disp,
    loadMatrix, loadMatrix', saveMatrix,
    latexFormat,
    dispf, disps, dispcf, format,
    dispDots, dispBlanks, dispShort,
-- * Element conversion
    Convert(..),
    roundVector,
    fromInt,toInt,fromZ,toZ,
    -- * Misc
    arctan2,
    separable,
    fromArray2D,
    module Data.Complex,
    Mod,
    Vector, Matrix, GMatrix, nRows, nCols

) where

import Internal.Vector
import Internal.Vectorized
import Internal.Matrix hiding (size)
import Internal.Element
import Internal.IO
import Internal.Numeric
import Internal.Container
import Internal.Util hiding ((&))
import Data.Complex
import Internal.Sparse
import Internal.Modular


{-# LANGUAGE TypeApplications #-}

module MinkowskiScalarWave where

import qualified Data.Vector.Unboxed as U
import Data.VectorSpace
import qualified Numeric.LinearAlgebra as L
import SSSFN

--------------------------------------------------------------------------------

-- | Scalar wave in double-null coordinates:
-- > (- dt^2 + dx^2) u = 0
-- > t = v + u   dt = dv + du
-- > x = v - u   dx = dv - du
-- > (- (dv + du)^2 + (dv - du)^2) u = 0
-- > - 4 dv du u = 0
dalembert :: (Eq a, L.Field a, U.Unbox a) => Op a
dalembert = -4 *^ (derivV * derivU)

incomingU :: forall a. (L.Container L.Vector a, Fractional a, Ord a) => Op a
incomingU = omap isneg (bndU @a)
  where
    isneg a = fromIntegral (fromEnum (a < 0))

incomingV :: forall a. (L.Container L.Vector a, Fractional a, Ord a) => Op a
incomingV = omap isneg (bndV @a)
  where
    isneg a = fromIntegral (fromEnum (a < 0))

originUV :: forall a. (L.Container L.Vector a, Fractional a, Ord a) => Op a
originUV = ozipWith (*) incomingU incomingV

potential :: forall a. (L.Container L.Vector a, Floating a) => Grid a
potential = gzipWith (\u v -> 0) (coordU @a) (coordV @a)

incoming :: forall a. (L.Container L.Vector a, Floating a) => Grid a
incoming = gsample bcond

--------------------------------------------------------------------------------

op :: (L.Field a, Ord a, U.Unbox a) => Op a
op = intr * dalembert + bndu * incomingU + bndv * incomingV + orig * originUV
  where
    bndu = incomingU - originUV
    bndv = incomingV - originUV
    orig = originUV
    intr = 1 - (bndu + bndv + orig)

rhs :: (Floating a, L.Numeric a, Ord a) => Grid a
rhs =
  intr #> potential
    ^+^ bndu #> incoming
    ^+^ bndv #> incoming
    ^+^ orig #> incoming
  where
    bndu = incomingU - originUV
    bndv = incomingV - originUV
    orig = originUV
    intr = 1 - (bndu + bndv + orig)

--------------------------------------------------------------------------------

bcond :: Floating a => (a, a) -> a
bcond (u, v) = cos (pi / 4 * t) * cos (pi / 4 * x)
  where
    t = v + u
    x = v - u

--------------------------------------------------------------------------------

main :: IO ()
main =
  do
    putStrLn "1+1D Scalar Field in the Minkowski spacetime in double-Null coordinates"
    -- putStrLn $ "intr " ++ show (intr :: Op Double)
    -- putStrLn $ "bndu " ++ show (bndu :: Op Double)
    -- putStrLn $ "bndv " ++ show (bndv :: Op Double)
    -- putStrLn $ "orig " ++ show (orig :: Op Double)
    -- putStrLn $ "dal " ++ show (omap approx (dalembert @Double))
    -- putStrLn $ "op " ++ show (omap approx (op @Double))
    -- let inc = incoming @Double
    -- putStrLn $ "inc " ++ show (gmap approx inc)
    -- putStrLn $ "rhs " ++ show (gmap approx (rhs @Double))
    -- let ires = dalembert #> inc
    -- putStrLn $ "ires " ++ show (gmap approx ires)
    -- putStrLn $ "|ires| " ++ show (gmaxabs ires)
    -- let jres = op #> inc ^-^ rhs
    -- putStrLn $ "jres " ++ show (gmap approx jres)
    -- putStrLn $ "|jres| " ++ show (gmaxabs jres)

    let sol = solve op (rhs @Double)
    putStrLn $ "sol " ++ show (gmap approx sol)
    let res = op #> sol ^-^ rhs
    putStrLn $ "res " ++ show (gmap approx res)
    putStrLn $ "|res| " ++ show (gmaxabs res)
    let err = sol ^-^ incoming
    putStrLn $ "err " ++ show (gmap approx err)
    putStrLn $ "|err| " ++ show (gmaxabs err)
    putStrLn "Done."
  where
    approx :: RealFrac a => a -> a
    approx x = fromInteger (round (1.0e+4 * x)) / 1.0e+4
-- intr = 1 - (incomingU + incomingV - originUV)
-- bndu = incomingU - originUV
-- bndv = incomingV - originUV
-- orig = originUV

--------------------------------------------------------------------------------

-- quantum measurements: the quantum system measures the state of the
-- classical system

-- Beatrice Bonga: conserved current from varying boundary term from
-- lagrangian <https://arxiv.org/abs/1911.04514>