#!/usr/bin/env python3

import argparse
import json
import os

import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit
from tqdm import tqdm

plt.rcParams["figure.constrained_layout.use"] = True


MAXIMUM_TIME = 500
LARGE_OBJECT_VOLUME = 587.0 * 233.0 * 220.0


def make_arg_parser():
    parser = argparse.ArgumentParser(
        description="compute the time one heuristics take to reach the results of another one"
    )
    parser.add_argument(
        "--figure",
        help="name of the file to save",
        default="time_to_reach.png",
    )
    parser.add_argument(
        "-f",
        "--files",
        nargs=2,
        help="Two data files",
    )
    parser.add_argument(
        "-l",
        "--labels",
        nargs=2,
        help="Two labels",
    )
    args = parser.parse_args()
    return args


def make_history(file_path) -> list:
    history = []
    with open(file_path, "r") as file:
        data = json.load(file)
    for instance_set in data:
        for instance_number in data[instance_set]:
            current_max = 0
            instance_data = data[instance_set][instance_number]["appendix"]
            solution_history = instance_data["solution_history"]
            for entry in solution_history:
                time_instant = float(entry["time"])
                volume_usage_percent = (
                    100.0 * float(entry["volume_usage"]) / LARGE_OBJECT_VOLUME
                )
                if current_max < volume_usage_percent:
                    current_max = volume_usage_percent
                    history.append(
                        {
                            "set": int(instance_set),
                            "instance": int(instance_number),
                            "instant": time_instant,
                            "value": volume_usage_percent,
                        }
                    )
    history.sort(key=lambda entry: (entry["instant"], entry["set"], entry["instance"]))
    return history


def make_initial_status(history) -> dict:
    initial_status = {}
    for h in history:
        instance_set = h["set"]
        instance_number = h["instance"]
        if instance_set not in initial_status:
            initial_status[instance_set] = {}
        initial_status[instance_set][instance_number] = 0
    return initial_status


# def compute_average(entry_set: dict) -> float:
#     n = 0
#     v = 0
#     for entry in entry_set:
#         n += 1
#         v += entry_set[entry]
#     v = v / n
#     return v


def compute_average_general(current: dict[dict], skip=set()) -> float:
    n = 0
    v = 0
    for instance_set in current:
        if instance_set in skip:
            continue
        for instance_number in current[instance_set]:
            n += 1
            v += current[instance_set][instance_number]
    v = v / n
    return v


def make_time_series_general(history: list) -> tuple[list, list]:
    # group entries in the same instant
    time = []
    average = []
    current = make_initial_status(history)
    previous_instant = -100000.0  # -inf
    for i, entry in enumerate(history):
        entry = history[i]
        instant = entry["instant"]
        value = entry["value"]
        if previous_instant != instant:
            time.append(instant)
            average.append(compute_average_general(current, skip={0}))
        instance_set = entry["set"]
        instance_number = entry["instance"]
        current[instance_set][instance_number] = value
        previous_instant = instant
    time.append(instant)
    average.append(compute_average_general(current, skip={0}))
    return time, average


# def make_time_series(history: list, target_instance_set: int) -> tuple[list, list]:
#     # group entries in the same instant
#     time = []
#     average = []
#     current = make_initial_status(history)
#     if target_instance_set not in current.keys():
#         return [], []
#     current = current[target_instance_set]
#     previous_instant = -100000.0  # -inf
#     for i in range(len(history)):
#         entry = history[i]
#         instant = entry["instant"]
#         instance_set = entry["set"]
#         instance_number = entry["instance"]
#         if instance_set == target_instance_set:
#             value = entry["value"]
#             if previous_instant != instant:
#                 time.append(instant)
#                 average.append(compute_average(current))
#             current[instance_number] = value
#             previous_instant = instant
#     time.append(instant)
#     average.append(compute_average(current))
#     return time, average


def build_datasets(args):
    datasets = []
    for i in tqdm(range(len(args.files)), desc="build datasets"):
        file_path = str(args.files[i])
        label = str(args.labels[i])
        history = make_history(file_path)
        time, average = make_time_series_general(history)
        datasets.append(
            {
                "label": label,
                "time": np.array(time),
                "average": np.array(average),
            }
        )
    return datasets


def find_first_time(v1, t2, v2):
    """
    Finds the first time instant from t2 where the corresponding value in v2
    is greater than or equal to each value in v1.

    Args:
        v1 (list of float): A list of increasing values corresponding to t1.
        t2 (list of float): A list of time instants corresponding to v2.
        v2 (list of float): A list of increasing values corresponding to t2.

    Returns:
        list of float: A list where each element is the first time instant in t2
                      where v2 is greater than or equal to the corresponding v1 value.
    """
    T = []
    j = 0  # Pointer for t2 and v2

    for i in range(len(v1)):
        # Find the first v2[j] >= v1[i]
        while j < len(v2) and v2[j] < v1[i]:
            j += 1
        # If j is within bounds, append the corresponding time from t2, else append None
        if j < len(t2):
            T.append(t2[j])
        else:
            break  # No value in v2 satisfies the condition

    T = np.array(T)
    T = T[T != np.array(None)]
    return T


def init_figure():
    A5 = np.array([210, 148])
    MM_TO_INCH = 0.0393700787
    size = 0.7 * MM_TO_INCH * A5
    figure = plt.figure(
        figsize=size,  # A5 paper size in inches
        dpi=300,
        layout="constrained",
    )
    return figure


def plot_dataset(figure, n_columns, column, dataset_reference, dataset, time_limit=None):
    axes = figure.add_subplot(1, n_columns, column)
    # data
    t1 = dataset_reference["time"]
    v1 = dataset_reference["average"]
    if time_limit is not None:
        v1 = v1[t1 < time_limit]
        t1 = t1[t1 < time_limit]
    t2 = dataset["time"]
    v2 = dataset["average"]
    T = find_first_time(v1, t2, v2)
    t1 = t1[: T.shape[0]]
    axes.plot(t1[: T.shape[0]], T, color="k", linestyle="-", label="experiment")
    # curve fit
    linear_model = lambda x, a: a * x
    param, _ = curve_fit(linear_model, t1, T)
    alpha = param[0]
    T_fit = linear_model(t1, alpha)
    axes.plot(t1, T_fit, linestyle="--", color="gray", label=f"f(t) = {alpha:.1f} · t")
    # description
    axes.set_ylabel(f"Time required by {dataset['label']} [s]")
    axes.set_xlabel(f"Time limit given to {dataset_reference['label']} [s]")
    axes.grid(True, which="both")
    axes.legend(loc="best")
    return


def save_figure(figure_file_path, figure):
    figure.savefig(figure_file_path)
    print(f"Figure saved to file: {os.path.realpath(figure_file_path)}")
    return


if __name__ == "__main__":
    args = make_arg_parser()
    datasets = build_datasets(args)
    figure = init_figure()
    figure.suptitle(
        f"Time for {datasets[1]['label']} to reach the result of {datasets[0]['label']}"
    )
    plot_dataset(figure, 2, 1, *datasets)
    plot_dataset(figure, 2, 2, *datasets, 10)
    save_figure(args.figure, figure)
