scripts/visualize.py

"""
Visualization and analysis (Person E).

Tasks: training curves (TensorBoard or training_summary.json), experiment comparison plots,
cross-attention heatmaps, translation-example Markdown.

Examples:
    python scripts/visualize.py --task training_curves --log-dir outputs/exp1/logs
    python scripts/visualize.py --task training_curves --summary-json checkpoints/training_summary.json
    python scripts/visualize.py --task comparison --results-dir outputs/experiments
    python scripts/visualize.py --task attention --config configs/default_config.yaml \\
        --checkpoint checkpoints/best_model.pt --src "Hello ." --tgt "Hi there ."
    python scripts/visualize.py --task examples --config configs/default_config.yaml \\
        --checkpoint checkpoints/best_model.pt --pairs-json result/translation_pairs.example.json
"""

from __future__ import annotations

import argparse
import json
import logging
import math
import sys
from pathlib import Path
from typing import Any, Optional

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import yaml
from omegaconf import OmegaConf

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))

REPO_ROOT = Path(__file__).resolve().parent.parent

logger = logging.getLogger(__name__)

plt.rcParams["axes.unicode_minus"] = False


def _format_metrics_table(labels: list[str], keys: list[str], metrics_map: dict[str, list[float]]) -> str:
    """Space-padded table (no tabs; Matplotlib renders tabs poorly)."""
    header = ["Experiment"] + [k.upper() for k in keys]
    rows: list[list[str]] = [header]
    for i, lab in enumerate(labels):
        rows.append(
            [str(lab)]
            + [f"{metrics_map[k][i]:.4f}" if i < len(metrics_map[k]) else "-" for k in keys]
        )
    ncols = len(header)
    widths = [max(len(rows[r][c]) for r in range(len(rows))) for c in range(ncols)]
    out_lines = []
    for row in rows:
        out_lines.append("  ".join(row[c].ljust(widths[c]) for c in range(ncols)))
    return "\n".join(out_lines)


def _read_training_summary(path: Path) -> dict[str, Any]:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _read_tensorboard_scalars(log_dir: Path) -> dict[str, tuple[list[int], list[float]]]:
    try:
        from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
    except ImportError as e:
        raise ImportError(
            "TensorBoard is required for --log-dir. Install with: pip install tensorboard"
        ) from e

    series: dict[str, tuple[list[int], list[float]]] = {}
    log_dir = Path(log_dir)
    if not log_dir.exists():
        return series

    ea = EventAccumulator(str(log_dir), size_guidance={"scalars": 0})
    ea.Reload()
    for tag in ea.Tags().get("scalars", []):
        events = ea.Scalars(tag)
        steps = [e.step for e in events]
        vals = [e.value for e in events]
        series[tag] = (steps, vals)
    return series


def plot_training_curves(
    log_dir: Optional[str] = None,
    output_path: str = "outputs/training_curves.png",
    summary_json: Optional[str] = None,
) -> Path:
    """Plot training curves from training_summary.json or TensorBoard scalars."""
    out = Path(output_path)
    out.parent.mkdir(parents=True, exist_ok=True)

    fig, axes = plt.subplots(2, 2, figsize=(12, 8))

    if summary_json:
        summ_path = Path(summary_json)
        if not summ_path.is_file():
            raise FileNotFoundError(f"training_summary.json not found: {summ_path}")
        data = _read_training_summary(summ_path)
        epochs = list(range(1, len(data.get("train_loss_history", [])) + 1))
        tl = data.get("train_loss_history", [])
        axes[0, 0].plot(epochs, tl, marker="o")
        axes[0, 0].set_title("Train Loss (per epoch)")
        axes[0, 0].set_xlabel("Epoch")
        axes[0, 0].set_ylabel("Loss")
        axes[0, 0].grid(True, alpha=0.3)

        vm = data.get("val_metrics_history", [])
        if vm:
            val_loss = [m.get("val_loss", float("nan")) for m in vm]
            axes[0, 1].plot(range(1, len(val_loss) + 1), val_loss, marker="o", color="tab:orange")
            axes[0, 1].set_title("Val Loss")
            axes[0, 1].set_xlabel("Epoch")
            axes[0, 1].grid(True, alpha=0.3)

            bleu = [m.get("bleu") for m in vm if isinstance(m.get("bleu"), (int, float))]
            if bleu:
                axes[1, 0].plot(range(1, len(bleu) + 1), bleu, marker="o", color="tab:green")
                axes[1, 0].set_title("BLEU (validation)")
                axes[1, 0].set_xlabel("Epoch")
                axes[1, 0].grid(True, alpha=0.3)
            else:
                axes[1, 0].text(0.5, 0.5, "No BLEU in validation logs", ha="center", va="center")
                axes[1, 0].axis("off")

        axes[1, 1].text(
            0.1,
            0.5,
            f"best_epoch: {data.get('best_epoch')}\n"
            f"metric: {data.get('metric_name')}\n"
            f"best: {data.get('best_metric')}\n"
            f"steps: {data.get('total_steps')}",
            fontsize=11,
            va="center",
        )
        axes[1, 1].axis("off")
        axes[1, 1].set_title("Summary")

    elif log_dir:
        series = _read_tensorboard_scalars(Path(log_dir))
        if not series:
            raise RuntimeError(f"No TensorBoard scalar events under {log_dir!r}")

        def plot_tag(ax, tag: str, title: str):
            if tag not in series:
                return
            steps, vals = series[tag]
            ax.plot(steps, vals)
            ax.set_title(title)
            ax.set_xlabel("Step")
            ax.grid(True, alpha=0.3)

        plot_tag(axes[0, 0], "Loss/train_step", "Train Loss (step)")
        plot_tag(axes[0, 1], "Loss/train", "Train Loss (epoch)")
        plot_tag(axes[1, 0], "Metrics/bleu", "BLEU")
        plot_tag(axes[1, 1], "LR/step", "Learning Rate")
    else:
        raise ValueError("Provide either --summary-json or --log-dir")

    fig.suptitle("EasyTranslate Training Curves", fontsize=14)
    fig.tight_layout()
    fig.savefig(out, dpi=150)
    plt.close(fig)
    logger.info("Saved training curves: %s", out)
    return out


def _collect_experiment_metrics(results_dir: Path) -> tuple[list[str], dict[str, list[float]]]:
    """Load metrics from experiments_summary.json or per-run evaluation_results.json under subdirs."""
    results_dir = Path(results_dir)
    labels: list[str] = []
    metrics_map: dict[str, list[float]] = {}

    direct = results_dir / "evaluation_results.json"
    if direct.is_file():
        labels.append(results_dir.name or "single")
        with open(direct, "r", encoding="utf-8") as f:
            m = json.load(f)
        for k, v in m.items():
            if isinstance(v, (int, float)) and not isinstance(v, bool):
                metrics_map.setdefault(k, []).append(float(v))
        return labels, metrics_map

    summary_file = results_dir / "experiments_summary.json"

    if summary_file.is_file():
        with open(summary_file, "r", encoding="utf-8") as f:
            rows = json.load(f)
        for row in rows:
            name = row.get("name", "unknown")
            labels.append(name)
            m = row.get("metrics") or {}
            for k, v in m.items():
                if isinstance(v, (int, float)) and not isinstance(v, bool):
                    metrics_map.setdefault(k, []).append(float(v))
        return labels, metrics_map

    for sub in sorted(results_dir.iterdir()):
        if not sub.is_dir():
            continue
        ev = sub / "evaluation_results.json"
        if not ev.is_file():
            continue
        labels.append(sub.name)
        with open(ev, "r", encoding="utf-8") as f:
            m = json.load(f)
        for k, v in m.items():
            if isinstance(v, (int, float)) and not isinstance(v, bool):
                metrics_map.setdefault(k, []).append(float(v))

    if len(labels) != len(next(iter(metrics_map.values()), [])) and metrics_map:
        # Metric length mismatch across runs: keep rows; plotting filters by available keys.
        pass
    return labels, metrics_map


def plot_experiment_comparison(
    results_dir: str,
    output_path: str = "outputs/experiment_comparison.png",
) -> Path:
    """Bar chart for BLEU / COMET / chrF / etc., plus a small text table."""
    out = Path(output_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    labels, metrics_map = _collect_experiment_metrics(Path(results_dir))
    if not labels:
        raise RuntimeError(
            f"No experiments_summary.json or */evaluation_results.json under {results_dir}"
        )

    preferred = ["bleu", "comet", "chrf", "ter"]
    keys = [k for k in preferred if k in metrics_map and len(metrics_map[k]) == len(labels)]
    if not keys:
        keys = [k for k, vals in metrics_map.items() if len(vals) == len(labels)]

    if not keys:
        raise RuntimeError(
            "No numeric metric columns aligned with each experiment; check evaluation_results.json"
        )

    n = len(keys)
    fig, axes = plt.subplots(1, max(n, 1), figsize=(4 * max(n, 1), 4))
    if n == 1:
        axes = [axes]

    for ax, key in zip(axes, keys):
        vals = metrics_map[key][: len(labels)]
        x = np.arange(len(labels))
        ax.bar(x, vals, color="steelblue")
        ax.set_xticks(x)
        ax.set_xticklabels(labels, rotation=25, ha="right")
        ax.set_title(key.upper())
        ax.grid(True, axis="y", alpha=0.3)

    table_text = _format_metrics_table(labels, keys, metrics_map)
    fig.subplots_adjust(bottom=0.28)
    fig.text(0.04, 0.02, table_text, fontsize=9, va="bottom", ha="left")
    fig.suptitle("Experiment Comparison", fontsize=14)
    fig.tight_layout()
    fig.savefig(out, dpi=150, bbox_inches="tight", pad_inches=0.25)
    plt.close(fig)
    logger.info("Saved experiment comparison plot: %s", out)
    return out


def _cross_attention_weight_matrix(
    attn_module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    memory_key_padding_mask: Optional[torch.BoolTensor],
) -> torch.Tensor:
    """Scaled dot-product attention weights [B, L_q, L_k], head-mean (for Flash / standard MHAttention)."""
    B, L_q, _ = query.shape
    L_k = key.shape[1]
    nhead = attn_module.nhead
    d_k = attn_module.d_k

    Q = attn_module.q_proj(query).view(B, L_q, nhead, d_k).transpose(1, 2)
    K = attn_module.k_proj(key).view(B, L_k, nhead, d_k).transpose(1, 2)
    if getattr(attn_module, "rope", None) is not None and attn_module.rope is not None:
        Q, K = attn_module.rope.apply_rotary_pos_emb(Q, K)

    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    if memory_key_padding_mask is not None:
        scores = scores.masked_fill(
            memory_key_padding_mask.unsqueeze(1).unsqueeze(2),
            float("-inf"),
        )
    w = torch.softmax(scores, dim=-1).mean(dim=1)
    return w[0]


def visualize_attention(
    model: torch.nn.Module,
    src_text: str,
    tgt_text: str,
    tokenizer,
    output_path: str = "outputs/attention_map.png",
    layer_idx: int = -1,
) -> Path:
    """
    Cross-attention alignment heatmap for the last (or chosen) decoder layer.
    Only models with ``decoder.layers[*].multihead_attn`` (e.g. TransformerTranslationModel).
    """
    from easytranslate.model.transformer import TransformerTranslationModel

    if not isinstance(model, TransformerTranslationModel):
        raise TypeError("visualize_attention only supports TransformerTranslationModel")

    device = next(model.parameters()).device
    model.eval()

    src_ids_list = tokenizer.encode(src_text, add_special_tokens=True)
    tgt_ids_list = tokenizer.encode(tgt_text, add_special_tokens=True)
    if len(tgt_ids_list) < 2:
        raise ValueError("target sequence too short for teacher-forcing visualization")

    teacher_tgt = tgt_ids_list[:-1]
    src_ids = torch.tensor([src_ids_list], dtype=torch.long, device=device)
    tgt_in = torch.tensor([teacher_tgt], dtype=torch.long, device=device)

    pad_id = model.pad_id
    src_padding = src_ids.eq(pad_id)
    tgt_padding = tgt_in.eq(pad_id)

    captured: dict[str, Any] = {}
    layer = model.decoder.layers[layer_idx]

    def _hook_layer_kw(m, args, kwargs, output):
        tgt_side, memory = args[0], args[1]
        mem_pad = kwargs.get("memory_key_padding_mask")
        query = m.norm2(tgt_side)
        captured["weights"] = _cross_attention_weight_matrix(
            m.multihead_attn, query, memory, mem_pad
        )

    def _hook_layer_legacy(m, inp, output):
        tgt_side, memory = inp[0], inp[1]
        query = m.norm2(tgt_side)
        captured["weights"] = _cross_attention_weight_matrix(
            m.multihead_attn, query, memory, None
        )

    try:
        handle = layer.register_forward_hook(_hook_layer_kw, with_kwargs=True)
    except TypeError:
        handle = layer.register_forward_hook(_hook_layer_legacy)

    with torch.no_grad():
        logits = model(src_ids, tgt_in, src_padding, tgt_padding)

    handle.remove()
    if "weights" not in captured:
        raise RuntimeError("cross-attention hook did not run")

    w = captured["weights"].detach().float().cpu().numpy()
    _ = logits

    src_tokens = [tokenizer.decode([i]) for i in src_ids_list]
    tgt_tokens = [tokenizer.decode([i]) for i in teacher_tgt]

    fig, ax = plt.subplots(figsize=(max(8, w.shape[1] * 0.35), max(6, w.shape[0] * 0.35)))
    im = ax.imshow(w, cmap="viridis", aspect="auto")
    ax.set_xticks(range(len(src_tokens)))
    ax.set_yticks(range(len(tgt_tokens)))
    ax.set_xticklabels(src_tokens, rotation=45, ha="right", fontsize=8)
    ax.set_yticklabels(tgt_tokens, fontsize=8)
    ax.set_xlabel("Source")
    ax.set_ylabel("Target (teacher forcing)")
    ax.set_title("Cross-attention (last layer, heads mean)")
    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    out = Path(output_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    fig.tight_layout()
    fig.savefig(out, dpi=150)
    plt.close(fig)
    logger.info("Saved attention heatmap: %s", out)
    return out


def generate_translation_examples(
    evaluator,
    test_pairs: list[tuple[str, str]],
    output_path: str = "outputs/translation_examples.md",
) -> Path:
    """Write Markdown: source, reference, hypothesis, sentence BLEU and chrF."""
    from easytranslate.evaluation.metrics import compute_bleu, compute_chrf

    srcs = [p[0] for p in test_pairs]
    refs = [p[1] for p in test_pairs]
    hyps = evaluator.translate(srcs)

    def esc(t: str) -> str:
        return t.replace("|", "\\|").replace("\n", " ")

    lines = [
        "# Translation examples",
        "",
        "| # | Source | Reference | Hypothesis | sent-BLEU | sent-chrF |",
        "|---|--------|-----------|------------|-----------|-----------|",
    ]
    for i, (s, r, h) in enumerate(zip(srcs, refs, hyps), 1):
        sb = compute_bleu([h], [r])["bleu"]
        ch = compute_chrf([h], [r])["chrf"]
        lines.append(f"| {i} | {esc(s)} | {esc(r)} | {esc(h)} | {sb:.2f} | {ch:.2f} |")

    lines.append("")
    lines.append("> Sentence BLEU/chrF are indicative only (tokenization-dependent).")

    out = Path(output_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text("\n".join(lines), encoding="utf-8")
    logger.info("Wrote translation examples: %s", out)
    return out


def _load_model_for_visual(
    config_path: Path,
    checkpoint_path: Path,
) -> tuple[torch.nn.Module, Any, dict]:
    """Load scratch Transformer + tokenizer from YAML and checkpoint (prefers config inside checkpoint)."""
    with open(config_path, "r", encoding="utf-8") as f:
        file_cfg = yaml.safe_load(f)

    try:
        ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    except TypeError:
        ckpt = torch.load(checkpoint_path, map_location="cpu")

    cfg = ckpt.get("config")
    if cfg is None:
        cfg = file_cfg
    else:
        try:
            from omegaconf import DictConfig

            if isinstance(cfg, DictConfig):
                cfg = OmegaConf.to_container(cfg, resolve=True)
        except Exception:
            pass
        if not isinstance(cfg, dict):
            cfg = dict(cfg)

    from easytranslate.model.transformer import TransformerTranslationModel
    from easytranslate.data.tokenizer import build_tokenizer

    tok_cfg = cfg.get("tokenizer") or cfg.get("data", {}).get("tokenizer") or file_cfg.get("tokenizer") or {}
    try:
        tokenizer = build_tokenizer(tok_cfg)
    except ValueError as e:
        raise ValueError(
            "Cannot build tokenizer: set tokenizer.path in config or store a loadable tokenizer "
            "section in checkpoint['config']."
        ) from e

    mcfg = cfg.get("model", {}).get("transformer", {}) or file_cfg.get("model", {}).get("transformer", {})
    model = TransformerTranslationModel(
        src_vocab_size=tokenizer.vocab_size,
        tgt_vocab_size=tokenizer.vocab_size,
        d_model=int(mcfg.get("d_model", 512)),
        nhead=int(mcfg.get("nhead", 8)),
        num_encoder_layers=int(mcfg.get("num_encoder_layers", 6)),
        num_decoder_layers=int(mcfg.get("num_decoder_layers", 6)),
        dim_feedforward=int(mcfg.get("dim_feedforward", 2048)),
        dropout=float(mcfg.get("dropout", 0.1)),
        activation=str(mcfg.get("activation", "gelu")),
        max_seq_len=int(mcfg.get("max_seq_len", 512)),
        use_flash_attention=bool(mcfg.get("use_flash_attention", True)),
        use_rotary_embedding=bool(mcfg.get("use_rotary_embedding", True)),
        pre_norm=bool(mcfg.get("pre_norm", True)),
        pad_id=tokenizer.pad_token_id,
    )
    model.load_state_dict(ckpt["model_state_dict"])
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    return model, tokenizer, cfg


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="EasyTranslate visualization CLI")
    p.add_argument(
        "--task",
        choices=["training_curves", "comparison", "attention", "examples"],
        required=True,
    )
    p.add_argument("--log-dir", type=str, default=None)
    p.add_argument("--summary-json", type=str, default=None)
    p.add_argument("--results-dir", type=str, default=None)
    p.add_argument("--output", type=str, default=None)
    p.add_argument("--config", type=str, default="configs/default_config.yaml")
    p.add_argument("--checkpoint", type=str, default=None)
    p.add_argument("--src", type=str, default=None)
    p.add_argument("--tgt", type=str, default=None)
    p.add_argument("--pairs-json", type=str, default=None)
    return p.parse_args()


def main() -> None:
    logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
    args = _parse_args()

    if args.task == "training_curves":
        outp = args.output or "outputs/training_curves.png"
        plot_training_curves(
            log_dir=args.log_dir,
            output_path=outp,
            summary_json=args.summary_json,
        )
        print(f"OK: {outp}")

    elif args.task == "comparison":
        rd = args.results_dir or "outputs/experiments"
        outp = args.output or "outputs/experiment_comparison.png"
        plot_experiment_comparison(rd, outp)
        print(f"OK: {outp}")

    elif args.task == "attention":
        if not args.checkpoint or not args.src or not args.tgt:
            raise SystemExit("--task attention requires --checkpoint --src --tgt")
        cfg_p = (REPO_ROOT / args.config).resolve()
        ckpt_p = (REPO_ROOT / args.checkpoint).resolve()
        model, tokenizer, _ = _load_model_for_visual(cfg_p, ckpt_p)
        outp = args.output or "outputs/attention_map.png"
        visualize_attention(model, args.src, args.tgt, tokenizer, output_path=outp)
        print(f"OK: {outp}")

    elif args.task == "examples":
        if not args.checkpoint or not args.pairs_json:
            raise SystemExit("--task examples requires --checkpoint --pairs-json")
        cfg_p = (REPO_ROOT / args.config).resolve()
        ckpt_p = (REPO_ROOT / args.checkpoint).resolve()
        model, tokenizer, cfg = _load_model_for_visual(cfg_p, ckpt_p)

        from easytranslate.evaluation.evaluator import Evaluator

        evaluator = Evaluator(model=model, tokenizer=tokenizer, config=cfg)

        pairs_path = (REPO_ROOT / args.pairs_json).resolve()
        with open(pairs_path, "r", encoding="utf-8") as f:
            raw = json.load(f)
        pairs: list[tuple[str, str]] = []
        for item in raw:
            if isinstance(item, dict):
                pairs.append((item["src"], item["ref"]))
            else:
                pairs.append((item[0], item[1]))

        outp = args.output or "outputs/translation_examples.md"
        generate_translation_examples(evaluator, pairs, outp)
        print(f"OK: {outp}")


if __name__ == "__main__":
    main()