{-# 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)


-- 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


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


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


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


{-# LANGUAGE BangPatterns               #-}
{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE DeriveFoldable             #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE DeriveTraversable          #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE NoImplicitPrelude          #-}
{-# LANGUAGE Rank2Types                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE StandaloneDeriving         #-}

-----------------------------------------------------------------------------
-- |
-- A class for semirings (types with two binary operations, one commutative and one associative, and two respective identities), with various general-purpose instances.
--
-----------------------------------------------------------------------------

module Data.Semiring
  ( -- * Semiring typeclass
    Semiring(..)
  , (+)
  , (*)
  , (^)
  , foldMapP
  , foldMapT
  , sum
  , product
  , sum'
  , product'
  , isZero
  , isOne

    -- * Types
  , Add(..)
  , Mul(..)
  , WrappedNum(..)
  , Mod2(..)
#if defined(VERSION_containers)
  , IntSetOf(..)
  , IntMapOf(..)
#endif

    -- * Ring typeclass
  , Ring(..)
  , fromInteger
  , fromIntegral
  , minus
  , (-)
  ) where

import           Control.Applicative (Applicative(..), Const(..), liftA2)
import           Data.Bits (Bits)
import           Data.Bool (Bool(..), (||), (&&), otherwise)
import           Data.Coerce (Coercible, coerce)
import           Data.Complex (Complex(..))
import           Data.Eq (Eq(..))
import           Data.Fixed (Fixed, HasResolution)
import           Data.Foldable (Foldable(foldMap))
import qualified Data.Foldable as Foldable
import           Data.Function ((.), const, id)
#if defined(VERSION_unordered_containers) || defined(VERSION_containers)
import           Data.Function (flip)
#endif
import           Data.Functor (Functor(..))
#if MIN_VERSION_base(4,12,0)
import           Data.Functor.Contravariant (Predicate(..), Equivalence(..), Op(..))
#endif
import           Data.Functor.Identity (Identity(..))
#if defined(VERSION_unordered_containers)
import           Data.Hashable (Hashable)
import           Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import           Data.HashSet (HashSet)
import qualified Data.HashSet as HashSet
#endif
import           Data.Int (Int, Int8, Int16, Int32, Int64)
import           Data.Maybe (Maybe(..))
#if MIN_VERSION_base(4,12,0)
import           Data.Monoid (Ap(..))
#endif
#if defined(VERSION_containers)
#if MIN_VERSION_base(4,7,0)
import           Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import           Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet
#endif
import           Data.Map (Map)
import qualified Data.Map as Map
#endif
import           Data.Monoid (Monoid(..), Dual(..))
import           Data.Ord (Ord((<)), (>=))
import           Data.Ord (Down(..))
import           Data.Proxy (Proxy(..))
import           Data.Ratio (Ratio, Rational, (%))
import           Data.Semigroup.Compat (Semigroup(..))
#if defined(VERSION_containers)
import           Data.Set (Set)
import qualified Data.Set as Set
#endif
import           Data.Traversable (Traversable)
import           Data.Typeable (Typeable)
import           Data.Word (Word, Word8, Word16, Word32, Word64)
import           Foreign.C.Types
  (CChar, CClock, CDouble, CFloat, CInt,
   CIntMax, CIntPtr, CLLong, CLong,
   CPtrdiff, CSChar, CSUSeconds, CShort,
   CSigAtomic, CSize, CTime, CUChar, CUInt,
   CUIntMax, CUIntPtr, CULLong, CULong,
   CUSeconds, CUShort, CWchar)
import           Foreign.Ptr (IntPtr, WordPtr)
import           Foreign.Storable (Storable)
import           GHC.Enum (Enum, Bounded)
import           GHC.Err (error)
import           GHC.Float (Float, Double)
import           GHC.Generics (Generic,Generic1)
import           GHC.IO (IO)
import           GHC.Integer (Integer)
import qualified GHC.Num as Num
import           GHC.Read (Read)
import           GHC.Real (Integral, Fractional, Real, RealFrac)
import qualified GHC.Real as Real
import           GHC.Show (Show)
import           Numeric.Natural (Natural)

#ifdef mingw32_HOST_OS
#define HOST_OS_WINDOWS 1
#else
#define HOST_OS_WINDOWS 0
#endif

#if !HOST_OS_WINDOWS
import           System.Posix.Types
  (CCc, CDev, CGid, CIno, CMode, CNlink,
   COff, CPid, CRLim, CSpeed, CSsize,
   CTcflag, CUid, Fd)
#endif

infixl 7 *, `times`
infixl 6 +, `plus`, -, `minus`
infixr 8 ^

{--------------------------------------------------------------------
  Helpers
--------------------------------------------------------------------}

-- | Raise a number to a non-negative integral power.
-- If the power is negative, this will call 'error'.
{-# SPECIALISE [1] (^) ::
        Integer -> Integer -> Integer,
        Integer -> Int -> Integer,
        Int -> Int -> Int #-}
{-# INLINABLE [1] (^) #-} -- See note [Inlining (^)]
(^) :: (Semiring a, Integral b) => a -> b -> a
x ^ y
  | y < 0 = error "Data.Semiring.^: negative power"
  | y == 0 = one
  | otherwise = getMul (stimes y (Mul x))

{- Note [Inlining (^)]
   ~~~~~~~~~~~~~~~~~~~
   The INLINABLE pragma allows (^) to be specialised at its call sites.
   If it is called repeatedly at the same type, that can make a huge
   difference, because of those constants which can be repeatedly
   calculated.

   Currently the fromInteger calls are not floated because we get
             \d1 d2 x y -> blah
   after the gentle round of simplification.
-}

{- Rules for powers with known small exponent
    see Trac #5237
    For small exponents, (^) is inefficient compared to manually
    expanding the multiplication tree.
    Here, rules for the most common exponent types are given.
    The range of exponents for which rules are given is quite
    arbitrary and kept small to not unduly increase the number of rules.
    It might be desirable to have corresponding rules also for
    exponents of other types (e.g., Word), but it's doubtful they
    would fire, since the exponents of other types tend to get
    floated out before the rule has a chance to fire. (Why?)

    Note: Trying to save multiplication by sharing the square for
    exponents 4 and 5 does not save time, indeed, for Double, it is
    up to twice slower, so the rules contain flat sequences of
    multiplications.
-}

{-# RULES
"^0/Int" forall x. x ^ (0 :: Int) = one
"^1/Int" forall x. x ^ (1 :: Int) = let u = x in u
"^2/Int" forall x. x ^ (2 :: Int) = let u = x in u*u
"^3/Int" forall x. x ^ (3 :: Int) = let u = x in u*u*u
"^4/Int" forall x. x ^ (4 :: Int) = let u = x in u*u*u*u
"^5/Int" forall x. x ^ (5 :: Int) = let u = x in u*u*u*u*u
"^0/Integer" forall x. x ^ (0 :: Integer) = one
"^1/Integer" forall x. x ^ (1 :: Integer) = let u = x in u
"^2/Integer" forall x. x ^ (2 :: Integer) = let u = x in u*u
"^3/Integer" forall x. x ^ (3 :: Integer) = let u = x in u*u*u
"^4/Integer" forall x. x ^ (4 :: Integer) = let u = x in u*u*u*u
"^5/Integer" forall x. x ^ (5 :: Integer) = let u = x in u*u*u*u*u
  #-}

-- | Infix shorthand for 'plus'.
(+) :: Semiring a => a -> a -> a
(+) = plus
{-# INLINE (+) #-}

-- | Infix shorthand for 'times'.
(*) :: Semiring a => a -> a -> a
(*) = times
{-# INLINE (*) #-}

-- | Infix shorthand for 'minus'.
(-) :: Ring a => a -> a -> a
(-) = minus
{-# INLINE (-) #-}

-- | Map each element of the structure to a semiring, and combine the results
--   using 'plus'.
foldMapP :: (Foldable t, Semiring s) => (a -> s) -> t a -> s
foldMapP f = Foldable.foldr (plus  . f) zero
{-# INLINE foldMapP #-}

-- | Map each element of the structure to a semiring, and combine the results
--   using 'times'.
foldMapT :: (Foldable t, Semiring s) => (a -> s) -> t a -> s
foldMapT f = Foldable.foldr (times . f) one
{-# INLINE foldMapT #-}

infixr 9 #.
(#.) :: Coercible b c => (b -> c) -> (a -> b) -> a -> c
(#.) _ = coerce

-- | The 'sum' function computes the additive sum of the elements in a structure.
--   This function is lazy. For a strict version, see 'sum''.
sum :: (Foldable t, Semiring a) => t a -> a
sum = getAdd #. foldMap Add
{-# INLINE sum #-}

-- | The 'product' function computes the product of the elements in a structure.
--   This function is lazy. for a strict version, see 'product''.
product :: (Foldable t, Semiring a) => t a -> a
product = getMul #. foldMap Mul
{-# INLINE product #-}

-- | The 'sum'' function computes the additive sum of the elements in a structure.
--   This function is strict. For a lazy version, see 'sum'.
sum'  :: (Foldable t, Semiring a) => t a -> a
sum'  = Foldable.foldl' plus zero
{-# INLINE sum' #-}

-- | The 'product'' function computes the additive sum of the elements in a structure.
--   This function is strict. For a lazy version, see 'product'.
product' :: (Foldable t, Semiring a) => t a -> a
product' = Foldable.foldl' times one
{-# INLINE product' #-}

-- | Monoid under 'plus'. Analogous to 'Data.Monoid.Sum', but
--   uses the 'Semiring' constraint rather than 'Num.Num'.
newtype Add a = Add { getAdd :: a }
  deriving
    ( Bounded
    , Enum
    , Eq
    , Foldable
    , Fractional
    , Functor
    , Generic
    , Generic1
    , Num.Num
    , Ord
    , Read
    , Real
    , RealFrac
    , Show
    , Storable
    , Traversable
    , Typeable
    )

instance Semiring a => Semigroup (Add a) where
  Add a <> Add b = Add (a + b)
  stimes n (Add a) = Add (fromNatural (Real.fromIntegral n) * a)
  {-# INLINE (<>) #-}

instance Semiring a => Monoid (Add a) where
  mempty = Add zero
  mappend = (<>)
  {-# INLINE mempty #-}
  {-# INLINE mappend #-}

-- | This is an internal type, solely for purposes
-- of default implementation of 'fromNatural'.
newtype Add' a = Add' { getAdd' :: a }

instance Semiring a => Semigroup (Add' a) where
  Add' a <> Add' b = Add' (a + b)

-- | Monoid under 'times'. Analogous to 'Data.Monoid.Product', but
--   uses the 'Semiring' constraint rather than 'Num.Num'.
newtype Mul a = Mul { getMul :: a }
  deriving
    ( Bounded
    , Enum
    , Eq
    , Foldable
    , Fractional
    , Functor
    , Generic
    , Generic1
    , Num.Num
    , Ord
    , Read
    , Real
    , RealFrac
    , Show
    , Storable
    , Traversable
    , Typeable
    )

instance Semiring a => Semigroup (Mul a) where
  Mul a <> Mul b = Mul (a * b)
  {-# INLINE (<>) #-}

instance Semiring a => Monoid (Mul a) where
  mempty = Mul one
  mappend = (<>)
  {-# INLINE mempty #-}
  {-# INLINE mappend #-}

-- | Provide Semiring and Ring for an arbitrary 'Num.Num'. It is useful with GHC 8.6+'s DerivingVia extension.
newtype WrappedNum a = WrapNum { unwrapNum :: a }
  deriving
    ( Bounded
    , Enum
    , Eq
    , Foldable
    , Fractional
    , Functor
    , Generic
    , Generic1
    , Num.Num
    , Ord
    , Read
    , Real
    , RealFrac
    , Show
    , Storable
    , Traversable
    , Typeable
    , Bits
    )

instance Num.Num a => Semiring (WrappedNum a) where
  plus  = (Num.+)
  zero  = 0
  times = (Num.*)
  one   = 1
  fromNatural = Real.fromIntegral

instance Num.Num a => Ring (WrappedNum a) where
  negate = Num.negate

-- | 'Mod2' represents the integers mod 2.
--
--   It is useful in the computing of <https://en.wikipedia.org/wiki/Zhegalkin_polynomial Zhegalkin polynomials>.
newtype Mod2 = Mod2 { getMod2 :: Bool }
  deriving
    ( Bounded
    , Enum
    , Eq
    , Ord
    , Read
    , Show
    , Generic
    )

instance Semiring Mod2 where
  -- we inline the definition of 'xor'
  -- on Bools, since the instance did not exist until
  -- base-4.7.0.
  plus (Mod2 x) (Mod2 y) = Mod2 (x /= y)
  times (Mod2 x) (Mod2 y) = Mod2 (x && y)
  zero = Mod2 False
  one = Mod2 True

instance Ring Mod2 where
  negate = id
  {-# INLINE negate #-}

{--------------------------------------------------------------------
  Classes
--------------------------------------------------------------------}

-- | The class of semirings (types with two binary
-- operations and two respective identities). One
-- can think of a semiring as two monoids of the same
-- underlying type, with the first being commutative.
-- In the documentation, you will often see the first
-- monoid being referred to as @additive@, and the second
-- monoid being referred to as @multiplicative@, a typical
-- convention when talking about semirings.
--
-- For any type R with a 'Num.Num'
-- instance, the additive monoid is (R, 'Prelude.+', 0)
-- and the multiplicative monoid is (R, 'Prelude.*', 1).
--
-- For 'Prelude.Bool', the additive monoid is ('Prelude.Bool', 'Prelude.||', 'Prelude.False')
-- and the multiplicative monoid is ('Prelude.Bool', 'Prelude.&&', 'Prelude.True').
--
-- Instances should satisfy the following laws:
--
-- [/additive left identity/]
--     @'zero' '+' x = x@
-- [/additive right identity/]
--     @x '+' 'zero' = x@
-- [/additive associativity/]
--     @x '+' (y '+' z) = (x '+' y) '+' z@
-- [/additive commutativity/]
--     @x '+' y = y '+' x@
-- [/multiplicative left identity/]
--     @'one' '*' x = x@
-- [/multiplicative right identity/]
--     @x '*' 'one' = x@
-- [/multiplicative associativity/]
--     @x '*' (y '*' z) = (x '*' y) '*' z@
-- [/left-distributivity of '*' over '+'/]
--     @x '*' (y '+' z) = (x '*' y) '+' (x '*' z)@
-- [/right-distributivity of '*' over '+'/]
--     @(x '+' y) '*' z = (x '*' z) '+' (y '*' z)@
-- [/annihilation/]
--     @'zero' '*' x = x '*' 'zero' = 'zero'@

class Semiring a where
#if __GLASGOW_HASKELL__ >= 708
  {-# MINIMAL plus, times, (zero, one | fromNatural) #-}
#endif
  plus  :: a -> a -> a -- ^ Commutative Operation
  zero  :: a           -- ^ Commutative Unit
  zero = fromNatural 0
  times :: a -> a -> a -- ^ Associative Operation
  one   :: a           -- ^ Associative Unit
  one = fromNatural 1
  fromNatural :: Natural -> a -- ^ Homomorphism of additive semigroups
  fromNatural 0 = zero
  fromNatural n = getAdd' (stimes n (Add' one))

-- | The class of semirings with an additive inverse.
--
--     @'negate' a '+' a = 'zero'@

class Semiring a => Ring a where
#if __GLASGOW_HASKELL__ >= 708
  {-# MINIMAL negate #-}
#endif
  negate :: a -> a

-- | Subtract two 'Ring' values. For any type @R@ with
-- a 'Num.Num' instance, this is the same as '(Prelude.-)'.
--
--     @x `minus` y = x '+' 'negate' y@
minus :: Ring a => a -> a -> a
minus x y = x + negate y
{-# INLINE minus #-}

-- | Convert from integer to ring.
--
-- When @{-#@ @LANGUAGE RebindableSyntax #-}@ is enabled,
-- this function is used for desugaring integer literals.
-- This may be used to facilitate transition from 'Num.Num' to 'Ring':
-- no need to replace 0 and 1 with 'one' and 'zero'
-- or to cast numeric literals.
fromInteger :: Ring a => Integer -> a
fromInteger x
  | x >= 0    = fromNatural (Num.fromInteger x)
  | otherwise = negate (fromNatural (Num.fromInteger (Num.negate x)))
{-# INLINE fromInteger #-}

-- | Convert from integral to ring.
fromIntegral :: (Integral a, Ring b) => a -> b
fromIntegral x
  | x >= 0    = fromNatural (Real.fromIntegral x)
  | otherwise = negate (fromNatural (Real.fromIntegral (Num.negate x)))
{-# INLINE fromIntegral #-}

{--------------------------------------------------------------------
  Instances (base)
--------------------------------------------------------------------}

instance Semiring b => Semiring (a -> b) where
  plus f g    = \x -> f x `plus` g x
  zero        = const zero
  times f g   = \x -> f x `times` g x
  one         = const one
  fromNatural = const . fromNatural
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring b => Ring (a -> b) where
  negate f x = negate (f x)
  {-# INLINE negate #-}

instance Semiring () where
  plus _ _  = ()
  zero      = ()
  times _ _ = ()
  one       = ()
  fromNatural _ = ()
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring () where
  negate _ = ()
  {-# INLINE negate #-}

instance Semiring (Proxy a) where
  plus _ _  = Proxy
  zero      = Proxy
  times _ _ = Proxy
  one       = Proxy
  fromNatural _ = Proxy
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Semiring Bool where
  plus  = (||)
  zero  = False
  times = (&&)
  one   = True
  fromNatural 0 = False
  fromNatural _ = True
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Semiring a => Semiring (Maybe a) where
  zero  = Nothing
  one   = Just one

  plus Nothing y = y
  plus x Nothing = x
  plus (Just x) (Just y) = Just (plus x y)

  times Nothing _ = Nothing
  times _ Nothing = Nothing
  times (Just x) (Just y) = Just (times x y)

  fromNatural 0 = Nothing
  fromNatural n = Just (fromNatural n)
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Semiring a => Semiring (IO a) where
  zero  = pure zero
  one   = pure one
  plus  = liftA2 plus
  times = liftA2 times
  fromNatural = pure . fromNatural
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring a => Ring (IO a) where
  negate = fmap negate
  {-# INLINE negate #-}

instance Semiring a => Semiring (Dual a) where
  zero = Dual zero
  Dual x `plus` Dual y = Dual (y `plus` x)
  one = Dual one
  Dual x `times` Dual y = Dual (y `times` x)
  fromNatural = Dual . fromNatural
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring a => Ring (Dual a) where
  negate (Dual x) = Dual (negate x)
  {-# INLINE negate #-}

instance Semiring a => Semiring (Const a b) where
  zero = Const zero
  one  = Const one
  plus  (Const x) (Const y) = Const (x `plus`  y)
  times (Const x) (Const y) = Const (x `times` y)
  fromNatural = Const . fromNatural
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring a => Ring (Const a b) where
  negate (Const x) = Const (negate x)
  {-# INLINE negate #-}

-- | This instance can suffer due to floating point arithmetic.
instance Ring a => Semiring (Complex a) where
  zero = zero :+ zero
  one  = one  :+ zero
  plus  (x :+ y) (x' :+ y') = plus x x' :+ plus y y'
  times (x :+ y) (x' :+ y')
    = (x * x' - (y * y')) :+ (x * y' + y * x')
  fromNatural n = fromNatural n :+ zero
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance Ring a => Ring (Complex a) where
  negate (x :+ y) = negate x :+ negate y
  {-# INLINE negate #-}

#if MIN_VERSION_base(4,12,0)
instance (Semiring a, Applicative f) => Semiring (Ap f a) where
  zero  = pure zero
  one   = pure one
  plus  = liftA2 plus
  times = liftA2 times
  fromNatural = pure . fromNatural
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

instance (Ring a, Applicative f) => Ring (Ap f a) where
  negate = fmap negate
  {-# INLINE negate #-}
#endif

#if MIN_VERSION_base(4,12,0)
deriving instance Semiring (Predicate a)

deriving instance Semiring a => Semiring (Equivalence a)

deriving instance Semiring a => Semiring (Op a b)
deriving instance Ring a => Ring (Op a b)
#endif

#define deriveSemiring(ty)         \
instance Semiring (ty) where {     \
   zero  = 0                       \
;  one   = 1                       \
;  plus  x y = (Num.+) x y         \
;  times x y = (Num.*) x y         \
;  fromNatural = Real.fromIntegral \
;  {-# INLINE zero #-}             \
;  {-# INLINE one  #-}             \
;  {-# INLINE plus #-}             \
;  {-# INLINE times #-}            \
;  {-# INLINE fromNatural #-}      \
}

deriveSemiring(Int)
deriveSemiring(Int8)
deriveSemiring(Int16)
deriveSemiring(Int32)
deriveSemiring(Int64)
deriveSemiring(Integer)
deriveSemiring(Word)
deriveSemiring(Word8)
deriveSemiring(Word16)
deriveSemiring(Word32)
deriveSemiring(Word64)
deriveSemiring(Float)
deriveSemiring(Double)
deriveSemiring(CUIntMax)
deriveSemiring(CIntMax)
deriveSemiring(CUIntPtr)
deriveSemiring(CIntPtr)
deriveSemiring(CSUSeconds)
deriveSemiring(CUSeconds)
deriveSemiring(CTime)
deriveSemiring(CClock)
deriveSemiring(CSigAtomic)
deriveSemiring(CWchar)
deriveSemiring(CSize)
deriveSemiring(CPtrdiff)
deriveSemiring(CDouble)
deriveSemiring(CFloat)
deriveSemiring(CULLong)
deriveSemiring(CLLong)
deriveSemiring(CULong)
deriveSemiring(CLong)
deriveSemiring(CUInt)
deriveSemiring(CInt)
deriveSemiring(CUShort)
deriveSemiring(CShort)
deriveSemiring(CUChar)
deriveSemiring(CSChar)
deriveSemiring(CChar)
deriveSemiring(IntPtr)
deriveSemiring(WordPtr)

#if !HOST_OS_WINDOWS
deriveSemiring(CCc)
deriveSemiring(CDev)
deriveSemiring(CGid)
deriveSemiring(CIno)
deriveSemiring(CMode)
deriveSemiring(CNlink)
deriveSemiring(COff)
deriveSemiring(CPid)
deriveSemiring(CRLim)
deriveSemiring(CSpeed)
deriveSemiring(CSsize)
deriveSemiring(CTcflag)
deriveSemiring(CUid)
deriveSemiring(Fd)
#endif

deriveSemiring(Natural)

instance Integral a => Semiring (Ratio a) where
  {-# SPECIALIZE instance Semiring Rational #-}
  zero  = 0 % 1
  one   = 1 % 1
  plus  = (Num.+)
  times = (Num.*)
  fromNatural n = Real.fromIntegral n % 1
  {-# INLINE zero  #-}
  {-# INLINE one   #-}
  {-# INLINE plus  #-}
  {-# INLINE times #-}
  {-# INLINE fromNatural #-}
deriving instance Semiring a => Semiring (Identity a)
deriving instance Semiring a => Semiring (Down a)
instance HasResolution a => Semiring (Fixed a) where
  zero  = 0
  one   = 1
  plus  = (Num.+)
  times = (Num.*)
  fromNatural = Real.fromIntegral
  {-# INLINE zero  #-}
  {-# INLINE one   #-}
  {-# INLINE plus  #-}
  {-# INLINE times #-}
  {-# INLINE fromNatural #-}

#define deriveRing(ty)          \
instance Ring (ty) where {      \
  negate = Num.negate           \
; {-# INLINE negate #-}         \
}

deriveRing(Int)
deriveRing(Int8)
deriveRing(Int16)
deriveRing(Int32)
deriveRing(Int64)
deriveRing(Integer)
deriveRing(Word)
deriveRing(Word8)
deriveRing(Word16)
deriveRing(Word32)
deriveRing(Word64)
deriveRing(Float)
deriveRing(Double)
deriveRing(CUIntMax)
deriveRing(CIntMax)
deriveRing(CUIntPtr)
deriveRing(CIntPtr)
deriveRing(CSUSeconds)
deriveRing(CUSeconds)
deriveRing(CTime)
deriveRing(CClock)
deriveRing(CSigAtomic)
deriveRing(CWchar)
deriveRing(CSize)
deriveRing(CPtrdiff)
deriveRing(CDouble)
deriveRing(CFloat)
deriveRing(CULLong)
deriveRing(CLLong)
deriveRing(CULong)
deriveRing(CLong)
deriveRing(CUInt)
deriveRing(CInt)
deriveRing(CUShort)
deriveRing(CShort)
deriveRing(CUChar)
deriveRing(CSChar)
deriveRing(CChar)
deriveRing(IntPtr)
deriveRing(WordPtr)

#if !HOST_OS_WINDOWS
deriveRing(CCc)
deriveRing(CDev)
deriveRing(CGid)
deriveRing(CIno)
deriveRing(CMode)
deriveRing(CNlink)
deriveRing(COff)
deriveRing(CPid)
deriveRing(CRLim)
deriveRing(CSpeed)
deriveRing(CSsize)
deriveRing(CTcflag)
deriveRing(CUid)
deriveRing(Fd)
#endif

instance Integral a => Ring (Ratio a) where
  negate = Num.negate
  {-# INLINE negate #-}

deriving instance Ring a => Ring (Down a)
deriving instance Ring a => Ring (Identity a)
instance HasResolution a => Ring (Fixed a) where
  negate = Num.negate
  {-# INLINE negate #-}

{--------------------------------------------------------------------
  Instances (containers)
--------------------------------------------------------------------}

#if defined(VERSION_containers)

-- | The multiplication laws are satisfied for
--   any underlying 'Monoid', so we require a
--   'Monoid' constraint instead of a 'Semiring'
--   constraint since 'times' can use
--   the context of either.
instance (Ord a, Monoid a) => Semiring (Set a) where
  zero  = Set.empty
  one   = Set.singleton mempty
  plus  = Set.union
  times xs ys = Foldable.foldMap (flip Set.map ys . mappend) xs
  fromNatural 0 = zero
  fromNatural _ = one
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

-- | Wrapper to mimic 'Set' ('Data.Semigroup.Sum' 'Int'),
-- 'Set' ('Data.Semigroup.Product' 'Int'), etc.,
-- while having a more efficient underlying representation.
newtype IntSetOf a = IntSetOf { getIntSet :: IntSet }
  deriving
    ( Eq
    , Generic
    , Generic1
    , Ord
    , Read
    , Show
    , Typeable
    , Semigroup
    , Monoid
    )

instance (Coercible Int a, Monoid a) => Semiring (IntSetOf a) where
  zero  = coerce IntSet.empty
  one   = coerce IntSet.singleton (mempty :: a)
  plus  = coerce IntSet.union
  xs `times` ys
    = coerce IntSet.fromList
        [ mappend k l
        | k :: a <- coerce IntSet.toList xs
        , l :: a <- coerce IntSet.toList ys
        ]
  fromNatural 0 = zero
  fromNatural _ = one
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

-- | The multiplication laws are satisfied for
--   any underlying 'Monoid' as the key type,
--   so we require a 'Monoid' constraint instead of
--   a 'Semiring' constraint since 'times' can use
--   the context of either.
instance (Ord k, Monoid k, Semiring v) => Semiring (Map k v) where
  zero = Map.empty
  one  = Map.singleton mempty one
  plus = Map.unionWith (+)
  xs `times` ys
    = Map.fromListWith (+)
        [ (mappend k l, v * u)
        | (k,v) <- Map.toList xs
        , (l,u) <- Map.toList ys
        ]
  fromNatural 0 = zero
  fromNatural n = Map.singleton mempty (fromNatural n)
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

-- | Wrapper to mimic 'Map' ('Data.Semigroup.Sum' 'Int') v,
-- 'Map' ('Data.Semigroup.Product' 'Int') v, etc.,
-- while having a more efficient underlying representation.
newtype IntMapOf k v = IntMapOf { getIntMap :: IntMap v }
  deriving
    ( Eq
    , Generic
    , Generic1
    , Ord
    , Read
    , Show
    , Typeable
    , Semigroup
    , Monoid
    )

instance (Coercible Int k, Monoid k, Semiring v) => Semiring (IntMapOf k v) where
  zero = coerce (IntMap.empty :: IntMap v)
  one  = coerce (IntMap.singleton :: Int -> v -> IntMap v) (mempty :: k) (one :: v)
  plus = coerce (IntMap.unionWith (+) :: IntMap v -> IntMap v -> IntMap v)
  xs `times` ys
    = coerce (IntMap.fromListWith (+) :: [(Int, v)] -> IntMap v)
        [ (mappend k l, v * u)
        | (k :: k, v :: v) <- coerce (IntMap.toList :: IntMap v -> [(Int, v)]) xs
        , (l :: k, u :: v) <- coerce (IntMap.toList :: IntMap v -> [(Int, v)]) ys
        ]
  fromNatural 0 = zero
  fromNatural n = coerce (IntMap.singleton :: Int -> v -> IntMap v) (mempty :: k) (fromNatural n :: v)
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

#endif

{--------------------------------------------------------------------
  Instances (unordered-containers)
--------------------------------------------------------------------}

#if defined(VERSION_unordered_containers)

-- | The multiplication laws are satisfied for
--   any underlying 'Monoid', so we require a
--   'Monoid' constraint instead of a 'Semiring'
--   constraint since 'times' can use
--   the context of either.
instance (Eq a, Hashable a, Monoid a) => Semiring (HashSet a) where
  zero = HashSet.empty
  one  = HashSet.singleton mempty
  plus = HashSet.union
  times xs ys = Foldable.foldMap (flip HashSet.map ys . mappend) xs
  fromNatural 0 = zero
  fromNatural _ = one
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}

-- | The multiplication laws are satisfied for
--   any underlying 'Monoid' as the key type,
--   so we require a 'Monoid' constraint instead of
--   a 'Semiring' constraint since 'times' can use
--   the context of either.
instance (Eq k, Hashable k, Monoid k, Semiring v) => Semiring (HashMap k v) where
  zero = HashMap.empty
  one  = HashMap.singleton mempty one
  plus = HashMap.unionWith (+)
  xs `times` ys
    = HashMap.fromListWith (+)
        [ (mappend k l, v * u)
        | (k,v) <- HashMap.toList xs
        , (l,u) <- HashMap.toList ys
        ]
  fromNatural 0 = zero
  fromNatural n = HashMap.singleton mempty (fromNatural n)
  {-# INLINE plus  #-}
  {-# INLINE zero  #-}
  {-# INLINE times #-}
  {-# INLINE one   #-}
  {-# INLINE fromNatural #-}
#endif

-- | Is the value 'zero'?
isZero :: (Eq a, Semiring a) => a -> Bool
isZero x = x == zero
{-# INLINEABLE isZero #-}

-- | Is the value 'one'?
isOne :: (Eq a, Semiring a) => a -> Bool
isOne x = x == one
{-# INLINEABLE isOne #-}