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

def boltzmann(x, x0, k):
    return 8.0/(1.0+np.exp(-k*(x-x0)))

if __name__ == "__main__":
    n = 50
    xmin = -18.0
    xmax = 18.0
    x0 = 2.0
    k = 0.25
    sigma = 0.6
    rng = np.random.RandomState(15281)
    x = np.linspace(xmin, xmax, n)
    y = boltzmann(x, x0, k) + sigma*rng.randn(len(x))
    xx = np.linspace(xmin, xmax, 200)
    yy = boltzmann(xx, x0, k)

    fig = plt.figure()
    spec = gridspec.GridSpec(nrows=1, ncols=2, left=0.1, wspace=0.4)
    ax1 = fig.add_subplot(spec[0, 0])
    show_spines(ax1, 'lb')
    ax1.plot(xx, yy, colors['red'], lw=2)
    ax1.scatter(x, y, c=colors['blue'], edgecolor='white', s=50)
    set_xlabel(ax1, 'Hair deflection', 'nm')
    set_ylabel(ax1, 'Conductance', 'nS')
    ax1.set_xlim(-20, 20)
    ax1.set_ylim(-1.5, 9.5)
    ax1.set_xticks(np.arange(-20.0, 21.0, 10.0))
    ax1.set_yticks(np.arange(0, 9, 2))
    
    ax2 = fig.add_subplot(spec[0, 1])
    show_spines(ax2, 'lb')
    xg = np.linspace(-3.0, 3.01, 200)
    yg = st.norm.pdf(xg, 0.0, sigma)
    ax2.plot(xg, yg, colors['red'], lw=2)
    bw = 0.25
    h, b = np.histogram(y-boltzmann(x, x0, k), np.arange(-3.0, 3.01, bw))
    ax2.bar(b[:-1], h/np.sum(h)/(b[1]-b[0]), fc=colors['yellow'], width=bar_fac*bw, align='edge')
    set_xlabel(ax2, 'Residuals', 'nS')
    set_ylabel(ax2, 'Pdf', '1/nS')
    ax2.set_xlim(-2.5, 2.5)
    ax2.set_ylim(0, 0.75)
    ax2.set_yticks(np.arange(0, 0.75, 0.2))

    fig.savefig("staticnonlinearity.pdf")
    plt.close()