__author__ = "Marie E. Rognes (meg@simula.no)"
__copyright__ = "Copyright (C) 2012 Marie Rognes"
__license__  = "GNU LGPL version 3 or any later version"

"""
This module contains functions for solving the nonlinear equation

  f(y) := y + 0.1 y^3 - 1 = 0

using the bisection method or Newton's method.
"""

def bisection(f, a, b, tolerance):

    print "Bisection method:"
    print "a = %g, b = %g, tolerance = %g" % (a, b, tolerance)

    assert (f(a)*f(b) < 0), "Input does not satisfy ansatz!"

    c = 0.5*(a + b)
    k = 1

    points = [c, ]
    values = [f(c), ]

    while (abs(f(c)) > tolerance):

        if f(a)*f(c) < 0:
            b = c
        else:
            a = c
        c = 0.5*(a + b)

        points += [c]
        values += [f(c)]
        k += 1

    return points, values

def newton(f, df, y, tolerance):
    print "Newton's method:"
    print "y = %g, tolerance = %g" % (y, tolerance)

    k = 1
    c = y

    points = [c, ]
    values = [f(c), ]
    while (abs(f(c)) > tolerance):
        c = c - f(c)/df(c)

        points += [c]
        values += [f(c)]
        k += 1

    return points, values


def print_table(output_format, points, values):

    if output_format == "plain":
        for (k, c) in enumerate(points):
            print "k = %d,\tc = %.10f,\tf(c) = %.10f" % (k, c, values[k])
    elif output_format == "latex":
        print "\\begin{center}"
        print "\\begin{tabular}{ccc}"
        print "\\toprule"
        print "$k$ & $y_k$ & $f(y_k)$ \\\\"
        print "\\midrule"
        for (k, c) in enumerate(points):
            print "%d & %.10g &  %.10g \\\\" % (k, c, values[k])
        print "\\bottomrule"
        print "\\end{tabular}"
        print "\\end{center}"

    else:
        print "Unknown output format: %s" % (output_format)


if __name__ == "__main__":

    # Define f (want f(x) == 0)
    def f(x):
        return x + 0.1*x**3 - 1.

    # The derivative of f
    def df(x):
        return 1 + 3*0.1*x**2

    # Plot f versus y for different ranges of y
    import pylab
    ys = pylab.linspace(-6, 6, 1000)
    fs = [f(y) for y in ys]
    pylab.plot(ys, fs)
    pylab.plot(ys, [0 for y in ys])
    pylab.grid(True)
    pylab.xlabel("y")
    pylab.ylabel("f")
    pylab.title("f versus y on [-6, 6]")
    pylab.savefig("f_vs_y.pdf")

    pylab.figure()
    ys = pylab.linspace(0, 2, 1000)
    fs = [f(y) for y in ys]
    pylab.plot(ys, fs)
    pylab.plot(ys, [0 for y in ys])
    pylab.grid(True)
    pylab.xlabel("y")
    pylab.ylabel("f")
    pylab.title("f versus y on [0, 2]")
    pylab.savefig("f_vs_y_zoom.pdf")

    # Compute zeros (solutions) using the bisection method
    points, values = bisection(f, 0., 2., 1.e-8)
    print_table("latex", points, values)

    # Compute zeros using Newton's method
    points, values = newton(f, df, 2, 1.e-8)
    print_table("latex", points, values)

    #pylab.show()