content: {-# OPTIONS_GHC -Wall #-}
  {-# LANGUAGE TypeFamilies, FlexibleContexts #-}
  
  -- | Algorithm 4.2 of Burden and Faires, 5th edition
  
  module Physics.Learn.AdaptiveQuadrature
  --    ( adaptiveQuad
  --    )
      where
  
  import Data.VectorSpace
      ( VectorSpace
      , InnerSpace
      , Scalar
      , (^+^)
      , (^-^)
      , (*^)
      , magnitude
      , sumV
      )
  
  -- | Simplest, most elegant implementation.
  --   Evaluates function at same spot multiple times.
  adaptiveQuad :: Double              -- ^ tolerance
               -> Double              -- ^ lower limit a
               -> Double              -- ^ upper limit b
               -> (Double -> Double)  -- ^ function f
               -> Double              -- ^ definite integral
  adaptiveQuad tol a b f
      = let s0 = simpson a b f
            m  = (a + b) / 2
            s1a = simpson a m f
            s1b = simpson m b f
        in if abs (s1a + s1b - s0) < 10 * tol
           then s1a + s1b
           else adaptiveQuad (tol/2) a m f + adaptiveQuad (tol/2) m b f
  
  simpson :: Double              -- ^ lower limit a
          -> Double              -- ^ upper limit b
          -> (Double -> Double)  -- ^ function f
          -> Double              -- ^ Simpson approximation
  simpson a b f = (b - a) / 6 * (f a + 4 * f ((a + b) / 2) + f b)
  
  -- | Version of adaptiveQuad for vectors.
  --   Evaluates function at same spot multiple times.
  adaptiveQuadVec :: (InnerSpace v, Scalar v ~ Double) =>
                     Double         -- ^ tolerance
                  -> Double         -- ^ lower limit a
                  -> Double         -- ^ upper limit b
                  -> (Double -> v)  -- ^ function f
                  -> v              -- ^ definite integral
  adaptiveQuadVec tol a b f
      = let s0 = simpsonVec a b f
            m  = (a + b) / 2
            s1a = simpsonVec a m f
            s1b = simpsonVec m b f
        in if magnitude (s1a ^+^ s1b ^-^ s0) < 10 * tol
           then s1a ^+^ s1b
           else adaptiveQuadVec (tol/2) a m f ^+^ adaptiveQuadVec (tol/2) m b f
  
  -- | Version of simpson for vectors.
  simpsonVec :: (VectorSpace v, Scalar v ~ Double) =>
                Double         -- ^ lower limit a
             -> Double         -- ^ upper limit b
             -> (Double -> v)  -- ^ function f
             -> v              -- ^ Simpson approximation
  simpsonVec a b f = ((b - a) / 6) *^ (f a ^+^ 4 *^ f ((a + b) / 2) ^+^ f b)
  
  -- | Burden and Faires, Example 2 on page 197
  example2f :: Double -> Double
  example2f x = (100 / x**2) * sin (10 / x)
  
  example2integral :: Double
  example2integral = adaptiveQuad 1e-4 1 3 example2f
  
  -- *AdaptiveQuadrature> example2integral 
  -- -1.426014810049443
  
  -- | Does no function evaluations itself.
  simpleSimpson :: Double              -- ^ lower limit a
                -> Double              -- ^ upper limit b
                -> Double              -- ^ value f(a)
                -> Double              -- ^ value f((a+b)/2)
                -> Double              -- ^ value f(b)
                -> Double              -- ^ Simpson approximation
  simpleSimpson a b fa fm fb = (b - a) / 6 * (fa + 4 * fm + fb)
  
  -- The workhorse of the adaptive Simpson method.
  -- Called by adaptiveSimpson
  adaptiveSimpsonStep :: Double              -- ^ tolerance
                      -> Double              -- ^ lower limit a
                      -> Double              -- ^ upper limit b
                      -> (Double -> Double)  -- ^ function f
                      -> Double              -- ^ value f(a)
                      -> Double              -- ^ value f((a+b)/2)
                      -> Double              -- ^ value f(b)
                      -> Double              -- ^ definite integral
  adaptiveSimpsonStep tol a b f fa fm fb
      = let s0 = simpleSimpson a b fa fm fb
            m  = (a + b) / 2
            am = (a + m) / 2
            mb = (m + b) / 2
            fam = f am
            fmb = f mb
            s1a = simpleSimpson a m fa fam fm
            s1b = simpleSimpson m b fm fmb fb
        in if abs (s1a + s1b - s0) < 10 * tol
           then s1a + s1b
           else adaptiveSimpsonStep (tol/2) a m f fa fam fm + adaptiveSimpsonStep (tol/2) m b f fm fmb fb
  
  -- | This version is more efficient in that it does not
  --   repeat function evaluations.
  adaptiveSimpson :: Double              -- ^ tolerance
                  -> Double              -- ^ lower limit a
                  -> Double              -- ^ upper limit b
                  -> (Double -> Double)  -- ^ function f
                  -> Double              -- ^ definite integral
  adaptiveSimpson tol a b f
      = let fa = f a
            m = (a + b) / 2
            fm = f m
            fb = f b
        in adaptiveSimpsonStep tol a b f fa fm fb
  
  -- | Does no function evaluations itself.
  --   For vector functions.
  simpleSimpsonVec :: (VectorSpace v, Fractional (Scalar v)) =>
                      Scalar v  -- ^ lower limit a
                   -> Scalar v  -- ^ upper limit b
                   -> v         -- ^ value f(a)
                   -> v         -- ^ value f((a+b)/2)
                   -> v         -- ^ value f(b)
                   -> v         -- ^ Simpson approximation
  simpleSimpsonVec a b fa fm fb = ((b - a) / 6) *^ (fa ^+^ 4 *^ fm ^+^ fb)
  
  ------------------------------------------
  -- Resource-limited adaptive quadrature --
  ------------------------------------------
  
  {-
  Want a version that gives an error estimate, and can be used by
  a scheduler for a resource-limited adaptive algorithm.
  We won't achieve a desired precision, but rather we'll use
  a fixed amount of resources in the best way possible.
  
  I think we'll need to create a data structure to hold the results
  of evaluations so far so that they can be fed to the next step
  if necessary.
  
  -- | This version does not repeat function evaluations.
  --   It provides an error estimate.
  
  
  -}
  
  -- data EvPair v = EvPair (Scalar v) v
  
  data SimpInterval3 v = SI3 { prLo    :: (Scalar v, v)
                             , prMi    :: (Scalar v, v)
                             , prHi    :: (Scalar v, v)
                             , intEst3 :: v
                             }
  
  data SimpInterval5 v = SI5 { pr0       :: (Scalar v, v)
                             , pr1       :: (Scalar v, v)
                             , pr2       :: (Scalar v, v)
                             , pr3       :: (Scalar v, v)
                             , pr4       :: (Scalar v, v)
                             , intEst012 :: v
                             , intEst234 :: v
                             , intEst024 :: v
                             , integralEst :: v  -- sum of intEst012 and intEst234
                             , errorEst  :: Scalar v
                             }
  
  divideInterval :: SimpInterval5 v -> (SimpInterval3 v, SimpInterval3 v)
  divideInterval (SI5 xy0 xy1 xy2 xy3 xy4 ie012 ie234 _ie024 _ _)
      = (SI3 xy0 xy1 xy2 ie012, SI3 xy2 xy3 xy4 ie234)
  
  refineInterval :: (InnerSpace v , Floating (Scalar v)) =>
                    (Scalar v -> v)
                 -> SimpInterval3 v
                 -> SimpInterval5 v
  refineInterval f (SI3 (x0,y0) (x2,y2) (x4,y4) ie024)
      = let x1 = (x0 + x2) / 2
            x3 = (x2 + x4) / 2
            y1 = f x1
            y3 = f x3
            ie012 = simpleSimpsonVec x0 x2 y0 y1 y2
            ie234 = simpleSimpsonVec x2 x4 y2 y3 y4
            ie = ie012 ^+^ ie234
            errEst = 1/10 * magnitude (ie ^-^ ie024)  -- 1/10 instead of 1/15
        in SI5 (x0,y0) (x1,y1) (x2,y2) (x3,y3) (x4,y4) ie012 ie234 ie024 ie errEst
  
  divideWorstInterval :: (InnerSpace v, Ord (Scalar v), Floating (Scalar v)) =>
                         (Scalar v -> v)
                      -> [SimpInterval5 v]
                      -> [SimpInterval5 v]
  divideWorstInterval _ [] = error "divideWorstInterval should never have been called on an empty list"
  divideWorstInterval f (si:sis)
      = let (si3a,si3b) = divideInterval si
            si5a = refineInterval f si3a
            si5b = refineInterval f si3b
        in insertSorted si5a $ insertSorted si5b sis
  
  insertSorted :: Ord (Scalar v) =>
                  SimpInterval5 v
               -> [SimpInterval5 v]
               -> [SimpInterval5 v]
  insertSorted si5 [] = [si5]
  insertSorted si5 (si:sis) = if errorEst si5 > errorEst si
                              then si5:si:sis
                              else si:insertSorted si5 sis
  
  adaptiveSimpEvalLimit :: (InnerSpace v, Ord (Scalar v), Floating (Scalar v)) =>
                           Int              -- ^ approximate number of function evals
                        -> Scalar v         -- ^ lower limit
                        -> Scalar v         -- ^ upper limit
                        -> (Scalar v -> v)  -- ^ scalar or vector function
                        -> v                -- ^ approximate integral
  adaptiveSimpEvalLimit n a b f
      = let m = (a + b) / 2
            fa = f a
            fm = f m
            fb = f b
            ie = simpleSimpsonVec a b fa fm fb
            si3 = SI3 (a,fa) (m,fm) (b,fb) ie
            si5 = refineInterval f si3
        in sumV $ map integralEst $ last $ take (div n 4) $ iterate (divideWorstInterval f) [si5]
  
  {-
  data SimpsonInterval5 v = SI5 { pLo         :: Scalar v
                                , pHi         :: Scalar v
                                , fLo         :: v
                                , fLM         :: v
                                , fM          :: v
                                , fMH         :: v
                                , fHi         :: v
                                , integralEst :: v
                                , errorEst    :: Scalar v
                                }
  -}
  
  -------------------------------
  -- Two-Dimensional integrals --
  -------------------------------
  
  adaptiveQuad2D :: Double              -- ^ tolerance
                 -> Double              -- ^ lower limit x_0
                 -> Double              -- ^ upper limit x_1
                 -> (Double -> Double)  -- ^ lower limit y_0(x)
                 -> (Double -> Double)  -- ^ upper limit y_1(x)
                 -> (Double -> Double -> Double)  -- ^ function f
                 -> Double              -- ^ definite integral
  adaptiveQuad2D tol x0 x1 y0 y1 f
      = let f1 x = adaptiveQuad tol' (y0 x) (y1 x) (f x)
            tol' = tol / abs (x1 - x0)
        in adaptiveQuad tol x0 x1 f1
  
  aq2dTest :: Double -> Double
  aq2dTest tol = adaptiveQuad2D tol (-1) 1 (\y -> -sqrt(1 - y**2)) (\y -> sqrt(1-y**2)) (\_ _ -> 1)
  
  adaptiveSimpson2D :: Double              -- ^ tolerance
                    -> Double              -- ^ lower limit x_0
                    -> Double              -- ^ upper limit x_1
                    -> (Double -> Double)  -- ^ lower limit y_0(x)
                    -> (Double -> Double)  -- ^ upper limit y_1(x)
                    -> (Double -> Double -> Double)  -- ^ function f
                    -> Double              -- ^ definite integral
  adaptiveSimpson2D tol x0 x1 y0 y1 f
      = let f1 x = adaptiveSimpson tol' (y0 x) (y1 x) (f x)
            tol' = tol / abs (x1 - x0)
        in adaptiveSimpson tol x0 x1 f1
  
  adaptiveSimpson3D :: Double              -- ^ tolerance
                    -> Double              -- ^ lower limit x_0
                    -> Double              -- ^ upper limit x_1
                    -> (Double -> Double)  -- ^ lower limit y_0(x)
                    -> (Double -> Double)  -- ^ upper limit y_1(x)
                    -> (Double -> Double -> Double)  -- ^ lower limit z_0(x,y)
                    -> (Double -> Double -> Double)  -- ^ upper limit z_1(x,y)
                    -> (Double -> Double -> Double -> Double)  -- ^ function f
                    -> Double              -- ^ definite integral
  adaptiveSimpson3D tol x0 x1 y0 y1 z0 z1 f
      = let f1 x = adaptiveSimpson2D tol' (y0 x) (y1 x) (z0 x) (z1 x) (f x)
            tol' = tol / abs (x1 - x0)
        in adaptiveSimpson tol x0 x1 f1
  
  as3dTest :: Double -> Double
  as3dTest tol = adaptiveSimpson3D tol (-1) 1
                 (\y -> -sqrt(1 - y**2)) (\y -> sqrt(1-y**2))
                 (\x y -> -sqrt(1 - x**2 - y**2)) (\x y -> sqrt(1 - x**2 - y**2))
                 (\_ _ _ -> 1)
  
  