#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Python implementation of Romberg's method for computing 1D integrals,
using a dictionary to store the intermediate results.

@reference: https://en.wikipedia.org/wiki/Romberg%27s_method#Method
@date: Sun Apr 05 15:00:27 2015.
@author: Lilian Besson for CS101 course at Mahindra Ecole Centrale 2015.
@licence: MIT Licence (http://lbesson.mit-license.org).
"""

from __future__ import division
# No issue with integer division ! / is now the mathematically correct /.

from math import fsum


# %% Romberg algorithm, with chosen values for (n, m)

def romberg(f, a, b, n, m, verb=False):
    r""" Compute the R(n, m) value inductively, to approximate \int_a^n f(x) dx.

    - The integrand f must be of class C^(2n+2) for the method to be correct.
    - Note: a memoization trick is too hard to set-up here, as this value R(n, m) depends on f, a and b.
    - Time complexity is O(n × m).
    - Memory complexity is also O(n × m) (using a dictionary to store all the value R(n, m)).
    """
    assert n >= 0
    assert m >= 0
    assert n >= m >= 1
    assert a <= b
    # First value :
    r = { (0, 0): 0.5 * (b - a) * (f(b) + f(a)) }

    # On side of the triangle
    for i in xrange(1, n+1):
        h_i = (b - a) / (2**i)
        xsamples = [ a + ((2*k - 1) * h_i) for k in xrange(1, 1+2**(i-1)) ]
#        if verb:
#            print "For i = {}, h_i is {}, and so we have {} samples.".format(i, h_i, len(xsamples))
        r[(i, 0)] = (0.5 * r[(i-1, 0)]) + h_i * fsum([ f(x) for x in xsamples ])

    # All the other values
    for j in xrange(1, m+1):
        for i in xrange(j, n+1):
            try:
                r[(i, j)] = ( ((4**j) * r[(i, j-1)]) - r[(i-1, j-1)]) / float( (4**j) - 1)
                if verb:
                    print "R({}, {}) = {}".format(i, j, r[(i, j)])
            except:
                print i, j, "was an issue !"

    # Computing the error. FIXME not perfect
    h_n = (b - a) / float(2**n)
    error = h_n ** (2*m+2)
    return r[(n, m)], error


# %% Romberg algorithm, with chosen precision

def romberg_prec(f, a, b, precision=1e-8, verb=False):
    r""" Compute the Romberg interpolation values inductively, to approximate \int_a^n f(x) dx, with a precision at least equal to *precision*.

    - The integrand f must be of class C^(2n+2) for the method to be correct.
    - Note: a memoization trick is too hard to set-up here, as this value R(n, m) depends on f, a and b.
    - Time complexity is O(n × m) for (n, m) the first values such that abs(R(n, m-1), R(n-1, m-1)) < precision .
    - Memory complexity is also O(n × m) (using a dictionary to store all the value R(n, m)).
    - Reference is the French wikipedia page: https://fr.wikipedia.org/wiki/M%C3%A9thode_de_Romberg#Crit.C3.A8re_d.E2.80.99arr.C3.AAt
    """
    # FIXME: do it!


# %% Examples
if __name__ == '__main__':
    from math import sin, cos, tan, log, exp, pi, sqrt

    # Example from https://en.wikipedia.org/wiki/Romberg%27s_method#Example
    f = lambda x: (2.0 / sqrt(pi)) * exp(-x**2)
    erf1, error = romberg(f, 0, 1, 5, 5)
    print r"\int_0^1 exp(-x^2) dx \simeq {}".format(erf1)
    print "The error is less than", error

    # \int_0^pi sin(x) dx is 2
    area, error = romberg(sin, 0, pi, 5, 5)
    print r"\int_0^pi sin(x) dx \simeq {}".format(area)
    print "The error is less than", error


    n = m = 5
    a = 0
    b = 1

    def f1(x):
        """ Quadrant of a circle."""
        return (1-x**2)**0.5

    print "For the function f1: x --> (1-x**2)**0.5, its integral from", a, "to", b, "is:"
    int1 = romberg(f1, a, b, n, m)
    print "approximatively", int1, "with a Romberg method with n =", n, "and m =", m

    def f2(x):
        """ Upward parabola."""
        return x**2

    print "For the function f2: x --> x**2, its integral from", a, "to", b, "is:"
    int2 = romberg(f2, a, b, n, m)
    print int2, "(expected value is 1/3)"
    print "approximatively with a Romberg method with n =", n, "and m =", m

    f3 = sin
    b = pi
    print "For the function f3: x --> sin(x), its integral from", a, "to", b, "is:"
    int3 = romberg(f3, a, b, n, m)
    print int3, "(expected value is 2)"
    print "approximatively with a Romberg method with n =", n, "and m =", m

    f4 = exp
    n, m = 5, 5
    a, b = -4, 19
    print "For the function f4: x --> exp(x), its integral from", a, "to", b, "is:"
    int4 = romberg(f4, a, b, n, m)
    print int4, "(expected value is exp(19) - exp(-4) ~=", f4(b) - f4(a), ")"
    print "approximatively with a Romberg method with n =", n, "and m =", m

    n, m = 10, 10
    a, b = -4, 19
    print "For the function f4: x --> exp(x), its integral from", a, "to", b, "is:"
    int4 = romberg(f4, a, b, n, m)
    # 1000 loops, best of 3: 558 µs per loop
    print int4, "(expected value is exp(19) - exp(-4) ~=", f4(b) - f4(a), ")"
    print "approximatively with a Romberg method with n =", n, "and m =", m

    a, b = -1000, 20
    print "For the function f4: x --> exp(x), its integral from", a, "to", b, "is:"
    int4 = romberg(f4, a, b, n, m)
    # 100 loops, best of 3: 3.1 ms per loop
    print int4, "(expected value is exp(20) - exp(-1000) ~=", f4(b) - f4(a), ")"
    print "approximatively with a Romberg method with n =", n, "and m =", m