import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize

nt = np.load('ardice_stability_0004.npy')

#plt.plot(nt[:, 0, :].T-nt[:,0,:].min(), nt[:, 1,:].T, 'g')

def model(theta, t):
    A, B, tau = theta
    return B + A * np.exp(-t/tau)

xmin = 2
xmax = 40

for led in range(16):
    plt.subplot(211)
    tstart = np.argmax(nt[led,1,1:]-nt[led,1,:-1])
    print(tstart)
    x = nt[led, 0, :]-nt[led,0,tstart+1]
    y = nt[led, 1,:]
    l = plt.plot(x, y)[0]

    goods = (x > xmin) & (x<xmax)
    def chi2(theta):
        return y[goods] - model(theta, x[goods])

    theta = scipy.optimize.leastsq(chi2, [1e-7, 1e-5, 10.])
    plt.plot(x[goods], model(theta[0], x[goods]), ls='--', color=l.get_color())

    plt.subplot(212)
    plt.plot(x[x>0], (y[x>0] / model(theta[0], x[x>0]))-1, color=l.get_color())
plt.show()