import os, uuid, warnings, math, tempfile
from pathlib import Path
from typing import List, Tuple

warnings.filterwarnings("ignore")

def _ensure_deps():
    try:
        import mediapipe, fpdf  
    except ImportError:
        os.system("pip install --quiet --upgrade mediapipe fpdf")

_ensure_deps()

import cv2
import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from fpdf import FPDF
import mediapipe as mp
from facenet_pytorch import InceptionResnetV1, MTCNN
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from torchvision import transforms
from transformers import AutoImageProcessor, AutoModelForImageClassification
from torchcam.methods import GradCAM as TCGradCAM
from captum.attr import Saliency
from skimage.feature import graycomatrix, graycoprops
import matplotlib.pyplot as plt
import pandas as pd
import spaces

plt.set_loglevel("ERROR")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

_face_det = MTCNN(select_largest=False, post_process=False, device=device).eval().to(device)
_df_model = InceptionResnetV1(pretrained="vggface2", classify=True, num_classes=1, device=device)
_df_model.load_state_dict(torch.load("resnet_inception.pth", map_location="cpu")["model_state_dict"])
_df_model.to(device).eval()
_df_cam = GradCAM(_df_model, target_layers=[_df_model.block8.branch1[-1]],
                  use_cuda=device.type == "cuda")

def _get_layer(model, name: str):
    mods = dict(model.named_modules())
    return mods.get(name) or next(m for n, m in mods.items() if n.endswith(name))

BIN_ID = "haywoodsloan/ai-image-detector-deploy"
_bin_proc = AutoImageProcessor.from_pretrained(BIN_ID)
_bin_mod  = AutoModelForImageClassification.from_pretrained(BIN_ID).to(device).eval()
_CAM_LAYER_BIN = "encoder.layers.3.blocks.1.layernorm_after"
_bin_cam = TCGradCAM(_bin_mod, target_layer=_get_layer(_bin_mod, _CAM_LAYER_BIN))

_susy_mod = torch.jit.load("SuSy.pt").to(device).eval()
_GEN_CLASSES = ["Stable Diffusion 1.x", "DALL·E 3",
                "MJ V5/V6", "Stable Diffusion XL", "MJ V1/V2"]
_PATCH, _TOP = 224, 5
_to_tensor = transforms.ToTensor()
_to_gray   = transforms.Compose([transforms.PILToTensor(), transforms.Grayscale()])

_calib_df_slope, _calib_df_inter = 1.0, 0.0
_calib_ai_slope, _calib_ai_inter = 1.0, 0.0


def _calibrate_df(p: float) -> float:   
    return p

def _calibrate_ai(p: float) -> float:
    return p

UNCERTAIN_GAP = 0.10
MIN_FRAMES, MAX_SAMPLES = 4, 20

def _extract_landmarks(rgb: np.ndarray) -> Tuple[np.ndarray, np.ndarray | None]:
    mesh = mp.solutions.face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1)
    res  = mesh.process(rgb); mesh.close()
    if not res.multi_face_landmarks:
        return rgb, None
    h, w, _ = rgb.shape
    out = rgb.copy()
    for lm in res.multi_face_landmarks[0].landmark:
        cx, cy = int(lm.x * w), int(lm.y * h)
        cv2.circle(out, (cx, cy), 1, (0, 255, 0), -1)
    return out, None

def _overlay_cam(cam, base):
    if torch.is_tensor(cam):  
        cam = cam.detach().cpu().numpy()

    cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-6)
    heat = Image.fromarray(
        (plt.cm.jet(cam)[:, :, :3] * 255).astype(np.uint8)
    ).resize((base.shape[1], base.shape[0]), Image.BICUBIC)

    return Image.blend(
        Image.fromarray(base).convert("RGBA"),
        heat.convert("RGBA"),
        alpha=0.45,
    )

def _render_pdf(title: str, verdict: str, conf: dict, pages: List[Image.Image]) -> str:
    out = Path(f"/tmp/report_{uuid.uuid4().hex}.pdf")
    pdf = FPDF(); pdf.set_auto_page_break(True, 15); pdf.add_page()
    pdf.set_font("Helvetica", size=14); pdf.cell(0, 10, title, ln=True, align="C")
    pdf.ln(4); pdf.set_font("Helvetica", size=12)
    pdf.multi_cell(0, 6, f"Verdict: {verdict}\n"
                         f"Confidence -> Real {conf['real']:.3f}  Fake {conf['fake']:.3f}")
    for idx, img in enumerate(pages):
        pdf.ln(4); pdf.set_font("Helvetica", size=11)
        pdf.cell(0, 6, f"Figure {idx+1}", ln=True)

        tmp = Path(tempfile.mktemp(suffix=".jpg"))
        img.convert("RGB").save(tmp, format="JPEG")
        pdf.image(str(tmp), x=10, w=90)
        tmp.unlink(missing_ok=True)
    pdf.output(out)
    return str(out)


def _susy_cam(tensor: torch.Tensor, class_idx: int) -> np.ndarray:
    sal = Saliency(_susy_mod)
    grad = sal.attribute(tensor, target=class_idx).abs().mean(1, keepdim=True)
    return grad.squeeze().detach().cpu().numpy()

@spaces.GPU
def _susy_predict(img: Image.Image):
    w, h = img.size
    npx, npy = max(1, w // _PATCH), max(1, h // _PATCH)
    patches  = np.zeros((npx * npy, _PATCH, _PATCH, 3), dtype=np.uint8)
    for i in range(npx):
        for j in range(npy):
            x, y = i * _PATCH, j * _PATCH
            patches[i*npy + j] = np.array(img.crop((x, y, x+_PATCH, y+_PATCH))
                                              .resize((_PATCH, _PATCH)))
    contrasts = []
    for p in patches:
        g = _to_gray(Image.fromarray(p)).squeeze(0).numpy()
        glcm = graycomatrix(g, [5], [0], 256, symmetric=True, normed=True)
        contrasts.append(graycoprops(glcm, "contrast")[0, 0])
    idx   = np.argsort(contrasts)[::-1][:_TOP]
    tens  = torch.from_numpy(patches[idx].transpose(0, 3, 1, 2)).float() / 255.0
    with torch.no_grad():
        probs = _susy_mod(tens.to(device)).softmax(-1).mean(0).cpu().numpy()[1:]
    return dict(zip(_GEN_CLASSES, probs))

def _fuse(p_ai: float, p_df: float) -> float:
    return 1 - (1 - p_ai) * (1 - p_df)

def _verdict(p: float) -> str:
    return "uncertain" if abs(p - 0.5) <= UNCERTAIN_GAP else ("fake" if p > 0.5 else "real")

@spaces.GPU
def _predict_image(pil: Image.Image):
    gallery: List[Image.Image] = []

    try:
        face = _face_det(pil)
    except Exception:
        face = None
    if face is not None:
        ft = F.interpolate(face.unsqueeze(0), (256, 256), mode="bilinear",
                           align_corners=False).float() / 255.0
        p_df_raw = torch.sigmoid(_df_model(ft.to(device))).item()
        p_df = _calibrate_df(p_df_raw)
        crop_np = (ft.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
        cam_df  = _df_cam(ft, [ClassifierOutputTarget(0)])[0]
        gallery.append(_overlay_cam(cam_df, crop_np))
        gallery.append(Image.fromarray(_extract_landmarks(
            cv2.cvtColor(np.array(pil), cv2.COLOR_BGR2RGB))[0]))
    else:
        p_df = 0.5

    inp_bin = _bin_proc(images=pil, return_tensors="pt").to(device)
    logits  = _bin_mod(**inp_bin).logits.softmax(-1)[0]
    p_ai_raw = logits[0].item()
    p_ai = _calibrate_ai(p_ai_raw)
    winner_idx = 0 if p_ai_raw >= logits[1].item() else 1
    inp_bin_h = {k: v.clone().detach().requires_grad_(True) for k, v in inp_bin.items()}
    cam_bin = _bin_cam(winner_idx, scores=_bin_mod(**inp_bin_h).logits)[0]
    gallery.append(_overlay_cam(cam_bin, np.array(pil)))

    bar_plot = gr.update(visible=False)
    if p_ai_raw > logits[1].item():
        gen_probs = _susy_predict(pil)
        bar_plot  = gr.update(value=pd.DataFrame(gen_probs.items(), columns=["class", "prob"]),
                              visible=True)
        susy_in   = _to_tensor(pil.resize((224, 224))).unsqueeze(0).to(device)
        g_idx     = _susy_mod(susy_in)[0, 1:].argmax().item() + 1
        cam_susy  = _susy_cam(susy_in, g_idx)
        gallery.append(_overlay_cam(cam_susy, np.array(pil)))

    # Fusion
    p_final = _fuse(p_ai, p_df)
    verdict = _verdict(p_final)
    conf    = {"real": round(1-p_final, 4), "fake": round(p_final, 4)}
    pdf     = _render_pdf("Unified Detector", verdict, conf, gallery[:3])

    return verdict, conf, gallery, bar_plot, pdf

def _sample_idx(n):  
    return list(range(n)) if n <= MAX_SAMPLES else np.linspace(0, n-1, MAX_SAMPLES, dtype=int)

@spaces.GPU
def _predict_video(path: str):
    cap = cv2.VideoCapture(path); total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 1
    probs, frames = [], []
    for i in _sample_idx(total):
        cap.set(cv2.CAP_PROP_POS_FRAMES, i)
        ok, frm = cap.read()
        if not ok:
            continue
        pil = Image.fromarray(cv2.cvtColor(frm, cv2.COLOR_BGR2RGB))
        verdict, conf, _, _, _ = _predict_image(pil)
        probs.append(conf["fake"])
        if len(frames) < MIN_FRAMES:
            frames.append(Image.fromarray(frm))
    cap.release()
    if not probs:
        blank = Image.new("RGB", (256, 256))
        return "No frames analysed", {"real": 0, "fake": 0}, [blank]

    p_final = float(np.mean(probs))
    return _verdict(p_final), {"real": round(1-p_final, 4), "fake": round(p_final, 4)}, frames

_css = "footer{visibility:hidden!important}.logo,#logo{display:none!important}"

with gr.Blocks(css=_css, title="AI-Fake & Deepfake Analyser") as demo:
    gr.Markdown("""
    ## Deepfake detector
    Upload an **image** or a short **video**.  
    The app fuses two complementary models, then shows heat-maps & a PDF report.Made by Code Alchemists
    Which is (Brijesh Khanoolkar, Shreeya Dessai, Slevin Rodrigues , Rafan Khan)
    """)

    with gr.Tab("Image"):
        with gr.Row():
            with gr.Column(scale=1):
                img_in = gr.Image(label="Upload image", type="pil")
                btn_i  = gr.Button("Analyze")
            with gr.Column(scale=2):
                txt_v  = gr.Textbox(label="Verdict", interactive=False)
                lbl_c  = gr.Label(label="Confidence")
        gal   = gr.Gallery(label="Explanations", columns=3, height=320)
        bar   = gr.BarPlot(x="class", y="prob", title="Likely generator",
                           y_label="probability", visible=False)
        pdf_f = gr.File(label="Download PDF report")

        btn_i.click(_predict_image, img_in, [txt_v, lbl_c, gal, bar, pdf_f])

    with gr.Tab("Video"):
        with gr.Row():
            with gr.Column(scale=1):
                vid_in = gr.Video(label="Upload MP4/AVI", format="mp4")
                btn_v  = gr.Button("Analyze")
            with gr.Column(scale=2):
                txt_vv = gr.Textbox(label="Verdict", interactive=False)
                lbl_cv = gr.Label(label="Confidence")
        gal_v = gr.Gallery(label="Sample frames", columns=4, height=240)

        btn_v.click(_predict_video, vid_in, [txt_vv, lbl_cv, gal_v])

demo.launch(share=True, show_api=False)