#!/usr/bin/env python3

from pathlib import Path
from dataclasses import dataclass
from simulation import (
    Simulation,
    ConfigTokens,
    get_argparser,
    get_options,
    simulation_results,
)

import matplotlib.pyplot as plt

SCRIPT_DIRECTORY = Path(__file__).parent.resolve()
DEFAULT_BASE_CONFIG = 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_tokens = ConfigTokens(
                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_tokens))


parser = get_argparser()
options = get_options(parser.parse_args())

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

df = simulation_results(options, simulations)

# Generate interesting plots!
for name, data in df.group_by("refresh_policy", "scheduler_buffer_size", maintain_order=True):
    plt.figure(figsize=(8.0, 12.0))

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

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

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

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