# -*- 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})


# %% 

MAIN_EXP_folder = 'C:/RE_qubit_TS/20260429_CW_PLE_dilute_sample2/50_pi_test/'
# MAIN_EXP_folder = 'P:/Tianyang/Data_Thibaut_2026/20260505/03_Bfield_dependence/'
os.chdir(MAIN_EXP_folder)

timestampsuffix = '_05051456.npy'

# mw_scan = np.load('MWfreq_x_axis_pm100-3000MHz_per10MHz'+ timestampsuffix)
wl_scan = np.load('wl_scan_x_axis_1530.00_1530.80_speed0.0100_CWGatedD'+ timestampsuffix)

raw_count_sum = np.zeros(len(wl_scan))
actual_count_sum = np.zeros(len(wl_scan))
for ii in range(1):
    raw_count_rep = np.load('raw_counts_B0Gs_rep{:d}'.format(ii) + timestampsuffix) 
    raw_count_sum += raw_count_rep 
    actual_count_rep = np.load('actual_counts_B0Gs_rep{:d}'.format(ii) + timestampsuffix) 
    actual_count_sum += actual_count_rep 


# actual_counts_wl = np.load('actual_counts_B0Gs_rep0_'+file_name[-12:])
# raw_counts_wl = np.load('raw_counts_B0Gs_rep0_'+file_name[-12:])


wn_scan = 1e7 / wl_scan 


# %%

plt.figure(100, figsize=[9,6], dpi=100)
plt.clf()
plt.plot(1e7/wl_scan, raw_count_sum, '.-b')
plt.plot(1e7/wl_scan, actual_count_sum, '.-r')
plt.grid()
plt.tight_layout()
plt.show()


# %% 

from scipy.optimize import curve_fit

def _1Lorentzian(x, x_c, area, width, y_offset):
#     y_c = y_offset + 2*area/width/np.pi    # peak value 
    return ( y_offset + (2*area/np.pi)*(width/(4*(x-x_c)**2 + width**2)))


(popt, pcov, infodict, mesg, ier) = curve_fit(_1Lorentzian, wn_scan, actual_count_rep, 
                                              p0 = [6534.34, 1, 1, 1], 
                                              bounds = ([6534.0, 0, 0, 0],
                                                        [6535.0, np.inf, np.inf, np.inf]),
                                              full_output=True)


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

plt.plot(wn_scan, actual_count_rep, '.-', label='experiment')  
plt.plot(wn_scan, _1Lorentzian(wn_scan, *popt), '.-', label='fitting')  

plt.grid()
plt.legend(loc=1)

plt.title('B = 0, $\\nu_c$ = {:.6f} cm-1'.format(popt[0]) )
plt.xlabel('Recorded wavenumber (cm$^{-1}$)')
plt.ylabel('Photon Counting (cps) ')

plt.tight_layout()
plt.show()

print(popt[0])    # 6534.344133644639 cm-1


wn_PLE_c = 6534.22767272914    # with precise fitting
wn_FTIR_c = 6534.356431
wn_scan = 1e7 / wl_scan - wn_PLE_c + wn_FTIR_c   # unit: [cm-1]