#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
Implementation of a Monte-Carlo method for a 2D integral.

A 3D plot is displayed while the volume is computed.

Note: thanks to Sreekar K (B3 group) for his initial program (which was already almost perfect!).


@date: Sat Mar 18 12:21:29 2015.
@author: Lilian Besson for CS101 course at Mahindra Ecole Centrale 2015.
@licence: MIT Licence (http://lbesson.mit-license.org).
"""


import random
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


# %% 2D Monte-Carlo:

def int_2d(f, a, b, g, h, n=1000):
    r""" Compute an approximation of the volume $\iint_{a < x < b, g(x) < y < h(x)} f(x, y) dx dy$.

    This Monte-Carlo method starts by computing an acceptable range for y values ([ymin, ymax]) then for z values ([zmin, zmax]).
    Then it will take n random points, uniformly taken in the 3D box: [a, b] x [ymin, ymax] x [zmin, zmax].

    These random samples are then used to evaluate the volume delimited by the 3D analytic surface z = f(x, y).
    """
    y_max = max([max(h(i), g(i)) for i in np.linspace(a, b, n)])
    y_min = min([min(h(i), g(i)) for i in np.linspace(a, b, n)])

    ax = plt.axes(projection = "3d")

    X = np.linspace(a, b, n/10)
    Y = np.linspace(y_min, y_max, n/10)
    # Plot the integration domain on the xy-plane
    ax.plot(X, g(X), X*0, linewidth=3, color='black')
    ax.plot(X, h(X), X*0, linewidth=3, color='magenta')

    X, Y = np.meshgrid(X, Y)
    Z = f(X, Y)
    plt.xlabel("Axis x")
    plt.ylabel("Axis y")
    plane = ax.plot_surface(X, Y, Z*0, linewidth=0)

    z_min = 1.02 * min(0, np.min(Z))
    z_max = 1.02 * max(0, np.max(Z))

    print "Sampling domain for x: a, b =", a, b
    print "Sampling domain for y: ymin, ymax =", y_min, y_max, "with the constraints g(x) <= y <= h(x)."
    print "Sampling domain for z: zmin, zmax =", z_min, z_max

    n_pts = 0
    lines = {}

    for _ in xrange(n):
        x = random.uniform(a, b)
        y = random.uniform(y_min, y_max)
        z = random.uniform(z_min, z_max)

        if min(g(x), h(x)) <= y <= max(g(x), h(x)):
            if 0 <= z <= f(x, y):
                n_pts += 1
                lines[0] = ax.plot([x], [y], [z], "go")
            elif f(x, y) <= z < 0:
                n_pts -= 1
                lines[1] = ax.plot([x], [y], [z], "mo")
            else:
                pass
        elif z < 0:
            # Below the curve (negatively below!)
            lines[2] = ax.plot([x], [y], [z], "c+", alpha=0.4)
        else:
            # Not below!
            lines[3] = ax.plot([x], [y], [z], "k+", alpha=0.4)


#    # Write the legend: FIXME do it!
##    import matplotlib.patches as mpatches
#
#    labels = [r'Between xy-plane and the surface ($z \geq 0$)', 'Between the surface and xy-plane  ($z < 0$)', 'Below the surface and the xy-plane', 'Above the surface and the xy-plane']
#
#    listlines = []
#    listlabels = []
#    for i, linei in lines.items():
#        if lines[i]:
##            red_patch = mpatches.Patch(color='red', label='The red data')
#            listlines.append(linei)
#            listlabels.append(labels[i])
#
##    ax.legend(listlines, listlabels)

    # Finally we plot the surface
    ax.plot_surface(X, Y, Z, cmap='hot', linewidth=0, alpha=0.8)

    # And we compute the volume
    box_volume = (b-a)*(y_max-y_min)*(z_max-z_min)
    observed_probability = float(n_pts)/n
    volume = observed_probability * box_volume
    plt.title("Monte-Carlo method, with {} points: approximated volume is {}.".format(n, volume))

#    plt.savefig("Monte-Carlo_method_for_computing_a_2D_integral.png", dpi=70)
    return volume


# %% One example with a simple function
if __name__ == '__main__':

    a, b = -4, 4

    def f(x, y):
        return x**2 + y**2

    def g(x):
        return -x

    def h(x):
        return x

    n = 2000

    print "For f(x, y) = x**2 - y**2"
    print "Domain for x is: [a, b] for a =", a, "b =", b
    print "Domain for y is: [g(x), h(x)] with g(x) = -x, h(x) = x"
    print "With", n, "points, we compute the volume as:"

    volume = int_2d(f, a, b, g, h, n)
    print volume