#!/usr/bin/env python3
"""Flow field visualization — vorticity panels and error maps.

Works on pre-saved vorticity PNGs or raw Macroscopic data.
All panels use unified [-0.03, 0.03] vorticity range.

Usage:
    conda run -n pycuda_3_10 python viz_flow.py [--scene re100]
"""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path
from typing import List, Optional, Tuple

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

_OUT_BASE = Path(__file__).resolve().parent / "output"


def log(msg: str) -> None:
    print(f"[viz_flow] {msg}", flush=True)


def _find_vorticity_files(scene_dir: Path) -> dict:
    """Scan scene output dir for vorticity images. Returns dict keyed by label."""
    files = {}
    for name in ["vorticity_target", "vorticity_controlled", "vorticity_zero",
                 "vorticity_uncontrolled"]:
        png = scene_dir / f"{name}.png"
        if png.exists():
            files[name.replace("vorticity_", "")] = png
    return files


def make_comparison_panel(scene_name: str, train_dir: Optional[Path],
                          repro_dir: Optional[Path], out_path: Path) -> None:
    """Create a multi-panel vorticity comparison for one scene.

    Layout: rows = [repro, train], cols = [target, controlled, zero]
    """
    panels: List[Tuple[str, Path]] = []

    if repro_dir and repro_dir.exists():
        rep_files = _find_vorticity_files(repro_dir)
        for label, p in rep_files.items():
            panels.append((f"Repro-{label}", p))

    if train_dir and train_dir.exists():
        trn_files = _find_vorticity_files(train_dir)
        for label, p in trn_files.items():
            panels.append((f"Train-{label}", p))

    if not panels:
        log(f"  {scene_name}: no vorticity images found, skipping")
        return

    n = len(panels)
    cols = min(3, n)
    rows = (n + cols - 1) // cols

    fig, axes = plt.subplots(rows, cols, figsize=(6 * cols, 4.5 * rows))
    if rows * cols == 1:
        axes = np.array([[axes]])
    elif rows == 1:
        axes = axes.reshape(1, -1)
    elif cols == 1:
        axes = axes.reshape(-1, 1)

    for idx, (label, img_path) in enumerate(panels):
        r, c = divmod(idx, cols)
        ax = axes[r, c]
        ax.imshow(plt.imread(str(img_path)))
        ax.set_title(label, fontsize=10)
        ax.axis("off")

    # Hide unused axes
    for idx in range(len(panels), rows * cols):
        r, c = divmod(idx, cols)
        axes[r, c].axis("off")

    fig.suptitle(f"{scene_name} - Vorticity Comparison", fontsize=14)
    fig.tight_layout()
    fig.savefig(out_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    log(f"  {scene_name}: {n} panels saved to {out_path.name}")


def make_side_by_side(scene_name: str, train_dir: Optional[Path],
                      repro_dir: Optional[Path], out_path: Path) -> None:
    """Simple side-by-side: repro controlled | train controlled."""
    images = []
    labels = []
    for label, d in [("Reproduce", repro_dir), ("Train", train_dir)]:
        if d is None or not d.exists():
            continue
        for name in ["controlled", "zero"]:
            png = d / f"vorticity_{name}.png"
            if png.exists():
                images.append(png)
                labels.append(f"{label} ({name})")
                break  # take first available

    if len(images) < 2:
        return

    fig, axes = plt.subplots(1, len(images), figsize=(6 * len(images), 4.5))
    if len(images) == 1:
        axes = [axes]
    for ax, img_path, label in zip(axes, images, labels):
        ax.imshow(plt.imread(str(img_path)))
        ax.set_title(label, fontsize=11)
        ax.axis("off")
    fig.suptitle(f"{scene_name}", fontsize=13)
    fig.tight_layout()
    fig.savefig(out_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    log(f"  {scene_name}: side-by-side saved")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> int:
    parser = argparse.ArgumentParser(description="Vorticity visualization")
    parser.add_argument("--scene", type=str, default=None,
                        help="Single scene to process")
    parser.add_argument("--side", type=str, default="all",
                        choices=["all", "train", "reproduce"],
                        help="Which pipeline to render")
    args = parser.parse_args()

    report_dir = _OUT_BASE / "reports" / "vorticity_panels"
    report_dir.mkdir(parents=True, exist_ok=True)

    # Scene names to process
    train_scenes = ["re100_karman", "transfer_re60", "transfer_re200",
                    "transfer_re400", "illusion_075L", "illusion_1L", "illusion_15L"]
    repro_scenes = ["karman_cloak", "steady_cloak", "illusion_075L",
                    "illusion_1L", "illusion_15L", "vortex_lamb", "vortex_taylor"]

    # Map train to repro for overlap
    overlap_map = {
        "re100_karman": "karman_cloak",
        "illusion_075L": "illusion_075L",
        "illusion_1L": "illusion_1L",
        "illusion_15L": "illusion_15L",
    }

    if args.side in ("all", "train"):
        for scene in train_scenes:
            if args.scene and scene != args.scene and not scene.startswith(args.scene):
                continue
            train_dir = _OUT_BASE / "train" / scene
            repro_dir = _OUT_BASE / "reproduce" / overlap_map.get(scene, "")
            make_comparison_panel(scene, train_dir,
                                  repro_dir if repro_dir.exists() else None,
                                  report_dir / f"vorticity_{scene}.png")

    if args.side in ("all", "reproduce"):
        for scene in repro_scenes:
            if args.scene and scene != args.scene:
                continue
            if scene in overlap_map.values():
                continue  # already covered above
            repro_dir = _OUT_BASE / "reproduce" / scene
            make_comparison_panel(scene, None, repro_dir,
                                  report_dir / f"vorticity_{scene}.png")

    # Summary: side-by-side for overlapping scenes
    for scene, rep_key in overlap_map.items():
        if args.scene and scene != args.scene and not scene.startswith(args.scene):
            continue
        train_dir = _OUT_BASE / "train" / scene
        repro_dir = _OUT_BASE / "reproduce" / rep_key
        make_side_by_side(scene, train_dir, repro_dir,
                          report_dir / f"vorticity_compare_{scene}.png")

    log(f"Reports in {report_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
