# Cal Wing
# Sep 2024

import numpy as np

from nptdms import TdmsFile

from makeGraph import makeGraph, pltKeyClose, UQ_COLOURS as UQC

# Photo Diode Things
SHOT_PD = "x2s4111", 1, "PD", np.timedelta64(1146, 'us') #np.timedelta64(1146800, 'ns')

# PCB Things
SHOT_PCB_1 = "x2s3667", 0.0148 * 1000, "PCB", np.timedelta64(1643, 'us')
SHOT_PCB_2 = "x2s3668", 0.0148 * 1000, "PCB", np.timedelta64(1648600, 'ns')

# Shot DATA FILE
DATA_FILE_FMT = './data/{0}/databox/{0}.tdms'

TIME_OFFSET = np.timedelta64(-10, 'us'), np.timedelta64(15, 'us') # us

def main():

    graphData = {
            "title": "Shock response time of PCBs",
            "xLabel": "Time ($\\mu$s)",
            "yLabel": "PCB Voltage Reading (mV)",
            "grid": True,
            "xLim": (-10, 15),
            "plots": []
        }

    for df, scale, name, zero_p in [SHOT_PD, SHOT_PCB_1, SHOT_PCB_2]:
        file_path = DATA_FILE_FMT.format(df)
        data = TdmsFile.read(file_path)

        channels = data.groups()[0].channels()

        for channel in channels:
            time = (channels[0][:] - channels[0][0]) # Convert to sec

            time_range_i = [0, 0]

            for i, value in enumerate(time):
                if value >= zero_p + TIME_OFFSET[0]:
                    time_range_i[0] = i
                    break
            
            for i, value in enumerate(time):
                if value >= zero_p + TIME_OFFSET[1]:
                    time_range_i[1] = i
                    break
                
            
            a = time > time[time_range_i[0]]
            b = time < time[time_range_i[1]]
            time_range = np.logical_and(a, b)
            time_2 = time[time_range][:] - time[time_range][0] + TIME_OFFSET[0]

            #print(time_range, a, b, time_range_i)

            if channel.name == "st2" and df is not SHOT_PD[0]:
                graphData["plots"].append({
                    "x": time_2,
                    "y": channel[time_range] * scale,
                    "label": f"{df} - {name}"
                })

            if channel.name == "st2" and df is SHOT_PD[0] and False:
                graphData["plots"].append({
                    "x": time,
                    "y": channel[:] * scale,
                    "label": f"{df} - {name}"
                })
 

    makeGraph(graphData, showPlot=True, doProgramBlock=True)



if __name__ == '__main__':
    main()

    #pltKeyClose()