"""M1: Forensics-pretrained backbone transfer test on AF data.

Strategy: load CLIP ViT-L/14 (forensics literature shows CLIP image features are
unexpectedly strong on AI-image-detection because they encode high-level texture
statistics), train a tiny LINEAR probe on AF train videos (per-second binary
labels), then evaluate per-video gap on held-out AF videos.

This is the simplest meaningful "external pretrained backbone + small head"
verifier. If it works, M2 is just: replace linear probe with temporal head and
push gap higher. If it fails, fall back to from-scratch training.

Compare against current Qwen2.5-VL ForgeryHead baseline (per-video gap = +0.009,
global AUC = 0.65). Target for M1 go signal: gap > 0.05, AUC > 0.70.
"""
import os
import random
import sys
import time

import decord
import numpy as np
import torch
import torch.nn.functional as F
from decord import VideoReader, cpu
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from src.open_r1.data_loader import TRAIN_GENERATORS, build_examples
from src.open_r1.forgery_head import frame_labels_from_segments

# --- Config ---
N_VIDEOS = 200            # 80/20 split → ~160 train / ~40 test
ANNOT = "/mnt/local-fast/zhangt/annot/annot"
VROOT = "/mnt/local-fast/zhangt/video"
MODEL_ID = "openai/clip-vit-large-patch14"
DEVICE = "cuda:0"
SEED = 42
SAMPLE_FPS = 1.0          # 1 frame / sec → matches fps_to_groups=1.0


def decode_video_at_1fps(video_path: str, duration: float):
    """Return list of PIL images, one per second, at native resolution."""
    vr = VideoReader(video_path, ctx=cpu(0))
    fps_video = vr.get_avg_fps()
    n_secs = max(1, int(duration))
    idxs = []
    for sec in range(n_secs):
        idx = min(int(sec * fps_video), len(vr) - 1)
        idxs.append(idx)
    frames = vr.get_batch(idxs).asnumpy()  # (T, H, W, 3) uint8
    pil = [Image.fromarray(f) for f in frames]
    return pil, len(idxs)


def main():
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)

    print(f"Loading CLIP {MODEL_ID} ...", flush=True)
    t0 = time.time()
    clip = CLIPModel.from_pretrained(MODEL_ID, torch_dtype=torch.float32).to(DEVICE).eval()
    proc = CLIPProcessor.from_pretrained(MODEL_ID)
    print(f"  loaded in {time.time()-t0:.1f}s. hidden={clip.config.projection_dim}", flush=True)

    print("Building AF train examples ...", flush=True)
    examples = build_examples(
        annot_dir=ANNOT, video_root=VROOT, generators=TRAIN_GENERATORS,
        split_prefix="train", preprocessed_data_path=None, require_video_exists=True,
    )
    random.shuffle(examples)
    examples = examples[:N_VIDEOS]
    print(f"  sampled {len(examples)} videos", flush=True)

    # --- Extract per-frame CLIP features ---
    all_feats, all_labels, gens, vids = [], [], [], []
    t0 = time.time()
    for i, ex in enumerate(examples, 1):
        try:
            pil_imgs, n_secs = decode_video_at_1fps(ex["video_path"], ex["durations"])
        except Exception as e:
            print(f"  [skip] {ex['video_path']}: {e}", flush=True)
            continue
        # Batch through CLIP image encoder
        with torch.no_grad():
            inputs = proc(images=pil_imgs, return_tensors="pt").to(DEVICE)
            feats = clip.get_image_features(**inputs)   # (T, 768)
            feats = F.normalize(feats, dim=-1)          # L2 normalise (CLIP convention)
        lbls = frame_labels_from_segments(
            ex["solution"], n_secs, fps_to_groups=SAMPLE_FPS
        ).numpy()
        all_feats.append(feats.cpu().numpy().astype(np.float32))
        all_labels.append(lbls.astype(np.float32))
        gens.append(ex["generator"])
        vids.append(ex["video_path"])

        if i % 20 == 0:
            print(f"  [{i}/{len(examples)}] elapsed={time.time()-t0:.0f}s", flush=True)

    # --- Train/test split (80/20 over VIDEOS, not frames) ---
    n_train = int(0.8 * len(all_feats))
    X_tr = np.concatenate(all_feats[:n_train], axis=0)   # (N_tr_frames, 768)
    y_tr = np.concatenate(all_labels[:n_train], axis=0)  # (N_tr_frames,)
    test_feats = all_feats[n_train:]
    test_labels = all_labels[n_train:]
    test_gens = gens[n_train:]
    print(f"\ntrain: {X_tr.shape[0]} frames ({(y_tr>0.5).sum()} pos / {(y_tr<0.5).sum()} neg)", flush=True)
    print(f"test:  {len(test_feats)} videos, {sum(len(x) for x in test_feats)} frames", flush=True)

    # --- Linear probe ---
    Xt = torch.tensor(X_tr, dtype=torch.float32, device=DEVICE)
    yt = torch.tensor(y_tr, dtype=torch.float32, device=DEVICE)
    probe = torch.nn.Linear(Xt.shape[1], 1).to(DEVICE)
    opt = torch.optim.AdamW(probe.parameters(), lr=1e-2, weight_decay=1e-4)
    pos_weight = torch.tensor([(yt < 0.5).sum().item() / max(1, (yt > 0.5).sum().item())]).to(DEVICE)
    print(f"BCE pos_weight={pos_weight.item():.3f}", flush=True)

    for epoch in range(100):
        logits = probe(Xt).squeeze(-1)
        loss = F.binary_cross_entropy_with_logits(logits, yt, pos_weight=pos_weight)
        opt.zero_grad(); loss.backward(); opt.step()
        if (epoch + 1) % 20 == 0:
            pred = (logits.sigmoid() > 0.5).float()
            acc = (pred == yt).float().mean()
            print(f"  epoch {epoch+1:3d}  loss={loss.item():.4f}  train_acc={acc.item():.3f}", flush=True)

    # --- Eval per-video ---
    probe.eval()
    per_video_gap = []
    per_gen = {}
    all_test_scores, all_test_labels = [], []
    with torch.no_grad():
        for feats, lbls, g in zip(test_feats, test_labels, test_gens):
            logits = probe(torch.tensor(feats, device=DEVICE)).squeeze(-1).cpu().numpy()
            scores = 1.0 / (1.0 + np.exp(-logits))
            all_test_scores.append(scores)
            all_test_labels.append(lbls)
            if lbls.any() and not lbls.all():
                m_in = float(scores[lbls > 0.5].mean())
                m_out = float(scores[lbls < 0.5].mean())
                per_video_gap.append(m_in - m_out)
                per_gen.setdefault(g, []).append((m_in, m_out))

    # --- Report ---
    print("\n========== CLIP LINEAR-PROBE — TRANSFER ON AF ==========")
    if all_test_scores:
        S = np.concatenate(all_test_scores)
        Y = np.concatenate(all_test_labels)
        n_pos, n_neg = int((Y > 0.5).sum()), int((Y < 0.5).sum())
        print(f"test frames: {len(S)}  pos={n_pos} neg={n_neg}")
        print(f"global score POS={S[Y>0.5].mean():.3f}  NEG={S[Y<0.5].mean():.3f}  gap={S[Y>0.5].mean()-S[Y<0.5].mean():+.3f}")
        # AUC via Mann-Whitney (sub-sample if large)
        pos_s = S[Y > 0.5]; neg_s = S[Y < 0.5]
        if len(pos_s) > 4000 or len(neg_s) > 4000:
            rng = np.random.default_rng(SEED)
            pos_s = rng.choice(pos_s, size=min(len(pos_s), 4000), replace=False)
            neg_s = rng.choice(neg_s, size=min(len(neg_s), 4000), replace=False)
        cmp = (pos_s[:, None] > neg_s[None, :]).astype(float)
        eq = (pos_s[:, None] == neg_s[None, :]).astype(float) * 0.5
        auc = (cmp + eq).mean()
        print(f"global AUC (sampled cmp): {auc:.3f}")

    if per_video_gap:
        arr = np.array(per_video_gap)
        print(f"\nper-video gap (in_GT - out_GT) over {len(arr)} videos:")
        for q in [0, 10, 25, 50, 75, 90, 100]:
            print(f"  p{q:3d} = {np.percentile(arr, q):+.3f}")
        print(f"  mean = {arr.mean():+.3f}  std = {arr.std():.3f}")
        print(f"  gap > 0.05 : {(arr > 0.05).mean():.2%}")
        print(f"  gap > 0.10 : {(arr > 0.10).mean():.2%}")
        print(f"  gap > 0.15 : {(arr > 0.15).mean():.2%}")

    if per_gen:
        print("\nper-generator (test split only):")
        print(f"  {'gen':<12} {'n':>4} {'pos':>6} {'neg':>6} {'gap':>6}")
        for g in sorted(per_gen.keys()):
            pairs = per_gen[g]
            mp = np.mean([p[0] for p in pairs])
            mn = np.mean([p[1] for p in pairs])
            print(f"  {g:<12} {len(pairs):>4} {mp:>6.3f} {mn:>6.3f} {mp-mn:>+6.3f}")

    # --- Baseline reminder ---
    print("\n--- Baseline (Qwen ForgeryHead, from head_sanity): "
          "per-video gap = +0.009, global AUC = 0.650 ---")
    print("Target for M1 go: gap mean > +0.05, AUC > 0.70.")


if __name__ == "__main__":
    main()