import matplotlib.pyplot as plt
import numpy as np
from plotstyle import *

def create_data():
    # wikipedia:
    # Generally, males vary in total length from 250 to 390 cm and
    # weigh between 90 and 306 kg
    c = 6
    x = np.arange(2.2, 3.9, 0.05)
    y = c * x**3.0
    rng = np.random.RandomState(32281)
    noise = rng.randn(len(x))*50
    y += noise
    return x, y, c


def gradient_descent(x, y):
    n = 20
    dc = 0.01
    eps = 0.0001
    cc = 1.1
    cs = []
    mses = []
    for k in range(n):
        m0 = np.mean((y-(cc*x**3.0))**2.0)
        m1 = np.mean((y-((cc+dc)*x**3.0))**2.0)
        dmdc = (m1 - m0)/dc
        cs.append(cc)
        mses.append(m0)
        cc -= eps*dmdc
    return cs, mses

    
def plot_mse(ax, x, y, c, cs):
    ms = np.zeros(len(cs))
    for i, cc in enumerate(cs):
        ms[i] = np.mean((y-(cc*x**3.0))**2.0)
    ccs = np.linspace(0.5, 10.0, 200)
    mses = np.zeros(len(ccs))
    for i, cc in enumerate(ccs):
        mses[i] = np.mean((y-(cc*x**3.0))**2.0)
        
    ax.plot(ccs, mses, zorder=10, **lsAm)
    ax.plot(cs[:12], ms[:12], zorder=20, **psB)
    ax.plot(cs[-1], ms[-1], zorder=30, **psC)
    for i in range(4):
        ax.annotate('',
                    xy=(cs[i+1]+0.2, ms[i+1]), xycoords='data',
                    xytext=(cs[i]+0.3, ms[i]+200), textcoords='data', ha='left',
                    arrowprops=dict(arrowstyle="->", relpos=(0.0,0.0),
                    connectionstyle="angle3,angleA=10,angleB=70") )
    ax.set_xlabel('c')
    ax.set_ylabel('Mean squared error')
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 25000)
    ax.set_xticks(np.arange(0.0, 10.1, 2.0))
    ax.set_yticks(np.arange(0, 30001, 10000))

        
def plot_descent(ax, cs, mses):
    ax.plot(np.arange(len(mses))+1, mses, **lpsBm)
    ax.set_xlabel('Iteration')
    #ax.set_ylabel('Mean squared error')
    ax.set_xlim(0, 12.5)
    ax.set_ylim(0, 25000)
    ax.set_xticks(np.arange(0.0, 12.1, 2.0))
    ax.set_yticks(np.arange(0, 30001, 10000))
    ax.set_yticklabels([])


if __name__ == "__main__":
    x, y, c = create_data()
    cs, mses = gradient_descent(x, y)
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=cm_size(figure_width, 1.1*figure_height))
    fig.subplots_adjust(**adjust_fs(left=8.0, right=1.2))
    plot_mse(ax1, x, y, c, cs)
    plot_descent(ax2, cs, mses)
    fig.savefig("cubicmse.pdf")
    plt.close()