import matplotlib.pyplot as plt
import numpy as np
from IPython import embed

def create_data():
    m = 0.75
    n=  -30
    x = np.arange(0.,101., 2.5)
    y = m * x + n;
    noise = np.random.randn(len(x))*15
    y += noise
    return x, y, m, n

    
def plot_data(x, y):
    plt.xkcd()
    plt.scatter(x, y, marker='o', color='dodgerblue', s=40)
    plt.xlabel("Input x")
    plt.ylabel("Output y")
    plt.xlim([-2.5, 102.5])
    ax = plt.gca()
    ax.spines["right"].set_visible(False)
    ax.spines["top"].set_visible(False)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')
    ax.xaxis.linewidth=1.5
    ax.yaxis.linewidth=1.5
    ax.tick_params(direction="out", width=1.25)
    ax.tick_params(direction="out", width=1.25)
    fig = plt.gcf()
    fig.set_facecolor("white")
    fig.set_size_inches(3., 3.)
    fig.savefig("figures/lin_regress.pdf")
    plt.close()
    

def plot_data_slopes(x, y, m, n):
    plt.xkcd()
    plt.scatter(x, y, marker='o', color='dodgerblue', s=40)
    for i in np.linspace(m/4, m*1.5, 5):
        plt.plot(x, i*x+n, color="r", lw=2)

    plt.xlabel("Input x")
    plt.ylabel("Output y")
    plt.xlim([-2.5, 102.5])
    ax = plt.gca()
    ax.spines["right"].set_visible(False)
    ax.spines["top"].set_visible(False)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')
    ax.xaxis.linewidth=1.5
    ax.yaxis.linewidth=1.5
    ax.tick_params(direction="out", width=1.25)
    ax.tick_params(direction="out", width=1.25)
    fig = plt.gcf()
    fig.set_facecolor("white")
    fig.set_size_inches(3., 3.)
    fig.savefig("figures/lin_regress_slope.pdf")
    plt.close()


def plot_data_intercepts(x, y, m, n):
    plt.xkcd()
    plt.scatter(x, y, marker='o', color='dodgerblue', s=40)
    for i in np.linspace(n-n/2, n+n/2, 5):
        plt.plot(x, m * x + i, color="r", lw=2)

    plt.xlabel("Input x")
    plt.ylabel("Output y")
    plt.xlim([-2.5, 102.5])
    ax = plt.gca()
    ax.spines["right"].set_visible(False)
    ax.spines["top"].set_visible(False)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')
    ax.xaxis.linewidth=1.5
    ax.yaxis.linewidth=1.5
    ax.tick_params(direction="out", width=1.25)
    ax.tick_params(direction="out", width=1.25)
    fig = plt.gcf()
    fig.set_facecolor("white")
    fig.set_size_inches(3., 3.)
    fig.savefig("figures/lin_regress_intercept.pdf")
    plt.close()


if __name__ == "__main__":
    x, y, m, n = create_data()
    plot_data(x,y)
    plot_data_slopes(x,y,m,n)
    plot_data_intercepts(x,y,m,n)