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

def fit_plot(Vbias, a1, a2, a3, a4):
    return np.exp(a1 * np.tanh(0.08*Vbias+a2)/a3) - a4

yes = (['y', 'Y', '', 'yes', 'Yes'])
no = (['n', 'N', 'no', 'No'])

shotnostart = 17739
toPlot = np.load("toPlot.npy") #toPlot je numpy 3D pole se osami: x=čas(0-4000), y=napětí(0) nebo proud(1), z=číslo měření
dataLoop = np.load("npzcache/loop_voltage" + str(shotnostart) + ".npz")
loop_voltage = dataLoop['data']


x = np.linspace(-100, 100, 1000)

timePlot = 2000 #modify iterator and range to change output
while timePlot in range(2000,2001):
    while True:

        try:

            xdata = toPlot[timePlot, 0, :]
            y = toPlot[timePlot, 1, :]
            ydata = y + 0.00009 * np.random.normal(size=len(xdata))

            initial_guess = [1.02,-0.45,7,-0.9]
            po, po_cov = spo.curve_fit(fit_plot, xdata, ydata, initial_guess)
            #ff=po[0]*np.exp(po[1] * np.tanh(po[2]*x+po[3])/po[4])-po[5]

            #print(po)
            #print(po_cov)

            ax1 = plt.subplot2grid((5,1),(0,0), rowspan=3)
            ax2 = plt.subplot2grid((6,3),(4,0))

            ax1.axis([-60, 60, -0.1, 0.6])
            ax1.grid()
            ax1.axvline(x=0, ls='-.', color='black')
            ax1.axhline(y=0, ls='-.', color='black')
            ax1.text(-60, -0.435, '$t = <' + str((timePlot * 10) - 10) + ';' + str((timePlot * 10) + 10) + '> [us]$', fontsize=12, color='grey')
            ax1.text(-60, 0.62, 'I = e^(' + str(round(po[0],2)) + '* tanh(0.08 * Vbias + ' + str(round(po[1],2)) + ')/ ' + str(round(po[2],1)) + ') + ' + str(round(po[3],2)), fontsize=10, color='grey')


            ax1.scatter(toPlot[timePlot, 0, :], toPlot[timePlot, 1, :], marker='x', color='red')
            ax1.plot(x, fit_plot(x, po[0], po[1], po[2], po[3]))

            ax2.plot(range(0, 40000), loop_voltage, color='blue')
            ax2.axis([6000, 40000, -1, 22])

            ax2.axvline(x=(timePlot * 10), ls='-.', color='red')
            cur_axes = plt.gca()
            cur_axes.axes.get_xaxis().set_ticklabels([])
            cur_axes.axes.get_yaxis().set_ticklabels([])
            plt.savefig('output/' + str(timePlot) + '.png')

            gc.collect()
            timePlot = timePlot + 1

        except RuntimeError:
            print("adaai")
            continue
        break