import matplotlib.pyplot as plt
import numpy as np
from matplotlib.transforms import Bbox

def create_data():
    m = 0.75
    n=  -40
    x = np.arange(5.,115., 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.scatter(x, y, marker='o', color='dodgerblue', s=40)
    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.tick_params(direction="out", width=1.25)
    ax.tick_params(direction="out", width=1.25)
    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))
    fig = plt.gcf()
    fig.set_facecolor("white")
    fig.set_size_inches(3., 3.)
    plt.tight_layout()
    fig.savefig("lin_regress.pdf")
    plt.close()
    

def plot_data_slopes(x, y, m, n):
    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.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.tick_params(direction="out", width=1.25)
    ax.tick_params(direction="out", width=1.25)
    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))
    fig = plt.gcf()
    fig.set_facecolor("white")
    fig.set_size_inches(3., 3.)
    plt.tight_layout()
    fig.savefig("lin_regress_slope.pdf")
    plt.close()


def plot_data_intercepts(x, y, m, n):
    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.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.tick_params(direction="out", width=1.25)
    ax.tick_params(direction="out", width=1.25)
    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))
    fig = plt.gcf()
    fig.set_facecolor("white")
    fig.set_size_inches(3., 3.)
    plt.tight_layout()
    fig.savefig("lin_regress_intercept.pdf")
    plt.close()


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