"""
VIPER — HuggingFace Spaces App
Deepfake detection via identity-anchored CLIP representations.

Runs on CPU (free tier). Inference: ~30s per video.
Auto-downloads model checkpoint from HuggingFace Hub.
"""

import os
import cv2
import numpy as np
import torch
import torch.nn as nn
import gradio as gr
import tempfile
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from PIL import Image
from torchvision import transforms as T
from scipy.fft import dctn
from scipy.special import rel_entr
from huggingface_hub import hf_hub_download

# ── Device (CPU for free Spaces) ──────────────────────────────
DEVICE = "cpu"

# ── CLIP preprocessing ────────────────────────────────────────
CLIP_TF = T.Compose([
    T.Resize(224), T.CenterCrop(224),
    T.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
                std=[0.26862954, 0.26130258, 0.27577711]),
])

# ── Model definition (same as training) ───────────────────────
class VIPERv3(nn.Module):
    def __init__(self, clip_visual, dropout=0.4):
        super().__init__()
        self.clip = clip_visual
        for p in self.clip.parameters():
            p.requires_grad = False
        self.head = nn.Sequential(
            nn.Linear(784, 512), nn.BatchNorm1d(512), nn.ReLU(), nn.Dropout(dropout),
            nn.Linear(512, 128), nn.BatchNorm1d(128), nn.ReLU(), nn.Dropout(dropout * 0.5),
            nn.Linear(128, 1),
        )
    def forward(self, crops, hand):
        B, T_, C, H, W = crops.shape
        with torch.no_grad():
            embs = self.clip(crops.view(B * T_, C, H, W))
        embs = embs.view(B, T_, -1).mean(dim=1)
        return self.head(torch.cat([embs.float(), hand], dim=1)).squeeze(-1)

# ── Face detection (InsightFace) ──────────────────────────────
from insightface.app import FaceAnalysis

_face_app = None
def get_face_app():
    global _face_app
    if _face_app is None:
        _face_app = FaceAnalysis(name="buffalo_sc",
                                  providers=["CPUExecutionProvider"])
        _face_app.prepare(ctx_id=-1, det_size=(320, 320))
    return _face_app

# ── Load model (downloads from Hub on first run) ──────────────
_model = None
def get_model():
    global _model
    if _model is not None:
        return _model

    import open_clip
    clip_model, _, _ = open_clip.create_model_and_transforms(
        "ViT-L-14", pretrained="openai"
    )
    clip_model = clip_model.to(DEVICE).eval()

    model = VIPERv3(clip_model.visual, dropout=0.4).to(DEVICE)

    # Download checkpoint from your model repo
    try:
        ckpt_path = hf_hub_download(
            repo_id="rxbinsingh/VIPER",
            filename="viper_best_v3_clip.pt",
        )
        state = torch.load(ckpt_path, map_location=DEVICE)
        model.load_state_dict(state)
        print("✓ Checkpoint loaded from HuggingFace Hub")
    except Exception as e:
        print(f"⚠ No checkpoint found ({e}). Running with untrained head.")

    model.eval()
    _model = model
    return model

# ── Preprocessing functions ───────────────────────────────────
def extract_faces_from_video(video_path, num_frames=16):
    """Extract face crops and ArcFace embeddings from video."""
    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if total < 4:
        cap.release()
        return None, None

    indices = np.linspace(int(total * 0.05), int(total * 0.95), num_frames, dtype=int)
    app = get_face_app()

    crops, embeddings = [], []
    for idx in indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
        ret, frame = cap.read()
        if not ret:
            continue
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        faces = app.get(rgb)
        if not faces:
            continue
        face = max(faces, key=lambda f: (f.bbox[2]-f.bbox[0])*(f.bbox[3]-f.bbox[1]))
        x1, y1, x2, y2 = [int(v) for v in face.bbox]
        pad_x, pad_y = int((x2-x1)*0.2), int((y2-y1)*0.2)
        h, w = frame.shape[:2]
        x1, y1 = max(0, x1-pad_x), max(0, y1-pad_y)
        x2, y2 = min(w, x2+pad_x), min(h, y2+pad_y)
        crop = cv2.resize(frame[y1:y2, x1:x2], (224, 224))
        crops.append(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
        embeddings.append(face.normed_embedding)

    cap.release()
    if len(crops) < 4:
        return None, None
    return crops, embeddings


def compute_hand_features(crops, embeddings):
    """Compute GIR + TFR analytical features (16-dim)."""
    # ArcFace anchor
    embs = np.stack(embeddings[:8])
    norms = np.linalg.norm(embs, axis=1)
    weights = np.exp(norms) / np.sum(np.exp(norms))
    anchor = np.sum(weights[:, None] * embs, axis=0)
    anchor = anchor / (np.linalg.norm(anchor) + 1e-8)

    # GIR
    gir_seq = []
    for emb in embeddings:
        emb_n = emb / (np.linalg.norm(emb) + 1e-8)
        gir_seq.append(1.0 - float(np.dot(emb_n, anchor)))

    # TFR (DCT)
    def dct_profile(crop, n_bins=64):
        gray = cv2.cvtColor(crop, cv2.COLOR_RGB2GRAY).astype(np.float32) / 255.0
        dct = dctn(gray, norm="ortho")
        mag = np.abs(dct)
        H, W = mag.shape
        cy, cx = H//2, W//2
        y_idx, x_idx = np.mgrid[0:H, 0:W]
        radius = np.sqrt((y_idx-cy)**2 + (x_idx-cx)**2)
        max_r = np.sqrt(cy**2 + cx**2)
        edges = np.linspace(0, max_r, n_bins+1)
        profile = np.zeros(n_bins, dtype=np.float32)
        for i in range(n_bins):
            mask = (radius >= edges[i]) & (radius < edges[i+1])
            if mask.sum() > 0:
                profile[i] = mag[mask].mean()
        total = profile.sum()
        return profile / total if total > 0 else profile

    anchor_profile = np.mean([dct_profile(c) for c in crops[:8]], axis=0) + 1e-8
    anchor_profile = anchor_profile / anchor_profile.sum()

    tfr_seq = []
    for crop in crops:
        p = dct_profile(crop) + 1e-8
        p = p / p.sum()
        tfr_seq.append(float(np.sum(rel_entr(p, anchor_profile))))

    # Stats
    gir_arr = np.array(gir_seq)
    tfr_arr = np.array(tfr_seq)
    gir_stats = [float(gir_arr.mean()), float(gir_arr.std()), float(np.mean(gir_arr > gir_arr.mean() + 2*gir_arr.std()))]
    tfr_stats = [float(tfr_arr.mean()), float(tfr_arr.std()), float(np.mean(tfr_arr > tfr_arr.mean() + 2*tfr_arr.std()))]

    hand = gir_stats + tfr_stats + [0.0]*4 + [min(1.0, len(embeddings)/8.0), 0.0, 0.0, 0.0, 0.0, 0.0]
    return np.array(hand, dtype=np.float32), gir_seq, tfr_seq

# ── Detection function ────────────────────────────────────────
def detect_deepfake(video_path):
    if video_path is None:
        return "Upload a video to analyze.", None

    # Extract faces
    crops, embeddings = extract_faces_from_video(video_path, num_frames=16)
    if crops is None:
        return "Could not detect faces in this video. Try a video with a clear face.", None

    # Hand features
    hand_feats, gir_seq, tfr_seq = compute_hand_features(crops, embeddings)

    # CLIP inference
    model = get_model()
    base_tf = T.ToTensor()
    tensors = [CLIP_TF(base_tf(Image.fromarray(c))) for c in crops[:16]]
    while len(tensors) < 16:
        tensors.append(tensors[-1])

    crops_t = torch.stack(tensors[:16]).unsqueeze(0).to(DEVICE)
    hand_t = torch.tensor(hand_feats, dtype=torch.float32).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        l1 = model(crops_t, hand_t)
        l2 = model(torch.flip(crops_t, [-1]), hand_t)
        prob = torch.sigmoid((l1 + l2) / 2).item()

    prediction = "FAKE" if prob > 0.65 else "REAL"
    confidence = prob if prediction == "FAKE" else (1 - prob)

    # Result text
    emoji = "🔴 FAKE DETECTED" if prediction == "FAKE" else "🟢 REAL VIDEO"
    result = f"""## {emoji}

**Confidence:** {confidence*100:.1f}%
**VIPER Score:** {prob:.4f} *(>0.5 = Fake)*
**Frames Analyzed:** {len(crops)}

---

### Displacement Reaction Analysis

`AB + C → AC + B`

| Signal | Value | Status |
|--------|-------|--------|
| GIR (Identity) | {hand_feats[0]:.4f} | {'⚠️ Elevated' if hand_feats[0] > 0.35 else '✅ Normal'} |
| TFR (Texture) | {hand_feats[3]:.4f} | {'⚠️ Elevated' if hand_feats[3] > 0.08 else '✅ Normal'} |

{'**The identity anchor failed to bond** — synthetic face displaced.**' if prediction == 'FAKE' else '**Identity anchor bonded successfully** — authentic video confirmed.'}
"""

    # Plot reaction curve
    fig, axes = plt.subplots(1, 2, figsize=(10, 3.5))
    color = "#d62728" if prediction == "FAKE" else "#2ca02c"
    for ax, seq, title, thresh in zip(axes, [gir_seq, tfr_seq],
        ["GIR (Identity Distance)", "TFR (Texture Divergence)"], [0.35, 0.08]):
        frames = list(range(len(seq)))
        ax.plot(frames, seq, color=color, linewidth=2, label="Residual")
        ax.axhline(thresh, color="gray", linestyle="--", linewidth=1.2, label=f"Threshold")
        ax.fill_between(frames, seq, thresh, where=[s>thresh for s in seq], alpha=0.2, color=color)
        ax.set_title(title, fontsize=10, fontweight="bold")
        ax.set_xlabel("Frame"); ax.set_ylabel("Residual")
        ax.legend(fontsize=8); ax.grid(True, alpha=0.3)
    fig.suptitle(f"VIPER Reaction Curve — {prediction}", fontsize=12, fontweight="bold", color=color)
    plt.tight_layout()
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    plt.savefig(tmp.name, dpi=120, bbox_inches="tight")
    plt.close()

    return result, tmp.name


# ── Gradio UI ─────────────────────────────────────────────────
with gr.Blocks(title="VIPER — Deepfake Detector", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🐍 VIPER — Deepfake Detection
    **Video Identity Perturbation and Extraction Residual**

    Upload any video to check if it contains a deepfake face.

    *Inspired by displacement reactions: AB + C → AC + B*

    **AUC: 0.991 | Accuracy: 95.2% | Detects face-swap & expression-transfer**
    """)

    with gr.Row():
        with gr.Column(scale=1):
            video_input = gr.Video(label="Upload Video")
            detect_btn = gr.Button("🔍 Analyze Video", variant="primary", size="lg")
            gr.Markdown("*Processing takes ~30s on CPU*")
        with gr.Column(scale=1):
            result_output = gr.Markdown(label="Result")
            plot_output = gr.Image(label="Reaction Curve", type="filepath")

    detect_btn.click(fn=detect_deepfake, inputs=[video_input], outputs=[result_output, plot_output])

    gr.Markdown("""
    ---
    **Robin Singh** · Bennett University · 2025
    | [GitHub](https://github.com/rxbinsingh/VIPER)
    | [Paper](https://www.researchgate.net/profile/Robin-Singh-61)
    | [HuggingFace](https://huggingface.co/rxbinsingh/VIPER)
    """)

if __name__ == "__main__":
    demo.launch()