# -*- coding: utf-8 -*-
"""
Created on Wed Feb  5 13:08:12 2025

@author: shen_t2
"""


import os
# os.chdir(os.path.abspath(os.path.dirname(__file__)))

import numpy as np
import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 14})

from scipy.fftpack import fft, ifft, fftfreq 


# %% 

MAIN_EXP_folder = 'C:/RE_qubit_TS/202503_Pulse_echo_sample2/02_thicker_Cu_plate/19_echo_halfpi_population/run1_2us_noEOM_avg100/'
os.chdir(MAIN_EXP_folder)

timestampsuffix = '_04071904.npy'

time_scan = np.load('half_pi_x_axis_300-1500ns_Pulse_echo_halfpi' + timestampsuffix)    # in [us]


raw_count_sum = np.zeros(len(time_scan))
actual_count_sum = np.zeros(len(time_scan))

total_average = 55

for ii in range(total_average):
    raw_count_rep = np.load('raw_counts_B680Gs_rep{:d}_tau500ns_count3ms'.format(ii) + timestampsuffix) 
    raw_count_sum += raw_count_rep
    
    actual_count_rep = np.load('actual_counts_B680Gs_rep{:d}_tau500ns_count3ms'.format(ii) + timestampsuffix) 
    actual_count_sum += actual_count_rep


raw_count = raw_count_sum / total_average
actual_count = actual_count_sum / total_average



# %% FFT

time_step = 1/409.6    # [us]
fs = 1e6 / time_step    # Hz

x, y = time_scan, actual_count

fft_x = fftfreq(len(x)) * fs
fft_y = fft(y)


# %%


plt.figure(1330, figsize=[9,6], dpi=100)
plt.clf()

# plt.plot(time_scan, raw_count, '.-b', label= 'raw counts' )
plt.plot(time_scan, actual_count, '.-r', label='Experiment' )

plt.grid()
plt.legend(loc=1)
# plt.title('Rabi measurement')
plt.xlabel('Excitation time (us)')
plt.ylabel('Actual photon counts (cps)')
plt.tight_layout()

# plt.savefig( MAIN_EXP_folder + 'fitted_T1_{:.0f}ms{:}.jpg'.format(popt[0], timestampsuffix[:-4]) )
plt.show()



# %%

plt.figure(1331, figsize=[9,6], dpi=100)
plt.clf()

# plt.plot(time_scan, raw_count, '.-b', label= 'raw counts' )
# plt.plot(time_scan, actual_count, '.-r', label='Experiment' )

# plt.plot(fft_x, abs(fft_y), '.-b', label='FFT')
plt.plot(fft_x[:len(fft_x)//2], abs(fft_y)[:len(fft_x)//2] / len(x) * 2, '.-r', label='FFT')

plt.grid()
plt.legend(loc=1)
# plt.title('Rabi measurement')
plt.xlabel('Frequency (Hz)')
plt.ylabel('FFT')
plt.yscale('log')
plt.tight_layout()

# plt.savefig( MAIN_EXP_folder + 'fitted_T1_{:.0f}ms{:}.jpg'.format(popt[0], timestampsuffix[:-4]) )
plt.show()