#!/usr/bin/env python3

import os
import matplotlib.pyplot as plt

from dataclasses import dataclass
from simulation import (
    Simulation,
    Configuration,
    get_options_from_args,
    simulation_results,
)

script_directory = os.path.dirname(os.path.abspath(__file__))
base_config = os.path.join(script_directory, "example_ddr3_simulations.json")

simulations: list[Simulation] = []


@dataclass(frozen=True)
class SchedulerBuffer:
    buffer: str
    size: int


refresh_policies = ["NoRefresh", "AllBank"]
patterns = ["sequential", "random"]
scheduler_buffers = [
    SchedulerBuffer("Bankwise", 8),
    SchedulerBuffer("Bankwise", 16),
]

for refresh_policy in refresh_policies:
    for scheduler_buffer in scheduler_buffers:
        for pattern in patterns:
            PAGE_POLICY = "Open" if pattern == "sequential" else "Closed"

            config = Configuration(
                f"{refresh_policy}_{scheduler_buffer.buffer}_{scheduler_buffer.size}_{pattern}",
                {
                    "refresh_policy": refresh_policy,
                    "pattern": pattern,
                    "page_policy": PAGE_POLICY,
                    "scheduler_buffer": scheduler_buffer.buffer,
                    "scheduler_buffer_size": scheduler_buffer.size,
                },
            )

            simulations.append(Simulation(config))


options = get_options_from_args()

if options.base_config is None:
    options.base_config = base_config

dataframe = simulation_results(options, simulations)

# Generate interesting plots!
for refresh_policy in refresh_policies:
    refresh_policy_filter = dataframe["refresh_policy"] == refresh_policy

    for scheduler_buffer in scheduler_buffers:
        scheduler_buffer_filter = (
            dataframe["scheduler_buffer"] == scheduler_buffer.buffer
        )
        scheduler_buffer_size_filter = (
            dataframe["scheduler_buffer_size"] == scheduler_buffer.size
        )

        filters = (
            refresh_policy_filter
            & scheduler_buffer_filter
            & scheduler_buffer_size_filter
        )

        TEMP_DATAFRAME = dataframe[filters]

        for pattern in patterns:
            pattern_filter = dataframe["pattern"] == pattern

            plt.figure(figsize=(8.0, 12.0))

            filters = (
                refresh_policy_filter
                & scheduler_buffer_filter
                & scheduler_buffer_size_filter
            )

            TEMP_DATAFRAME = dataframe[filters]

            plt.bar(
                "pattern",
                "bandwidth",
                data=TEMP_DATAFRAME,
            )

        # Plot MAX line
        plt.bar(
            "MAX",
            "max_bandwidth",
            data=TEMP_DATAFRAME,
            color="darkgray",
        )

        title = f"{refresh_policy} {scheduler_buffer.buffer} {scheduler_buffer.size}"

        plt.title(title)
        plt.axis(ymin=0, ymax=120)
        plt.xlabel("Access Pattern")
        plt.ylabel("Bandwidth [Gb/s]")

        plt.savefig(f"{options.out_dir}/{title}.png")
        plt.close()