# -*- coding: utf-8 -*-
"""
Created on Thu Nov  7 15:38:33 2024

@author: shen_t2
"""


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 ext_samp_clk_trig_count_task, gen_trig_src_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 = "_CWCD"


# %% Experiment parameters 


from Global_Experiment_Parameters import * 

MAIN_EXP_folder = 'C:/RE_qubit_TS/202411_CW_PLE_3_g/{:}{:}_BzScan_auto_{:}_{:}/'.format(now.strftime("%Y%m%d"), EXP_TYPE, laser_pol, ODFilter)
# os.chdir(MAIN_EXP_folder)


total_average = 1


EMagnet_poles_gap = "small"    # "small" or "large"
Bext_field_scan_positive = np.concatenate((np.arange(50, 2000, 50), 
                                           np.arange(2000, 3000, 200), 
                                           np.arange(3000, 5000, 500), 
                                           np.arange(5000, 8000, 1000)))
Bext_field_scan = np.concatenate( ( -Bext_field_scan_positive[::-1], 
                                   np.array([0]), 
                                   Bext_field_scan_positive ) )    # Gauss, always keep increasing (from - to +)


wide_scan_output_channel = 79    # 50 = piezo V // 79 = CTL wl //
# wide_scan_start_wl = 1530.00    # nm
wide_scan_trigger_threshold = 1530.00    # nm
wide_scan_end_wl = 1530.40    # nm
wide_scan_speed = 0.0010    # nm/s

wide_scan_start_wl = wide_scan_trigger_threshold - 15 * wide_scan_speed


DAQ_counting_time = 0.1    # s

DAQ_counting_rate = 1 / DAQ_counting_time    # Hz
DAQ_record_time_est = (wide_scan_end_wl - wide_scan_trigger_threshold) / wide_scan_speed    # s
DAQ_samps_per_chan = int(DAQ_record_time_est / DAQ_counting_time)
# DAQ_loop_per_chan = 3


SAPD_dead_time = 15 * 1e-6    # s



# %% Devices initialization 

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


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

dlcpro = set_widescan_para(Laser_DLCPRO_IP, 
                           wide_scan_output_channel, 
                           wide_scan_start_wl,
                           wide_scan_end_wl,
                           wide_scan_speed, 
                           wide_scan_trigger_threshold)

dlcpro.open()
print("Wide Scan Status:", dlcpro.laser1.wide_scan.state.get(), dlcpro.laser1.wide_scan.state_txt.get())


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

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


task_count = ext_samp_clk_trig_count_task(DAQ_counting_rate, DAQ_samps_per_chan)
task_clock = gen_trig_src_task(DAQ_counting_rate)




# %% 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")


wl_x_axis = np.linspace(wide_scan_trigger_threshold, wide_scan_end_wl, DAQ_samps_per_chan)
np.save( MAIN_EXP_folder + 'wl_scan_x_axis_{:.2f}_{:.2f}_speed{:.4f}'.format(wide_scan_trigger_threshold, wide_scan_end_wl, wide_scan_speed) + timestempID, 
        wl_x_axis[1:])

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



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 +
    
    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)
        dlcpro.laser1.wide_scan.start()
        
        task_clock.start()
        task_count.start()
        
        print("------------------------------")
        print("Start: Exp. {} Gs, Loop {}.".format(B_field, ii))
        
        raw_counts_accumulated = task_count.read(nidaqmx.constants.READ_ALL_AVAILABLE, nidaqmx.constants.WAIT_INFINITELY)
        
        if len(raw_counts_accumulated) == DAQ_samps_per_chan:
            # print("------------------------------")
            print("End: no errors.")
            
            
        task_count.stop()
        task_clock.stop()
        
        
        
        # Data processing and saving 
        
        raw_counts_accumulated = np.array(raw_counts_accumulated)
        raw_counts = raw_counts_accumulated[1:] - raw_counts_accumulated[:-1]
        np.save( MAIN_EXP_folder + 'raw_counts_B{:d}Gs_rep{:d}'.format(B_field, ii) + 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) + timestempID, actual_counts)
        print("Data saved.")
        
        raw_counts_all.append(raw_counts)
        
        
        #######################################
        # real-time plot  
        plt.figure(1, figsize=[9,6], dpi=100)
        plt.clf()
        plt.plot(wl_x_axis[1:], actual_counts, '-r', label='loop {:}'.format(ii) )
        plt.grid()
        plt.legend(loc=1)
        # plt.title(plot_title)
        plt.xlabel('Wavelength (nm)')
        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, timestempID) )
        plt.show()
        plt.pause(0.1)
        
        
    raw_counts_all = np.array(raw_counts_all)
    np.save( MAIN_EXP_folder + 'raw_counts_All_Bfields_rep{:d}'.format(ii) + 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) + timestempID, actual_counts_all)



# %% Completed  

set_Bfield_Current(0, EMagnet_PowerSupply, EMagnet_PSUPPLY_IP, EMagnet_PSUPPLY_PORT)

dlcpro.laser1.wide_scan.trigger.output_enabled.set(False)
dlcpro.close()

task_count.close()
task_clock.close()


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



# %% END