#!/usr/bin/env python3
"""Evaluate MAVT 3D (triplane) reconstruction + understanding quality.

Recon metrics (per-plane + mean):
  PSNR  (per-plane, higher is better)
  SSIM  (per-plane, higher is better)
  LPIPS (per-plane, AlexNet, lower is better)
  FID   (Inception-V3 features over all 3 planes concatenated, lower is better)

Understanding metric:
  cos_sim_teacher : mean cosine similarity between MAVT.semantic and frozen
                    SigLIP2 teacher's pooler_output, fed on the XY plane
                    (the "natural-image" proxy used during training).

Pipeline mirrors eval_image.py / eval_video.py: load Lightning ckpt,
pre-create cd_split poolers found in the ckpt, then run forward over
UniversalThreeDDataset and accumulate metrics.

Usage:
  PYTHONPATH=src .venv/bin/python eval_threed.py \\
      --ckpt checkpoints/stage3/balanced/mavt-stage3-balanced-step=0050000-val/loss=0.2500.ckpt \\
      --threed_root dataset/universal_3d \\
      --max_objects 512 \\
      --output eval_threed.json
"""
from __future__ import annotations

import argparse
import inspect
import json
from pathlib import Path
from typing import Dict, List

import torch
from torch.utils.data import DataLoader, Subset
from torchmetrics.image import StructuralSimilarityIndexMeasure
from torchmetrics.image.fid import FrechetInceptionDistance
from torchmetrics.image.psnr import PeakSignalNoiseRatio
from torchvision.utils import make_grid

import lpips

from mavt.training.lightning_module import MAVTLightningModule
from mavt.data.datasets import UniversalThreeDDataset
from mavt.data.datamodule import _collate


PLANE_NAMES = ('oxoy', 'oxoz', 'oyoz')  # front, top, side


def _to_unit(x: torch.Tensor) -> torch.Tensor:
    """[-1, 1] → [0, 1]."""
    return (x.clamp(-1.0, 1.0) + 1.0) * 0.5


def _plane_strip(planes: torch.Tensor) -> torch.Tensor:
    """(B, 3, 3, H, W) → (3, 3*H, B*W) tensor suitable for make_grid.

    Stacks 3 planes vertically per object; concatenates objects horizontally.
    """
    B, P, C, H, W = planes.shape
    # rearrange to (B, P*C, H, W) where order = oxoy_RGB, oxoz_RGB, oyoz_RGB
    planes = planes.reshape(B, P * C, H, W)
    return planes


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpt', required=True, help='Lightning .ckpt path')
    ap.add_argument('--threed_root', required=True,
                    help='Root directory with 3d_objects/renders/<id>/{oxoy,oxoz,oyoz}.png')
    ap.add_argument('--output', default='eval_threed.json')
    ap.add_argument('--max_objects', type=int, default=512,
                    help='Cap total objects evaluated (None = all)')
    ap.add_argument('--triplane_res', type=int, default=256)
    ap.add_argument('--batch_size', type=int, default=8)
    ap.add_argument('--num_workers', type=int, default=4)
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    ap.add_argument('--lpips_chunk', type=int, default=8,
                    help='Sub-batch size for LPIPS to control memory (per plane)')
    ap.add_argument('--fid_feature', type=int, default=2048,
                    choices=[64, 192, 768, 2048])
    ap.add_argument('--semantic', action=argparse.BooleanOptionalAction, default=True,
                    help='Compute cosine similarity to frozen SigLIP2 teacher (XY proxy)')
    ap.add_argument('--save_samples', type=int, default=4,
                    help='Save this many GT-vs-recon comparison PNGs')
    args = ap.parse_args()

    device = torch.device(args.device)
    torch.backends.cudnn.benchmark = True

    # --- Model: pre-create cd_split poolers from ckpt before loading -------
    print(f'[eval-threed] loading checkpoint: {args.ckpt}')
    ckpt = torch.load(args.ckpt, map_location='cpu', weights_only=False)
    raw_hp = dict(ckpt.get('hyper_parameters', {}))
    state = ckpt.get('state_dict', {})

    valid = set(inspect.signature(MAVTLightningModule.__init__).parameters)
    hparams = {k: v for k, v in raw_hp.items() if k in valid}
    module = MAVTLightningModule(**hparams)

    pooler_combos = set()
    for k in state.keys():
        if k.startswith('model.cd_split._content_poolers.'):
            shape = k.split('.')[3]
            if '_' in shape and all(s.isdigit() for s in shape.split('_')):
                a, b = shape.split('_')
                pooler_combos.add((int(a), int(b)))
    # Threed is not in active_modalities for stage 1/2 → must inject the
    # expected combo (N=3*S²//patch²) so the param groups are populated.
    S = args.triplane_res
    patch = int(hparams.get('patch_size', 16))
    N_threed = 3 * (S // patch) * (S // patch)
    threed_c = max(1, int(N_threed * 0.35))
    threed_d = max(1, int(N_threed * 0.25))
    module.model.cd_split.prepare_poolers(threed_c, threed_d)
    pooler_combos.add((threed_c, threed_d))
    print(f'[eval-threed] pre-created poolers for combos: {sorted(pooler_combos)}')

    missing, unexpected = module.load_state_dict(state, strict=False)
    real_missing = [k for k in missing if not k.startswith('semantic_teacher.')]
    print(f'[eval-threed] load: {len(real_missing)} missing (excl. teacher), '
          f'{len(unexpected)} unexpected')
    if real_missing:
        print(f'[eval-threed]   missing sample: {real_missing[:5]}')
    if unexpected:
        print(f'[eval-threed]   unexpected sample: {unexpected[:5]}')

    module.eval().to(device)

    # --- Data ---------------------------------------------------------------
    ds = UniversalThreeDDataset(args.threed_root, resolution=args.triplane_res)
    if args.max_objects and args.max_objects < len(ds):
        ds = Subset(ds, list(range(args.max_objects)))
    print(f'[eval-threed] {len(ds)} objects in eval set')

    loader = DataLoader(
        ds, batch_size=args.batch_size, shuffle=False,
        num_workers=args.num_workers, pin_memory=(device.type == 'cuda'),
        collate_fn=_collate, drop_last=False,
    )

    # --- Reconstruction metrics (per-plane + aggregate) ---------------------
    psnr_per_plane = [
        PeakSignalNoiseRatio(data_range=1.0).to(device) for _ in range(3)
    ]
    ssim_per_plane = [
        StructuralSimilarityIndexMeasure(data_range=1.0).to(device) for _ in range(3)
    ]
    lpips_per_plane_sum = [0.0, 0.0, 0.0]
    lpips_per_plane_n = [0, 0, 0]
    fid_metric = FrechetInceptionDistance(
        feature=args.fid_feature, normalize=True,
    ).to(device)
    lpips_fn = lpips.LPIPS(net='alex', verbose=False).to(device).eval()

    # --- Understanding metric -----------------------------------------------
    teacher = None
    teacher_input_size = 224
    if args.semantic:
        teacher_name = hparams.get('siglip2_model_name', 'google/siglip2-base-patch16-224')
        print(f'[eval-threed] loading semantic teacher: {teacher_name}')
        from transformers import AutoModel
        siglip = AutoModel.from_pretrained(teacher_name)
        teacher = siglip.vision_model.to(device).eval()
        for p in teacher.parameters():
            p.requires_grad_(False)
        try:
            teacher_input_size = int(siglip.config.vision_config.image_size)
        except AttributeError:
            teacher_input_size = 224
        print(f'[eval-threed] teacher input size: {teacher_input_size}')
    cos_sim_sum, cos_sim_n = 0.0, 0

    autocast_dtype = torch.bfloat16 if device.type == 'cuda' else torch.float32
    saved = 0
    out_dir = Path(args.output).with_suffix('')
    if args.save_samples > 0:
        out_dir.mkdir(parents=True, exist_ok=True)

    for bi, batch in enumerate(loader):
        x = batch['data'].to(device, non_blocking=True)  # (B, 3, 3, H, W) in [-1, 1]

        with torch.no_grad(), torch.amp.autocast(
                device_type=device.type, dtype=autocast_dtype, enabled=device.type == 'cuda'):
            out = module.model(x, 'threed', decode=True)
        recon = out.reconstruction.float().clamp(-1.0, 1.0)  # (B, 3, 3, H, W)

        rec01 = _to_unit(recon)
        tgt01 = _to_unit(x)

        # Per-plane metrics
        for p in range(3):
            psnr_per_plane[p].update(rec01[:, p], tgt01[:, p])
            ssim_per_plane[p].update(rec01[:, p], tgt01[:, p])
            # LPIPS per plane (treat each plane as an independent image)
            for s in range(0, recon.shape[0], args.lpips_chunk):
                d = lpips_fn(
                    recon[s:s + args.lpips_chunk, p],
                    x[s:s + args.lpips_chunk, p],
                )
                lpips_per_plane_sum[p] += d.sum().item()
                lpips_per_plane_n[p] += d.numel()

        # FID over all 3 planes concatenated as separate images (B*3 images)
        B = rec01.shape[0]
        flat_real = tgt01.reshape(B * 3, 3, args.triplane_res, args.triplane_res)
        flat_fake = rec01.reshape(B * 3, 3, args.triplane_res, args.triplane_res)
        fid_metric.update(flat_real, real=True)
        fid_metric.update(flat_fake, real=False)

        # Understanding: XY plane (index 0) as proxy for SigLIP2
        if teacher is not None:
            with torch.no_grad(), torch.amp.autocast(
                    device_type=device.type, dtype=autocast_dtype, enabled=device.type == 'cuda'):
                xy = x[:, 0]  # (B, 3, H, W)
                if xy.shape[-1] != teacher_input_size:
                    teacher_in = torch.nn.functional.interpolate(
                        xy, size=teacher_input_size, mode='bilinear', align_corners=False)
                else:
                    teacher_in = xy
                t_emb = teacher(pixel_values=teacher_in).pooler_output
                m_emb = out.semantic.float()
            cos = torch.nn.functional.cosine_similarity(
                m_emb.float(), t_emb.float(), dim=-1)
            cos_sim_sum += cos.sum().item()
            cos_sim_n += cos.numel()

        # Save sample visualizations
        if saved < args.save_samples:
            for i in range(min(args.save_samples - saved, x.shape[0])):
                # 3 planes stacked vertically for GT vs recon
                pair = torch.cat([
                    _plane_strip(tgt01[i:i + 1].cpu()),
                    _plane_strip(rec01[i:i + 1].cpu()),
                ], dim=2)  # concat vertically: GT on top, recon on bottom
                obj_id = batch['id'][i] if 'id' in batch else f'idx_{bi * args.batch_size + i}'
                # pair shape: (1, 9, H, W) → make_grid to image
                grid = make_grid(pair[0], nrow=3, padding=2, pad_value=1.0)
                from PIL import Image
                arr = (grid.clamp(0, 1).permute(1, 2, 0).numpy() * 255).astype('uint8')
                Image.fromarray(arr).save(out_dir / f'sample_{saved:03d}_{obj_id[:16]}.png')
                saved += 1
                if saved >= args.save_samples:
                    break

        if (bi + 1) % 5 == 0 or (bi + 1) == len(loader):
            cos_str = f'  cos_sim={cos_sim_sum / max(1, cos_sim_n):.4f}' if cos_sim_n else ''
            print(f'[eval-threed] {bi + 1}/{len(loader)} batches  '
                  f'PSNR_xy={psnr_per_plane[0].compute().item():.3f}  '
                  f'PSNR_xz={psnr_per_plane[1].compute().item():.3f}  '
                  f'PSNR_yz={psnr_per_plane[2].compute().item():.3f}{cos_str}')

    fid = float(fid_metric.compute().item())

    psnr_vals = [float(m.compute().item()) for m in psnr_per_plane]
    ssim_vals = [float(m.compute().item()) for m in ssim_per_plane]
    lpips_vals = [
        lpips_per_plane_sum[p] / max(1, lpips_per_plane_n[p])
        for p in range(3)
    ]

    results = {
        'ckpt': args.ckpt,
        'threed_root': args.threed_root,
        'n_objects': len(ds),
        'triplane_res': args.triplane_res,
        'psnr_xy': psnr_vals[0],
        'psnr_xz': psnr_vals[1],
        'psnr_yz': psnr_vals[2],
        'psnr_mean': sum(psnr_vals) / 3,
        'ssim_xy': ssim_vals[0],
        'ssim_xz': ssim_vals[1],
        'ssim_yz': ssim_vals[2],
        'ssim_mean': sum(ssim_vals) / 3,
        'lpips_alex_xy': lpips_vals[0],
        'lpips_alex_xz': lpips_vals[1],
        'lpips_alex_yz': lpips_vals[2],
        'lpips_alex_mean': sum(lpips_vals) / 3,
        'fid_inception': fid,
        'cos_sim_teacher': cos_sim_sum / cos_sim_n if cos_sim_n else None,
        'fid_feature_dim': args.fid_feature,
    }
    print(json.dumps(results, indent=2))
    Path(args.output).write_text(json.dumps(results, indent=2))
    print(f'[eval-threed] wrote {args.output}')
    if saved > 0:
        print(f'[eval-threed] wrote {saved} sample PNGs to {out_dir}/')


if __name__ == '__main__':
    main()