#!/usr/bin/env python3
"""Diagnostic for the BPN supervision-target fix (E1, tum_fr1_desk).

Checks whether, for the 47 "nonsharp" frames where loss_blur was active
during training, the learned BPN kernel actually moved the rendered image
TOWARDS the RAW (motion-blurred) source frame:

    PSNR(render_sharp, RAW)   -- baseline: how far the sharp 3D render
                                 already is from the blurry photo
    PSNR(BPN(render_sharp), RAW) -- after applying the learned kernel

If the kernel learned something meaningful, BPN(render_sharp) should be
CLOSER to RAW (higher PSNR / lower L1) than render_sharp is. Also reports
mask.mean() (how much blending BPN applies) and both renders' distance to
the EVSSM target, for context.

Run with: conda run -n trigsplat --cwd <repo> python <this_script>
"""
import os, sys, json, glob
import numpy as np
import cv2
import torch
from skimage.metrics import structural_similarity as ssim_fn
from skimage.metrics import peak_signal_noise_ratio as psnr_fn

REPO_TRI = "/srv2/szha0669/blur_slam_exp/repos/triangle-splatting"
sys.path.insert(0, REPO_TRI)
os.chdir(REPO_TRI)

from scene import Scene, TriangleModel
from triangle_renderer import render
from arguments import ModelParams, PipelineParams, get_combined_args
from train_bpn import build_bpn_modules, apply_bpn_blur
from argparse import ArgumentParser

BASE = "/home/szha0669/storage/blur_slam_exp"
SCENE_NAME = "tum_fr1_desk"
ITERATION = 30000
MODEL_PATH = f"{BASE}/outputs/trigsplat_i2slam_gtall_30k_colmappose_sparsedepth_E1/{SCENE_NAME}"
DATA = f"{BASE}/data/i2slam_trigsplat/tum_fr1_desk_abl1"
SHARP_JSON = f"{BASE}/outputs/logs/tum_fr1desk_sharp_frames.json"
RAW_GLOB = f"{BASE}/data/TUM_RGBD/rgbd_dataset_freiburg1_desk/rgb/*.png"
STRIDE = 10

sys.argv = ["diag.py", "-m", MODEL_PATH, "-s", DATA, "--images", "images", "-r", "2",
            "--iteration", str(ITERATION)]
parser = ArgumentParser()
model = ModelParams(parser, sentinel=True)
pipeline = PipelineParams(parser)
parser.add_argument("--iteration", default=-1, type=int)
args = get_combined_args(parser)
dataset = model.extract(args)
pipe = pipeline.extract(args)

bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")

triangles = TriangleModel(dataset.sh_degree)
scene = Scene(args=dataset, triangles=triangles, init_opacity=None, init_size=None,
               nb_points=None, set_sigma=None, no_dome=False,
               load_iteration=args.iteration, shuffle=False)

train_cameras = scene.getTrainCameras()
n_cams = len(train_cameras)
cam_index_by_name = {cam.image_name: idx for idx, cam in enumerate(train_cameras)}
h, w = int(scene.orig_h), int(scene.orig_w)

# cfg_args on disk only stores the ModelParams group (not train_bpn.py's custom
# BPN args), so the kernel sizes used for E1 training are hardcoded here.
ks1, ks2, ks3, ks_ss = 5, 9, 21, 21
mlp_ms, mlp_ss = build_bpn_modules(n_cams, h, w, ks1, ks2, ks3, ks_ss, args)
bpn_ckpt = torch.load(os.path.join(MODEL_PATH, f"bpn_{ITERATION}.pth"), map_location="cuda")
mlp_ms.load_state_dict(bpn_ckpt["mlp_ms"])
mlp_ss.load_state_dict(bpn_ckpt["mlp_ss"])
mlp_ms.eval()
mlp_ss.eval()

bpn = {
    "mlp_ms": mlp_ms, "mlp_ss": mlp_ss,
    "ks1": ks1, "ks2": ks2, "ks3": ks3, "ks_ss": ks_ss,
    "blur_chunk_rows": 96,
    "no_curriculum": getattr(args, "no_bpn_curriculum", False),
}

sharp_set = set(json.load(open(SHARP_JSON)))
raw_paths = sorted(glob.glob(RAW_GLOB))


def to_np(t):
    return (torch.clamp(t, 0, 1).detach().cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8)


rows = []
with torch.no_grad():
    for cam in train_cameras:
        if cam.image_name in sharp_set:
            continue
        render_pkg = render(cam, triangles, pipe, background)
        render_sharp = render_pkg["render"]
        depth = render_pkg.get("surf_depth", torch.zeros_like(render_sharp[:1]))
        cam_idx = cam_index_by_name[cam.image_name]
        blur_image, mask, _ = apply_bpn_blur(render_sharp, depth, cam_idx, ITERATION, bpn)

        raw_idx = int(cam.image_name) * STRIDE
        raw_img = cv2.imread(raw_paths[raw_idx])
        raw_img = cv2.cvtColor(raw_img, cv2.COLOR_BGR2RGB)
        raw_img = cv2.resize(raw_img, (cam.image_width, cam.image_height), interpolation=cv2.INTER_LANCZOS4)

        evssm = to_np(cam.original_image)
        sharp_np = to_np(render_sharp)
        blur_np = to_np(blur_image)

        row = dict(
            frame=cam.image_name,
            mask_mean=float(mask.mean().item()),
            psnr_sharp_raw=float(psnr_fn(raw_img, sharp_np, data_range=255)),
            psnr_blur_raw=float(psnr_fn(raw_img, blur_np, data_range=255)),
            ssim_sharp_raw=float(ssim_fn(raw_img, sharp_np, channel_axis=2, data_range=255)),
            ssim_blur_raw=float(ssim_fn(raw_img, blur_np, channel_axis=2, data_range=255)),
            psnr_sharp_evssm=float(psnr_fn(evssm, sharp_np, data_range=255)),
            psnr_blur_evssm=float(psnr_fn(evssm, blur_np, data_range=255)),
            l1_sharp_blur=float(np.abs(sharp_np.astype(np.float32) - blur_np.astype(np.float32)).mean()),
        )
        rows.append(row)
        print(f"{row['frame']}  mask={row['mask_mean']:.4f}  "
              f"PSNR(sharp,RAW)={row['psnr_sharp_raw']:.3f}  PSNR(blur,RAW)={row['psnr_blur_raw']:.3f}  "
              f"d_raw={row['psnr_blur_raw']-row['psnr_sharp_raw']:+.3f}  "
              f"PSNR(sharp,EVSSM)={row['psnr_sharp_evssm']:.3f}  PSNR(blur,EVSSM)={row['psnr_blur_evssm']:.3f}  "
              f"|sharp-blur|={row['l1_sharp_blur']:.3f}")

n = len(rows)
agg = {k: float(np.mean([r[k] for r in rows])) for k in rows[0] if k != "frame"}
print(f"\n=== aggregate over {n} nonsharp frames ===")
for k, v in agg.items():
    print(f"  {k:20s} = {v:.4f}")
print(f"\nDelta PSNR vs RAW (blur - sharp): {agg['psnr_blur_raw'] - agg['psnr_sharp_raw']:+.4f} dB "
      f"({'kernel moved render TOWARDS raw blur' if agg['psnr_blur_raw'] > agg['psnr_sharp_raw'] else 'kernel did NOT help match raw blur'})")

out = f"{BASE}/outputs/logs/diag_e1_bpn_kernel_vs_raw.json"
json.dump({"per_frame": rows, "aggregate": agg}, open(out, "w"), indent=1)
print(f"\nSaved -> {out}")