import h5py
import matplotlib.pyplot as plt
import numpy as np
import sys
import re
import ast
from scipy.optimize import curve_fit

# Solo levanto algunos experimentos
ALL_FILES = """000000704-SingleLine.h5
000000702-SingleLine.h5
000000706-SingleLine.h5
000000708-SingleLine.h5
000000712-SingleLine.h5
000000714-SingleLine.h5
000000716-SingleLine.h5
000000718-SingleLine.h5"""

BINW = 20e-9
T0 = 1.15e-6

def expo(T, tau, N0, C):
    global T0
    return N0*np.exp(-(T-T0)/tau) + C

def pow_from_amp(amp):
    """Paso de amplitud urukul a potencia medida por Nico"""
    # Forma altamente ineficiente de hacer esto, pero me salio asi
    amplitudes_UV = np.array([0.08, 0.10, 0.12, 0.14, 0.16, 0.18, 0.20, 0.22, 0.24, 0.26, 0.28, 0.30])
    assert amp in amplitudes_UV
    potencias_UV = np.array([2, 5, 11, 20, 32, 47, 67, 86, 105, 120, 134, 144])
    return potencias_UV[np.where(amplitudes_UV == amp)][0]

amplitudes_UV = np.array([0.08, 0.10, 0.12, 0.14, 0.16, 0.18, 0.20, 0.22, 0.24, 0.26, 0.28, 0.30])
potencias_UV = np.array([2, 5, 11, 20, 32, 47, 67, 86, 105, 120, 134, 144])
plt.plot(amplitudes_UV, potencias_UV, 'ko-', lw=0.2)
plt.xlabel("Amplitud Urukul")
plt.ylabel("Potencia /uW")
plt.grid()
#%%

## Mostrar corte de los histos:
# fig, ax = plt.subplots()
# ax.axvline(T0, color='k')

fig0, [ax0, ax1_a] = plt.subplots(1, 2)
ax1_b = ax1_a.twinx()

allamps = np.array([])
allpows = np.array([])
alltaus = np.array([])
allN0 = np.array([])

for i, fname in enumerate(ALL_FILES.split()):
    data = h5py.File(fname, 'r') # Leo el h5: Recordar que nuestros datos estan en 'datasets'

    # Aca hago algo repugnante para poder levantar los strings que dejamos
    # que además tenian un error de tipeo al final. Esto no deberá ser necesario
    # cuando se solucione el error este del guardado.
    laser_UV = ast.literal_eval(data['datasets']['laser_UV'][()].decode().replace('"}', '}'))

    counts = np.array(data['datasets']['counts'])
    bines = np.arange(counts.min(), counts.max()+BINW, BINW)

    ## Mostrar corte de los histos
    # ax.hist(counts, bines[bines<3e-6], histtype='step', align='mid', color=f"C{i}")

    heigs, binsf  = np.histogram(counts, bines[bines>T0])
    ax0.step(binsf[:-1], heigs, label=f"AMP: {laser_UV['amp']}", where='mid',
            color=f"C{i}", lw=0.5, alpha=0.4)

    popt, pcov = curve_fit(expo, binsf[:-1], heigs, p0=(2e-6, 400, 300))
    print(popt) # tau, N0, C
    ax0.plot(binsf, expo(binsf, *popt), label=f"tau: {popt[0]}",
            color=f"C{i}", ls='-', lw=1, zorder=99)

    allamps = np.append(allamps, laser_UV['amp'])
    laser_pow = pow_from_amp(laser_UV['amp'])
    allpows = np.append(allpows, laser_pow)
    alltaus = np.append(alltaus, popt[0])
    allN0 = np.append(allN0, popt[1])

    ax1_a.plot(laser_pow , popt[0], 'o', color=f"C{i}", ms=5, )
    ax1_b.plot(laser_pow, popt[1], '^', color=f"C{i}", ms=7, )

ax1_a.plot(allpows, alltaus, 'k-', lw=0.2, zorder=0)
ax1_b.plot(allpows, allN0, 'k-', lw=0.2, zorder=0)


# plt.annotate(f"bin: {BINW}", (0,5e-5, 700), fontsize=14)
ax0.set_xlabel("Tiempo")
ax0.set_ylabel("Cuentas")

ax1_a.set_xlabel("Potencia [uW]")
ax1_a.set_ylabel("Tau (circulo)")
ax1_b.set_ylabel("Alturas (triang)")

ax1_a.grid(alpha=0.3)
# plt.show()

plt.figure()
plt.plot(allpows, allN0/alltaus, 'ko-', lw=0.2)
plt.grid(alpha=0.3)
plt.xlabel("Potencia [uW]")
plt.ylabel("Alturas/Tau")

plt.show()
input()