import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as mt
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_gradient(ax, x, y, c):
    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)
    cmin = ccs[np.argmin(mses)]
    gradient = np.diff(mses)/(ccs[1]-ccs[0])
        
    ax.plot([cmin, cmin], [-10000, 10000], **lsSpine)
    ax.plot([ccs[0], ccs[-1]], [0, 0], **lsSpine)
    ax.plot(ccs[:-1], gradient, **lsBm)
    ax.set_xlabel('c')
    ax.set_ylabel('Derivative')
    ax.set_xlim(0, 10)
    ax.set_ylim(-10000, 10000)
    ax.set_xticks(np.arange(0.0, 10.1, 2.0))
    ax.set_yticks(np.arange(-10000, 10001, 10000))
    ax.set_yticklabels(['', '0', ''])

    
def plot_mse(ax, x, y, c):
    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)
    cmin = ccs[np.argmin(mses)]
    gradient = np.diff(mses)/(ccs[1]-ccs[0])

    ay = 1500.0
    asB = dict(arrowprops=dict(arrowstyle="->", shrinkA=0, shrinkB=0,
                               color=lsB['color'], lw=2))
    ax.annotate('', xy=(3.0, ay), xytext=(1.0, ay), **asB)
    ax.annotate('', xy=(5.0, ay), xytext=(3.8, ay), **asB)
    ax.annotate('', xy=(6.2, ay), xytext=(7.4, ay), **asB)
    ax.annotate('', xy=(8.0, ay), xytext=(10.0, ay), **asB)
    ax.plot([cmin, cmin], [0, 30000], **lsSpine)
    ax.plot(ccs, mses, zorder=10, **lsAm)
    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))
    ax.set_yticklabels(['0', '', '', ''])


if __name__ == "__main__":
    x, y, c = create_data()
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=cm_size(figure_width, 1.1*figure_height))
    fig.subplots_adjust(wspace=0.5, **adjust_fs(left=5.0, right=1.2))
    plot_gradient(ax1, x, y, c)
    plot_mse(ax2, x, y, c)
    fig.savefig("cubicgradient.pdf")
    plt.close()