plot_emax_boxplot.py

from __future__ import annotations

import argparse
import re
from collections import defaultdict
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

import matplotlib.pyplot as plt
import numpy as np
import torch

METHOD_ORDER = ["cdfquant", "uniform_renorm", "uniform_norenorm", "repq_vit"]
REQUIRED_METHOD_ORDER = ["cdfquant", "uniform_renorm", "uniform_norenorm"]
OPTIONAL_METHOD_ORDER = [m for m in METHOD_ORDER if m not in REQUIRED_METHOD_ORDER]
ALL_MODE_SUFFIX = METHOD_ORDER + ["fp"]
METHOD_TICK = {
    "cdfquant": "CDFQuant",
    "uniform_renorm": "Uniform-renorm",
    "uniform_norenorm": "Uniform-norenorm",
    "repq_vit": "RepQ-ViT",
}
METHOD_COLOR = {
    # Pastel palette requested by user (4 boxes).
    "cdfquant": "#E8F0FE",         # light blue
    "uniform_renorm": "#D4EDDA",   # light green
    "uniform_norenorm": "#FFF3CD", # light yellow
    "repq_vit": "#F8D7DA",         # light red
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Paper-style E_max box/violin plots")
    parser.add_argument("--workdir", type=str, default=".")
    parser.add_argument("--glob", type=str, default="emax*.pt", help="Glob for E_max .pt files")
    parser.add_argument("--model", type=str, default=None, help="Optional exact model-name filter")
    parser.add_argument("--bits", type=int, nargs="+", default=[3, 4])
    parser.add_argument("--layers", type=int, nargs="+", default=[3, 6, 9])
    parser.add_argument(
        "--allow-incomplete-modes",
        action="store_true",
        default=False,
        help="Allow incomplete mode groups. Default behavior requires complete "
        "groups containing cdfquant/uniform_renorm/uniform_norenorm.",
    )
    parser.add_argument("--kind", type=str, choices=["box", "violin"], default="box")
    parser.add_argument("--fig-width", type=float, default=7.0, help="Figure width in inches")
    parser.add_argument("--fig-height", type=float, default=5.2, help="Figure height in inches")
    parser.add_argument(
        "--ymax-override",
        type=str,
        nargs="*",
        default=[],
        help="Optional per-bit y-max override, e.g. '3:0.20 4:0.12'. "
        "Useful for a zoomed readability figure.",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="emax_boxplot_b34.pdf",
        help="Output PDF path. If multiple models are found, suffix _<model>.pdf is added.",
    )
    return parser.parse_args()


def _safe_name(name: str) -> str:
    name = name.strip().lower().replace(" ", "_")
    return re.sub(r"[^a-z0-9_.-]+", "_", name)


def _normalize_mode(mode: str) -> str:
    norm = mode.strip().lower().replace("-", "_")
    alias_map = {
        "cdfquant": "cdfquant",
        "uniform_renorm": "uniform_renorm",
        "uniform_renormalize": "uniform_renorm",
        "uniform_renormalized": "uniform_renorm",
        "uniform_norenorm": "uniform_norenorm",
        "uniform_no_renorm": "uniform_norenorm",
        "uniform_no_renormalize": "uniform_norenorm",
        "repq": "repq_vit",
        "repq_vit": "repq_vit",
        "repqvit": "repq_vit",
    }
    return alias_map.get(norm, norm)


def _infer_mode_from_stem(stem: str) -> str:
    s = stem.lower().replace("-", "_")
    for mode in ALL_MODE_SUFFIX:
        suffix = f"_{mode}"
        if s.endswith(suffix):
            return mode
    for token in ("repq_vit", "repqvit", "repq"):
        if token in s:
            return "repq_vit"
    return ""


def _setup_ieee_style() -> None:
    plt.rcParams.update(
        {
            "font.size": 9,
            "axes.labelsize": 9,
            "axes.titlesize": 9,
            "xtick.labelsize": 8,
            "ytick.labelsize": 8,
            "font.family": "DejaVu Sans",
            "figure.facecolor": "white",
            "axes.facecolor": "white",
            "axes.linewidth": 0.8,
            "grid.linewidth": 0.6,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )


def _theoretical_bound(bits: int) -> float:
    return 1.0 / (2.0 * ((2 ** bits) - 1))


def parse_ymax_override(items: List[str]) -> Dict[int, float]:
    override: Dict[int, float] = {}
    for item in items:
        if ":" not in item:
            raise ValueError(f"Invalid --ymax-override item '{item}', expected <bit>:<value>")
        bit_str, value_str = item.split(":", 1)
        bit = int(bit_str.strip())
        value = float(value_str.strip())
        if value <= 0:
            raise ValueError(f"Invalid y-max value {value} for bit {bit}, must be > 0")
        override[bit] = value
    return override


def _load_pt(path: Path) -> Dict[str, object] | None:
    try:
        obj = torch.load(path, map_location="cpu")
    except Exception as exc:
        print(f"[Warn] Skip unreadable file {path}: {exc}")
        return None
    if not isinstance(obj, dict):
        print(f"[Warn] Skip unexpected format {path}")
        return None
    return obj


def _infer_run_key(path: Path) -> str:
    stem = path.stem
    for mode in ALL_MODE_SUFFIX:
        suffix = f"_{mode}"
        if stem.endswith(suffix):
            return stem[: -len(suffix)]
    return stem


def collect_model_distributions(
    paths: Iterable[Path],
    bits_filter: List[int],
    layers_filter: List[int],
    model_filter: str | None,
    require_complete_modes: bool,
) -> Dict[str, Dict[Tuple[int, int, str], np.ndarray]]:
    grouped_files: Dict[
        Tuple[str, int, str], List[Tuple[Path, Dict[str, object]]]
    ] = defaultdict(list)

    for path in paths:
        obj = _load_pt(path)
        if obj is None:
            continue

        model = str(obj.get("model", "unknown"))
        mode = _normalize_mode(str(obj.get("mode", "")))
        if mode not in METHOD_ORDER:
            mode = _infer_mode_from_stem(path.stem)
        bits = int(obj.get("bits", -1))

        if model_filter is not None and model != model_filter:
            continue
        if bits not in bits_filter:
            continue
        if mode not in METHOD_ORDER:
            continue
        run_key = _infer_run_key(path)
        grouped_files[(model, bits, run_key)].append((path, obj))

    valid_items: List[Tuple[Path, Dict[str, object]]] = []
    pending_optional: List[Tuple[Tuple[str, int, str], List[Tuple[Path, Dict[str, object]]]]] = []
    complete_pairs: set[Tuple[str, int]] = set()
    skipped_groups = 0
    for (model, bits, run_key), file_objs in grouped_files.items():
        modes_in_group: set[str] = set()
        for path, obj in file_objs:
            mode = _normalize_mode(str(obj.get("mode", "")))
            if mode not in METHOD_ORDER:
                mode = _infer_mode_from_stem(path.stem)
            if mode in METHOD_ORDER:
                modes_in_group.add(mode)

        if require_complete_modes and not set(REQUIRED_METHOD_ORDER).issubset(modes_in_group):
            if modes_in_group and modes_in_group.issubset(set(OPTIONAL_METHOD_ORDER)):
                pending_optional.append(((model, bits, run_key), file_objs))
            else:
                skipped_groups += 1
            continue

        valid_items.extend(file_objs)
        complete_pairs.add((model, bits))

    if require_complete_modes:
        for (model, bits, _run_key), file_objs in pending_optional:
            if (model, bits) in complete_pairs:
                valid_items.extend(file_objs)

    if require_complete_modes and skipped_groups > 0:
        print(
            f"[Info] Skipped {skipped_groups} incomplete file groups "
            "(missing one or more required modes)."
        )

    raw: Dict[str, Dict[Tuple[int, int, str], List[np.ndarray]]] = defaultdict(
        lambda: defaultdict(list)
    )
    for path, obj in valid_items:
        model = str(obj.get("model", "unknown"))
        mode = _normalize_mode(str(obj.get("mode", "")))
        if mode not in METHOD_ORDER:
            mode = _infer_mode_from_stem(path.stem)
        if mode not in METHOD_ORDER:
            continue
        bits = int(obj.get("bits", -1))
        layers = obj.get("layers", {})
        if not isinstance(layers, dict):
            continue

        for layer_key, value in layers.items():
            try:
                layer = int(layer_key)
            except Exception:
                continue
            if layer not in layers_filter:
                continue
            if not torch.is_tensor(value):
                continue

            arr = value.detach().to(torch.float32).reshape(-1).cpu().numpy()
            if arr.size == 0:
                continue
            raw[model][(bits, layer, mode)].append(arr)

    merged: Dict[str, Dict[Tuple[int, int, str], np.ndarray]] = {}
    for model, by_key in raw.items():
        merged[model] = {}
        for key, arrays in by_key.items():
            merged[model][key] = np.concatenate(arrays, axis=0)

    return merged


def _plot_one_axis(
    ax: plt.Axes,
    distributions: Dict[Tuple[int, int, str], np.ndarray],
    bit: int,
    layer: int,
    kind: str,
    methods: List[str],
) -> None:
    methods = methods or METHOD_ORDER
    values: List[np.ndarray] = []
    positions: List[int] = []

    for i, method in enumerate(methods, start=1):
        arr = distributions.get((bit, layer, method))
        if arr is None or arr.size == 0:
            continue
        values.append(arr)
        positions.append(i)

    ax.set_axisbelow(True)
    ax.set_xticks(range(1, len(methods) + 1))
    ax.set_xticklabels([METHOD_TICK[m] for m in methods], rotation=18, ha="right")
    ax.grid(axis="y", color="#ECECEC", linewidth=0.7, alpha=1.0)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_color("#6E6E6E")
    ax.spines["bottom"].set_color("#6E6E6E")

    bound = _theoretical_bound(bit)
    ax.axhline(bound, color="#C00000", linestyle=(0, (4, 2)), linewidth=1.1)
    ax.text(
        0.03,
        0.96,
        f"$\\Delta/2$={bound:.4f}",
        transform=ax.transAxes,
        ha="left",
        va="top",
        fontsize=8,
        color="red",
    )

    if not values:
        ax.text(0.5, 0.5, "No data", transform=ax.transAxes, ha="center", va="center", color="#666666")
        return

    if kind == "box":
        bp = ax.boxplot(
            values,
            positions=positions,
            widths=0.58,
            patch_artist=True,
            showfliers=False,
            whis=(5, 95),
            medianprops={"color": "#FFFFFF", "linewidth": 1.4},
            whiskerprops={"linewidth": 0.95, "color": "#666666"},
            capprops={"linewidth": 0.95, "color": "#666666"},
            showmeans=True,
            meanprops={
                "marker": "D",
                "markersize": 3.0,
                "markerfacecolor": "#202020",
                "markeredgecolor": "#FFFFFF",
                "markeredgewidth": 0.35,
            },
        )
        for patch, pos in zip(bp["boxes"], positions):
            method = methods[pos - 1]
            patch.set_facecolor(METHOD_COLOR[method])
            patch.set_edgecolor("#2F2F2F")
            patch.set_linewidth(0.9)
            patch.set_alpha(0.92)
    else:
        vp = ax.violinplot(
            values,
            positions=positions,
            widths=0.8,
            showmeans=False,
            showmedians=True,
            showextrema=False,
        )
        for body, pos in zip(vp["bodies"], positions):
            method = methods[pos - 1]
            body.set_facecolor(METHOD_COLOR[method])
            body.set_edgecolor("#2F2F2F")
            body.set_alpha(0.88)
        if "cmedians" in vp:
            vp["cmedians"].set_color("#FFFFFF")
            vp["cmedians"].set_linewidth(1.2)


def _row_ymax(
    distributions: Dict[Tuple[int, int, str], np.ndarray],
    bit: int,
    layers: List[int],
    override: Dict[int, float],
    methods: List[str],
) -> float:
    if bit in override:
        return override[bit]

    row_arrays: List[np.ndarray] = []
    for layer in layers:
        for mode in methods:
            arr = distributions.get((bit, layer, mode))
            if arr is not None and arr.size > 0:
                row_arrays.append(arr)

    bound = _theoretical_bound(bit)
    if not row_arrays:
        return max(bound * 1.4, 0.1)

    concat = np.concatenate(row_arrays, axis=0)
    q995 = float(np.quantile(concat, 0.995))
    ymax = max(q995 * 1.10, bound * 1.35, 0.08)
    return min(ymax, 1.0)


def make_figure(
    model: str,
    distributions: Dict[Tuple[int, int, str], np.ndarray],
    bits: List[int],
    layers: List[int],
    kind: str,
    fig_width: float,
    fig_height: float,
    ymax_override: Dict[int, float],
    output: Path,
) -> None:
    _setup_ieee_style()
    methods = [
        m
        for m in METHOD_ORDER
        if any(
            (distributions.get((bit, layer, m)) is not None)
            and (distributions[(bit, layer, m)].size > 0)
            for bit in bits
            for layer in layers
        )
    ]
    if not methods:
        methods = METHOD_ORDER

    nrows = len(bits)
    ncols = len(layers)
    fig, axes = plt.subplots(nrows, ncols, figsize=(fig_width, fig_height), squeeze=False)

    for r, bit in enumerate(bits):
        ymax = _row_ymax(
            distributions=distributions,
            bit=bit,
            layers=layers,
            override=ymax_override,
            methods=methods,
        )
        for c, layer in enumerate(layers):
            ax = axes[r][c]
            _plot_one_axis(
                ax,
                distributions=distributions,
                bit=bit,
                layer=layer,
                kind=kind,
                methods=methods,
            )
            ax.set_ylim(0.0, ymax)

            if r == 0:
                ax.set_title(f"Layer {layer}")
            if c == 0:
                ax.set_ylabel(f"{bit}-bit\n$E_{{max}}$")

    legend_handles = [
        plt.Line2D(
            [0],
            [0],
            color=METHOD_COLOR[m],
            lw=6,
            solid_capstyle="butt",
            label=METHOD_TICK[m],
        )
        for m in methods
    ]
    fig.legend(
        handles=legend_handles,
        ncol=max(1, len(methods)),
        loc="upper center",
        frameon=False,
        bbox_to_anchor=(0.5, 1.01),
        columnspacing=1.2,
        handlelength=1.7,
    )
    fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.95])
    fig.savefig(output, bbox_inches="tight")
    plt.close(fig)
    print(f"[Saved] {output}")


def main() -> None:
    args = parse_args()
    workdir = Path(args.workdir).resolve()
    out_base = Path(args.output).resolve()
    ymax_override = parse_ymax_override(args.ymax_override)
    require_complete_modes = not args.allow_incomplete_modes

    paths = sorted(workdir.glob(args.glob))
    print(f"[Info] Found {len(paths)} files with pattern '{args.glob}' in {workdir}")

    model_data = collect_model_distributions(
        paths=paths,
        bits_filter=args.bits,
        layers_filter=args.layers,
        model_filter=args.model,
        require_complete_modes=require_complete_modes,
    )

    if not model_data:
        print("[Error] No matched E_max data found. Check --glob/--model/--bits/--layers.")
        return

    models = sorted(model_data.keys())
    print(f"[Info] Models found: {', '.join(models)}")

    if len(models) == 1:
        model = models[0]
        make_figure(
            model=model,
            distributions=model_data[model],
            bits=args.bits,
            layers=args.layers,
            kind=args.kind,
            fig_width=args.fig_width,
            fig_height=args.fig_height,
            ymax_override=ymax_override,
            output=out_base,
        )
        return

    # Multiple models: export one PDF per model to avoid mixed, messy distributions.
    for model in models:
        out_path = out_base.with_name(f"{out_base.stem}_{_safe_name(model)}{out_base.suffix}")
        make_figure(
            model=model,
            distributions=model_data[model],
            bits=args.bits,
            layers=args.layers,
            kind=args.kind,
            fig_width=args.fig_width,
            fig_height=args.fig_height,
            ymax_override=ymax_override,
            output=out_path,
        )


if __name__ == "__main__":
    main()