#!/usr/bin/env python3
"""Visualize the global/high INS encoder branches on a RenderMe view.

The images are activation diagnostics. They show spatial response, PCA
projections, and representative channels; they do not assign semantic labels
to the features.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch
import torch.nn.functional as F
from torch import Tensor
from torchvision.utils import make_grid, save_image

REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
if str(REPOSITORY_ROOT) not in sys.path:
    sys.path.insert(0, str(REPOSITORY_ROOT))

from datasets.renderme_ins_dataset import RenderMeINSViewDataset
from scripts.evaluate_ins_stage1_renderme import (
    build_model_from_checkpoint,
    read_checkpoint,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Visualize global and high-resolution INS branch activations."
    )
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--data-root", type=Path, default=None)
    parser.add_argument("--frame-index", type=int, default=None)
    parser.add_argument("--camera-id", default=None)
    parser.add_argument("--source-camera-id", default=None)
    parser.add_argument("--resolution", type=int, default=512)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--amp", choices=("none", "bf16"), default="none")
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("/workspace/cpfs-data/tmp/ins_branch_viz"),
    )
    parser.add_argument("--top-channels", type=int, default=12)
    parser.add_argument("--save-features", action="store_true")
    args = parser.parse_args()
    if not args.checkpoint.is_file():
        raise FileNotFoundError(args.checkpoint)
    if args.resolution <= 0 or args.top_channels <= 0:
        raise ValueError("resolution and top-channels must be positive.")
    return args


def robust_normalize(value: Tensor) -> Tensor:
    value = value.float()
    low = torch.quantile(value, 0.01)
    high = torch.quantile(value, 0.99)
    return ((value - low) / (high - low).clamp_min(1e-8)).clamp(0.0, 1.0)


def feature_mean_map(features: Tensor) -> Tensor:
    value = robust_normalize(features.abs().mean(dim=1, keepdim=True))
    return value.expand(-1, 3, -1, -1)


def feature_pca_rgb(features: Tensor) -> Tensor:
    features = features.float()
    batch, channels, height, width = features.shape
    if batch != 1:
        raise ValueError("Branch visualization expects batch size one.")
    samples = features[0].permute(1, 2, 0).reshape(height * width, channels)
    samples = samples - samples.mean(dim=0, keepdim=True)
    _, _, vectors = torch.pca_lowrank(samples, q=3, center=False, niter=2)
    projected = samples @ vectors[:, :3]
    projected = projected.reshape(height, width, 3).permute(2, 0, 1).unsqueeze(0)
    return robust_normalize(projected)


def top_channel_grid(features: Tensor, count: int) -> tuple[Tensor, list[int]]:
    features = features.float()
    count = min(count, features.shape[1])
    scores = features.abs().flatten(2).mean(dim=2)[0]
    indices = torch.topk(scores, k=count).indices.tolist()
    maps = [
        robust_normalize(features[:, index : index + 1]).expand(-1, 3, -1, -1)
        for index in indices
    ]
    return make_grid(torch.cat(maps, dim=0), nrow=min(4, count), padding=2), indices


def gradient_energy(features: Tensor) -> Tensor:
    features = features.float()
    horizontal = (features[:, :, :, 1:] - features[:, :, :, :-1]).abs().mean()
    vertical = (features[:, :, 1:, :] - features[:, :, :-1, :]).abs().mean()
    return 0.5 * (horizontal + vertical)


def capture_branches(model) -> tuple[dict[str, Tensor], list[object]]:
    captured: dict[str, Tensor] = {}

    def make_hook(name: str):
        def hook(_module, _inputs, output):
            captured[name] = output.detach().float()

        return hook

    handles = [
        model.encoder.f_encoder.register_forward_hook(make_hook("global")),
        model.encoder.high_encoder.register_forward_hook(make_hook("high")),
        model.encoder.final_encoder.register_forward_hook(make_hook("fused")),
    ]
    return captured, handles


def resolve_data_config(args: argparse.Namespace, checkpoint: dict) -> dict:
    saved = checkpoint.get("args", {})
    if not isinstance(saved, dict):
        saved = {}
    root = args.data_root or Path(
        saved.get(
            "data_root",
            "/workspace/cpfs-data/data/renderme_all/0026/0026_e0",
        )
    )
    frame_index = int(saved.get("frame_index", 0)) if args.frame_index is None else args.frame_index
    source = args.source_camera_id
    if source is None:
        source = saved.get("source_camera_id")
    canonical_radius = float(saved.get("canonical_radius", 2.7))
    canonical_focus = tuple(float(v) for v in saved.get("canonical_focus", (0.0, 0.0, 0.2)))
    background = tuple(float(v) for v in saved.get("background", (1.0, 1.0, 1.0)))
    return {
        "root": root,
        "frame_index": frame_index,
        "source": None if source is None else str(source),
        "canonical_radius": canonical_radius,
        "canonical_focus": canonical_focus,
        "background": background,
    }


@torch.inference_mode()
def visualize(args: argparse.Namespace) -> dict[str, object]:
    device = torch.device(args.device)
    if device.type != "cuda" or not torch.cuda.is_available():
        raise RuntimeError("Branch visualization requires CUDA.")
    torch.cuda.set_device(device)
    checkpoint = read_checkpoint(args.checkpoint)
    model = build_model_from_checkpoint(checkpoint, args.resolution, device)
    config = resolve_data_config(args, checkpoint)

    # The geometric center is resolved by the dataset when source is omitted.
    camera_id = args.camera_id or config["source"]
    dataset = RenderMeINSViewDataset(
        root=config["root"],
        frame_index=config["frame_index"],
        image_size=512,
        camera_ids=None if camera_id is None else [camera_id],
        source_camera_id=config["source"],
        canonical_radius=config["canonical_radius"],
        canonical_focus=config["canonical_focus"],
        background=config["background"],
    )
    camera_id = str(camera_id or dataset.source_camera_id)
    sample = dataset.get_view(camera_id)
    image = sample["image"].unsqueeze(0).to(device)
    captured, handles = capture_branches(model)
    with torch.autocast(
        device_type="cuda",
        dtype=torch.bfloat16,
        enabled=args.amp == "bf16",
    ):
        triplanes = model.encoder(image)
    for handle in handles:
        handle.remove()
    if set(captured) != {"global", "high", "fused"}:
        raise RuntimeError(f"Missing branch activations: {sorted(captured)}")

    args.output_dir.mkdir(parents=True, exist_ok=True)
    source = ((image.float() + 1.0) * 0.5).clamp(0.0, 1.0)
    global_features = captured["global"]
    high_features = captured["high"]
    fused_features = captured["fused"]
    global_mean = feature_mean_map(global_features)
    high_mean = feature_mean_map(high_features)
    global_pca = feature_pca_rgb(global_features)
    high_pca = feature_pca_rgb(high_features)
    fused_mean = feature_mean_map(fused_features)
    global_channels, global_indices = top_channel_grid(global_features, args.top_channels)
    high_channels, high_indices = top_channel_grid(high_features, args.top_channels)

    source_small = F.interpolate(
        source, size=global_mean.shape[-2:], mode="bilinear", align_corners=False
    )
    overview = make_grid(
        torch.cat(
            [source_small, global_mean, global_pca, high_mean, high_pca, fused_mean],
            dim=0,
        ),
        nrow=3,
        padding=3,
    )
    overview_path = args.output_dir / "branch_overview.png"
    global_path = args.output_dir / "global_top_channels.png"
    high_path = args.output_dir / "high_top_channels.png"
    save_image(overview, overview_path)
    save_image(global_channels, global_path)
    save_image(high_channels, high_path)

    stats: dict[str, object] = {
        "checkpoint": str(args.checkpoint.resolve()),
        "camera_id": camera_id,
        "global_shape": list(global_features.shape),
        "high_shape": list(high_features.shape),
        "fused_shape": list(fused_features.shape),
        "triplane_shape": list(triplanes.shape),
        "global_top_channels": global_indices,
        "high_top_channels": high_indices,
        "global_mean_abs": float(global_features.abs().mean()),
        "high_mean_abs": float(high_features.abs().mean()),
        "global_gradient_energy": float(gradient_energy(global_features)),
        "high_gradient_energy": float(gradient_energy(high_features)),
        "global_spatial_std": float(global_features.flatten(2).std(dim=2).mean()),
        "high_spatial_std": float(high_features.flatten(2).std(dim=2).mean()),
        "outputs": {
            "overview": str(overview_path),
            "global_channels": str(global_path),
            "high_channels": str(high_path),
        },
    }
    if args.save_features:
        feature_path = args.output_dir / "branch_features.pt"
        torch.save(
            {
                "global": global_features.cpu(),
                "high": high_features.cpu(),
                "fused": fused_features.cpu(),
            },
            feature_path,
        )
        stats["outputs"]["features"] = str(feature_path)
    with (args.output_dir / "stats.json").open("w", encoding="utf-8") as file:
        json.dump(stats, file, indent=2)
        file.write("\n")
    print(json.dumps(stats, indent=2), flush=True)
    return stats


def main() -> int:
    visualize(parse_args())
    return 0


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