#!/usr/bin/env python3
"""
Visualize calibration drift across grouped splits such as train / case / clip.

The script is designed for split-based calibration summary tables like
`calib_profiling_results/train_calib_table.csv`, where each row belongs to a
group in the `split` column and contains calibration parameters as numeric
columns.

Outputs:
  1. `split_parameter_panels.png`
     Per-parameter strip plots with the baseline group's mean ± std overlay.
  2. `group_zscore_heatmap.png`
     Mean parameter drift of each non-baseline group relative to the baseline.
  3. `group_drift_ranking.png`
     Ranked horizontal bar charts of the strongest drifts per group.
  4. `parameter_pair_relationships.png`
     Selected 2D parameter relationships colored by split.
  5. `group_comparison_report.txt`
     Compact text summary with counts and top drifts.

Usage:
    python3 tools/scripts_for_gt/profiling/visualize_calib_group_comparison.py \
        --input calib_profiling_results/train_calib_table.csv \
        --output-dir calib_profiling_results/group_comparison
"""

from __future__ import annotations

import argparse
import csv
import math
from collections import defaultdict
from pathlib import Path
from typing import Dict, Iterable, List, Sequence

try:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
except ImportError as exc:
    raise SystemExit(
        "matplotlib is required for visualization. "
        "Install the project plotting dependencies first, for example:\n"
        "  pip install -r requirements.txt"
    ) from exc

try:
    import numpy as np
except ImportError as exc:
    raise SystemExit(
        "numpy is required for visualization. "
        "Install the project dependencies first, for example:\n"
        "  pip install -r requirements.txt"
    ) from exc


DEFAULT_GROUP_ORDER = ("train", "case", "clip")
DEFAULT_PAIR_SPECS = (
    ("focal_u", "focal_v"),
    ("cu", "cv"),
    ("yaw", "vanish_x"),
    ("pitch", "vanish_y"),
    ("pos_x", "pos_z"),
    ("roll", "pos_z"),
)
PARAMETER_GROUPS = {
    "Intrinsics": ("focal_u", "focal_v", "cu", "cv"),
    "Angles": ("pitch", "yaw", "roll"),
    "Position": ("pos_x", "pos_y", "pos_z"),
    "Vanishing Point": ("vanish_x", "vanish_y"),
}
PALETTE = {
    "train": "#1f77b4",
    "case": "#d62728",
    "clip": "#2ca02c",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Visualize calibration drift across grouped splits."
    )
    parser.add_argument(
        "--input",
        type=str,
        default="calib_profiling_results/train_calib_table.csv",
        help="Input CSV with split-based calibration rows.",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="calib_profiling_results/group_comparison",
        help="Directory where plots and report will be written.",
    )
    parser.add_argument(
        "--split-column",
        type=str,
        default="split",
        help="Column containing group labels.",
    )
    parser.add_argument(
        "--id-column",
        type=str,
        default="vehicle/date",
        help="Column containing row identifiers for optional annotations.",
    )
    parser.add_argument(
        "--baseline-split",
        type=str,
        default="train",
        help="Baseline split used for drift comparison.",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=8,
        help="How many parameters to keep in each drift ranking panel.",
    )
    parser.add_argument(
        "--figure-dpi",
        type=int,
        default=180,
        help="Saved figure DPI.",
    )
    parser.add_argument(
        "--annotate-pairs",
        action="store_true",
        help="Annotate pair plots with row identifiers when the table is small.",
    )
    return parser.parse_args()


def _to_float(value: str) -> float:
    if value is None:
        return math.nan
    text = str(value).strip()
    if not text:
        return math.nan
    try:
        return float(text)
    except ValueError:
        return math.nan


def _mean(values: Sequence[float]) -> float:
    if not values:
        return math.nan
    return float(sum(values) / len(values))


def _std(values: Sequence[float]) -> float:
    if not values:
        return math.nan
    mu = _mean(values)
    return float(math.sqrt(sum((v - mu) ** 2 for v in values) / len(values)))


def ordered_splits(available: Iterable[str], baseline_split: str) -> List[str]:
    available_unique = list(dict.fromkeys(available))
    ordered = [s for s in DEFAULT_GROUP_ORDER if s in available_unique]
    if baseline_split in available_unique and baseline_split not in ordered:
        ordered.insert(0, baseline_split)
    ordered.extend(s for s in available_unique if s not in ordered)
    return ordered


def ordered_parameters(columns: Sequence[str]) -> List[str]:
    seen = set()
    ordered: List[str] = []
    for group_cols in PARAMETER_GROUPS.values():
        for col in group_cols:
            if col in columns and col not in seen:
                ordered.append(col)
                seen.add(col)
    for col in columns:
        if col not in seen:
            ordered.append(col)
            seen.add(col)
    return ordered


def load_table(
    csv_path: Path, split_col: str, id_col: str
) -> tuple[List[dict], List[str], List[str]]:
    with csv_path.open("r", newline="") as f:
        reader = csv.DictReader(f)
        raw_rows = list(reader)
        if not raw_rows:
            raise ValueError(f"No rows found in {csv_path}")

    columns = list(raw_rows[0].keys())
    for required in (split_col, id_col):
        if required not in columns:
            raise KeyError(
                f"Required column '{required}' not found in {csv_path}. "
                f"Available columns: {columns}"
            )

    candidate_numeric = [c for c in columns if c not in (split_col, id_col)]
    numeric_cols = []
    for col in candidate_numeric:
        if any(not math.isnan(_to_float(row.get(col, ""))) for row in raw_rows):
            numeric_cols.append(col)

    parsed_rows: List[dict] = []
    for row in raw_rows:
        parsed = {
            split_col: row.get(split_col, "").strip(),
            id_col: row.get(id_col, "").strip(),
        }
        for col in numeric_cols:
            parsed[col] = _to_float(row.get(col, ""))
        parsed_rows.append(parsed)

    splits = ordered_splits(
        [row[split_col] for row in parsed_rows if row[split_col]], baseline_split=""
    )
    return parsed_rows, splits, ordered_parameters(numeric_cols)


def build_group_values(
    rows: Sequence[dict], splits: Sequence[str], split_col: str, params: Sequence[str]
) -> Dict[str, Dict[str, List[float]]]:
    grouped: Dict[str, Dict[str, List[float]]] = {
        split: {param: [] for param in params} for split in splits
    }
    for row in rows:
        split = row.get(split_col, "")
        if split not in grouped:
            continue
        for param in params:
            value = row.get(param, math.nan)
            if not math.isnan(value):
                grouped[split][param].append(value)
    return grouped


def compute_baseline_stats(
    grouped_values: Dict[str, Dict[str, List[float]]],
    baseline_split: str,
    params: Sequence[str],
) -> Dict[str, dict]:
    baseline = grouped_values.get(baseline_split)
    if baseline is None:
        raise KeyError(
            f"Baseline split '{baseline_split}' not found in input data. "
            f"Available splits: {list(grouped_values.keys())}"
        )

    stats = {}
    for param in params:
        values = baseline[param]
        stats[param] = {
            "count": len(values),
            "mean": _mean(values),
            "std": _std(values),
            "min": min(values) if values else math.nan,
            "max": max(values) if values else math.nan,
        }
    return stats


def compute_group_means(
    grouped_values: Dict[str, Dict[str, List[float]]], params: Sequence[str]
) -> Dict[str, Dict[str, float]]:
    means: Dict[str, Dict[str, float]] = {}
    for split, param_map in grouped_values.items():
        means[split] = {param: _mean(param_map[param]) for param in params}
    return means


def compute_group_zscores(
    group_means: Dict[str, Dict[str, float]],
    baseline_stats: Dict[str, dict],
    splits: Sequence[str],
    params: Sequence[str],
) -> Dict[str, Dict[str, float]]:
    zscores: Dict[str, Dict[str, float]] = {}
    for split in splits:
        zscores[split] = {}
        for param in params:
            baseline_mean = baseline_stats[param]["mean"]
            baseline_std = baseline_stats[param]["std"]
            group_mean = group_means[split][param]
            if math.isnan(group_mean) or math.isnan(baseline_mean):
                zscores[split][param] = math.nan
            elif baseline_std == 0 or math.isnan(baseline_std):
                zscores[split][param] = 0.0 if group_mean == baseline_mean else math.nan
            else:
                zscores[split][param] = (group_mean - baseline_mean) / baseline_std
    return zscores


def format_value(value: float) -> str:
    return "nan" if math.isnan(value) else f"{value:.3f}"


def short_label(text: str, max_len: int = 28) -> str:
    if len(text) <= max_len:
        return text
    return text[: max_len - 3] + "..."


def plot_split_parameter_panels(
    rows: Sequence[dict],
    splits: Sequence[str],
    params: Sequence[str],
    split_col: str,
    baseline_stats: Dict[str, dict],
    output_path: Path,
    dpi: int,
) -> None:
    n_params = len(params)
    ncols = 4
    nrows = math.ceil(n_params / ncols)
    fig, axes = plt.subplots(
        nrows,
        ncols,
        figsize=(4.3 * ncols, 3.4 * nrows),
        constrained_layout=True,
    )
    axes = np.atleast_1d(axes).reshape(nrows, ncols)
    rng = np.random.default_rng(20260318)

    for idx, param in enumerate(params):
        ax = axes[idx // ncols][idx % ncols]
        train_mean = baseline_stats[param]["mean"]
        train_std = baseline_stats[param]["std"]

        if not math.isnan(train_mean):
            ax.axhline(train_mean, color="#444444", linewidth=1.2, linestyle="--")
            if not math.isnan(train_std) and train_std > 0:
                ax.axhspan(
                    train_mean - train_std,
                    train_mean + train_std,
                    color="#444444",
                    alpha=0.10,
                )

        for x_idx, split in enumerate(splits):
            values = [
                row[param]
                for row in rows
                if row.get(split_col) == split and not math.isnan(row.get(param, math.nan))
            ]
            if not values:
                continue
            x = x_idx + rng.uniform(-0.12, 0.12, size=len(values))
            color = PALETTE.get(split, f"C{x_idx}")
            ax.scatter(
                x,
                values,
                s=42,
                alpha=0.85,
                linewidth=0.5,
                edgecolors="white",
                color=color,
                zorder=3,
            )
            ax.scatter(
                [x_idx],
                [_mean(values)],
                marker="D",
                s=52,
                color=color,
                edgecolors="black",
                linewidth=0.7,
                zorder=4,
            )

        ax.set_title(param, fontsize=11, pad=8)
        ax.set_xticks(range(len(splits)))
        ax.set_xticklabels(splits, rotation=20)
        ax.grid(axis="y", linestyle=":", alpha=0.35)
        ax.set_axisbelow(True)

    for idx in range(n_params, nrows * ncols):
        axes[idx // ncols][idx % ncols].axis("off")

    fig.suptitle(
        "Calibration Parameter Comparison by Split\n"
        "Dashed line and shaded band show baseline mean ± std",
        fontsize=14,
    )
    fig.savefig(output_path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)


def plot_group_zscore_heatmap(
    splits: Sequence[str],
    params: Sequence[str],
    zscores: Dict[str, Dict[str, float]],
    baseline_split: str,
    output_path: Path,
    dpi: int,
) -> None:
    target_splits = [split for split in splits if split != baseline_split]
    if not target_splits:
        return

    matrix = np.array(
        [[zscores[split][param] for param in params] for split in target_splits],
        dtype=float,
    )
    finite_vals = matrix[np.isfinite(matrix)]
    vmax = max(1.0, float(np.max(np.abs(finite_vals)))) if finite_vals.size else 1.0

    fig_w = max(10.5, 0.9 * len(params))
    fig_h = max(3.6, 1.3 * len(target_splits))
    fig, ax = plt.subplots(figsize=(fig_w, fig_h), constrained_layout=True)

    im = ax.imshow(matrix, cmap="coolwarm", vmin=-vmax, vmax=vmax, aspect="auto")
    ax.set_xticks(np.arange(len(params)))
    ax.set_xticklabels(params, rotation=35, ha="right")
    ax.set_yticks(np.arange(len(target_splits)))
    ax.set_yticklabels(target_splits)
    ax.set_title(
        f"Mean Parameter Drift vs {baseline_split} baseline (z-score units)", fontsize=13
    )

    for i in range(len(target_splits)):
        for j in range(len(params)):
            value = matrix[i, j]
            text = "nan" if not np.isfinite(value) else f"{value:.2f}"
            color = "white" if np.isfinite(value) and abs(value) > vmax * 0.55 else "black"
            ax.text(j, i, text, ha="center", va="center", fontsize=8, color=color)

    cbar = fig.colorbar(im, ax=ax, shrink=0.9)
    cbar.set_label("z-score of group mean relative to baseline", rotation=90)
    fig.savefig(output_path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)


def plot_group_drift_ranking(
    splits: Sequence[str],
    params: Sequence[str],
    zscores: Dict[str, Dict[str, float]],
    baseline_split: str,
    output_path: Path,
    dpi: int,
    top_k: int,
) -> None:
    target_splits = [split for split in splits if split != baseline_split]
    if not target_splits:
        return

    fig, axes = plt.subplots(
        len(target_splits),
        1,
        figsize=(10.5, 3.1 * len(target_splits)),
        constrained_layout=True,
    )
    axes = np.atleast_1d(axes)

    for ax, split in zip(axes, target_splits):
        scored = []
        for param in params:
            value = zscores[split][param]
            if math.isnan(value):
                continue
            scored.append((param, value))
        scored.sort(key=lambda item: abs(item[1]), reverse=True)
        scored = scored[:top_k]

        if not scored:
            ax.axis("off")
            continue

        labels = [param for param, _ in scored][::-1]
        values = [value for _, value in scored][::-1]
        colors = ["#d62728" if value > 0 else "#1f77b4" for value in values]

        ax.barh(labels, values, color=colors, alpha=0.88)
        ax.axvline(0.0, color="#333333", linewidth=1.0)
        ax.grid(axis="x", linestyle=":", alpha=0.35)
        ax.set_title(f"Top parameter drift for split: {split}", fontsize=12)
        ax.set_xlabel("Mean drift relative to baseline (z-score units)")

        for y_idx, value in enumerate(values):
            text_x = value + 0.05 if value >= 0 else value - 0.05
            align = "left" if value >= 0 else "right"
            ax.text(text_x, y_idx, f"{value:.2f}", va="center", ha=align, fontsize=9)

    fig.suptitle("Ranked group drift relative to baseline", fontsize=14)
    fig.savefig(output_path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)


def plot_parameter_pair_relationships(
    rows: Sequence[dict],
    splits: Sequence[str],
    split_col: str,
    id_col: str,
    output_path: Path,
    dpi: int,
    annotate_pairs: bool,
) -> None:
    available_pairs = []
    for x_col, y_col in DEFAULT_PAIR_SPECS:
        if all(x_col in row and y_col in row for row in rows):
            if any(not math.isnan(row.get(x_col, math.nan)) for row in rows) and any(
                not math.isnan(row.get(y_col, math.nan)) for row in rows
            ):
                available_pairs.append((x_col, y_col))

    if not available_pairs:
        return

    ncols = 2
    nrows = math.ceil(len(available_pairs) / ncols)
    fig, axes = plt.subplots(
        nrows,
        ncols,
        figsize=(6.0 * ncols, 4.8 * nrows),
        constrained_layout=True,
    )
    axes = np.atleast_1d(axes).reshape(nrows, ncols)

    should_annotate = annotate_pairs or len(rows) <= 24

    for idx, (x_col, y_col) in enumerate(available_pairs):
        ax = axes[idx // ncols][idx % ncols]
        for split_idx, split in enumerate(splits):
            split_rows = [
                row
                for row in rows
                if row.get(split_col) == split
                and not math.isnan(row.get(x_col, math.nan))
                and not math.isnan(row.get(y_col, math.nan))
            ]
            if not split_rows:
                continue

            xs = [row[x_col] for row in split_rows]
            ys = [row[y_col] for row in split_rows]
            ax.scatter(
                xs,
                ys,
                s=54,
                alpha=0.88,
                linewidth=0.6,
                edgecolors="white",
                color=PALETTE.get(split, f"C{split_idx}"),
                label=split,
            )

            if should_annotate:
                for row in split_rows:
                    label = short_label(row.get(id_col, ""), max_len=18)
                    ax.annotate(
                        label,
                        (row[x_col], row[y_col]),
                        textcoords="offset points",
                        xytext=(4, 4),
                        fontsize=7,
                        alpha=0.80,
                    )

        ax.set_title(f"{x_col} vs {y_col}", fontsize=11)
        ax.set_xlabel(x_col)
        ax.set_ylabel(y_col)
        ax.grid(linestyle=":", alpha=0.35)

    for idx in range(len(available_pairs), nrows * ncols):
        axes[idx // ncols][idx % ncols].axis("off")

    handles, labels = axes[0][0].get_legend_handles_labels()
    if handles:
        fig.legend(handles, labels, loc="upper center", ncol=min(len(labels), 4), frameon=False)

    fig.suptitle("Selected parameter relationships by split", fontsize=14)
    fig.savefig(output_path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)


def write_report(
    output_path: Path,
    csv_path: Path,
    splits: Sequence[str],
    rows: Sequence[dict],
    split_col: str,
    params: Sequence[str],
    baseline_split: str,
    grouped_values: Dict[str, Dict[str, List[float]]],
    baseline_stats: Dict[str, dict],
    group_means: Dict[str, Dict[str, float]],
    zscores: Dict[str, Dict[str, float]],
    top_k: int,
) -> None:
    counts = defaultdict(int)
    for row in rows:
        counts[row.get(split_col, "")] += 1

    lines = []
    lines.append("=" * 72)
    lines.append("CALIBRATION GROUP COMPARISON REPORT")
    lines.append("=" * 72)
    lines.append(f"Input CSV               : {csv_path}")
    lines.append(f"Baseline split          : {baseline_split}")
    lines.append("")
    lines.append("Split counts:")
    for split in splits:
        lines.append(f"  {split:<24} {counts.get(split, 0)}")
    lines.append("")
    lines.append("Baseline summary:")
    for param in params:
        stat = baseline_stats[param]
        lines.append(
            f"  {param:<24} n={stat['count']:<3d} "
            f"mean={format_value(stat['mean'])} std={format_value(stat['std'])} "
            f"range=[{format_value(stat['min'])}, {format_value(stat['max'])}]"
        )

    for split in splits:
        if split == baseline_split:
            continue
        lines.append("")
        lines.append(f"Top drift vs baseline for split '{split}':")
        ranked = []
        for param in params:
            z_value = zscores[split][param]
            group_mean = group_means[split][param]
            base_mean = baseline_stats[param]["mean"]
            if math.isnan(z_value) or math.isnan(group_mean) or math.isnan(base_mean):
                continue
            ranked.append((param, z_value, group_mean - base_mean))
        ranked.sort(key=lambda item: abs(item[1]), reverse=True)
        for param, z_value, delta in ranked[:top_k]:
            lines.append(
                f"  {param:<24} delta={delta:+.3f}   z={z_value:+.2f}"
            )

    output_path.write_text("\n".join(lines) + "\n")


def main() -> None:
    args = parse_args()
    csv_path = Path(args.input)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    rows, file_splits, params = load_table(
        csv_path=csv_path,
        split_col=args.split_column,
        id_col=args.id_column,
    )
    splits = ordered_splits(file_splits, args.baseline_split)
    grouped_values = build_group_values(rows, splits, args.split_column, params)
    baseline_stats = compute_baseline_stats(grouped_values, args.baseline_split, params)
    group_means = compute_group_means(grouped_values, params)
    zscores = compute_group_zscores(group_means, baseline_stats, splits, params)

    plot_split_parameter_panels(
        rows=rows,
        splits=splits,
        params=params,
        split_col=args.split_column,
        baseline_stats=baseline_stats,
        output_path=output_dir / "split_parameter_panels.png",
        dpi=args.figure_dpi,
    )
    plot_group_zscore_heatmap(
        splits=splits,
        params=params,
        zscores=zscores,
        baseline_split=args.baseline_split,
        output_path=output_dir / "group_zscore_heatmap.png",
        dpi=args.figure_dpi,
    )
    plot_group_drift_ranking(
        splits=splits,
        params=params,
        zscores=zscores,
        baseline_split=args.baseline_split,
        output_path=output_dir / "group_drift_ranking.png",
        dpi=args.figure_dpi,
        top_k=args.top_k,
    )
    plot_parameter_pair_relationships(
        rows=rows,
        splits=splits,
        split_col=args.split_column,
        id_col=args.id_column,
        output_path=output_dir / "parameter_pair_relationships.png",
        dpi=args.figure_dpi,
        annotate_pairs=args.annotate_pairs,
    )
    write_report(
        output_path=output_dir / "group_comparison_report.txt",
        csv_path=csv_path,
        splits=splits,
        rows=rows,
        split_col=args.split_column,
        params=params,
        baseline_split=args.baseline_split,
        grouped_values=grouped_values,
        baseline_stats=baseline_stats,
        group_means=group_means,
        zscores=zscores,
        top_k=args.top_k,
    )

    print("Saved outputs to:", output_dir)
    print(" - split_parameter_panels.png")
    print(" - group_zscore_heatmap.png")
    print(" - group_drift_ranking.png")
    print(" - parameter_pair_relationships.png")
    print(" - group_comparison_report.txt")


if __name__ == "__main__":
    main()
