/** Routines for numerical integration.

    See:
    https://en.wikipedia.org/wiki/Simpson%27s_rule
    https://en.wikipedia.org/wiki/Adaptive_Simpson%27s_method

    https://web.mit.edu/10.001/Web/Course_Notes/Integration_Notes/numerical_integration.pdf

    See examples and tests in NumericalIntegrationTest.frink
*/

/** Numerically integrate function f from a to b. */
NIntegrate[f, a, b, eps=1e-14] :=
{
   return EngelenIntegrator.NIntegrate[f, a, b, eps, 6]
   //NIntegrateAdaptiveSimpson[f, a, b, eps, 30, undef]
}

/** Numerically integrate function f[x, data] from a to b.  data is auxiliary
    data (e.g. constant parameters) that will be passed to the function. */
NIntegrateData[f, a, b, data, eps=1e-14, steps=6] :=
{
   return EngelenIntegrator.NIntegrateData[f, a, b, data, eps, steps]
   //return NIntegrateAdaptiveSimpson[f, a, b, eps, 30, data]
}

/** Integrate the function from a to b with the specified number of steps. */
NIntegrateCompositeTrapezoid[f, a, b, steps, data] :=
{
   if data == undef
   {
      fa = f[a]
      fb = f[b]
   } else
   {
      fa = f[a, data]
      fb = f[b, data]
   }
   
   h = (b-a)/steps
   integral = (fa + fb) / 2   // TODO: Make this work with dates
   xi = a
   for i = 1 to steps-1
   {
      xi = xi + h
      if data == undef
         fxi = f[xi]
      else
         fxi = f[xi, data]
      
      integral = integral + fxi
   }

   return integral * h
}


/** Integrate the function from a to b with the specified number of steps using
    the composite Simpson's 1/3 rule.  Simpson's rule is exact when the
    integrand is a polynomial of degree 3 or lower.

    See:
    https://en.wikipedia.org/wiki/Simpson%27s_rule#Composite_Simpson's_1/3_rule

    The error is proportional to
    -1/180 h^4 (b-a) D[f, 4]

    that is, proportional to the 4th derivative of f.  h = (b-a)/steps
*/
NIntegrateCompositeSimpson13[f, a, b, steps, data] :=
{
   // Make sure steps is even
   if steps mod 2 == 1
      steps = steps + 1
   
   h = (b-a)/steps
   if data == undef
   {
      fa = f[a]
      fb = f[b]
   } else
   {
      fa = f[a, data]
      fb = f[b, data]
   }
   integral = fa + fb
   xi = a

   integral1 = 0 integral
   for i = 1 to steps div 2
   {
      x2iminus1 = a + (2i-1) h
      if data == undef
         ff = f[x2iminus1]
      else
         ff = f[x2iminus1, data]
      integral1 = integral1 + ff
   }

   integral2 = 0 integral
   for i = 1 to steps div 2 - 1
   {
      x2i = a + (2 i) h
      if data == undef
         ff = f[x2i]
      else
         ff = f[x2i, data]

      integral2 = integral2 + ff
   }

   integral = integral + 4 integral1 + 2 integral2

   return 1/3 h integral
}

/** Adaptive Simpson's method.

    see https://en.wikipedia.org/wiki/Adaptive_Simpson%27s_method

    f: function to integrate
    a, b: limits of integration
    epsilon: maximum error
    maxRecDepth: maximum recursion depth
*/
NIntegrateAdaptiveSimpson[f, a, b, epsilon, maxRecDepth, data] :=
{
   h = b - a
   if h == 0 h
      return 0 h

   if data == undef
   {
      fa = f[a]
      fb = f[b]
      fm = f[(a + b)/2]
   } else
   {
      fa = f[a, data]
      fb = f[b, data]
      fm = f[(a + b)/2, data]
   }
   
   S = (h/6)*(fa + 4*fm + fb)
   eps = epsilon factorUnits[b]@1 factorUnits[fb]@1 // Make units come out right
   
   return adaptiveSimpsonRec[f, a, b, eps, S, fa, fb, fm, maxRecDepth, data]
}

/** Adaptive Simpson's Rule, private recursive function */
adaptiveSimpsonRec[f, a, b, eps, whole, fa, fb, fm, rec, data] :=
{
    m  = (a + b)/2
    h  = (b - a)/2
    lm = (a + m)/2
    rm = (m + b)/2

    // serious numerical trouble: it won't converge
    if ((eps/2 == eps) or (a == lm))
    {
       println["Serious numerical trouble in integration.  eps is $eps"]
       return whole;
    }

    if data == undef
    {
       flm = f[lm]
       frm = f[rm]
    } else
    {
       flm = f[lm, data]
       frm = f[rm, data]
    }       
    
    left  = (h/6) * (fa + 4*flm + fm)
    right = (h/6) * (fm + 4*frm + fb)
    delta = left + right - whole

    if (rec <= 0)
    {
       println["Recursive integration limit too shallow between $a and $b."]
    }

    //    println["Got here.  delta is $delta and eps is $eps"]
    
    // Lyness 1969 + Richardson extrapolation; see article
    if rec <= 0 or abs[delta] <= 15*eps
       return left + right + delta/15
    
    return adaptiveSimpsonRec[f, a, m, eps/2, left,  fa, fm, flm, rec-1, data] +
           adaptiveSimpsonRec[f, m, b, eps/2, right, fm, fb, frm, rec-1, data]
}

/** Convenience method to call the class below. */
NIntegrateTanh[f, a, b, eps=1e-14, debug=false] :=
{
   TanhSinhIntegrator.NIntegrate[f, a, b, eps, debug]
}

/** Numerical integrator using the tanh-sinh quadrature.

    See:
    Takahasi, Hidetosi; Mori, Masatake (1974),
    "Double Exponential Formulas for Numerical Integration",
    Publications of the Research Institute for Mathematical Sciences,
    9 (3): 721–741 
    https://www.ems-ph.org/journals/show_pdf.php?issn=0034-5318&vol=9&iss=3&rank=12

    The algorithm has been additionally improved by Michashki & Mosig (2016):
    https://www.ingentaconnect.com/content/tandf/jew/2016/00000030/00000003/art00001;jsessionid=55mdimh9e060a.x-ic-live-03

    and Van Engelen (2022):
    https://www.genivia.com/files/qthsh.pdf
    (see especially page 24 if you're going to implement it with all the
    improvements or, better yet, page 33 if you want to implement infinite
    endpoints.)

    This implementation is based on:
    https://github.com/Proektsoftbg/Numerical/blob/main/Numerical/Integrator/TanhSinh.cs
*/
class TanhSinhIntegrator
{
   // Precalculated abcissae and weights from the Van Engelen implementation
   class var n = 11
   class var m = [ 6, 7, 13, 26, 53, 106, 212, 423, 846, 1693, 3385 ]
   class var r = new array[n]
   class var w = new array[n]
   class var initialized = TanhSinhIntegrator.initialize[]

   /** Numerically integrate function f from x1 to x2 and the specified
       precision. */
   class NIntegrate[f, x1, x2, precision = 1e-14, debug=false] :=
   {
      // Friendly correction to limits of integration
      signCorrection = 1
      if x1 > x2
         [x1,x2,signCorrection] = [x2, x1, -1]   // Flip sign
      
      ic = 1    // Iteration count
      c = (x1 + x2) / 2
      d = (x2 - x1) / 2    // Half-width of bounds of integration
      s = f[c]
      err = undef
      i = 0
      eps = clamp[precision * 0.1, 1e-15, 1e-8]
      tol = 10 precision
      do
      {
         q = 0
         p = 0
         fp = 0
         fm = 0
         j = 0
         
         do
         {
            x = r@i@j * d
            if (x1 + x > x1)
            {
               y = f[x1 + x]
               ic = ic + 1
               fp = y
            }

            if (x2 - x < x2)
            {
               y = f[x2 - x]
               ic = ic + 1
               fm = y
            }

            q = w@i@j * (fp + fm)
            p = p + q
            j = j + 1
         } while (abs[q] > eps * abs[p] and j < m@i)

         err = 2 s
         s = s + p
         err = abs[err - s]
         i = i + 1
      } while (err > tol * abs[s] and i < n)

      // When the integral is smaller than precision, the absolute error is
      // evaluated
      if abs[s] > 1
         err = err / abs[s]

      if debug
         println["err is $err, iteration count is $ic"]
      
      return d * s * 2^(1-i) signCorrection
   }

   /** Initialize the abcissae and weights */
   class initialize[] :=
   {
      h = 2
      for i = 0 to n-1
      {
         h = h / 2
         eh = exp[h]
         t = eh
         r@i = new array[m@i]
         w@i = new array[m@i]
         if i > 0
            eh = eh * eh

         for j= 0 to m@i - 1
         {
            t1 = 1 / t
            u = exp[t1 - t]
            u1 = 1 / (1 + u)
            d = 2 * u * u1
            if d == 0
               break
            r@i@j = d
            w@i@j = (t1 + t) * d * u1
            t = t * eh
         }
      }
      
      return true
   }
}

/** This class integrates numerically using the advanced algorithm in

    Improving the Double Exponential Quadrature Tanh-Sinh, Sinh-Sinh and
    Exp-Sinh Formulas, Dr. Robert A. van Engelen, Genivia Labs,

    https://www.genivia.com/files/qthsh.pdf
*/
class EngelenIntegrator
{
   class var debug = false
   
   // integrate function f, range a..b, max steps, error tolerance eps
   class NIntegrate[f, a, b, eps=1e-14, steps=6, data=undef] :=
   {
      tol = 10 eps
      c = 0
      d = 1
      h = 2
      sign = 1
      v = undef
      
      k = 0
      mode = 0    // Tanh-Sinh = 0, Exp-Sinh = 1, Sinh-Sinh = 2
      if isfinite[a] and isfinite[b] and b < a  // Friendly swap of bounds
      {
         v = b
         [a,b,sign] = [b,a,-1]
      }

      if isfinite[a] and isfinite[b]
      {
         // Both a and b are finite, mode 0
         c = (a+b)/2
         d = (b-a)/2
         v = c
      } else
      if isfinite[a]
      {
         // a is finite, b is +infinity
         mode = 1 // Exp-Sinh
         // alternatively d = exp_sinh_opt_d(f, a, eps, d)
         d = exp_sinh_opt_d[f, a, eps, d, data]
         c = a
         v = a+d
      } else
      if isfinite[b]
      {
         // a is -infinity, b is finite
         mode = 1 // Exp-Sinh
         //d = -d   // alternatively d = exp_sinh_opt_d(f, b, eps, -d)
         d = exp_sinh_opt_d[f, b, eps, -d, data]
         sign = -sign
         c = b
         v = b+d
      } else
      {
         // Both a and b are infinite
         mode = 2 // Sinh-Sinh
         v = 0
      }

      // println["got here, mode=$mode, v=$v"]
      if data == undef
         s = applySafe[f, v]
      else
         s = applySafe[f, [v,data]]
      
      // println["and here, s=$s"]
      if ! isfinite[s]
      {
         println["Error: this integral cannot be peformed because its center point $v is infinite."]
         // Try to do it as the sum of two one-sided integrals.
         return EngelenIntegrator.NIntegrate[f, a, v, eps, steps, data] + EngelenIntegrator.NIntegrate[f, v, b, eps, steps, data]
      }
      
      do
      {
         p = 0
         fp = 0
         fm = 0
         h = h/2
         t = eh = exp[h]
         if k > 0
            eh = eh * eh
         if mode == 0
         {
            // Tanh-Sinh
            do
            {
               u = exp[1/t-t]      // = exp(-2*sinh(j*h)) = 1/exp(sinh(j*h))^2
               r = 2*u/(1+u)       // = 1 - tanh(sinh(j*h))
               w = (t+1/t)*r/(1+u) // = cosh(j*h)/cosh(sinh(j*h))^2
               x = d*r
               
               if (a+x > a)
               {
                  // if too close to a then reuse previous fp
                  arg = a+x
                  if data == undef
                     y = applySafe[f, arg]
                  else
                     y = applySafe[f, [arg, data]]
                  
                  if isfinite[y]
                     fp = y // if f(x) is finite, add to local sum
               }

               if b-x < b
               {
                  // if too close to a then reuse previous fp
                  arg = b-x
                  if data == undef
                     y = applySafe[f, arg]
                  else
                     y = applySafe[f, [arg, data]]
                  
                  if isfinite[y]
                     fm = y // if f(x) is finite, add to local sum
               }
               //println["w=$w, fp=$fp, fm=$fm"]
               q = w*(fp+fm)
               p = p+q
               t = t * eh
               //println["First loop, q=$q, p=$p"]
            } while abs[q] > eps*abs[p]
         } else
         {
            t = t / 2
            do
            {
               r = exp[t-(1/4)/t] // = exp(sinh(j*h))
               w = r
               q = 0
               if mode == 1
               {
                  // Exp-Sinh
                  x = c + d/r
                  if (x == c) // if x hit the finite endpoint then break
                     break
                  if data == undef
                     y = applySafe[f, x]
                  else
                     y = applySafe[f, [x, data]]
                  
                  if isfinite[y] // if f(x) is finite, add to local sum
                     q = q + y/w
               } else
               {
                  // Sinh-Sinh
                  r = (r-1/r)/2 // = sinh(sinh(j*h))
                  w = (w+1/w)/2 // = cosh(sinh(j*h))
                  x = c - d*r

                  if data == undef
                     y = applySafe[f, x]
                  else
                     y = applySafe[f, [x, data]]
                  
                  if isfinite[y] // if f(x) is finite, add to local sum
                     q = q + y*w
               }

               x = c + d*r

               if data == undef
                  y = applySafe[f, x]
               else
                  y = applySafe[f, [x, data]]
               
               if isfinite[y] // if f(x) is finite, add to local sum
                  q = q + y*w
               q = q * (t+(1/4)/t) // q *= cosh(j*h)
               p = p + q
               t = t * eh
               //println["q=$q, p=$p"]
            } while abs[q] > eps*abs[p]
         }
         v = s - p
         s = s + p
         k = k + 1
         //println["v=$v, s=$s"]
      } while abs[v] > tol*abs[s] and k <= steps

      err = abs[v]/(abs[s]+eps)   // This is estimated error
      if debug
      {
         println["err is $err"]
         if err > eps
            println["Warning: integration error $err is greater than epsilon $eps.  Integration may not have converged.  Consider reducing epsilon or taking more steps?"]
      }
      return sign d s h // result with estimated relative error e
   }

   class NIntegrateData[f, a, b, data, eps=1e-14, steps=6] :=
   {
      return NIntegrate[f, a, b, eps, steps, data]
   }

   // Returns optimized Exp-Sinh integral split point d
   class exp_sinh_opt_d[f, a, eps, d, data] :=
   {
      if data == undef
      {
         fa = applySafe[f, a + d/2]
         fb = applySafe[f, a + d*2]
      } else
      {
         fa = applySafe[f, [a + d/2, data]]
         fb = applySafe[f, [a + d*2, data]]
      }
      
      h2 = fa - fb * 4
      
      i = 1
      j = 32    // j=32 is optimal to find r
      if isfinite[h2] and abs[h2] > 1e-5            // if |h2| > 2^-16
      {
         s = 0
         lr = 2
         do
         {
            // find max j such that fl and fr are finite
            j = j / 2
            r = shiftLeft[1, i + j]

            if data == undef
            {
               fl = applySafe[f, a + d/r]
               fr = applySafe[f, a + d*r] r^2
            } else
            {
               fl = applySafe[f, [a + d/r, data]]
               fr = applySafe[f, [a + d*r, data]] r^2
            }
            
            h = fl - fr
         } while j > 1 and !isfinite[h]
         
         if j > 1 and isfinite[h] and realSignum[h] != realSignum[h2]
         {
            lfl = fl // last fl=f(a+d/r)
            lfr = fr // last fr=f(a+d*r)*r*r
            do
            {
               // bisect in 4 iterations
               j = j / 2
               r = shiftLeft[1, i + j]

               if data == undef
               {
                  fl = applySafe[f, a + d/r]
                  fr = applySafe[f, a + d*r] r^2
               } else
               {
                  fl = applySafe[f, [a + d/r, data]]
                  fr = applySafe[f, [a + d*r, data]] r^2
               }
               
               h = fl - fr
               if isfinite[h]
               {
                  s = s + abs[h] // sum |h| to remove noisy cases
                  if realSignum[h] == realSignum[h2]
                     i = i + j // search right half
                  else
                  {
                     // search left half
                     lfl = fl   // record last fl=f(a+d/r)
                     lfr = fr   // record last fl=f(a+d*r)*r*r
                     lr = r     // record last r
                  }
               }
            } while j > 1

            if s > eps
            {
               // if sum of |h| > eps
               h = lfl - lfr // use last fl and fr before the sign change
               r = lr        // use last r before the sign change
               if (h != 0)   // if last diff != 0, back up r by one step
                  r = r / 2

               if abs[lfl] < abs[lfr]
                  d = d / r  // move d closer to the finite endpoint
               else
                  d = d * r  // move d closer to the infinite endpoint
            }
         }
      }
//      println["optimized d is $d"]
      return d
   }
   
   // Determines if a result is infinite for the purposes of integration.
   class isfinite[expr] :=
   {
      if isUnit[expr]
         return true
      else
         return false
   }
}