#!/usr/bin/env python3

import argparse
import itertools
import json
import math
import os
import sys

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from tqdm import tqdm

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

MAXIMUM_TIME = 500
XTICK_MAJOR_STRIDE = 250
XTICK_MINOR_STRIDE = 50
YTICK_MAJOR_STRIDE = 1.0
YTICK_MINOR_STRIDE = 0.5

LARGE_OBJECT_VOLUME = 587.0 * 233.0 * 220.0

LINE_STYLE_LIST = [
    "solid",
    "dashed",
    "dotted",
    "dashdot",
]


def make_arg_parser():
    parser = argparse.ArgumentParser(
        description="make time series analysis of the results"
    )
    parser.add_argument(
        "--single",
        help="make a single plot, with only the average for instance sets 1-15",
        action=argparse.BooleanOptionalAction,
        default=False,
    )
    parser.add_argument(
        "-f",
        "--figure_file_path",
        help="name of the file to save",
        default="time_serie_analysis.jpg",
    )
    parser.add_argument(
        "files_and_labels",
        nargs="+",
        help="Pairs of data files and their labels (e.g., file1.json label1 file2.json label2)",
    )
    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 (time_instant <= MAXIMUM_TIME) and (
                    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):
    if len(args.files_and_labels) % 2 != 0:
        parser.error(
            "You must provide a list of files, each one with a label.\nExample: file_X.json label-X file_Y.json label-Y"
        )
    datasets = []
    for i in tqdm(range(0, len(args.files_and_labels), 2), desc="build datasets"):
        file_path = str(args.files_and_labels[i])
        label = str(args.files_and_labels[i + 1])
        linestyle = LINE_STYLE_LIST[(i // 2) % len(LINE_STYLE_LIST)]
        history = make_history(file_path)
        general = make_time_series_general(history)
        specific = {
            instance_set: make_time_series(history, instance_set)
            for instance_set in range(15 + 1)
        }
        datasets.append(
            {
                "label": label,
                "linestyle": linestyle,
                "general": general,
                "specific": specific,
            }
        )
    return datasets


def closest(value, multiple, f):
    return int(f(value / multiple)) * multiple


def initialize_single_figure():
    A6 = np.array([148, 2*105//3])  # landscape
    MM_TO_INCH = 0.0393700787
    size = MM_TO_INCH * A6
    figure = plt.figure(
        figsize=size,
        dpi=300,
        layout="constrained",
    )
    composed_figure = figure
    composed_figure_axes = composed_figure.subplots(1, 2)
    return {
        "figure": figure,
        "composed_figure": figure,
        "composed_figure_axes": composed_figure_axes,
    }


def initialize_figure():
    A4 = np.array([210, 297])
    MM_TO_INCH = 0.0393700787
    size = 0.7 * MM_TO_INCH * A4
    figure = plt.figure(
        figsize=size,  # A4 paper size in inches
        dpi=300,
        layout="constrained",
    )
    composed_figure, separated_figure = figure.subfigures(
        nrows=2,
        ncols=1,
        height_ratios=[1, 2],
    )
    composed_figure.suptitle("Average volume usage of sets 1-15")
    separated_figure.suptitle("Average volume usage by instance set")
    composed_figure_axes = composed_figure.subplots(1, 2)
    separated_figure_axes = separated_figure.subplots(
        nrows=4,
        ncols=4,
        sharex=True,
        sharey=True,
        squeeze=False,
    )
    return {
        "figure": figure,
        "composed_figure": composed_figure,
        "separated_figure": separated_figure,
        "composed_figure_axes": composed_figure_axes,
        "separated_figure_axes": separated_figure_axes,
    }


def plot_single_dataset(figure_data, dataset):
    composed_figure = figure_data["composed_figure"]
    linestyle = dataset["linestyle"]
    label = dataset["label"]
    time, average = dataset["general"]
    xlim = [0, closest(max(time), XTICK_MAJOR_STRIDE, math.ceil)]
    ylim = [94.5, 97.5]
    xlabel = "time [s]"
    ylabel = "volume usage [%]"
    xtick_major = np.arange(
        closest(min(time), XTICK_MAJOR_STRIDE, math.floor),
        closest(max(time), XTICK_MAJOR_STRIDE, math.ceil) + 1,
        XTICK_MAJOR_STRIDE,
    )
    xtick_minor = np.arange(
        closest(min(time), XTICK_MINOR_STRIDE, math.floor),
        closest(max(time), XTICK_MINOR_STRIDE, math.ceil) + 1,
        XTICK_MINOR_STRIDE,
    )
    ytick_major = [95, 96, 97]
    ytick_minor = [95.5, 96.5, 97.5]
    # composed
    axes = figure_data["composed_figure_axes"][0]
    axes.plot(time, average, color="k", linestyle=linestyle, label=label)
    axes.set_xticks(xtick_major, minor=False)
    axes.set_xticks(xtick_minor, minor=True)
    axes.set_yticks(ytick_major, minor=False)
    axes.set_yticks(ytick_minor, minor=True)
    axes.set_xlim(xlim)
    axes.set_ylim(ylim)
    axes.set_ylabel(ylabel)
    axes.set_xlabel(xlabel)
    axes.grid(True, which="both")
    axes.legend(loc="lower right")
    # focus
    focus_xlim = [0, 10]
    axes = figure_data["composed_figure_axes"][1]
    axes.set_title(f"{focus_xlim[0]} to {focus_xlim[1]} seconds")
    axes.plot(time, average, color="k", linestyle=linestyle, label=label)
    axes.set_xticks(np.arange(0, 11, 2), minor=False)
    axes.set_xticks(np.arange(0, 11, 1), minor=True)
    axes.set_yticks(ytick_major, minor=False)
    axes.set_yticks(ytick_minor, minor=True)
    axes.set_xlim(focus_xlim)
    axes.set_ylim(ylim)
    axes.set_ylabel(ylabel)
    axes.set_xlabel(xlabel)
    axes.grid(True, which="both")
    axes.legend(loc="best")
    return


def plot_dataset(figure_data, dataset):
    composed_figure = figure_data["composed_figure"]
    separated_figure = figure_data["separated_figure"]
    linestyle = dataset["linestyle"]
    label = dataset["label"]
    time, average = dataset["general"]
    xlim = [0, closest(max(time), XTICK_MAJOR_STRIDE, math.ceil)]
    ylim = [94.5, 97.5]
    xlabel = "time [s]"
    ylabel = "volume usage [%]"
    xtick_major = np.arange(
        closest(min(time), XTICK_MAJOR_STRIDE, math.floor),
        closest(max(time), XTICK_MAJOR_STRIDE, math.ceil) + 1,
        XTICK_MAJOR_STRIDE,
    )
    xtick_minor = np.arange(
        closest(min(time), XTICK_MINOR_STRIDE, math.floor),
        closest(max(time), XTICK_MINOR_STRIDE, math.ceil) + 1,
        XTICK_MINOR_STRIDE,
    )
    ytick_major = [95, 96, 97]
    ytick_minor = [95.5, 96.5, 97.5]
    # composed
    axes = figure_data["composed_figure_axes"][0]
    axes.set_title(f"{xlim[0]} to {xlim[1]} seconds")
    axes.plot(time, average, color="k", linestyle=linestyle, label=label)
    axes.set_xticks(xtick_major, minor=False)
    axes.set_xticks(xtick_minor, minor=True)
    axes.set_yticks(ytick_major, minor=False)
    axes.set_yticks(ytick_minor, minor=True)
    axes.set_xlim(xlim)
    axes.set_ylim(ylim)
    axes.set_ylabel(ylabel)
    axes.set_xlabel(xlabel)
    axes.grid(True, which="both")
    axes.legend(loc="lower right")
    # focus
    focus_xlim = [0, 10]
    axes = figure_data["composed_figure_axes"][1]
    axes.set_title(f"{focus_xlim[0]} to {focus_xlim[1]} seconds")
    axes.plot(time, average, color="k", linestyle=linestyle, label=label)
    axes.set_xticks(np.arange(0, 11, 2), minor=False)
    axes.set_xticks(np.arange(0, 11, 1), minor=True)
    axes.set_yticks(ytick_major, minor=False)
    axes.set_yticks(ytick_minor, minor=True)
    axes.set_xlim(focus_xlim)
    axes.set_ylim(ylim)
    axes.set_ylabel(ylabel)
    axes.set_xlabel(xlabel)
    axes.grid(True, which="both")
    axes.legend(loc="best")
    # instance sets
    axes = figure_data["separated_figure_axes"]
    indices = [(i, j) for i in range(len(axes)) for j in range(len(axes[i]))]
    for i, j in indices:
        instance_set_number = i * len(axes) + j
        time, average = dataset["specific"][instance_set_number]
        axes[i][j].set_title(f"{instance_set_number}")
        axes[i][j].plot(time, average, color="k", linestyle=linestyle, label=label)
        axes[i][j].set_xticks(xtick_major, minor=False)
        axes[i][j].set_xticks(xtick_minor, minor=True)
        axes[i][j].set_yticks(ytick_major, minor=False)
        axes[i][j].set_yticks(ytick_minor, minor=True)
        axes[i][j].set_xlim(xlim)
        axes[i][j].set_ylim(ylim)
        axes[i][j].grid(True, which="both")
        if j == 0 and ((i == len(axes) - 1) or (i == 0)):
            axes[i][j].set_ylabel(ylabel)
        if i == len(axes) - 1:
            axes[i][j].set_xlabel(xlabel)
    return


def save_figure(figure_file_path, figure_data: dict):
    figure_data["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)
    if args.single:
        figure_data = initialize_single_figure()
        for dataset in tqdm(datasets, desc="plot datasets"):
            plot_single_dataset(figure_data, dataset)
    else:
        figure_data = initialize_figure()
        for dataset in tqdm(datasets, desc="plot datasets"):
            plot_dataset(figure_data, dataset)
    save_figure(args.figure_file_path, figure_data)
