#!/usr/bin/env python3
"""
Evaluate CoMoGaussian rendered test views across pipelines.
Reads TEST/ renders produced by train.py, computes PSNR/SSIM/LPIPS,
and saves visual comparison grids.

Pipeline dirs expected:
  outputs/<pipeline>/scannet/<scene>/TEST/img_7000_NNN.png   (renders)
  outputs/<pipeline>/scannet/<scene>/TEST/GT_NNN.png         (GT, saved at first test iter)
  outputs/<pipeline>/scannet/<scene>/psnr.txt                (inline metrics from training)

Outputs:
  outputs/eval/comogaussian/per_scene/<scene>_metrics.json
  outputs/eval/comogaussian/per_scene/<scene>_grid.jpg
  outputs/eval/comogaussian/summary.json
  outputs/eval/comogaussian/summary.md

Usage:
  python eval_comogaussian_renders.py
  python eval_comogaussian_renders.py --pipelines turtle_comogaussian evssm_comogaussian
  python eval_comogaussian_renders.py --scenes scene0000_00 --iter 7000
"""
import argparse, json, re
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw

BASE = Path(__file__).parent.parent

DEFAULT_PIPELINES = ["blur_comogaussian", "turtle_comogaussian", "evssm_comogaussian"]


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--out-root",   default=str(BASE / "outputs"))
    p.add_argument("--pipelines",  nargs="*", default=None)
    p.add_argument("--scenes",     nargs="*", default=None)
    p.add_argument("--eval-dir",   default=str(BASE / "outputs/eval/comogaussian"))
    p.add_argument("--iter",       type=int, default=7000)
    p.add_argument("--n-vis",      type=int, default=4)
    return p.parse_args()


def psnr_np(img1, img2):
    mse = np.mean((img1.astype(np.float64) - img2.astype(np.float64)) ** 2)
    if mse < 1e-10:
        return 100.0
    return 20.0 * np.log10(255.0 / np.sqrt(mse))


def ssim_np(img1, img2):
    from skimage.metrics import structural_similarity
    return float(structural_similarity(img1, img2, channel_axis=2, data_range=255))


def lpips_score(img1_t, img2_t, lpips_fn):
    import torch
    return float(lpips_fn(img1_t * 2 - 1, img2_t * 2 - 1).item())


def to_tensor(arr):
    import torch
    return torch.from_numpy(arr.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0).cuda()


def parse_psnr_txt(psnr_path):
    """Parse psnr.txt written by CoMoGaussian train.py."""
    result = {}
    try:
        with open(psnr_path) as f:
            for line in f:
                m = re.search(r'ITER (\d+).*?(PSNR|SSIM|LPIPS)\s+([\d.]+)', line)
                if m:
                    it, metric, val = int(m.group(1)), m.group(2), float(m.group(3))
                    if it not in result:
                        result[it] = {}
                    result[it][metric] = val
    except Exception:
        pass
    return result


def make_label(text, w, h=26, bg=(30, 30, 30), fg=(255, 255, 255)):
    img = np.full((h, w, 3), bg, dtype=np.uint8)
    pil = Image.fromarray(img)
    ImageDraw.Draw(pil).text((6, 4), text, fill=fg)
    return np.array(pil)


def load_test_renders(model_dir, iteration, scale_factor=1):
    """
    Load renders and GTs from render.py output:
      <model>/test/ours_<iter>/test_preds_<scale>/NNNNN.png
      <model>/test/ours_<iter>/gt_<scale>/NNNNN.png
    Falls back to train.py inline output:
      <model>/TEST/img_<iter>_NNN.png  +  GT_NNN.png
    """
    # render.py output (preferred)
    render_path = Path(model_dir) / "test" / f"ours_{iteration}" / f"test_preds_{scale_factor}"
    gt_path     = Path(model_dir) / "test" / f"ours_{iteration}" / f"gt_{scale_factor}"
    if render_path.exists() and gt_path.exists():
        renders = [np.array(Image.open(f).convert("RGB")) for f in sorted(render_path.glob("*.png")) if "depth" not in f.name]
        gts     = [np.array(Image.open(f).convert("RGB")) for f in sorted(gt_path.glob("*.png"))]
        return renders, gts

    # fallback: train.py inline TEST/ output
    test_dir = Path(model_dir) / "TEST"
    if test_dir.exists():
        renders = [np.array(Image.open(f).convert("RGB")) for f in sorted(test_dir.glob(f"img_{iteration}_*.png"))]
        gts     = [np.array(Image.open(f).convert("RGB")) for f in sorted(test_dir.glob("GT_*.png"))]
        return renders, gts

    return [], []


def load_results_json(model_dir):
    """Load metrics.py output results.json: {method: {SSIM, PSNR, LPIPS}}"""
    p = Path(model_dir) / "results.json"
    if p.exists():
        with open(p) as f:
            return json.load(f)
    return {}


def evaluate_scene_pipeline(scene, pipeline, out_root, iteration, lpips_fn):
    model_dir = Path(out_root) / pipeline / "scannet" / scene
    if not model_dir.exists():
        return None

    renders, gts = load_test_renders(model_dir, iteration)
    if not renders:
        print(f"    [{pipeline}] no renders at iter {iteration}")
        # fall back to psnr.txt
        psnr_data = parse_psnr_txt(model_dir / "psnr.txt")
        if psnr_data and iteration in psnr_data:
            return {"source": "psnr.txt", **psnr_data[iteration],
                    "n_frames": 0, "renders": [], "gts": []}
        return None

    import torch
    metrics_per_frame = []
    for r, g in zip(renders, gts):
        g_resized = np.array(Image.fromarray(g).resize((r.shape[1], r.shape[0])))
        p = psnr_np(r, g_resized)
        s = ssim_np(r, g_resized)
        lp = None
        if lpips_fn is not None:
            try:
                lp = lpips_score(to_tensor(r), to_tensor(g_resized), lpips_fn)
            except Exception:
                pass
        metrics_per_frame.append({"psnr": round(p, 4), "ssim": round(s, 4),
                                   "lpips": round(lp, 4) if lp is not None else None})

    avg_psnr = float(np.mean([x["psnr"] for x in metrics_per_frame]))
    avg_ssim = float(np.mean([x["ssim"] for x in metrics_per_frame]))
    lpips_vals = [x["lpips"] for x in metrics_per_frame if x["lpips"] is not None]
    avg_lpips = float(np.mean(lpips_vals)) if lpips_vals else None

    return {
        "source": "renders",
        "PSNR": round(avg_psnr, 4),
        "SSIM": round(avg_ssim, 4),
        "LPIPS": round(avg_lpips, 4) if avg_lpips is not None else None,
        "n_frames": len(renders),
        "per_frame": metrics_per_frame,
        "renders": renders,
        "gts": gts,
    }


def save_comparison_grid(pipeline_renders, out_path, n_vis=4):
    """pipeline_renders: dict pipeline -> {"renders": [...], "gts": [...]}"""
    pipelines = list(pipeline_renders.keys())
    if not pipelines:
        return

    # use GT from first available pipeline
    gts = None
    for p in pipelines:
        if pipeline_renders[p]["gts"]:
            gts = pipeline_renders[p]["gts"]
            break
    if gts is None:
        return

    n_frames = min(len(gts), min(
        len(pipeline_renders[p]["renders"]) for p in pipelines
        if pipeline_renders[p]["renders"]
    ) if any(pipeline_renders[p]["renders"] for p in pipelines) else 0)

    if n_frames == 0:
        return

    idxs = np.linspace(0, n_frames - 1, min(n_vis, n_frames), dtype=int)
    rows = []
    for idx in idxs:
        cols = []
        for p in pipelines:
            rlist = pipeline_renders[p]["renders"]
            if rlist and idx < len(rlist):
                img = rlist[idx]
                label = make_label(f"{p}", img.shape[1])
                cols.append(np.vstack([label, img]))
        # GT column
        gt = gts[idx]
        label = make_label("GT", gt.shape[1])
        cols.append(np.vstack([label, gt]))
        if cols:
            rows.append(np.hstack(cols))
    if rows:
        grid = np.vstack(rows)
        Image.fromarray(grid).save(out_path, quality=92)


def main():
    args = parse_args()

    try:
        import torch, lpips as lpips_lib
        lpips_fn = lpips_lib.LPIPS(net='alex').cuda()
    except Exception as e:
        print(f"[warn] LPIPS unavailable ({e}), skipping")
        lpips_fn = None

    out_root = Path(args.out_root)
    pipelines = args.pipelines or DEFAULT_PIPELINES

    # find available scenes
    if args.scenes:
        scenes = args.scenes
    else:
        scene_set = set()
        for pl in pipelines:
            d = out_root / pl / "scannet"
            if d.exists():
                scene_set.update(p.name for p in d.iterdir() if p.is_dir())
        scenes = sorted(scene_set)

    print(f"Pipelines: {pipelines}")
    print(f"Scenes:    {len(scenes)}")

    eval_dir = Path(args.eval_dir)
    (eval_dir / "per_scene").mkdir(parents=True, exist_ok=True)

    all_metrics = {}
    for scene in scenes:
        print(f"\n[{scene}]")
        scene_metrics = {}
        pipeline_vis = {}

        for pl in pipelines:
            result = evaluate_scene_pipeline(scene, pl, str(out_root), args.iter, lpips_fn)
            if result is None:
                continue
            p_str = f"PSNR={result.get('PSNR', '?'):.4f}" if isinstance(result.get('PSNR'), float) else "PSNR=?"
            s_str = f"SSIM={result.get('SSIM', '?'):.4f}" if isinstance(result.get('SSIM'), float) else "SSIM=?"
            print(f"  {pl:30s}  {p_str}  {s_str}  n={result.get('n_frames',0)}")
            scene_metrics[pl] = {k: v for k, v in result.items()
                                  if k not in ("renders", "gts", "per_frame")}
            pipeline_vis[pl] = {"renders": result.get("renders", []),
                                 "gts": result.get("gts", [])}

        all_metrics[scene] = scene_metrics

        # save per-scene JSON
        with open(eval_dir / "per_scene" / f"{scene}_metrics.json", "w") as f:
            json.dump({"scene": scene, "iter": args.iter, "metrics": scene_metrics}, f, indent=2)

        # visual comparison grid
        if any(pipeline_vis[p]["renders"] for p in pipeline_vis):
            save_comparison_grid(pipeline_vis, eval_dir / "per_scene" / f"{scene}_grid.jpg", n_vis=args.n_vis)
            print(f"  grid saved -> {eval_dir}/per_scene/{scene}_grid.jpg")

    # summary JSON
    summary = {}
    for pl in pipelines:
        psnrs = [all_metrics[s][pl]["PSNR"] for s in all_metrics
                 if pl in all_metrics[s] and isinstance(all_metrics[s][pl].get("PSNR"), float)]
        ssims = [all_metrics[s][pl]["SSIM"] for s in all_metrics
                 if pl in all_metrics[s] and isinstance(all_metrics[s][pl].get("SSIM"), float)]
        lpips_vals = [all_metrics[s][pl]["LPIPS"] for s in all_metrics
                      if pl in all_metrics[s] and isinstance(all_metrics[s][pl].get("LPIPS"), float)]
        if psnrs:
            summary[pl] = {
                "PSNR_mean": round(float(np.mean(psnrs)), 4),
                "SSIM_mean": round(float(np.mean(ssims)), 4),
                "LPIPS_mean": round(float(np.mean(lpips_vals)), 4) if lpips_vals else None,
                "n_scenes": len(psnrs),
            }

    with open(eval_dir / "summary.json", "w") as f:
        json.dump({"per_scene": all_metrics, "overall": summary}, f, indent=2)

    # markdown table
    lines = [
        "# CoMoGaussian Render Quality (Test Views)\n",
        f"Iteration: {args.iter}\n",
        "| Scene | " + " | ".join(f"{pl} PSNR" for pl in pipelines) +
        " | " + " | ".join(f"{pl} SSIM" for pl in pipelines) + " |",
        "|-------|" + "|".join(["------"] * len(pipelines)) + "|" +
        "|".join(["------"] * len(pipelines)) + "|",
    ]
    for scene in sorted(all_metrics):
        m = all_metrics[scene]
        psnr_cols = [f"{m[pl]['PSNR']:.4f}" if pl in m and isinstance(m[pl].get('PSNR'), float) else "—" for pl in pipelines]
        ssim_cols = [f"{m[pl]['SSIM']:.4f}" if pl in m and isinstance(m[pl].get('SSIM'), float) else "—" for pl in pipelines]
        lines.append(f"| {scene} | " + " | ".join(psnr_cols) + " | " + " | ".join(ssim_cols) + " |")

    lines.append(f"\n**Overall mean:**")
    for pl, s in summary.items():
        lp = f"  LPIPS={s['LPIPS_mean']:.4f}" if s.get("LPIPS_mean") else ""
        lines.append(f"- {pl}: PSNR={s['PSNR_mean']:.4f}  SSIM={s['SSIM_mean']:.4f}{lp}  (n={s['n_scenes']})")

    with open(eval_dir / "summary.md", "w") as f:
        f.write("\n".join(lines))

    print(f"\n{'='*60}")
    print("Summary:")
    for pl, s in summary.items():
        print(f"  {pl:30s}  PSNR={s['PSNR_mean']:.4f}  SSIM={s['SSIM_mean']:.4f}")
    print(f"\nOutputs -> {eval_dir}/")


if __name__ == "__main__":
    main()