import sys
import matplotlib.pyplot as plt
from pathlib import Path
from InquirerPy.base import Choice
from InquirerPy import inquirer

from utils import load_measured_data

default_dataset_dir = Path("../single_capture_test_results")

battery_thermal_limit = 45.0  # Celsius
case_thermal_limit = 41.0  # Celsius

def select_waveforms(dataset_directory = default_dataset_dir):
    """
    Select waveforms from a given dataset directory.

    Args:
        dataset_directory (Path): The directory containing the dataset.

    Returns:
        list: A list of selected waveforms.
    """

    if not dataset_directory.exists():
        print(f"Dataset directory {dataset_directory} does not exist.")
        return []

    # glob all .csv files in the directory
    all_csv_files = list(dataset_directory.glob("*.csv"))
    if not all_csv_files:
        print(f"No CSV files found in {dataset_directory}.")
        return []

    external_temp_files = list(dataset_directory.glob("external_temp.*.csv"))

    # Filter out external temperature files
    waveform_files = [f for f in all_csv_files if f not in external_temp_files]

    choices = []
    for waveform_file in waveform_files:
        time_id = waveform_file.stem.split(".")[1]

        ch = Choice(name=f"{waveform_file.name}", value={'waveform':waveform_file, 'external_temp': None})

        for temp_file in external_temp_files:
            if time_id in temp_file.stem:
                ch.name += f" (ext. temp available)"
                ch.value['external_temp'] = temp_file
                break

        choices.append(ch)

    try:
        selected = inquirer.fuzzy(
            message="Select one or more waveforms:",
            choices=choices,
            multiselect=True,
            instruction="(Use <tab> to select, <enter> to confirm)"
        ).execute()

    except KeyboardInterrupt:
        print("Selection cancelled by user.")
        sys.exit(0)

    except Exception as e:
        print(f"An error occurred during selection: {e}")
        sys.exit(1)

    return selected

def colored_region_plot(axis, time_vector, data_vector, mask, color='red', alpha=0.5):

    start = None
    in_region = False
    for i , val in enumerate(mask):
        if val and not in_region:
            start = i
            in_region = True
        elif not val and in_region:
            axis.plot(time_vector[start:i-1], data_vector[start:i-1], color=color, alpha=alpha)
            in_region = False

    if in_region:
        axis.plot(time_vector[start:(i-1)], data_vector[start:i-1], color=color, alpha=alpha)

def colored_region_box(axis, time_vector, mask, color='orange', alpha=0.5):

    start = None
    in_region = False
    for i , val in enumerate(mask):
        if val and not in_region:
            start = i
            in_region = True
        elif not val and in_region:
            axis.axvspan(time_vector[start], time_vector[i-1], color=color, alpha=alpha)
            in_region = False

    if in_region:
        axis.axvspan(time_vector[start], time_vector[-1], color=color, alpha=alpha)

def sec_to_min(time_vector):
    return (time_vector - time_vector[0]) / 60.0

def plot_temperature_profile(waveform_name, profile_data):

    fig, ax = plt.subplots(2)
    fig.canvas.manager.set_window_title(waveform_name)

    ax[0].plot(sec_to_min(profile_data.time), profile_data.battery_temp, color='green', label='battery temeperature')
    ax[0].axhline(y=battery_thermal_limit, color='green', linestyle='--')

    ax[0].plot(sec_to_min(profile_data.time), profile_data.pmic_die_temp, color='orange', label='pmic die temperature')

    colored_region_plot(
        ax[0],
        sec_to_min(profile_data.time),
        profile_data.battery_temp,
        profile_data.battery_temp > battery_thermal_limit,
        color='red',
        alpha=1)

    if profile_data.ext_temp is not None:
        ax[0].plot(sec_to_min(profile_data.ext_temp_time), profile_data.ext_temp, color='blue', label='case temperature', linestyle='--')
        ax[0].axhline(y=case_thermal_limit, color='blue', linestyle='--')

        colored_region_plot(
            ax[0],
            sec_to_min(profile_data.ext_temp_time),
            profile_data.ext_temp,
            profile_data.ext_temp > case_thermal_limit,
            color='red',
            alpha=1)


    ax[0].set_xlabel("Time (min)")
    ax[0].set_ylabel("Temperature (C)")
    ax[0].set_title("Temperature Profile: " + waveform_name)
    ax[0].set_xlim(left=sec_to_min(profile_data.time)[0], right=sec_to_min(profile_data.time)[-1])

    ax[0].legend()
    ax[0].grid(True)

    def min_to_hr(x): return x / 60.0
    def hr_to_min(x): return x * 60.0

    secax = ax[0].secondary_xaxis('top', functions=(min_to_hr, hr_to_min))
    secax.set_xlabel("Time (hours)")

    # Change background color according to charging state
    usb_charging_mask = (profile_data.usb == "USB_connected") & (abs(profile_data.battery_current) > 0)
    wlc_charging_mask = (profile_data.wlc == "WLC_connected") & ~usb_charging_mask & (abs(profile_data.battery_current) > 0)
    colored_region_box(ax[0], sec_to_min(profile_data.time), usb_charging_mask, color='blue', alpha=0.2)
    colored_region_box(ax[0], sec_to_min(profile_data.time), wlc_charging_mask, color='green', alpha=0.2)

    ax[1].plot(sec_to_min(profile_data.time), profile_data.battery_current, color='purple', label='battery current')
    ax[1].set_xlabel("Time (min)")
    ax[1].set_ylabel("Current (mA)")
    ax[1].set_xlim(left=sec_to_min(profile_data.time)[0], right=sec_to_min(profile_data.time)[-1])
    ax[1].grid(True)
    ax[1].legend()

    colored_region_box(ax[1], sec_to_min(profile_data.time), usb_charging_mask, color='blue', alpha=0.2)
    colored_region_box(ax[1], sec_to_min(profile_data.time), wlc_charging_mask, color='green', alpha=0.2)

def main():

    selected_waveforms = select_waveforms()

    for waveform in selected_waveforms:

        # Load data from files
        profile_data = load_measured_data(
            data_file_path=waveform['waveform'],
            extern_temp_file_path=waveform['external_temp']
        )

        plot_temperature_profile(waveform['waveform'].name, profile_data)

    # Plot graphs
    plt.show()

if __name__ == "__main__":
    main()