#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import os
import re
import subprocess
from pathlib import Path

os.environ.setdefault("MPLCONFIGDIR", str(Path("out") / ".matplotlib"))

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

STAT_RE = re.compile(
    r"STAT: program=(\w+) size=(\d+) depth=(\d+) min_size=(\d+) valid=(\d+) time=([\d.eE+-]+)"
)


def seed_for(size: int, repeat: int) -> int:
    return 13013 + size * 17 + repeat * 1_000_003


def run_once(binary: str, program: str, size: int, depth: int, min_size: int, seed: int) -> dict:
    cmd = [
        binary,
        "--size",
        str(size),
        "--depth",
        str(depth),
        "--min-size",
        str(min_size),
        "--seed",
        str(seed),
    ]

    proc = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)

    if proc.returncode != 0:
        raise RuntimeError("Command failed: " + " ".join(cmd) + "\n" + proc.stderr)

    match = STAT_RE.search(proc.stderr)
    if not match:
        raise RuntimeError("STAT line not found: " + proc.stderr)

    stat_program, stat_size, stat_depth, stat_min_size, valid, elapsed = match.groups()

    if stat_program != program:
        raise RuntimeError(f"expected program={program}, got {stat_program}")
    if valid != "1":
        raise RuntimeError("sort validation failed: " + proc.stderr)

    return {
        "program": program,
        "size": int(stat_size),
        "depth": int(stat_depth),
        "min_size": int(stat_min_size),
        "seed": seed,
        "time": float(elapsed),
    }


def write_csv(path: Path, rows: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    header = ["program", "size", "depth", "min_size", "seed", "time"]

    with path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=header)
        writer.writeheader()
        writer.writerows(rows)


def draw_graph(path: Path, rows: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)

    sizes = sorted({row["size"] for row in rows})
    programs = ["shm", "pipe"]
    markers = {"shm": "o", "pipe": "s"}
    labels = {"shm": "Разделяемая память", "pipe": "Каналы"}

    fig, (time_ax, speed_ax) = plt.subplots(1, 2, figsize=(15, 6))

    for program in programs:
        cur = [row for row in rows if row["program"] == program]
        cur.sort(key=lambda row: row["size"])
        time_ax.plot(
            [row["size"] for row in cur],
            [row["time"] for row in cur],
            marker=markers[program],
            label=labels[program],
        )

    for program in programs:
        cur = [row for row in rows if row["program"] == program]
        cur.sort(key=lambda row: row["size"])
        speed_ax.plot(
            [row["size"] for row in cur],
            [row["size"] / row["time"] if row["time"] > 0 else 0 for row in cur],
            marker=markers[program],
            label=labels[program],
        )

    time_ax.set_xscale("log", base=2)
    speed_ax.set_xscale("log", base=2)

    time_ax.set_xlabel("Размер списка")
    time_ax.set_ylabel("Время, секунды")
    time_ax.set_title("Время сортировки связного списка слиянием")
    time_ax.grid(True)
    time_ax.legend()

    speed_ax.set_xlabel("Размер списка")
    speed_ax.set_ylabel("Элементов в секунду")
    speed_ax.set_title("Скорость сортировки на одинаковых нагрузках")
    speed_ax.grid(True)
    speed_ax.legend()

    fig.suptitle("Разделяемая память и каналы: задача сортировки списка из лабораторной 1")
    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", default="out", help="output directory")
    parser.add_argument("--depth", type=int, default=3)
    parser.add_argument("--min-size", type=int, default=4096)
    parser.add_argument("--repeats", type=int, default=3)
    parser.add_argument(
        "--sizes",
        type=int,
        nargs="+",
        default=[2**power for power in range(9, 19)],
    )
    args = parser.parse_args()

    out_dir = Path(args.out)
    rows = []

    for size in args.sizes:
        for program, binary in [("shm", "./shm_sort"), ("pipe", "./pipe_sort")]:
            measurements = []

            for repeat in range(args.repeats):
                seed = seed_for(size, repeat)
                row = run_once(binary, program, size, args.depth, args.min_size, seed)
                measurements.append(row)
                print(
                    f"{program}: size={size} repeat={repeat + 1}/{args.repeats} "
                    f"time={row['time']:.6f}",
                    flush=True,
                )

            best = min(measurements, key=lambda row: row["time"])
            rows.append(best)

    write_csv(out_dir / "benchmark.csv", rows)
    draw_graph(out_dir / "speed_comparison.png", rows)
    print(f"Saved {out_dir / 'benchmark.csv'}")
    print(f"Saved {out_dir / 'speed_comparison.png'}")


if __name__ == "__main__":
    main()