import gradio as gr
from transformers import pipeline

# ---- Model registry (all are image-classification on HF) ----
# Sources:
# - prithivMLmods/Deep-Fake-Detector-v2-Model (labels: Realism / Deepfake)
# - Wvolf/ViT_Deepfake_Detection (real/fake)
# - yermandy/deepfake-detection (CLIP-encoder baseline for deepfake)
MODEL_REGISTRY = {
    "ViT Deepfake v2 (Prithiv)": "prithivMLmods/Deep-Fake-Detector-v2-Model",
    "ViT Deepfake (Wvolf)":      "Wvolf/ViT_Deepfake_Detection",
    "CLIP Deepfake (yermandy)":  "yermandy/deepfake-detection",
}

_pipes = {}

def _get_pipe(model_id: str):
    if model_id not in _pipes:
        _pipes[model_id] = pipeline("image-classification",
                                    model=model_id,
                                    device_map="auto")
    return _pipes[model_id]

def _fake_real_probs(result):
    # result: list[{'label': str, 'score': float}]
    fake, real = 0.0, 0.0
    for r in result:
        lbl = r["label"].strip().lower()
        s   = float(r["score"])
        if ("fake" in lbl) or ("deepfake" in lbl) or ("ai" in lbl):
            fake = max(fake, s)
        if ("real" in lbl) or ("realism" in lbl) or ("authentic" in lbl):
            real = max(real, s)
    if fake==0.0 and real==0.0:
        # fallback: take top-1 and mirror
        top = max(result, key=lambda x: x["score"])
        is_fake = ("fake" in top["label"].lower()) or ("deepfake" in top["label"].lower()) or ("ai" in top["label"].lower())
        if is_fake:
            fake, real = float(top["score"]), 1.0 - float(top["score"])
        else:
            real, fake = float(top["score"]), 1.0 - float(top["score"])
    # normalize to sum<=1 if both present
    s = fake + real
    if s > 1.0 and s > 0:
        fake, real = fake/s, real/s
    return fake, real

def predict(img, model_name, ensemble, top_k, threshold):
    model_ids = list(MODEL_REGISTRY.values()) if ensemble else [MODEL_REGISTRY[model_name]]

    agg_fake, agg_real, rows = [], [], []
    for mid in model_ids:
        pipe = _get_pipe(mid)
        out  = pipe(img, top_k=top_k)
        f, r = _fake_real_probs(out)
        agg_fake.append(f); agg_real.append(r)
        for item in out:
            rows.append([mid, item["label"], float(item["score"])])

    fake_prob = sum(agg_fake)/len(agg_fake)
    real_prob = sum(agg_real)/len(agg_real)
    pred = "FAKE" if fake_prob >= threshold else "REAL"

    # aggregate chart (Label expects {label: score})
    chart = {"FAKE": float(fake_prob), "REAL": float(real_prob)}
    # top-k table
    rows = sorted(rows, key=lambda x: x[2], reverse=True)[: top_k * len(model_ids)]
    return pred, chart, rows

with gr.Blocks(theme="soft") as demo:
    gr.Markdown("# 🔎 Deepfake Detector\nChoose a model or use an **Ensemble** for a more robust score.")
    with gr.Row():
        with gr.Column(scale=3):
            img = gr.Image(type="pil", label="Upload image (face works best)")
            with gr.Accordion("Settings", open=False):
                model_name = gr.Dropdown(list(MODEL_REGISTRY.keys()),
                                         value="ViT Deepfake v2 (Prithiv)",
                                         label="Backbone")
                ensemble   = gr.Checkbox(label="Ensemble (use all models)", value=False)
                top_k      = gr.Slider(1, 5, value=3, step=1, label="Top-k per model")
                threshold  = gr.Slider(0.1, 0.9, value=0.5, step=0.01, label="FAKE threshold")
            btn = gr.Button("Predict", variant="primary")
        with gr.Column(scale=2):
            pred  = gr.Label(label="Prediction (FAKE vs REAL)")
            chart = gr.Label(label="Aggregated probabilities")
            table = gr.Dataframe(headers=["model", "label", "score"], wrap=True)

    btn.click(predict, [img, model_name, ensemble, top_k, threshold], [pred, chart, table])

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