# -*- coding: utf-8 -*-
"""
Created on Tue Feb  4 23:35:57 2025

@author: shen_t2
"""

"""
! Remember to start RFSoC codes after running this file !
"""


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

import time
from datetime import datetime
now = datetime.now()
# timestempID = now.strftime("_%y%m%d%H%M%S")
timestempID = now.strftime("_%m%d%H%M")

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


import nidaqmx
from DAQCountingFunctions import arbitaryTTL_gated_count_task, correct_cps

# import PowerSupplyCaylarLib
# from EMagnetSetFields import EMagnet_connect_test, set_Bfield_Gauss, set_Bfield_Current

# from toptica.lasersdk.dlcpro.v2_6_0 import DLCpro, NetworkConnection, DeviceNotFoundError
# from LaserWideScanSettings import set_widescan_para

EXP_TYPE = "_Pulse_Rabi"


# %% Experiment parameters 

from Global_Experiment_Parameters import * 

EMagnet_poles_gap = "large"    # "small" 39 mm upto 1.3 T, or "large" 80 mm upto 0.8 T
EMagnet_initialization = True
EMagnet_initialization = False

total_average = 100


# Scanning excitation pulse width
time_init = 0.000            # in [us]
time_final = 1.500
time_step = 0.005
time_list = np.linspace(time_init, time_final, int((time_final-time_init)/time_step)+1)
# time_list = np.logspace(1, 32, num=50, endpoint=True, base=2.0) / 1e3
x_all = time_list 

DAQ_samps_per_chan = len(x_all)
# DAQ_loop_per_chan = 3

# AOM_ON_time = 30e-3    # s
DAQ_counting_time = 10e-3    # s 
resonance_freq = +1500    # in [MHz]



exp_notes = 'laser {:d} mA (with EOM), {:}, {:},\n+{:d} Gs, sample{:}, {:}'.format(laser_current, ODFilter, laser_pol, Bext_field_scan[0], sample, Temp_info)    # as plot title 
exp_notes += '\nDAQ counting window = {:.0f} ms'.format(DAQ_counting_time*1e3)
# exp_notes += '\nResonance frequency = {:.0f} MHz'.format(resonance_freq)
exp_notes += '\nwithout EOM, laser directly locked to 0-0 transition'

fn_suffix = '_count{:.0f}ms'.format(DAQ_counting_time*1e3)

MAIN_EXP_folder = 'C:/RE_qubit_TS/202503_Pulse_echo_sample2/02_thicker_Cu_plate/61_left3_left_splitting/run21_Rabi_smallApt_1500ns_noEOM_count10ms_avg{:d}/'.format(total_average)
# MAIN_EXP_folder += 'excitation_{:.1f}ms_count_{:.0f}ms/'.format(AOM_ON_time*1e3, DAQ_counting_time*1e3)
# os.chdir(MAIN_EXP_folder)



# %% Devices initialization 

# EMagnet_PowerSupply = PowerSupplyCaylarLib.CaylarPowerSupply(2)
# EMagnet_connect_test(EMagnet_PowerSupply, EMagnet_PSUPPLY_IP, EMagnet_PSUPPLY_PORT)


######################################

###   Laser NOT used! 


######################################

try:
    task_count.close()
    task_clock.close()
except:
    pass
finally:
    print("------------------------------")
    print("Tasks have been cleared.\n")


task_count = arbitaryTTL_gated_count_task(DAQ_samps_per_chan)



# %% Start experiments

try:
    os.mkdir(MAIN_EXP_folder)
except:
    pass
finally:
    print("==============================")
    print("==============================")
    print("==============================")
    print("Experiment folder has been created.\n")
    print("Experiment STARTS here.\n")


np.save( MAIN_EXP_folder + 'excitation_time_x_axis_{:.0f}-{:.0f}ms'.format(time_init * 1e-3, time_final * 1e-3) + EXP_TYPE + timestempID, 
        x_all)

# np.save( MAIN_EXP_folder + 'wl_scan_Bfields' + timestempID, Bext_field_scan)


raw_counts_sum = np.zeros(len(x_all))
actual_counts_sum = np.zeros(len(x_all))
ext_reminder = True

for ii in range(total_average):
    
    # raw_counts_all = []
    
    
    # Always initialize to -100 as minus saturation,
    # because the calibration follows the lower branch of the magnetic hysteresis. 
    # Then, scan from - to +
    
    if EMagnet_initialization:
        set_Bfield_Current(-100, EMagnet_PowerSupply, EMagnet_PSUPPLY_IP, EMagnet_PSUPPLY_PORT)
        time.sleep(5)
        
        
    for B_field in Bext_field_scan:
        
        # set_Bfield_Gauss(B_field, EMagnet_poles_gap, EMagnet_PowerSupply, EMagnet_PSUPPLY_IP, EMagnet_PSUPPLY_PORT)
        # time.sleep(5)
        
        task_count.start()
        
        print("------------------------------")
        print("Start: Exp. {} Gs, Loop {}.".format(B_field, ii))
        
        if ext_reminder == True:
            print("==============================")
            print("!!! Now, start RFSoC output !!!")
            print("!!! And, make sure AOM TTL input from RFSoC Pmod box !!!")
            print("==============================")
            ext_reminder = False
        
        raw_counts = task_count.read(nidaqmx.constants.READ_ALL_AVAILABLE, nidaqmx.constants.WAIT_INFINITELY)
        
        if len(raw_counts) == DAQ_samps_per_chan:
            # print("------------------------------")
            print("End: no errors.")
            
            
        task_count.stop()
        
        
        
        # Data processing and saving 
        
        raw_counts = np.array(raw_counts)
        np.save( MAIN_EXP_folder + 'raw_counts_B{:d}Gs_rep{:d}'.format(B_field, ii) + fn_suffix + timestempID, raw_counts)
        
        actual_counts = correct_cps(raw_counts, DAQ_counting_time, SAPD_dead_time)
        np.save( MAIN_EXP_folder + 'actual_counts_B{:d}Gs_rep{:d}'.format(B_field, ii) + fn_suffix + timestempID, actual_counts)
        print("Data saved.")
        
    
        raw_counts_sum += raw_counts
        actual_counts_sum += actual_counts
        
        # raw_counts_all.append(raw_counts)
        
        
        #######################################
        # real-time plot  
        plt.figure(32, figsize=[9,6], dpi=100)
        plt.clf()
        plt.plot(x_all * 1e-3, raw_counts_sum /(ii+1), '.-b', label='avg {:}, raw counts'.format(ii+1) )
        plt.plot(x_all * 1e-3, actual_counts_sum /(ii+1), '.-r', label='avg {:}, corrected'.format(ii+1) )
        plt.grid()
        plt.legend(loc=1)
        plt.title(exp_notes)
        plt.xlabel('Excitation time/ms')
        plt.ylabel('Actual photon counts (cps)')
        plt.tight_layout()
        plt.savefig( MAIN_EXP_folder + 'plot_counts_B{:d}Gs_rep{:d}{:}{:}.jpg'.format(B_field, ii, fn_suffix, timestempID) )
        plt.show()
        plt.pause(0.1)
        
        
    # if len(Bext_field_scan) > 1:
    #     raw_counts_all = np.array(raw_counts_all)
    #     np.save( MAIN_EXP_folder + 'raw_counts_All_Bfields_rep{:d}'.format(ii) + fn_suffix + timestempID, raw_counts_all)
        
    #     actual_counts_all = correct_cps(raw_counts_all, DAQ_counting_time, SAPD_dead_time)
    #     np.save( MAIN_EXP_folder + 'actual_counts_All_Bfields_rep{:d}'.format(ii) + fn_suffix + timestempID, actual_counts_all)


if total_average > 1:
    np.save( MAIN_EXP_folder + 'raw_counts_{:d}Averaged'.format(total_average) + EXP_TYPE + timestempID, raw_counts_sum / total_average)
    np.save( MAIN_EXP_folder + 'actual_counts_{:d}Averaged'.format(total_average) + EXP_TYPE + timestempID, actual_counts_sum / total_average)



# %% Completed  

# set_Bfield_Current(0, EMagnet_PowerSupply, EMagnet_PSUPPLY_IP, EMagnet_PSUPPLY_PORT)


task_count.close()


print("==============================")
print("ALL END. (NO ERRORS)\n")



# %% END


plt.figure(32, figsize=[9,6], dpi=100)
plt.clf()
plt.plot(x_all, raw_counts_sum /(ii+1), '.-b', label='avg {:}, raw counts'.format(ii+1) )
plt.plot(x_all, actual_counts_sum /(ii+1), '.-r', label='avg {:}, corrected'.format(ii+1) )
plt.grid()
plt.legend(loc=1)
plt.title(exp_notes)
plt.xlabel('Excitation time/us')
plt.ylabel('Actual photon counts (cps)')
# plt.xscale('log')
plt.tight_layout()
plt.savefig( MAIN_EXP_folder + 'plot_counts_B{:d}Gs_{:d}Averaged{:}{:}.jpg'.format(B_field, total_average, fn_suffix, timestempID) )
plt.show()