#!/usr/bin/env python3
"""Generate publication-ready Raptor vs Pimcomp comparison charts.

Usage:
    .venv/bin/python validation/tools/pim/pimcomp/compare/generate_plots.py

Dependencies: matplotlib (no pandas or seaborn required).
By default, reads results_comparison.csv beside this script and writes plots/
beside it, with outputs grouped in png/ and svg/ under the output directory.

YOLO models without Pimcomp measurements get separate Raptor-only charts.
Comparisons use vertical bars for ratios, diverging advantages, and absolute
values, with compiler colors and patterns for the Raptor pipeline stages.
Throughput-mode latency is the output interval, a bottleneck-stage proxy,
not end-to-end inference latency. Separate throughput plots use inferences/s.
Energy is reported per inference; power is averaged over the simulation.
Numerical results whose compiler did not pass functional validation have a red
status following their value label; missing Pimcomp
measurements appear as 'N/A'. The fixed compiler palette is shared by
all metrics. Relative plots use Pimcomp / Raptor for latency, power and energy,
and Raptor / Pimcomp for throughput, so >1 favors Raptor. Diverging plots put the winner's
factor above (Raptor) or below (Pimcomp) parity, with 1x at the zero line.
"""
from __future__ import annotations

import argparse
import csv
import math
import sys
from dataclasses import dataclass
from itertools import groupby
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Patch, Rectangle
from matplotlib.lines import Line2D
from matplotlib.ticker import FuncFormatter, MultipleLocator

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from plot_style import COLORS, HATCH_COLOR, TEXT_COLOR, configure_style, fit_value_labels, save_figure, value_label


COMPILER_COLORS = {"raptor": COLORS[0], "pimcomp": COLORS[2]}
STAGE_PATTERNS = {1: "", 2: "..", 4: "//", 8: "\\\\"}
ARCHITECTURES = ("arch-a", "arch-b")
MODELS = ("vgg8", "resnet18", "resnet34", "googlenet", "yolo11n")
CONFIGURATIONS = (("latency", 1), ("throughput", 2), ("throughput", 4), ("throughput", 8))
LINE_COLOR = "#A8B2C0"
VALIDATION_STYLES = {
    "PASS": ("#168C8B", "white", ""),
    "FAIL": ("#D18D32", HATCH_COLOR, ""),
    "SKIPPED": ("#E7DFF0", TEXT_COLOR, ""),
    "NA": ("#E8EDF1", TEXT_COLOR, ""),
}


@dataclass(frozen=True)
class Result:
    arch: str
    model: str
    mode: str
    pipeline: int
    pimcomp_schedule: str
    raptor_validation: str
    pimcomp_validation: str
    latency_speedup: float
    power: float
    energy: float
    raptor_latency: float
    pimcomp_latency: float
    raptor_throughput: float
    pimcomp_throughput: float
    raptor_power: float
    pimcomp_power: float
    raptor_energy: float
    pimcomp_energy: float

    @property
    def comparison_passed(self) -> bool:
        return self.raptor_validation == self.pimcomp_validation == "PASS"

    @property
    def throughput_speedup(self) -> float:
        return quotient(self.raptor_throughput, self.pimcomp_throughput)


def positive_float(text: str | None) -> float:
    try:
        value = float(text)
    except (TypeError, ValueError):
        return math.nan
    return value if math.isfinite(value) and value > 0 else math.nan


def quotient(numerator: float, denominator: float) -> float:
    if not (math.isfinite(numerator) and math.isfinite(denominator)):
        return math.nan
    return numerator / denominator


def validation_status(text: str) -> str:
    status = text.strip().upper()
    if status not in VALIDATION_STYLES:
        raise ValueError(f"Unknown functional-validation status: {text!r}")
    return status


def load_results(path: Path) -> list[Result]:
    with path.open(newline="", encoding="utf-8-sig") as handle:
        source = list(csv.DictReader(handle))
    rows = []
    expected = {(arch, model, mode, stage)
                for arch in ARCHITECTURES for model in MODELS
                for mode, stage in CONFIGURATIONS}
    seen = {}
    for line, row in enumerate(source, 2):
        try:
            arch, model, mode = (row[field].strip().lower() for field in ("arch", "model", "mode"))
            stage = int(row["raptor_pipeline"])
        except (KeyError, AttributeError, TypeError, ValueError) as exc:
            raise ValueError(f"Invalid configuration key at CSV row {line}: {exc}") from exc
        key = arch, model, mode, stage
        if key not in expected:
            raise ValueError(f"Unexpected configuration at CSV row {line}: {key}")
        if key in seen:
            raise ValueError(f"Duplicate configuration {key} at CSV rows {seen[key]} and {line}")
        seen[key] = line
        schedule = row["pimcomp_pipeline"].strip().lower()
        if schedule != ("element" if mode == "latency" else "batch"):
            raise ValueError(f"{key}: incompatible Pimcomp schedule {schedule!r}")
        rows.append(Result(
            arch=arch,
            model=model,
            mode=mode,
            pipeline=stage,
            pimcomp_schedule=schedule,
            raptor_validation=validation_status(row["raptor_functional_validation"]),
            pimcomp_validation=validation_status(row["pimcomp_functional_validation"]),
            latency_speedup=quotient(positive_float(row["pimcomp_latency_ms"]),
                                     positive_float(row["raptor_latency_ms"])),
            power=quotient(positive_float(row["pimcomp_power_mw"]),
                           positive_float(row["raptor_power_mw"])),
            energy=quotient(positive_float(row["pimcomp_energy_pj"]),
                            positive_float(row["raptor_energy_pj"])),
            raptor_latency=positive_float(row["raptor_latency_ms"]),
            pimcomp_latency=positive_float(row["pimcomp_latency_ms"]),
            raptor_throughput=positive_float(row["raptor_throughput_samples_s"]),
            pimcomp_throughput=positive_float(row["pimcomp_throughput_samples_s"]),
            raptor_power=positive_float(row["raptor_power_mw"]),
            pimcomp_power=positive_float(row["pimcomp_power_mw"]),
            raptor_energy=positive_float(row["raptor_energy_pj"]),
            pimcomp_energy=positive_float(row["pimcomp_energy_pj"]),
        ))
    missing = expected - seen.keys()
    if missing:
        raise ValueError(f"Missing {len(missing)} expected configurations. Examples: {sorted(missing)[:5]}")
    rows.sort(key=lambda r: (ARCHITECTURES.index(r.arch), MODELS.index(r.model), r.pipeline))
    absolute_columns(rows)
    return rows


def advantage_compiler(ratio: float) -> str | None:
    if ratio > 1:
        return "raptor"
    if ratio < 1:
        return "pimcomp"
    return None


def figure_legend(fig, results: list[Result], field: str,
                  absolute=False, raptor_only=False):
    compilers = ("raptor",) if raptor_only else ("raptor", "pimcomp")
    if absolute:
        measurements = [(row, compiler, getattr(row, f"{compiler}_{field}"))
                        for row, compiler in absolute_columns(results, raptor_only)]
    else:
        measurements = [(row, None, getattr(row, field)) for row in results]
    available = [(row, compiler, value) for row, compiler, value in measurements
                 if math.isfinite(value)]
    stages = {row.pipeline for row, compiler, _ in available if compiler != "pimcomp"}
    handles = []
    visible_compilers = {compiler if absolute else advantage_compiler(value)
                         for _, compiler, value in available} - {None}
    if len(visible_compilers) > 1:
        for compiler in compilers:
            handles.append(Patch(facecolor=COMPILER_COLORS[compiler],
                                 label=compiler.capitalize() + ("" if absolute else " advantage")))
    handles += [Patch(facecolor="white", edgecolor=TEXT_COLOR, hatch=STAGE_PATTERNS[depth],
                      label=f"{depth} stage" + (" (latency)" if depth == 1 else "s"))
                for depth in sorted(stages) if stages != {1}]
    styles = []
    if len(available) < len(measurements):
        styles.append(Patch(facecolor="white", edgecolor="#B2BAC5",
                            label="N/A: missing measurement"))
    handles += styles
    if handles:
        fig.legend(handles=handles, loc="outside upper center", ncol=(len(handles) + 1) // 2,
                   frameon=False, fontsize=9.5, handlelength=1.7, columnspacing=1.5,
                   handletextpad=0.6, labelspacing=0.4)


def absolute_columns(results: list[Result], raptor_only: bool = False) -> list[tuple[Result, str]]:
    columns = []
    for (arch, model), group in groupby(results, key=lambda row: (row.arch, row.model)):
        rows = sorted(group, key=lambda row: row.pipeline)
        for row in rows:
            columns.append((row, "raptor"))
            if row.mode == "latency" and not raptor_only:
                columns.append((row, "pimcomp"))
        throughput = [row for row in rows if row.mode == "throughput"]
        if throughput and not raptor_only:
            signatures = {
                (row.pimcomp_schedule, row.pimcomp_validation,
                 *(None if math.isnan(value) else value for value in (
                     row.pimcomp_latency, row.pimcomp_throughput, row.pimcomp_power, row.pimcomp_energy)))
                for row in throughput
            }
            if len(signatures) != 1:
                raise ValueError(f"{arch}/{model}: inconsistent shared Pimcomp batch results across pipelines")
            columns.append((throughput[0], "pimcomp"))
    return columns


def label_groups(ax, results: list[Result], positions: list[float] | None = None) -> None:
    if positions is None:
        positions = list(range(len(results)))
    centers, labels = [], []
    start = 0
    for index, (model, rows) in enumerate(groupby(results, key=lambda row: row.model)):
        end = start + len(list(rows))
        centers.append((positions[start] + positions[end - 1]) / 2)
        labels.append(model.upper())
        left = ax.get_xlim()[0] if start == 0 else (positions[start - 1] + positions[start]) / 2
        right = ax.get_xlim()[1] if end == len(results) else (positions[end - 1] + positions[end]) / 2
        if index % 2 == 0:
            ax.axvspan(left, right, facecolor="#F7F9FC", zorder=0)
        if index > 0:
            ax.axvline(left, color="#DDE3EB", linewidth=0.9)
        start = end
    ax.set_xticks(centers, labels, fontsize=9.5)


def diverging_ratio(ratio: float) -> float:
    return ratio - 1 if ratio >= 1 else 1 - 1 / ratio


def draw_validation(results: list[Result], output_dir: Path) -> None:
    columns = absolute_columns(results)
    configurations = list(dict.fromkeys((row.mode, row.pipeline if compiler == "raptor" else row.pimcomp_schedule,
                                        compiler) for row, compiler in columns))
    positions = []
    for index, (mode, _, _) in enumerate(configurations):
        gap = 0.4 if index and mode != configurations[index - 1][0] else 0
        positions.append(positions[-1] + 1 + gap if index else 0)
    configuration_x = dict(zip(configurations, positions))
    headers = [f"Raptor\n{stage} stages" if mode == "throughput" and compiler == "raptor"
               else "Pimcomp\nbatch" if mode == "throughput" else compiler.capitalize()
               for mode, stage, compiler in configurations]
    fig, axes = plt.subplots(1, 2, figsize=(14, 3.8), layout="constrained")
    for ax, arch in zip(axes, ARCHITECTURES):
        ax.set_xlim(positions[0] - 0.55, positions[-1] + 0.55)
        ax.set_ylim(len(MODELS) - 0.5, -0.5)
        centers, modes = [], []
        for mode, group in groupby(zip(configurations, positions), key=lambda item: item[0][0]):
            group_positions = [x for _, x in group]
            centers.append((group_positions[0] + group_positions[-1]) / 2)
            modes.append(mode.capitalize())
        ax.set_xticks(centers, modes, fontsize=10)
        ax.set_xticks(positions, headers, minor=True)
        ax.tick_params(axis="x", which="minor", top=True, labeltop=True,
                       bottom=False, labelbottom=False, length=0, pad=6, labelsize=9)
        ax.set_yticks(range(len(MODELS)), [model.upper() for model in MODELS], fontsize=9)
        ax.tick_params(axis="both", length=0, pad=6)
        architecture = "ARCHITECTURE A" if arch == "arch-a" else "ARCHITECTURE B"
        ax.set_xlabel(architecture, fontsize=11, weight="bold", labelpad=8)
        for row, compiler in columns:
            if row.arch != arch:
                continue
            key = (row.mode, row.pipeline if compiler == "raptor" else row.pimcomp_schedule, compiler)
            x, y = configuration_x[key], MODELS.index(row.model)
            status = getattr(row, f"{compiler}_validation")
            color, text_color, _ = VALIDATION_STYLES[status]
            ax.add_patch(Rectangle((x - 0.5, y - 0.42), 1, 0.84,
                                   facecolor=color, edgecolor="#657081", linewidth=0.5))
            ax.text(x, y, "N/A" if status == "NA" else status,
                    color=text_color, ha="center", va="center", fontsize=9)
    save_figure(fig, "validation_status", output_dir)


def draw_metric(results: list[Result], field: str, ylabel: str,
                basename: str, output_dir: Path, diverging: bool = False) -> None:
    data = [diverging_ratio(getattr(row, field)) if diverging else getattr(row, field)
            for row in results]
    available = [v for v in data if math.isfinite(v)]
    upper = max(0, max(available, default=1.0))
    lower = min(0, min(available, default=0)) if diverging else 0
    padding = max(upper - lower, 1.0) * 0.3
    y_max = math.ceil((upper + padding) * 2) / 2
    y_min = math.floor((lower - padding) * 2) / 2 if diverging else 0
    y_interval = 1.0 if y_max - y_min > 4.5 else 0.5

    fig, axes = plt.subplots(1, 2, sharey=True, figsize=(14, 5.2),
                             layout="constrained")
    figure_legend(fig, results, field)

    for ax, arch in zip(axes, ARCHITECTURES):
        selected = [row for row in results if row.arch == arch]
        ax.set_xlabel("ARCHITECTURE A" if arch == "arch-a" else "ARCHITECTURE B",
                      fontsize=11, weight="bold", labelpad=8)
        ax.set_axisbelow(True)
        ax.set_ylim(y_min, y_max)
        ax.set_xlim(-0.5, len(selected) - 0.5)
        label_groups(ax, selected)
        ax.tick_params(axis="y", labelleft=True, pad=5, length=0)
        ax.tick_params(axis="x", labelsize=9.5, length=0, pad=6)
        ax.yaxis.set_major_locator(MultipleLocator(y_interval))
        if diverging:
            ax.yaxis.set_major_formatter(FuncFormatter(lambda value, _: f"{abs(value) + 1:g}×"))
        ax.grid(axis="y", color="#E6EAF0", linewidth=0.8)
        ax.axhline(0 if diverging else 1, color="#536478", linewidth=1.2,
                   linestyle="--", zorder=3)
        ax.set_ylabel(ylabel, fontsize=10.5, labelpad=6)

        for position, row in enumerate(selected):
            ratio = getattr(row, field)
            if not math.isfinite(ratio):
                ax.text(position, (y_max - y_min) * 0.02, "N/A", ha="center", va="bottom",
                        fontsize=9.5, color="#85909F", fontstyle="italic")
                continue
            compiler = advantage_compiler(ratio)
            statuses = [status for status in (row.raptor_validation, row.pimcomp_validation) if status != "PASS"]
            status = "FAIL" if "FAIL" in statuses else next(iter(statuses), "PASS")
            pattern = STAGE_PATTERNS[row.pipeline]
            measurement = diverging_ratio(ratio) if diverging else ratio
            ax.bar(position, measurement, width=0.72,
                    color=COMPILER_COLORS.get(compiler, "#657081"),
                    hatch=pattern,
                    edgecolor=TEXT_COLOR if pattern else "none",
                    linewidth=0.7 if pattern else 0.0, zorder=2)
            direction = -1 if measurement < 0 else 1
            factor = max(ratio, 1 / ratio) if diverging else ratio
            value_label(ax, position, measurement + direction * (y_max - y_min) * 0.012,
                        f"{factor:.2f}×", status, direction, fontsize=9.1)

    fit_value_labels(fig)
    save_figure(fig, basename, output_dir)


def draw_absolute(results: list[Result], field: str, ylabel: str,
                  basename: str, output_dir: Path, scale: float = 1.0,
                  raptor_only: bool = False) -> None:
    columns = absolute_columns(results, raptor_only)
    available = [getattr(row, f"{compiler}_{field}") * scale
                 for row, compiler in columns
                 if math.isfinite(getattr(row, f"{compiler}_{field}"))]
    y_max = max(available, default=1.0) * 1.35
    fig, axes = plt.subplots(1, 2, sharey=True, figsize=(14, 5.2),
                             layout="constrained")
    figure_legend(fig, results, field, absolute=True, raptor_only=raptor_only)

    for ax, arch in zip(axes, ARCHITECTURES):
        selected = [(row, compiler) for row, compiler in columns if row.arch == arch]
        positions = []
        for index, (row, _) in enumerate(selected):
            if index == 0:
                positions.append(0.0)
                continue
            previous = selected[index - 1][0]
            gap = 0.8 if row.model != previous.model else 0.4 if row.mode != previous.mode else 0.0
            positions.append(positions[-1] + 1.0 + gap)
        ax.set_xlabel("ARCHITECTURE A" if arch == "arch-a" else "ARCHITECTURE B",
                      fontsize=11, weight="bold", labelpad=8)
        ax.set_axisbelow(True)
        ax.set_ylim(0, y_max)
        if positions:
            ax.set_xlim(positions[0] - 0.9, positions[-1] + 0.9)
        else:
            ax.set_xlim(-0.9, 0.9)
        label_groups(ax, [row for row, _ in selected], positions)
        if {row.mode for row, _ in selected} == {"latency", "throughput"}:
            centers, labels = [], []
            for (_, mode), group in groupby(enumerate(selected), key=lambda item: (item[1][0].model, item[1][0].mode)):
                group_positions = [positions[position] for position, _ in group]
                centers.append((group_positions[0] + group_positions[-1]) / 2)
                labels.append(mode.capitalize())
            ax.set_xticks(centers, labels, minor=True)
            ax.tick_params(axis="x", which="minor", length=0, pad=4, labelsize=8)
        ax.set_ylabel(ylabel, fontsize=10.5)
        if field == "throughput":
            ax.ticklabel_format(axis="y", style="plain", useOffset=False)
        ax.grid(axis="y", color="#E6EAF0", linewidth=0.8)
        ax.tick_params(axis="y", length=0, pad=6)
        ax.tick_params(axis="x", which="major", length=0,
                       pad=22 if {row.mode for row, _ in selected} == {"latency", "throughput"} else 6)

        for x, (row, compiler) in zip(positions, selected):
            measurement = getattr(row, f"{compiler}_{field}") * scale
            if not math.isfinite(measurement):
                ax.text(x, y_max * 0.02, "N/A", rotation=90, ha="center",
                        va="bottom", fontsize=8, color="#85909F")
                continue
            status = getattr(row, f"{compiler}_validation")
            pattern = STAGE_PATTERNS[row.pipeline] if compiler == "raptor" else ""
            ax.bar(x, measurement, width=1.0,
                   color=COMPILER_COLORS[compiler],
                   hatch=pattern,
                   edgecolor=TEXT_COLOR if pattern else "none",
                   linewidth=0.7 if pattern else 0.0, zorder=2)
            label = (f"{measurement:,.2f}".rstrip("0").rstrip(".")
                     if field == "throughput" else f"{measurement:.3g}")
            value_label(ax, x, measurement + y_max * 0.012, label, status)

    fit_value_labels(fig)
    save_figure(fig, basename, output_dir)


def draw_pipeline_effect(results: list[Result], output_dir: Path) -> None:
    rows = [row for row in results if row.mode == "throughput" or row.pipeline == 1]
    stages = sorted({row.pipeline for row in rows})
    panels = [(model, color) for model, color in zip(MODELS, COLORS)
              if any(row.model == model for row in rows)]
    lookup = {(row.arch, row.model, row.pipeline): row for row in rows}
    rates = {key: quotient(1000, row.raptor_latency) if row.mode == "latency" else row.raptor_throughput
             for key, row in lookup.items()}
    fig, axes = plt.subplots(len(ARCHITECTURES), len(panels), sharex=True, sharey="col",
                             figsize=(3 * len(panels), 7.2), squeeze=False, layout="constrained")
    for column, (model, color) in enumerate(panels):
        for index, arch in enumerate(ARCHITECTURES):
            ax = axes[index, column]
            keys = [(arch, model, stage) for stage in stages]
            values = [rates.get(key, math.nan) for key in keys]
            valid = [key in lookup and lookup[key].raptor_validation == "PASS" for key in keys]
            ax.plot(stages, values, color=color, linewidth=1.8)
            for stage, value, passed in zip(stages, values, valid):
                if math.isfinite(value):
                    marker = ("s" if stage == 1 else "o") if passed else "x"
                    ax.plot(stage, value, color=color, marker=marker,
                            markersize=6, linestyle="none")
            ax.set_xticks(stages)
            ax.set_xlabel(model.upper(), fontsize=10.5, weight="bold", labelpad=8)
            ax.grid(axis="y", color="#E6EAF0", linewidth=0.8)
            ax.tick_params(axis="both", length=0, pad=6, labelbottom=True, labelleft=True)
            ax.ticklabel_format(axis="y", style="plain", useOffset=False)
            if column == 0:
                ax.set_ylabel(f"{arch.replace('arch-', 'ARCHITECTURE ').upper()}\n"
                              "Throughput (inferences/s)", fontsize=10.5)
    for ax in axes.flat:
        ax.set_ylim(bottom=0)
    handles = [Line2D([], [], color="#657081", marker=marker, linestyle="none", label=label)
               for marker, label in (("s", "Derived as 1000ms / latency_ms"),
                                     ("o", "Throughput measurement"))]
    if any(row.raptor_validation != "PASS" and math.isfinite(rates[key]) for key, row in lookup.items()):
        handles.append(Line2D([], [], color="#424B58", marker="x", linestyle="none",
                              label="Raptor validation not PASS"))
    fig.legend(handles=handles, loc="outside upper center", ncol=len(handles), frameon=False, fontsize=9.5)
    fig.supxlabel("Raptor pipeline stages", fontsize=11)
    save_figure(fig, "raptor_pipeline_throughput", output_dir)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    script_dir = Path(__file__).resolve().parent
    parser.add_argument("--input", type=Path, default=script_dir / "results_comparison.csv")
    parser.add_argument("--output-dir", type=Path, default=script_dir / "plots")
    args = parser.parse_args()
    configure_style()
    results = load_results(args.input)
    valid = sum(row.comparison_passed for row in results)
    print(f"Results: {len(results)} configurations, {valid} valid comparisons, "
          f"{len(results)-valid} comparisons without two PASS statuses")
    draw_validation(results, args.output_dir)
    draw_pipeline_effect(results, args.output_dir)
    yolo_only_models = {
        model for model in MODELS if model.startswith("yolo")
        and not any(math.isfinite(getattr(row, f"pimcomp_{field}"))
                    for row in results if row.model == model
                    for field in ("latency", "power", "energy"))
    }
    comparisons = [row for row in results if row.model not in yolo_only_models]
    yolo_only = [row for row in results if row.model in yolo_only_models]
    for metric, name, label, mode in [
        ("latency_speedup", "longest_stage_latency", "Longest stage latency speedup (x)", None),
        ("latency_speedup", "latency", "Single-inference latency speedup (x)", "latency"),
        ("throughput_speedup", "throughput", "Throughput speedup (x)", "throughput"),
        ("power", "power", "Average power advantage (x)", None),
        ("energy", "energy", "Energy per inference advantage (x)", None),
    ]:
        selected = [row for row in comparisons if mode is None or row.mode == mode]
        draw_metric(selected, metric, label,
                    f"{name}_comparison", args.output_dir)
        draw_metric(selected, metric, label,
                    f"{name}_diverging", args.output_dir, diverging=True)
    for metric, name, label, scale, mode in [
        ("latency", "longest_stage_latency", "Longest stage latency (ms)", 1.0, None),
        ("latency", "latency", "Single-inference latency (ms)", 1.0, "latency"),
        ("power", "power", "Average power (mW)", 1.0, None),
        ("energy", "energy", "Energy per inference (mJ)", 1e-9, None),
        ("throughput", "throughput", "Throughput (inferences/s)", 1.0, "throughput"),
    ]:
        for suffix, group, raptor_only in (
            ("absolute", comparisons, False), ("yolo", yolo_only, True)
        ):
            selected = [row for row in group if mode is None or row.mode == mode]
            if not selected:
                continue
            draw_absolute(selected, metric, label,
                          f"{name}_{suffix}", args.output_dir, scale,
                          raptor_only=raptor_only)


if __name__ == "__main__":
    main()