# DESCRIPTION: Plot the number of wanted and expected trials
# EXPECTED FILES: get_next_trials.csv
# TEST_OUTPUT_MUST_CONTAIN: Got
# TEST_OUTPUT_MUST_CONTAIN: Requested

import argparse
import importlib.util
import os
import signal
import sys
import traceback
from typing import Any, Union

import matplotlib.pyplot as plt
import pandas as pd

from beartype import beartype

script_dir = os.path.dirname(os.path.realpath(__file__))
helpers_file = f"{script_dir}/.helpers.py"
spec = importlib.util.spec_from_file_location(
    name="helpers",
    location=helpers_file,
)
if spec is not None and spec.loader is not None:
    helpers = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(helpers)
else: # pragma: no cover
    raise ImportError(f"Could not load module from {helpers_file}")

signal.signal(signal.SIGINT, signal.SIG_DFL)

@beartype
def parse_log_file(args: Any, log_file_path: str) -> Union[pd.DataFrame, None]:
    try:
        data = pd.read_csv(log_file_path, header=None, names=['time', 'got', 'requested'])

        valid_time_mask = data['time'].apply(helpers.is_valid_time_format)
        if not valid_time_mask.all():
            if not os.environ.get("NO_NO_RESULT_ERROR"): # pragma: no cover
                helpers.log_error("Some rows have invalid time format and will be removed.")
        data = data[valid_time_mask]

        if "time" not in data:
            if not os.environ.get("NO_NO_RESULT_ERROR"): # pragma: no cover
                print("time could not be found in data")
            sys.exit(19)

        data['time'] = pd.to_datetime(data['time'], format='%Y-%m-%d %H:%M:%S')

        # Sort data by time
        data = data.sort_values(by='time')

        return data
    except FileNotFoundError: # pragma: no cover
        helpers.log_error(f"File '{log_file_path}' not found.")
        raise
    except AssertionError as e: # pragma: no cover
        helpers.log_error(str(e))
        raise
    except UnicodeDecodeError: # pragma: no cover
        if not os.environ.get("PLOT_TESTS"):
            print(f"{args.run_dir}/get_next_trials.csv seems to be invalid utf8.")
        sys.exit(7)
    except Exception as e: # pragma: no cover
        helpers.log_error(f"An unexpected error occurred: {e}")
        print(traceback.format_exc(), file=sys.stderr)
        raise

    return None # pragma: no cover

@beartype
def plot_trial_usage(args: Any, log_file_path: str) -> None:
    try:
        data = parse_log_file(args, log_file_path)

        plt.figure(figsize=(12, 6))

        # Plot 'got'
        if data is not None:
            plt.plot(data['time'], data['got'], label='Got', color='blue')

            # Plot 'requested'
            plt.plot(data['time'], data['requested'], label='Requested', color='orange')

            plt.xlabel('Time')
            plt.ylabel('Count')
            plt.title('Trials Usage Plot')
            plt.legend()

            plt.gcf().autofmt_xdate()  # Rotate and align the x labels

            plt.tight_layout()
            if args.save_to_file:
                fig = plt.figure(1)
                helpers.save_to_file(fig, args, plt)
            else: # pragma: no cover
                if not args.no_plt_show:
                    plt.show()
        else: # pragma: no cover
            helpers.log_error("Failed to get job data")
            sys.exit(8)
    except Exception as e: # pragma: no cover
        helpers.log_error(f"An error occurred while plotting: {e}")
        raise

@beartype
def main() -> None:
    parser = argparse.ArgumentParser(description='Plot trial usage from log file')
    parser.add_argument('--run_dir', type=str, help='Directory containing log file')
    parser.add_argument('--debug', action='store_true', help='Enable debug mode')

    parser.add_argument('--save_to_file', type=str, help='Save the plot to the specified file', default=None)
    parser.add_argument('--no_plt_show', help='Disable showing the plot', action='store_true', default=False)
    args = parser.parse_args()

    if args.debug: # pragma: no cover
        print(f"Debug mode enabled. Run directory: {args.run_dir}")

    if args.run_dir:
        log_file_path = os.path.join(args.run_dir, "get_next_trials.csv")
        if os.path.exists(log_file_path):
            try:
                plot_trial_usage(args, log_file_path)
            except Exception as e: # pragma: no cover
                helpers.log_error(f"main Error: {e}")
                sys.exit(3)
        else: # pragma: no cover
            helpers.log_error(f"File '{log_file_path}' does not exist.")

if __name__ == "__main__":
    main()
