# Cal Wing (c.wing@uq.net.au) - Oct 2024
# Thesis Graphing

import os

import numpy as np
import pandas as pd

import yaml

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)

# Load Data
DATA_PATH = "./data"
DATA_INFO = "_info.yaml"

data_to_load = [
    "x2s5823"
]

data = {}

for dp in data_to_load:
    data_path = f"{DATA_PATH}/{dp}/"
    data_info_path = data_path + DATA_INFO
    if not os.path.exists(data_info_path):
        print(f"[ERR] Could not find data info file: '{data_info_path}'")
        print(f"[WARN] Not Loading Data '{dp}'")
        continue
    
    with open(data_info_path, 'r') as file:
        # Load data info (Cal)
        dataInfo = yaml.safe_load(file)
        x2_shot = dataInfo["shot-info"]["name"]

        x2_tdms_data = TdmsFile.read(data_path + dataInfo["shot-info"]['tdms'])
        x2_channels = x2_tdms_data.groups()[0].channels()
        
        if dataInfo["probe-info"]["data-record"]["type"] == "scope":
            scope_data_path = data_path + dataInfo["probe-info"]["data-record"]["data"]
            scope_config_path = data_path + dataInfo["probe-info"]["data-record"]["config"]

            # Generate Headers
            with open(scope_data_path, 'r') as dfile:
                scope_header = []
                
                header_lines = []
                for i, line in enumerate(dfile):
                    if i > 1: break
                    header_lines.append(line.strip().split(","))
                
                for i, name in enumerate(header_lines[0]):
                    if name == "x-axis":
                        name = "Time"
                    
                    if header_lines[1][i] in ["second", "Volt"]:
                        outStr = f"{name} [{header_lines[1][i][0]}]"
                    else:
                        outStr = f"{name} [{header_lines[1][i]}]"

                    scope_header.append(outStr)
                      
            scope_data = pd.read_csv(scope_data_path, names=scope_header, skiprows=1)

        data[x2_shot] = {
            "info": dataInfo,
            "probes": scope_data,
            "x2": x2_channels,
            "x2-tdms": x2_tdms_data
        }
    
loaded_data = data.keys()

print("Loaded Data")