import csv
import os

import matplotlib
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter

# Ensure Matplotlib doesn't require a display (important for running in headless environments)
matplotlib.use("Agg")

# Constants
PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__))
RESULTS_DIR = os.path.join(PROJECT_ROOT, "results")
PLOT_DATA_FILE = os.path.join(RESULTS_DIR, "plot_data.csv")

# Create the results directory if it doesn't exist
os.makedirs(RESULTS_DIR, exist_ok=True)


def load_plot_data(file_path):
    """Loads plot data from a CSV file into a dictionary."""
    data = {}
    with open(file_path, "r") as f:
        reader = csv.DictReader(f)

        # Add a check to ensure fieldnames is not None
        if reader.fieldnames is None:
            raise ValueError("CSV file must have a header row.")

        for field in reader.fieldnames:
            data[field] = []

        for row in reader:
            for field in reader.fieldnames:
                value = row[field]
                if field == "Input Size":
                    data[field].append(int(value))
                elif "_Time" in field or "_Memory" in field:
                    data[field].append(float(value))
                else:
                    data[field].append(int(value))
    return data


# Formatter functions
def millions_formatter(x, pos):
    """Format numbers in millions."""
    return f"{x * 1e-6:.1f}"


# def thousands_formatter(x, pos):
#     """Format numbers in thousands."""
#     return f"{x * 1e-3:.1f}"


def seconds_formatter(x, pos):
    """Format time in seconds."""
    return f"{x:.2f}"


def milliseconds_formatter(x, pos):
    """Format time in milliseconds."""
    return f"{x * 1e3:.2f}"


def memory_mb_formatter(x, pos):
    """Format memory usage in MB."""
    return f"{x:.2f}"


def memory_kb_formatter(x, pos):
    """Format memory usage in KB."""
    return f"{x * 1e3:.2f}"


def plot_data(input_sizes, algorithms_data, plot_name, plot_title, y_label, formatter):
    """Generates and saves a plot for the given data."""
    plt.figure(figsize=(10, 6))
    markers = {
        "Quick Sort (Random Pivot)": "o",
        "Quick Sort (Deterministic Pivot)": "v",
        "Merge Sort": "s",
        "Heap Sort": "^",
    }

    for algorithm, data in algorithms_data.items():
        plt.plot(input_sizes, data, label=algorithm, marker=markers.get(algorithm, "o"))

    plt.title(plot_title)
    plt.xlabel("Input Size")
    plt.ylabel(y_label)
    plt.gca().yaxis.set_major_formatter(FuncFormatter(formatter))
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(os.path.join(RESULTS_DIR, plot_name), dpi=300)
    plt.close()


def main():
    # Load data
    data = load_plot_data(PLOT_DATA_FILE)
    input_sizes = data["Input Size"]

    # Split indices
    split_index = 5  # First 5 datasets (indices 0 to 4) correspond to input sizes from 10k to 100k

    # Define algorithms
    algorithms = [
        "Quick Sort (Random Pivot)",
        "Quick Sort (Deterministic Pivot)",
        "Merge Sort",
        "Heap Sort",
    ]

    # Plot comparisons
    comparisons_data_small = {alg: data[alg][:split_index] for alg in algorithms}
    comparisons_data_large = {alg: data[alg][split_index - 1 :] for alg in algorithms}

    plot_data(
        input_sizes[:split_index],
        comparisons_data_small,
        "comparisons_10000_to_100000.png",
        "Number of Comparisons vs. Input Size",
        "Number of Comparisons (Millions)",
        formatter=millions_formatter,
    )

    plot_data(
        input_sizes[split_index - 1 :],
        comparisons_data_large,
        "comparisons_100000_to_800000.png",
        "Number of Comparisons vs. Input Size",
        "Number of Comparisons (Millions)",
        formatter=millions_formatter,
    )

    # Plot running times
    times_algorithms = {
        "Quick Sort (Random Pivot)": data["QS_Random_Time"],
        "Quick Sort (Deterministic Pivot)": data["QS_Deterministic_Time"],
        "Merge Sort": data["MS_Time"],
        "Heap Sort": data["HS_Time"],
    }

    times_data_small = {
        alg: times_algorithms[alg][:split_index] for alg in times_algorithms
    }
    times_data_large = {
        alg: times_algorithms[alg][split_index - 1 :] for alg in times_algorithms
    }

    plot_data(
        input_sizes[:split_index],
        times_data_small,
        "time_10000_to_100000.png",
        "Sorting Time vs. Input Size",
        "Time (Milliseconds)",
        formatter=milliseconds_formatter,
    )

    plot_data(
        input_sizes[split_index - 1 :],
        times_data_large,
        "time_100000_to_800000.png",
        "Sorting Time vs. Input Size",
        "Time (Seconds)",
        formatter=seconds_formatter,
    )

    # Plot memory usage
    memory_algorithms = {
        "Quick Sort (Random Pivot)": data["QS_Random_Memory"],
        "Quick Sort (Deterministic Pivot)": data["QS_Deterministic_Memory"],
        "Merge Sort": data["MS_Memory"],
        "Heap Sort": data["HS_Memory"],
    }

    memory_data_small = {
        alg: memory_algorithms[alg][:split_index] for alg in memory_algorithms
    }
    memory_data_large = {
        alg: memory_algorithms[alg][split_index - 1 :] for alg in memory_algorithms
    }

    plot_data(
        input_sizes[:split_index],
        memory_data_small,
        "memory_10000_to_100000.png",
        "Memory Usage vs. Input Size",
        "Memory Usage (Kilobytes)",
        formatter=memory_kb_formatter,
    )

    plot_data(
        input_sizes[split_index - 1 :],
        memory_data_large,
        "memory_100000_to_800000.png",
        "Memory Usage vs. Input Size",
        "Memory Usage (Megabytes)",
        formatter=memory_mb_formatter,
    )


if __name__ == "__main__":
    main()