# Cal Wing
# Sep 2024

import os

import numpy as np

from nptdms import TdmsFile
from makeGraph import makeGraph, pltKeyClose, UQ_COLOURS as UQC

# Folder correction
#  Make sure the relevant folders folder exists
folders = ["./images"]
for folder in folders:
    if not os.path.isdir(folder): os.mkdir(folder)

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

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

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

TIME_OFFSET = np.timedelta64(-22, 'us'), np.timedelta64(33, 'us') # us

def main():

    graphData = {
            "title": "Shock response time of PCBs\nAt position st2",
            "xLabel": "Time ($\\mu$s)",
            "yLabel": "PCB Voltage Reading (mV)",
            "grid": True,
            "xLim": (-20, 20),
            "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}"
                })
    
    graphData['plots'].append(
        {
            "x":0, 
            "type":"axvLine", 
            #"label":"Shock", 
            "color": UQC["dark_grey"],
            "args": {
                "linestyle": "--",
                "zorder": -2
            }
        }
    )

    makeGraph(graphData, figSavePath="./images/{0}.png", showPlot=False, doProgramBlock=False)



if __name__ == '__main__':
    main()

    #pltKeyClose()