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

# Solo levanto algunos experimentos
ALL_FILES_SP = """000001523-SingleLine.h5"""

ALL_FILES_DP = """000001525-SingleLine.h5
000001526-SingleLine.h5
000001527-SingleLine.h5
000001528-SingleLine.h5
000001529-SingleLine.h5
000001530-SingleLine.h5"""

def Background_Risetime_AOM(t, bineo, bkgr_in, bkgr_end, rise_time, rise_time_init):
    """simulo un prendido de AOM"""
    i = 0
    j = 1
    Background = []
    while i < len(t):
        if t[i] < rise_time_init:
            Background.append(bkgr_in)
        else:
            rise_step = (bkgr_end-bkgr_in)*bineo/rise_time
            if t[i] < rise_time_init + rise_time:
                Background.append(bkgr_in + j*rise_step)
                j = j + 1
            else:
                Background.append(bkgr_end)
        i = i + 1
    return np.array(Background)


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.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([5, 11, 20, 32, 47, 67, 86, 105, 120, 134, 144])
    return potencias_UV[np.where(amplitudes_UV == amp)][0]

def SubstractUniformBackground(counts, times, timetarget):
    "a partir de time, deberia solo haber background"
    i = min(range(len(times)), key=lambda i: abs(times[i]-timetarget))
    countsbkg = counts[i:]
    background = int(np.mean(countsbkg))
    print(background)
    return background, np.array([c-background for c in counts])

def SubstractDynamicBackground(counts, times, rise_time_init, rise_time, time_initbkgr_target, time_endbkgr_target):
    "a partir de time, deberia solo haber background"

    i_initbkg = min(range(len(times)), key=lambda i: abs(times[i]-time_initbkgr_target))
    i_endbkg = min(range(len(times)), key=lambda i: abs(times[i]-time_endbkgr_target))
    countsbkg_init = counts[0:i_initbkg]
    countsbkg_end = counts[i_endbkg:]
    background_init = int(np.mean(countsbkg_init))
    background_end = int(np.mean(countsbkg_end))
    
    DynamicBkg = Background_Risetime_AOM(times, times[1]-times[0], background_init, background_end, rise_time, rise_time_init)
    
    Counts_Substracted_Bkg = np.array([counts[i] - DynamicBkg[i] for i in range(len(counts))])
    
    return Counts_Substracted_Bkg, DynamicBkg

amplitudes_UV = np.flip(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.flip(np.array([4, 10, 19, 32, 49, 71, 96, 125, 155, 183, 208, 229]))

amplitudes_IR = np.array([0.35, 0.30, 0.25, 0.20, 0.15, 0.10])
potencias_IR = np.array([255, 210, 151, 88, 37, 9])

"""
plt.figure()
plt.plot(amplitudes_UV, potencias_UV, 'ko-', lw=0.2)
plt.xlabel("Amplitud Urukul")
plt.ylabel("Potencia UV/uW")
plt.grid()

plt.figure()
plt.plot(amplitudes_IR, potencias_IR, 'bo-', lw=0.2)
plt.xlabel("Amplitud Urukul")
plt.ylabel("Potencia IR/uW")
plt.grid()
"""

#%%

## Mostrar corte de los histos:
# fig, ax = plt.subplots()
# ax.axvline(T0, color='k')
os.chdir('/home/oem/Documentos/Doctorado/Artiq/Repositorio/artiq_experiments/artiq_master/results/2021-07-02/17')


BINW = 25e-9
T0 = 0.01e-6
#T0 = 0.02e-6

figSP, axSP = plt.subplots()
#axSP = axSP.twinx()

for i, fname in enumerate(ALL_FILES_SP.split()):
    
    if i==0:
        data = h5py.File(fname, 'r') # Leo el h5: Recordar que nuestros datos estan en 'datasets'
        
        laser_UV_amp = data['datasets']['laser_UV_amp']
        laser_UV_freq = data['datasets']['laser_UV_freq']
           
        #print(laser_UV)
        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_SP, binsf_SP  = np.histogram(counts, bines[bines>T0])
        
        backgroundtime = 4e-6
        
        
        
        rise_time_init = 1.02
        rise_time = 0.11
        time_initbkgr_target = 0.6
        time_endbkgr_target = 4
        
        CountsSubstSP, DynamicBkg = SubstractDynamicBackground(heigs_SP, binsf_SP, rise_time_init*1e-6, rise_time*1e-6, time_initbkgr_target*1e-6, time_endbkgr_target*1e-6)
        #bckg, CountsSubstSP = SubstractUniformBackground(heigs_SP, binsf_SP[:-1], backgroundtime)

        #print('bckg DP:' + str(bckg))
        
        #axSP.plot([b*1e6 for b in binsf_SP], DynamicBkg, label='Rising AOM')            
       
        axSP.plot([b*1e6 for b in binsf_SP[:-1]], CountsSubstSP, label='Background substracted')
           
        axSP.plot([b*1e6 for b in binsf_SP[:-1]], heigs_SP, zorder=1)    
        axSP.set_title('SP Transition')
        axSP.set_xlabel('Time (us)')
        axSP.set_ylabel('Counts')
        #axSP.hlines(bckg, 4, 6, color='red', zorder=2, label='Background')
        axSP.axvline(0.97, color='blue', linestyle='--', linewidth=0.5, label='UV turns on')
        axSP.axvspan(0, 0.97, color='lightcyan')
        axSP.legend()
        
        TotalCountsDetected_SP = np.sum(CountsSubstSP[int(0.*len(heigs_SP)):int(0.6*len(heigs_SP))])

#%%
figDP, axDP = plt.subplots()


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

os.chdir('/home/oem/Documentos/Doctorado/Artiq/Repositorio/artiq_experiments/artiq_master/results/2021-07-02/18')        
for i, fname in enumerate(ALL_FILES_DP.split()):
    
    if i==0:
        data = h5py.File(fname, 'r') # Leo el h5: Recordar que nuestros datos estan en 'datasets'
        
        laser_IR_amp = data['datasets']['laser_IR_amp']
        laser_IR_freq = data['datasets']['laser_IR_freq']
           
        #print(laser_UV)
        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_DP, binsf_DP  = np.histogram(counts, bines[bines>T0])
        
        backgroundtime = 3e-6

        bckg, CountsSubstDP = SubstractUniformBackground(heigs_DP, binsf_SP[:-1], backgroundtime)
        
        print('bckg DP:' + str(bckg))
        
        axDP.plot([b*1e6 for b in binsf_DP[:-1]], heigs_DP, zorder=1)    
        axDP.plot([b*1e6 for b in binsf_DP[:-1]], CountsSubstDP, label='Background substracted')
        axDP.set_title('DP Transition')
        axDP.set_xlabel('Time (us)')
        axDP.set_ylabel('Counts')
        axDP.hlines(bckg, 4, 6, color='red', zorder=2, label='Background')
        axDP.axvline(0.64, color='red', linestyle='--', linewidth=0.5, label='IR turns on')
        axDP.axvspan(0.4, 0.64, color='papayawhip')
        axDP.legend()
        
        TotalCountsDetected_DP= np.sum(CountsSubstDP[0:int(0.2*len(heigs_DP))])
        
print(TotalCountsDetected_DP)        
print('BR: ', TotalCountsDetected_SP/TotalCountsDetected_DP)
#%%
figBOTH, axBOTH = plt.subplots()
axBOTH.plot([1e6*t+0.32 - 1 for t in binsf_DP[:-1]], SubstractUniformBackground(heigs_DP, binsf_DP[:-1], backgroundtime), label='DP transition')        
axBOTH.plot([1e6*t - 1 for t in binsf_SP[:-1]], SubstractUniformBackground(heigs_SP, binsf_SP[:-1], backgroundtime), label='SP transition')        
axBOTH.set_xlim(-0.1, 3)
axBOTH.set_xlabel('Time from pulse (us)')
axBOTH.set_ylabel('Counts')
axBOTH.set_title('PotUV: 229 uW, PotIR: 255 uW. Branching Fraction: 14.1')
axBOTH.legend()