# -*- coding: utf-8 -*-
"""
@author: Timon Schapeler and Niklas Lamberty

Script to load in the raw SNSPD data.
Plot of an example trace.
Code to filter out faulty traces, as described in the accompanying paper:
https://doi.org/10.48550/arXiv.2310.12471

"""

import numpy as np

# Attenuations (i.e., mean photon numbers) used in the experiment
dB = np.load('data/dB.npy')[:-1]
# These attenuations correlate to the following mean photon numbers per pulse
mpn = np.arange(0.5,5.5,0.5) #mean photon numbers from 0.5 to 5 in steps of 0.5
# Measurements where randomized according to the following array
dB_shuffled = np.load('data/db_shuffled.npy')


# Here we find out which indices in the raw data (Trace Nr.) belong to which mean photon number
dB_index_all = np.zeros((len(dB),200),dtype=int)
# This array is just to make sure whether all indices were found
# We expect 200 in the first element and 100 in all others, which equals the number of times we
# set the according attenuation in the experiment
dB_caught = np.zeros_like(dB,dtype=int)

for j in range(len(dB)):
    for i in range(len(dB_shuffled)):        
        if dB_shuffled[i] == dB[j]:
            dB_index_all[j][dB_caught[j]] = i  
            dB_caught[j] += 1     


#%% Plotting an example trace with the correct scaling factors

import matplotlib.pyplot as plt
x_inc = 0.78125e-2 #x axis (time) scaling factor
y_inc = 8.633e-3 #y axis (voltage) scaling factor
x_0 = 0 #x axis offset
y_0 = -1.84924e2*-1 #y axis offset

i = 0 #index of the trace to plot
n = 1000 #number of traces per file
l = 30000 #length of on trace

voltage = np.load('data/TracesNr0.npy')[i*l:i*l+l] * -1 #multiply by -1 to flip the trace 
time = np.arange(l)*x_inc + x_0 #time axis

fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.plot(time, voltage*y_inc + y_0, color='black', linewidth=1)
ax.set_ylabel('Voltage [mV]')
ax.set_xlabel('Time [ns]')
fig.tight_layout()
fig.show()

#%% Going through the entire data set to filter faulty traces (not really optimized)

# Array that will count the number of "empty" traces, i.e., traces where no detection event occurred
no_photon = np.zeros_like(dB)

for l in range(len(dB)):
    for j in range(dB_caught[l]):
        print(f'Attenuation {l+1}/{len(dB)} | Subset {j+1}/{dB_caught[l]}')
        mat = np.zeros((1000,19000),dtype=int)
        data=np.load('data/TracesNr'+str(dB_index_all[l][j])+'.npy')
        for i in range(1000):
            mat[i][:]=data[i*30000:(i)*30000+20000][1000:] #the 19000 elements contain the relevant SNSPD trace
        
        # Filtering the data
        mat = mat*-1
        faulty = [] #list containing the indices of the faulty traces
        for i in range(mat.shape[0]):
            if max(mat[i,:2000]) > -20000: #this cuts out detection events that happened just before the actual interesting time window
                faulty.append(i)
                
            elif max(mat[i,5000:]) > 0: #this cuts out detection events that happened after the actual interesting time window (non optimal extinction ratio)
                faulty.append(i)
            
            elif mat[i,3000] < 10000: #this cuts out empty traces (when no detection event happened), comment this out, if these events should be included
                faulty.append(i) 
                no_photon[l] += 1
                
        mat=np.delete(mat,faulty,axis=0)
        
        if len(mat)==0:
            break
        
        #################################################
        # Here goes whatever you want to do to the data #
        #################################################
        
        # break #remove this, if you want to go through the entire data set
    # break #remove this, if you want to go through the entire data set