#!/usr/bin/env python3
import os
import argparse
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path

def process_files(csv_files, fixed_param, fixed_value, max_value=0):
    """Combine and average simulation times by the varying parameter."""
    results = {}
    varying_param = "depth" if fixed_param == "width" else "width"

    for path in csv_files:
        df = pd.read_csv(path)
        df = df[df[fixed_param] == fixed_value].copy()
        if max_value > 0:
            df = df[df[varying_param] <= max_value].copy()

        df["total_time"] = (
            df["simulation_time"]
            + df["creation_time"]
            + df["setup_time"]
        )

        Q1 = df["total_time"].quantile(0.25)
        Q3 = df["total_time"].quantile(0.75)
        IQR = Q3 - Q1
        lower_bound = Q1 - 1.5 * IQR
        upper_bound = Q3 + 1.5 * IQR
        df = df[(df["total_time"] >= lower_bound) & (df["total_time"] <= upper_bound)]
        grouped = df.groupby(varying_param, as_index=False)["total_time"].mean()
        results[path.parent.name + "-" + path.stem] = grouped

    return results, varying_param

def plot_results(results, varying_param, fixed_param, fixed_value, name):
    plt.figure(figsize=(8, 6))
    for label, df in results.items():
        plt.plot(df[varying_param], df["total_time"], marker='o', label=label)

    plt.title(f"Average Simulation Time vs {varying_param.capitalize()} (fixed {fixed_param}={fixed_value})")
    plt.xlabel(varying_param.capitalize())
    plt.ylabel("Average Simulation Time (s)")
    plt.legend()
    plt.grid(True)

    path = "../Plots/"
    output_name = f"{name}_{fixed_param}_{fixed_value}.png"
    plt.tight_layout()
    plt.savefig(path + output_name, dpi=300)
    print(f"Saved plot as {output_name}")


def main():
    parser = argparse.ArgumentParser(description="Compare simulation results from multiple CSV files.")
    parser.add_argument("csv_files", nargs="+", type=Path, help="Paths to the result CSV files.")
    parser.add_argument("-n", "--name", default="DEVS", type=str, help="Name of the output file.")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("-w", "--width", type=int, help="Fixed width value.")
    group.add_argument("-d", "--depth", type=int, help="Fixed depth value.")
    parser.add_argument('-m', '--max', type=int, default=0, help='Maximum value of the varying parameter')

    args = parser.parse_args()
    if args.width is not None:
        fixed_param, fixed_value = "width", args.width
    else:
        fixed_param, fixed_value = "depth", args.depth

    results, varying_param = process_files(args.csv_files, fixed_param, fixed_value, args.max)
    plot_results(results, varying_param, fixed_param, fixed_value, args.name)

if __name__ == "__main__":
    main()