# -*- 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 pulse_gated_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 = "_CWGatedD"


# %% Experiment parameters 

from Global_Experiment_Parameters import * 

exp_notes = 'laser {:d} mA (w/o EOM), {:}, {:},\ndiff. Gs, sample{:}, {:}'.format(laser_current, ODFilter, laser_pol, sample, Temp_info)
exp_notes += '\nPiezo = 60V'
MAIN_EXP_folder = 'P:/Tianyang/Data_Thibaut_2026/20260505/03_Bfield_dependence/'
# os.chdir(MAIN_EXP_folder)


total_average = 1


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


Bext_field_scan = np.concatenate((np.arange(-500, 600, 10), 
                                  np.arange(600, 2020, 20)))    # Gauss, always keep increasing (from - to +)
Bext_field_scan = np.arange(-500, 2600, 50)
# Bext_field_scan = [+1000]


wide_scan_output_channel = 79    # 50 = piezo V // 79 = CTL wl //
# wide_scan_start_wl = 1530.0    # nm
wide_scan_trigger_threshold = 1530.00    # nm
wide_scan_end_wl = 1530.80   # nm
wide_scan_speed = 0.0200/2    # nm/s


# wide_scan_output_channel = 50    # 50 = piezo V // 79 = CTL wl //
# # wide_scan_start_wl = 0    # V
# wide_scan_trigger_threshold = 20    # V
# wide_scan_end_wl = 130    # V
# wide_scan_speed = 1/2    # V/s


wide_scan_start_wl = wide_scan_trigger_threshold - 15 * wide_scan_speed
if wide_scan_output_channel == 50:
    wide_scan_start_wl = wide_scan_trigger_threshold - 2 * wide_scan_speed
    wide_scan_start_wl = max(5, wide_scan_start_wl)


AOM_ON_time =  0.010    # s
DAQ_counting_time = 0.010    # s

DAQ_counting_rate = 1 / (AOM_ON_time + 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_rate)
# DAQ_loop_per_chan = 3



# %% 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 = pulse_gated_count_task(DAQ_samps_per_chan)
task_clock = gen_trig_src_task(DAQ_counting_rate, 
                               pulse_duty_cycle = AOM_ON_time * 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) + EXP_TYPE + timestempID, 
        wl_x_axis)

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 +
    
    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)
        dlcpro.laser1.wide_scan.start()
        
        task_clock.start()
        task_count.start()
        
        print("------------------------------")
        print("Start: Exp. {} Gs, Loop {}.".format(B_field, ii))
        
        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()
        task_clock.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) + 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(21, figsize=[9,6], dpi=100)
        plt.clf()
        plt.plot(wl_x_axis, raw_counts, '.-b', label='loop {:}, raw counts'.format(ii) )
        plt.plot(wl_x_axis, actual_counts, '.-r', label='loop {:}, corrected'.format(ii) )
        plt.grid()
        plt.legend(loc=1)
        plt.title(exp_notes)
        if wide_scan_output_channel == 79:
            plt.xlabel('Wavelength (nm)')
        if wide_scan_output_channel == 50:
            plt.xlabel('Piezo Voltage (V)')
        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


# plt.figure(2, figsize=[9,6], dpi=100)
# plt.clf()
# plt.plot(wl_x_axis, raw_counts, '.-b', label='loop {:}, raw counts'.format(ii) )
# plt.plot(wl_x_axis, actual_counts, '.-r', label='loop {:}, corrected'.format(ii) )
# plt.grid()
# plt.legend(loc=1)
# plt.title(exp_notes)
# plt.xlabel('Wavelength (nm)')
# plt.ylabel('Actual photon counts (cps)')
# plt.show()


# wl_scan = np.load('wl_scan_x_axis_1529.90_1530.50_speed0.0010_12031938.npy')
# actual_counts = np.load('actual_counts_B1000Gs_rep0_12031938.npy')
# raw_counts = np.load('raw_counts_B1000Gs_rep0_12031938.npy')
# plt.figure(3, figsize=[9,6], dpi=100)
# plt.clf()
# plt.plot(wl_scan, raw_counts, '.-b')
# plt.plot(wl_scan, actual_counts, '.-r')
# plt.grid()
# plt.show()