"""Cobweb plots
Adapted from https://scipython.com/blog/cobweb-plots/
"""

import numpy as np
import matplotlib
import matplotlib.pyplot as plt

# Use LaTeX throughout the figure for consistency
matplotlib.rc('font', **{'family': 'serif', 'serif': ['Computer Modern'], 'size': 16})
matplotlib.rc('text', usetex=True)
# Figure dpi
dpi = 72

def plot_cobweb(f, r, x0, nmax=40, xlim=(0,1), ylim=(0,1)):
    """Make a cobweb plot.

    Plot y = f(x; r) and y = x for 0 <= x <= 1, and illustrate the behaviour of
    iterating x = f(x) starting at x = x0. r is a parameter to the function.

    """
    xmin, xmax = xlim
    x = np.linspace(xmin, xmax, 500)
    fig = plt.figure(figsize=(600/dpi, 450/dpi), dpi=dpi)
    ax = fig.add_subplot(111)

    # Plot y = f(x) and y = x
    ax.plot(x, f(x, r), c='#444444', lw=2)
    ax.plot(x, x, c='#444444', lw=2)

    # Iterate x = f(x) for nmax steps, starting at (x0, 0).
    px, py = np.empty((2,nmax+1,2))
    px[0], py[0] = x0, 0
    for n in range(1, nmax, 2):
        px[n] = px[n-1]
        py[n] = f(px[n-1], r)
        px[n+1] = py[n]
        py[n+1] = py[n]

    # Plot the path traced out by the iteration.
    ax.plot(px, py, c='b', alpha=0.7)

    # Annotate and tidy the plot.
    ax.minorticks_on()
    ax.grid(which='minor', alpha=0.5)
    ax.grid(which='major', alpha=0.5)
    ax.set_ylim(*ylim)
    ax.set_aspect('equal')
    ax.set_xlabel('$x$')
    ax.set_ylabel(f.latex_label)
    ax.axvline(c='grey')
    ax.axhline(c='grey')
    ax.set_title('$x_0 = {:.1}, {} = {:.2}$'.format(x0, f.par_name, r))

    plt.savefig('{}_x0_{:.1}_r_{:.2}.png'.format(f.map_name, x0, r), dpi=dpi)

class AnnotatedFunction:
    """A small class representing a mathematical function.

    This class is callable so it acts like a Python function, but it also
    defines a string giving its latex representation.

    """

    def __init__(self, func, latex_label, map_name, par_name):
        self.func = func
        self.latex_label = latex_label
        self.map_name = map_name
        self.par_name = par_name

    def __call__(self, *args, **kwargs):
        return self.func(*args, **kwargs)

# Example:

map1 = AnnotatedFunction(lambda x, mu: mu + x * (1 - x), 'Mapa 1', 'mapa1', r'\mu')

plot_cobweb(map1, 0.16, 0.3, xlim=(0.25, 0.45), ylim=(0.25, 0.45))