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

def create_data():
    m = 0.75
    n=  -40
    x = np.concatenate( (np.arange(10.,110., 2.5), np.arange(0.,120., 2.0)) )
    y = m * x + n;
    rng = np.random.RandomState(37281)
    noise = rng.randn(len(x))*15
    y += noise
    return x, y, m, n

    
def plot_data(ax, x, y, m, n):
    ax.set_xlabel('Input x')
    ax.set_ylabel('Output y')
    ax.set_xlim(0, 120)
    ax.set_ylim(-80, 80)
    ax.set_xticks(np.arange(0,121, 40))
    ax.set_yticks(np.arange(-80,81, 40))
    ax.annotate('Error',
                xy=(x[34]+1, y[34]+15), xycoords='data',
                xytext=(80, -50), textcoords='data', ha='left',
                arrowprops=dict(arrowstyle="->", relpos=(0.9,1.0),
                connectionstyle="angle3,angleA=50,angleB=-30") )
    ax.plot(x[:40], y[:40], zorder=0, **psAm)
    inxs = [3, 13, 16, 19, 25, 34, 36]
    ax.plot(x[inxs], y[inxs], zorder=10, **psA)
    xx = np.asarray([2, 118])
    ax.plot(xx, m*xx+n, **lsBm)
    for i in inxs :
        xx = [x[i], x[i]]
        yy = [m*x[i]+n, y[i]]
        ax.plot(xx, yy, zorder=5, **lsDm)
    

def plot_error_hist(ax, x, y, m, n):
    ax.set_xlabel('Squared error')
    ax.set_ylabel('Frequency')
    bins = np.arange(0.0, 602.0, 50.0)
    ax.set_xlim(bins[0], bins[-1])
    ax.set_ylim(0, 35)
    ax.set_xticks(np.arange(bins[0], bins[-1], 100))
    ax.set_yticks(np.arange(0, 36, 10))
    errors = (y-(m*x+n))**2.0
    mls = np.mean(errors)
    ax.annotate('Mean\nsquared\nerror',
                xy=(mls, 0.5), xycoords='data',
                xytext=(350, 20), textcoords='data', ha='left',
                arrowprops=dict(arrowstyle="->", relpos=(0.0,0.2),
                connectionstyle="angle3,angleA=10,angleB=90") )
    ax.hist(errors, bins, **fsD)



if __name__ == "__main__":
    x, y, m, n = create_data()
    fig, axs = plt.subplots(1, 2)
    fig.subplots_adjust(**adjust_fs(fig, left=6.0))
    plot_data(axs[0], x, y, m, n)
    plot_error_hist(axs[1], x, y, m, n)
    fig.savefig("linear_least_squares.pdf")
    plt.close()