"""Fine-tune LiquidInbetween from RIFE's weights.

Each sample is two frames 2–12 steps apart at 15 fps (0.13–0.8 s: photos
taken moments apart), and a real frame somewhere between them as the target
at its own t. The held-out set is a different video entirely.

    python train.py <frames dir> <eval frames dir> <teacher.safetensors> <out dir> [iterations]
"""
import json, math, random, sys, time
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 load_teacher, student_from_teacher

frames_dir, eval_dir, teacher_path, out_dir = map(Path, sys.argv[1:5])
iterations = int(sys.argv[5]) if len(sys.argv) > 5 else 6000
out_dir.mkdir(parents=True, exist_ok=True)
device = torch.device("mps")
random.seed(7)
np.random.seed(7)
torch.manual_seed(7)


def load_set(d):
    index = json.loads((d / "index.json").read_text())
    return [(np.load(d / f"{v['name']}.npy", mmap_mode="r"), set(v["cuts"])) for v in index]


train_set, eval_set = load_set(frames_dir), load_set(eval_dir)


def triplet(videos, gap_range, size=None):
    """Two frames `gap` apart with no cut between them, and one in between."""
    while True:
        frames, cuts = random.choices(videos, weights=[len(v[0]) for v in videos])[0]
        gap = random.randint(*gap_range)
        if len(frames) <= gap:
            continue
        i = random.randrange(0, len(frames) - gap)
        if any(i < c <= i + gap for c in cuts):
            continue
        j = random.randint(1, gap - 1)
        a, m, b = frames[i], frames[i + j], frames[i + gap]
        t = j / gap
        if size:
            h, w = a.shape[:2]
            y, x = random.randrange(0, h - size + 1), random.randrange(0, w - size + 1)
            a, m, b = (f[y:y + size, x:x + size] for f in (a, m, b))
        return a, m, b, t


def to_tensor(arrays):
    return torch.from_numpy(np.stack(arrays)).permute(0, 3, 1, 2).float().div(255)


def batch(n, size=224):
    items = []
    for _ in range(n):
        a, m, b, t = triplet(train_set, (2, 12), size)
        if random.random() < 0.5:
            a, m, b = a[:, ::-1], m[:, ::-1], b[:, ::-1]
        if random.random() < 0.5:
            a, b, t = b, a, 1 - t
        items.append((np.ascontiguousarray(a), np.ascontiguousarray(m), np.ascontiguousarray(b), t))
    a, m, b = (to_tensor([x[k] for x in items]).to(device) for k in range(3))
    t = torch.tensor([x[3] for x in items], dtype=torch.float32, device=device).view(-1, 1, 1, 1)
    return a, m, b, t


def charbonnier(x, y):
    return torch.sqrt((x - y) ** 2 + 1e-6).mean()


def laplacian_loss(x, y, levels=4):
    loss = 0
    for _ in range(levels):
        xd, yd = F.avg_pool2d(x, 2), F.avg_pool2d(y, 2)
        loss = loss + (x - F.interpolate(xd, size=x.shape[-2:], mode="bilinear", align_corners=False)
                       - (y - F.interpolate(yd, size=y.shape[-2:], mode="bilinear", align_corners=False))).abs().mean()
        x, y = xd, yd
    return loss


teacher = load_teacher(teacher_path).to(device).eval()
student = student_from_teacher(load_teacher(teacher_path)).to(device)
for p in teacher.parameters():
    p.requires_grad_(False)

# Fixed held-out triplets: a different video, larger gaps, three moments.
random.seed(123)
held = []
while len(held) < 120:
    a, m, b, t = triplet(eval_set, (6, 10))
    if t in (0.25, 0.5, 0.75) or len(held) % 3 == 0:
        held.append((a, m, b, t))
random.seed(7)


def psnr(x, y):
    return (-10 * torch.log10(((x - y) ** 2).mean(dim=(1, 2, 3)) + 1e-10)).mean().item()


@torch.no_grad()
def evaluate(net, timestep):
    scores = {"blend": [], "model": []}
    for a, m, b, t in held:
        a_, m_, b_ = (to_tensor([np.ascontiguousarray(f)]).to(device) for f in (a, m, 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")
        if timestep:
            out = net(pad(a_), pad(b_), torch.tensor([[[[t]]]], device=device))
        else:
            # RIFE only makes the middle; reach other moments by halving recursively.
            def at(lo, hi, x0, x1, depth=0):
                mid = net(x0, x1)
                c = (lo + hi) / 2
                if abs(c - t) < 1e-6 or depth > 4:
                    return mid
                return at(lo, c, x0, mid, depth + 1) if t < c else at(c, hi, mid, x1, depth + 1)
            out = at(0.0, 1.0, pad(a_), pad(b_))
        out = out[..., :h, :w].clamp(0, 1)
        scores["model"].append(psnr(out, m_))
        scores["blend"].append(psnr(a_ * (1 - t) + b_ * t, m_))
    return {k: float(np.mean(v)) for k, v in scores.items()}


log = []
baseline = evaluate(teacher, timestep=False)
start = evaluate(student, timestep=True)
print(f"held-out PSNR: blend {baseline['blend']:.2f}, RIFE {baseline['model']:.2f}, ours at start {start['model']:.2f}", flush=True)
log.append({"iteration": 0, "blend": baseline["blend"], "rife": baseline["model"], "ours": start["model"]})

optimizer = torch.optim.AdamW(student.parameters(), lr=1e-4, weight_decay=1e-4)
best = start["model"]
began = time.time()
for it in range(1, iterations + 1):
    lr = 1e-5 + 0.5 * (1e-4 - 1e-5) * (1 + math.cos(math.pi * it / iterations))
    for g in optimizer.param_groups:
        g["lr"] = lr
    a, m, b, t = batch(8)
    outs = student(a, b, t, all_scales=True)
    loss = charbonnier(outs[-1], m) + 0.5 * laplacian_loss(outs[-1], m) + 0.25 * charbonnier(outs[-2], m)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
    optimizer.step()
    if it % 100 == 0:
        print(f"it {it} loss {loss.item():.4f} lr {lr:.2e} {(time.time() - began) / it:.2f}s/it", flush=True)
    if it % 500 == 0 or it == iterations:
        student.eval()
        score = evaluate(student, timestep=True)["model"]
        student.train()
        log.append({"iteration": it, "ours": score})
        print(f"== it {it}: held-out PSNR ours {score:.2f} (RIFE {baseline['model']:.2f}, blend {baseline['blend']:.2f})", flush=True)
        if score > best:
            best = score
            torch.save(student.state_dict(), out_dir / "liquid_inbetween.pt")
        (out_dir / "log.json").write_text(json.dumps(log, indent=1))
torch.save(student.state_dict(), out_dir / "liquid_inbetween_last.pt")
print("done; best held-out PSNR", best, flush=True)