# Quick Plotter
# Cal Wing Nov 2025

import os, shutil, time, sys
from pathlib import Path

import numpy as np

from makeGraph import makeGraph, UQ_COLOURS as UQC

import matplotlib.pyplot as plt
from matplotlib.widgets import CheckButtons
from matplotlib.widgets import RadioButtons

def main():
    args = sys.argv[1:]
    assert len(args) >= 1, "CSV File must be passed in!"

    data_path = Path(args[0])

    data = {}
    numeric_headers = []
    headers = np.genfromtxt(data_path, delimiter=",", skip_header=0, max_rows=1, dtype=str)
    for i, header in enumerate(headers):
        data[header] = np.genfromtxt(data_path, delimiter=",", skip_header=1, max_rows=None, dtype=float, usecols=(i,))

        if np.isnan(data[header][0]):
            data[header] = np.genfromtxt(data_path, delimiter=",", skip_header=1, max_rows=None, dtype=str, usecols=(i,))
        else:
            numeric_headers.append(header)

    data["Index"] = tuple(range(len(data[numeric_headers[0]])))

    fig, ax = plt.subplot_mosaic(
        [
            ['main', 'x_axis'],
            #['main', 'button_box'],
            ['main', 'y_plot'],
        ],
        width_ratios=[5, 1],
        layout='constrained',
    )

    radio_background = 'lightgoldenrodyellow'

    ax['x_axis'].set_facecolor(radio_background)
    radio = RadioButtons(ax['x_axis'], ["Index"] + numeric_headers)

    # ax['button_box'].set_facecolor(radio_background)
    # norm_check = CheckButtons(
    #     ax=ax['button_box'],
    #     labels=["Normalize Values"],
    #     actives=[False],
    # )

    ax['y_plot'].set_facecolor(radio_background)
    check = CheckButtons(
        ax=ax['y_plot'],
        labels=numeric_headers,
        actives=[(x in numeric_headers[1:2]) for x in numeric_headers],
    )

    def do_plot(foo):
        x_header = radio.value_selected
        x_axis = data[x_header]
        
        ax["main"].clear()
        ax["main"].grid()
        ax["main"].set_xlabel(x_header)
        ax["main"].set_title(f"{data_path.parent.name}\\{data_path.parent.name}\\{data_path.name}")
        for y_plot in check.get_checked_labels():
            y_data = data[y_plot]
            # if len(norm_check.get_checked_labels()) > 0:
            #     y_data = y_data.copy()
            #     y_data -= y_data.min()
            #     y_data /= y_data.max()

            ax["main"].plot(x_axis, y_data, label=y_plot)
        
        ax["main"].legend()
        
        fig.canvas.draw()
    
    radio.on_clicked(do_plot)
    # norm_check.on_clicked(do_plot)
    check.on_clicked(do_plot)

    do_plot(None)

    plt.show()