import numpy as np
import scipy.stats as st
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from plotstyle import *

fig = plt.figure()
spec = gridspec.GridSpec(nrows=1, ncols=2, wspace=0.3,
                         **adjust_fs(fig, left=5.5, top=1))
spec1 = gridspec.GridSpecFromSubplotSpec(1, 2, spec[0, 0], width_ratios=[3, 1], wspace=0.0)
spec2 = gridspec.GridSpecFromSubplotSpec(1, 2, spec[0, 1], width_ratios=[3, 1], wspace=0.0)

# the line:
slope = 2.0
xx = np.arange(0.0, 4.1, 0.1)
yy = slope*xx
# the data:
n = 40
rng = np.random.RandomState(5218)
sigma = 1.5
x = 4.0*rng.rand(n)
y = slope*x+rng.randn(n)*sigma
# fit:
slopef = np.sum(x*y)/np.sum(x*x)
yf = slopef*xx

# plot it:
ax = fig.add_subplot(spec1[0, 0])
ax.set_xticks(np.arange(0.0, 4.1))
ax.set_xlim(0.0, 4.2)
ax.set_ylim(-4.0, 12.0)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.plot(x, y, label='data', zorder=10, **psAm)
ax.plot(xx, yy, label='original', zorder=5, **lsB)
ax.plot(xx, yf, label='fit', zorder=7, **lsCm)
ax.legend(loc='upper left', bbox_to_anchor=(-0.05, 1.2))

ax = fig.add_subplot(spec1[0, 1])
ax.show_spines('l')
ax.set_xticks([])
ax.set_ylim(-4.0, 12.0)
ax.set_yticks([])
bins = np.arange(-4.0, 12.1, 0.75)
ax.hist(y, bins, orientation='horizontal', zorder=10, **fsA)

ax = fig.add_subplot(spec2[0, 0])
ax.set_xticks(np.arange(0.0, 4.1))
ax.set_xlim(0.0, 4.2)
ax.set_ylim(-4.0, 12.0)
ax.set_xlabel('x')
ax.set_ylabel('y - mx')
ax.plot(x, y - slopef*x, label='residuals', zorder=10, **psAm)
#ax.legend(loc='upper left', bbox_to_anchor=(0.0, 1.0))

ax = fig.add_subplot(spec2[0, 1])
ax.show_spines('l')
ax.set_xlim(0.0, 11.0)
ax.set_xticks([])
ax.set_ylim(-4.0, 12.0)
ax.set_yticks([])
r = y - slopef*x
ax.hist(r, bins, orientation='horizontal', zorder=10, **fsA)
gx = np.arange(-4.0, 12.1, 0.1)
gy = st.norm.pdf(gx, np.mean(r), np.std(r))
ax.plot(1.0+gy*29.0, gx, zorder=5, **lsBm)

plt.savefig('mlepropline.pdf')