# -*- coding: utf-8 -*-
"""final lip-sync.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1YrBzX2QVG77cfqTXjgBRvc2SIU9iDe5R
"""


import gradio as gr
import cv2
import torch
import torch.nn as nn
from torchvision import transforms
import librosa
import numpy as np
import os
import gc
import re
import requests
import base64
import random
from io import BytesIO
from PIL import Image
from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances
from sklearn.preprocessing import normalize
import open_clip
import laion_clap
import whisper
import timm  # ADDED: Required for your custom Xception Lip-Sync model

# Transformers & Models
from transformers import (
    ViTForImageClassification, ViTImageProcessor,
    VideoMAEImageProcessor, VideoMAEForVideoClassification,
    AutoFeatureExtractor, AutoModelForAudioClassification, AutoModelForAudioFrameClassification,
    Wav2Vec2Processor, Wav2Vec2Model
)
from retinaface import RetinaFace
import torch.nn.functional as F
from groq import Groq

# --- 1. SETUP & AUTH ---
print("🚀 Initializing Master Forensic System...")

# API Keys
# os.environ["HF_TOKEN"] = "YOUR_TOKEN_HERE"
OPENROUTER_API_KEY = "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"
GROQ_API_KEY = "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"

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
client = Groq(api_key=GROQ_API_KEY)
AST_MODEL_PATH = "./Final_Audio_Model"
LIP_MODEL_PATH = hf_hub_download(repo_id="aneela-pervez/My-Deepfake-Models", filename="best_model (3).pth")

def clear_memory():
    gc.collect()
    torch.cuda.empty_cache()

clear_memory()

# --- 2. FORENSIC PROMPTS ---
FORENSIC_PROMPTS = {
    "DeepFake": [
        "Check face resolution vs hair.", "Check chin blurring.", "Check eyebrow definition.",
        "Check blinking mechanics.", "Check skin tone mismatch.", "Check rectangular artifacts.",
        "Check wax-like skin.", "Check gaze alignment.", "Check teeth separation.",
        "Check double chin blur.", "Check glasses artifacts.", "Check face flickering.",
        "Check edge seams.", "Check flat lighting.", "Check texture repetition."
    ],
    "MotionAnomaly": [
        "Check for jittery head movements.", "Check if expressions freeze unnaturally.",
        "Check if eye blinking is too slow or fast.", "Check for morphing artifacts during movement.",
        "Check if the face slides over the background.", "Check mouth movement consistency."
    ]
}

# CLAP SEMANTIC PROMPTS (ADDED)
CLAP_REAL_PROMPTS = [
    "A person speaking naturally",
    "Human voice with natural breathing and emotion",
    "Clear natural human conversation"
]
CLAP_FAKE_PROMPTS = [
    "Robotic AI generated voice",
    "Monotone synthetic speech without emotion",
    "Deepfake voice clone with metallic resonance"
]

# --- 3. AUDIO FORENSIC CATEGORIES ---
FORENSIC_CATEGORIES_AUDIO = {
    "AUTHENTIC": ["Natural human speech", "Human voice with natural breathing", "Consistent room ambiance"],
    "AI_SYNTHETIC": ["AI voice clone robotic smoothness", "Synthetic neural vocoder", "Deepfake metallic resonance"],
    "MANIPULATED": ["Manually spliced audio", "Inconsistent stitching artifacts", "Micro-cuts and clicks"],
    "DISGUISED": ["Artificial pitch shifting", "Muffled voice mask identity", "Time-stretched artifacts"]
}
ALL_AUDIO_LABELS = [item for sublist in FORENSIC_CATEGORIES_AUDIO.values() for item in sublist]

# --- 4. MODEL LOADING (ALL EXPERTS) ---
def load_all_models():
    print("Loading All Forensic Models...")

    try:
        vit_proc = ViTImageProcessor.from_pretrained("dima806/deepfake_vs_real_image_detection", token=os.environ["HF_TOKEN"])
        vit_model = ViTForImageClassification.from_pretrained("dima806/deepfake_vs_real_image_detection", token=os.environ["HF_TOKEN"]).to(DEVICE).eval()

        vmae_path = "MCG-NJU/videomae-base-finetuned-kinetics"
        vmae_proc = VideoMAEImageProcessor.from_pretrained(vmae_path)
        vmae_model = VideoMAEForVideoClassification.from_pretrained(vmae_path).to(DEVICE).eval()

        clip_net, _, clip_tf = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
        clip_net.to(DEVICE).eval()

        w2v_proc = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
        w2v_model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base-960h").to(DEVICE).eval()

        whisper_mod = whisper.load_model("base", device=DEVICE)

        id_ext = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base-960h")
        id_mod = AutoModelForAudioFrameClassification.from_pretrained("facebook/wav2vec2-base-960h").to(DEVICE).eval()

        try:
            ast_ext = AutoFeatureExtractor.from_pretrained(AST_MODEL_PATH, local_files_only=True)
            ast_mod = AutoModelForAudioClassification.from_pretrained(AST_MODEL_PATH, local_files_only=True).to(DEVICE).eval()
            print("✅ Loaded Custom AST Model")
        except:
            print("⚠️ Custom AST not found, using standard MIT/ast-finetuned-audioset")
            ast_ext = AutoFeatureExtractor.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593")
            ast_mod = AutoModelForAudioClassification.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593").to(DEVICE).eval()

        clap_mod = laion_clap.CLAP_Module(enable_fusion=False, amodel='HTSAT-tiny').to(DEVICE)
        clap_mod.load_ckpt(model_id=1)

        lip_mod = timm.create_model('legacy_xception', pretrained=False, num_classes=2)
        if os.path.exists(LIP_MODEL_PATH):
            try:
                cp = torch.load(LIP_MODEL_PATH, map_location=DEVICE, weights_only=False)
                state = cp.get('model_state_dict', cp)
                lip_mod.load_state_dict(state, strict=False)
                print("✅ Custom Lip-Sync Model Loaded")
            except Exception as e:
                print(f"⚠️ Custom Lip-Sync Model load failed: {e}. Using generic weights.")
        else:
            print("⚠️ Custom Lip-Sync Model not found at path, using generic weights.")
        lip_mod.to(DEVICE).eval()

        print("✅ All Models Loaded Successfully")
        return {
            "vit": (vit_model, vit_proc),
            "vmae": (vmae_model, vmae_proc),
            "clip": (clip_net, clip_tf),
            "w2v": (w2v_model, w2v_proc),
            "whisper": whisper_mod,
            "id": (id_mod, id_ext),
            "ast": (ast_mod, ast_ext),
            "clap": clap_mod,
            "lip": lip_mod
        }

    except Exception as e:
        print(f"❌ Critical Model Load Error: {e}")
        return None

models = load_all_models()

# --- 5. CLIP TRANSFORM ---
clip_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
])

# --- 6. PROJECTION LAYER FOR LSE-NET ---
class LSEProjection(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(400, 768)
    def forward(self, x):
        return self.proj(x)

lse_proj = LSEProjection().to(DEVICE)

# --- 7. HELPER FUNCTIONS ---
def extract_faces_retina(frame_img):
    frame_np = np.array(frame_img)
    try: faces = RetinaFace.detect_faces(frame_np)
    except: return frame_img
    if not faces or isinstance(faces, tuple): return frame_img
    max_area = 0; best_face = frame_img
    for key in faces:
        identity = faces[key]
        facial_area = identity["facial_area"]
        x1, y1, x2, y2 = facial_area
        width = x2 - x1; height = y2 - y1
        area = width * height
        if area > max_area:
            max_area = area
            margin = int(width * 0.2)
            x1 = max(0, x1 - margin); y1 = max(0, y1 - margin)
            x2 = min(frame_np.shape[1], x2 + margin); y2 = min(frame_np.shape[0], y2 + margin)
            best_face = frame_img.crop((x1, y1, x2, y2))
    return best_face

def extract_16_frames(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened(): return []
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    indices = np.linspace(0, total-1, 16, dtype=int)
    frames = []
    for i in indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, i)
        ret, frame = cap.read()
        if ret: frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
    cap.release()
    while len(frames) < 16: frames.append(frames[-1])
    return frames[:16]

# --- 8. AUDIO ANALYSIS MODULE (UPDATED WITH CLAP PROMPTS) ---
def analyze_audio_fuzzy(audio_path):
    try:
        y, sr = librosa.load(audio_path, sr=16000)
        y_48, _ = librosa.load(audio_path, sr=48000) # Required for CLAP
        y = y / (np.max(np.abs(y)) + 1e-9)

        # A. Extract DNA (ID Match)
        id_mod, id_ext = models["id"]
        def get_dna(sig):
            inp = id_ext(sig, sampling_rate=16000, return_tensors="pt").to(DEVICE)
            with torch.no_grad():
                logits = id_mod(**inp).logits
                emb = torch.mean(logits, dim=1).cpu().numpy()
            return normalize(emb)

        segments = np.array_split(y, 3)
        seg_embs = [get_dna(s) for s in segments]
        cos_sim = cosine_similarity(seg_embs[0], seg_embs[-1])[0][0]

        # B. Feature Extraction (AST Texture)
        ast_mod, ast_ext = models["ast"]
        inp = ast_ext(y, sampling_rate=16000, return_tensors="pt").to(DEVICE)
        with torch.no_grad():
            ast_prob = torch.nn.functional.softmax(ast_mod(**inp).logits, dim=-1).cpu().numpy()[0][1]
            ast_score = ast_prob * 100

        # C. CLAP SEMANTIC PROMPT MATCHING (The Fix for MP4 Noise)
        clap_mod = models["clap"]
        audio_emb = torch.from_numpy(clap_mod.get_audio_embedding_from_data(x=[y_48])).to(DEVICE)
        real_emb = torch.from_numpy(clap_mod.get_text_embedding(CLAP_REAL_PROMPTS)).to(DEVICE)
        fake_emb = torch.from_numpy(clap_mod.get_text_embedding(CLAP_FAKE_PROMPTS)).to(DEVICE)

        sim_real = torch.nn.functional.cosine_similarity(audio_emb, real_emb).mean().item()
        sim_fake = torch.nn.functional.cosine_similarity(audio_emb, fake_emb).mean().item()

        # Normalize CLAP scores
        sim_real_norm = max(0, sim_real + 0.1)
        sim_fake_norm = max(0, sim_fake + 0.1)
        clap_fake_score = (sim_fake_norm / (sim_real_norm + sim_fake_norm)) * 100

        # D. Edits / Micro-cuts
        num_edits = len(np.where(np.diff(librosa.onset.onset_strength(y=y, sr=16000)) > 8.0)[0])
        edit_penalty = max(0, (num_edits - 15) * 3.0)

        # E. Hybrid Final Anomaly Score
        # Trusts CLAP Semantics (65%) more than AST Texture (35%) to ignore background noise
        final_anomaly_score = (clap_fake_score * 0.65) + (ast_score * 0.35) + edit_penalty
        final_anomaly_score = min(100.0, final_anomaly_score)

        metrics = f"CLAP Fake Match: {clap_fake_score:.1f}% | AST Texture: {ast_score:.1f}% | Cuts: {num_edits} | ID Match: {cos_sim:.2f}"

        # F. Verdict formatting expected by master_pipeline
        if final_anomaly_score > 70.0:
            verdict = "FAKE"
            reason = "AI Voice Detected (CLAP Semantic Match)"
        else:
            verdict = "REAL"
            reason = "Authentic Voice (CLAP Verified)"

        return verdict, final_anomaly_score, reason, metrics

    except Exception as e:
        return "ERROR", 0, str(e), ""

# --- 9. INTEGRATED LIP SYNC MODULE (LSE-NET APPROACH) ---
# Yahan Face Image model ki jagah, Audio-Visual Vector Binding lagai gayi hai!
def analyze_lip_sync(frames, audio_path):
    try:
        # 1. Audio Features (Wav2Vec2)
        speech, _ = librosa.load(audio_path, sr=16000)
        w2v_mod, w2v_proc = models["w2v"]

        a_inputs = w2v_proc(speech, return_tensors="pt", sampling_rate=16000).input_values.to(DEVICE)
        with torch.no_grad():
            audio_embeddings = w2v_mod(a_inputs).last_hidden_state.mean(dim=1)
            audio_embeddings = F.normalize(audio_embeddings, dim=-1) # Shape: [1, 768]

        # 2. Video Features (VideoMAE)
        vmae_mod, vmae_proc = models["vmae"]
        v_inputs = vmae_proc(list(frames), return_tensors="pt").to(DEVICE)

        with torch.no_grad():
            video_outputs = vmae_mod(**v_inputs).logits # Output size: 400
            video_embeddings = lse_proj(video_outputs)  # Projected to size: 768
            video_embeddings = F.normalize(video_embeddings, dim=-1) # Shape: [1, 768]

        # 3. Cross-Modal Binding Match (Cosine Similarity)
        sync_similarity = F.cosine_similarity(audio_embeddings, video_embeddings).item()

        # Calculate Anomaly Score (100% means terrible match, 0% means perfect sync)
        sync_error_score = (1 - max(0, sync_similarity)) * 100

        return sync_error_score

    except Exception as e:
        print(f"LSE-Net Lip-Sync Verification Failed: {e}")
        return 100.0 # Error ki surat mein isay suspicious declare kar dein

# --- 10. MASTER PIPELINE ---
def master_pipeline(video_path):
    clear_memory()
    if not video_path: return None, "No Video", "Error"

    # 1. Video Analysis
    frames = extract_16_frames(video_path)
    if not frames: return None, "Extraction Failed", "Error"

    # VideoMAE
    vmae_mod, vmae_proc = models["vmae"]
    v_in = vmae_proc(list(frames), return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        v_conf = torch.softmax(vmae_mod(**v_in).logits, dim=1).max().item() * 100
        video_score = 100 - v_conf # Low confidence = Anomaly

    # ViT (Visual Check on Faces)
    vit_mod, vit_proc = models["vit"]
    face = extract_faces_retina(frames[8]) # Check middle frame
    v_inp = vit_proc(images=face, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        vit_prob = vit_mod(**v_inp).logits.softmax(dim=1)[0][0].item() * 100 # 0=Fake

    # 2. Audio Analysis (Using CLAP Prompts)
    audio_verdict, audio_score, audio_reason, audio_metrics = analyze_audio_fuzzy(video_path)

    # 3. Lip Sync (Using LSE-Net Bridge: Audio-Visual Match)
    sync_score = analyze_lip_sync(frames, video_path)

    # --- 🧠 MASTER DECISION LOGIC (5 CASES) ---
    VIDEO_FAKE_THRESH = 75
    AUDIO_FAKE_THRESH = 75
    SYNC_BAD_THRESH = 65

    is_video_fake = (video_score > VIDEO_FAKE_THRESH) or (vit_prob > 50)
    is_audio_fake = (audio_verdict != "REAL")

    if not is_video_fake and not is_audio_fake:
        # Both Real -> Check Sync
        if sync_score > SYNC_BAD_THRESH:
            final_case = "CASE 1: Real Video + Real Audio + Lip Sync Issue"
            verdict_short = "SYNC_MISMATCH"
        else:
            final_case = "CASE 5: Real Video + Real Audio + Authentic"
            verdict_short = "AUTHENTIC"

    elif not is_video_fake and is_audio_fake:
        final_case = "CASE 2: Real Video + Fake Audio"
        verdict_short = "AUDIO_FAKE"

    elif is_video_fake and not is_audio_fake:
        final_case = "CASE 3: Fake Video + Real Audio"
        verdict_short = "VIDEO_FAKE"

    else: # Both Fake
        final_case = "CASE 4: Fake Video + Fake Audio"
        verdict_short = "FULL_DEEPFAKE"

    # --- CLAUDE FINAL JUDGE ---
    prompts = random.sample(FORENSIC_PROMPTS["DeepFake"], 3)
    prompt_text = (
        f"Role: Senior Forensic Judge. Analyze this case.\n"
        f"Data:\n"
        f"1. Video Anomaly: {video_score:.1f}% (High=Fake)\n"
        f"2. Audio Anomaly: {audio_score:.1f}% (Verdict: {audio_verdict})\n"
        f"3. Lip-Sync Error (Cross-Modal Expert): {sync_score:.1f}% (High=Tampered)\n"
        f"4. Detected Case: {final_case}\n\n"
        f"Task: Provide a professional forensic justification for determining '{final_case}'. "
        f"Mention specific artifacts derived from these scores. Keep it concise."
    )

    # Send Keyframe to Claude
    buffered = BytesIO()
    frames[8].save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue()).decode()

    payload = {
        "model": "anthropic/claude-3.5-sonnet",
        "messages": [{"role": "user", "content": [
            {"type": "text", "text": prompt_text},
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img_str}"}}
        ]}]
    }

    try:
        response = requests.post("https://openrouter.ai/api/v1/chat/completions",
                               headers={"Authorization": f"Bearer {OPENROUTER_API_KEY}"},
                               json=payload)
        claude_reasoning = response.json()['choices'][0]['message']['content']
    except:
        claude_reasoning = f"Automated Verdict: {final_case}. Metrics align with this classification."

    # Final Output Formatting
    summary = (
        f"# 🛡️ FINAL VERDICT: {verdict_short}\n"
        f"## 📂 Classification: {final_case}\n\n"
        f"### 🤖 Forensic Reasoning (Claude):\n{claude_reasoning}\n\n"
        f"### 📊 Detailed Metrics:\n"
        f"- **Lip-Sync Anomaly (LSE-Net):** {sync_score:.1f}% (Threshold: {SYNC_BAD_THRESH})\n"
        f"- **Video Score:** {video_score:.1f}% (Threshold: {VIDEO_FAKE_THRESH})\n"
        f"- **Audio Score:** {audio_score:.1f}% (Threshold: {AUDIO_FAKE_THRESH})\n"
        f"- **Audio Details:** {audio_metrics}\n"
        f"- **Video Artifacts:** Checked {', '.join(prompts)}"
    )

    return frames[8], summary, claude_reasoning

# --- 11. UI ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🏢 Master Forensic Lab (Video + Audio + LSE-Net)")
    gr.Markdown("Integrated system categorizing 5 distinct forensic scenarios using LSE-Net Audio-Visual Lip-Sync detection.")

    with gr.Row():
        vid_in = gr.Video(label="Upload Investigation Video")
        btn = gr.Button("🔍 Run Master Analysis", variant="primary")

    with gr.Row():
        img_out = gr.Image(label="Keyframe Analysis")
        txt_out = gr.Markdown(label="Master Report")

    reason_out = gr.Textbox(label="Judge Reasoning Log", lines=2)

    btn.click(master_pipeline, inputs=[vid_in], outputs=[img_out, txt_out, reason_out])

if __name__ == "__main__":
    demo.launch(share=True, debug=True)