%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

# use seaborn plotting defaults
# If this causes an error, you can comment it out.
import seaborn as sns
sns.set()

from fig_code import linear_data_sample
x, y, dy = linear_data_sample()
plt.errorbar(x, y, dy, fmt='o');

plt.plot(x, y, 'o');

def model(theta, x):
    b, m = theta
    return m * x + b

def least_abs_deviation(theta, x, y):
    return np.sum(abs(model(theta, x) - y))

def least_squares(theta, x, y):
    return np.sum((model(theta, x) - y) ** 2)

from scipy.optimize import fmin
theta_guess = [0, 1]

theta_LAD = fmin(least_abs_deviation, theta_guess, args=(x, y))
theta_LS = fmin(least_squares, theta_guess, args=(x, y))

xfit = np.linspace(0, 10)

plt.plot(x, y, 'o')
plt.plot(xfit, model(theta_LAD, xfit), label='Least Abs. Dev.')
plt.plot(xfit, model(theta_LS, xfit), label='Least Sq.')
plt.legend(loc='best');

ydiff = np.linspace(-4, 4, 1000)
plt.plot(ydiff, stats.norm.pdf(ydiff, 0, 1))
plt.xlabel('$y - y_{true}$', size=18)
plt.ylabel('$p(y|y_{true})$', size=18);

X = np.vstack([np.ones_like(x), x]).T

theta_hat = np.linalg.solve(np.dot(X.T, X),
                            np.dot(X.T, y))
print(theta_hat)
print(theta_LS)

xfit = np.linspace(0, 10)

plt.plot(x, y, 'o')
plt.plot(xfit, model(theta_hat, xfit), label='Least Sq.')
plt.legend(loc='best');

XT_Sigma_X = np.dot(X.T / dy ** 2, X)
XT_Sigma_y = np.dot(X.T / dy ** 2, y)
theta_LS_err = np.linalg.solve(XT_Sigma_X, XT_Sigma_y)

plt.errorbar(x, y, dy, fmt='o')
plt.plot(xfit, model(theta_LAD, xfit), label='Least Abs. Dev.')
plt.plot(xfit, model(theta_LS, xfit), label='Least Sq.')
plt.plot(xfit, model(theta_LS_err, xfit), label='Fit w/ errors')
plt.legend(loc='best');

from fig_code import linear_data_sample_big_errs
x, y, dy = linear_data_sample_big_errs()
plt.errorbar(x, y, dy, fmt='o');

theta_LAD = fmin(least_abs_deviation, theta_guess, args=(x, y))

X = np.vstack([np.ones_like(x), x]).T

theta_LS = np.linalg.solve(np.dot(X.T, X),
                           np.dot(X.T, y))

theta_LS_err = np.linalg.solve(np.dot(X.T / dy ** 2, X),
                               np.dot(X.T / dy ** 2, y))

plt.errorbar(x, y, dy, fmt='o')
plt.ylim(-20, 40)
plt.plot(xfit, model(theta_LAD, xfit), label='Least Abs. Dev.')
plt.plot(xfit, model(theta_LS, xfit), label='Least Sq.')
plt.plot(xfit, model(theta_LS_err, xfit), label='Least Sq. w/ errors')
plt.legend(loc='best');