"""Score in-betweens where things move: PSNR over pixels whose two endpoint
frames differ (dilated), on held-out triplets, for a blend, RIFE and ours."""
import json, random, sys
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
sys.path.insert(0, str(Path(__file__).parent))
from liquid_model import IFNet, load_teacher

eval_dir, teacher_path, ours_path = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3])
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
index = json.loads((eval_dir / "index.json").read_text())
videos = [(np.load(eval_dir / f"{v['name']}.npy", mmap_mode="r"), set(v["cuts"])) for v in index]
teacher = load_teacher(teacher_path).to(device).eval()
ours = IFNet(timestep=True).to(device).eval()
ours.load_state_dict(torch.load(ours_path, map_location=device))

random.seed(99)
samples = []
while len(samples) < 90:
    frames, cuts = random.choice(videos)
    gap = random.choice([6, 8, 10, 12])
    i = random.randrange(0, len(frames) - gap)
    if any(i < c <= i + gap for c in cuts):
        continue
    j = random.choice([gap // 4, gap // 2, 3 * gap // 4])
    samples.append((frames[i], frames[i + j], frames[i + gap], j / gap))

def tensor(f):
    return torch.from_numpy(np.ascontiguousarray(f)).permute(2, 0, 1)[None].float().div(255).to(device)

def rife_at(t, x0, x1, lo=0.0, hi=1.0, depth=0):
    mid = teacher(x0, x1)
    c = (lo + hi) / 2
    if abs(c - t) < 1e-6 or depth > 4:
        return mid
    return rife_at(t, x0, mid, lo, c, depth + 1) if t < c else rife_at(t, mid, x1, c, hi, depth + 1)

scores = {"blend": [], "rife": [], "ours": []}
with torch.no_grad():
    for a, m, b, t in samples:
        A, M, B = tensor(a), tensor(m), tensor(b)
        h, w = A.shape[-2:]
        H, W = (h + 31) // 32 * 32, (w + 31) // 32 * 32
        pad = lambda x: F.pad(x, (0, W - w, 0, H - h), mode="replicate")
        moving = ((A - B).abs().max(1, keepdim=True)[0] > 0.08).float()
        moving = (F.max_pool2d(moving, 9, 1, 4) > 0).float()
        if moving.mean() < 0.01:
            continue
        outs = {
            "blend": A * (1 - t) + B * t,
            "rife": rife_at(t, pad(A), pad(B))[..., :h, :w],
            "ours": ours(pad(A), pad(B), torch.tensor([[[[t]]]], device=device))[..., :h, :w],
        }
        for k, o in outs.items():
            mse = (((o.clamp(0, 1) - M) ** 2) * moving).sum() / (moving.sum() * 3)
            scores[k].append(-10 * torch.log10(mse + 1e-10).item())
print({k: round(float(np.mean(v)), 2) for k, v in scores.items()}, "on", len(scores["ours"]), "moving samples")