# -*- coding: utf-8 -*-
"""
Cleaned fake vs forged vs real audio module (Gradio UI & Notebook Magics removed)
"""

# Note: In pure Python files, install requirements via terminal, not in the code.
# pip install torch librosa numpy laion-clap groq openai-whisper transformers
# apt-get install -y ffmpeg

import os
from huggingface_hub import login
import torch
import librosa
import numpy as np
import numpy
# Naye PyTorch ko batana ke audio model ki weights safe hain
torch.serialization.add_safe_globals([numpy.core.multiarray.scalar])
import laion_clap
import whisper
from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances
from sklearn.preprocessing import normalize
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification, AutoModelForAudioFrameClassification
from groq import Groq

# --- FETCH SECRETS & LOGIN ---
# Hugging Face Secrets se keys fetch karna
HF_TOKEN = os.environ.get("HF_TOKEN")
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")

if HF_TOKEN:
    login(token=HF_TOKEN)
    print("✅ Login Successful! Ab aap models load kar sakti hain.")
else:
    print("⚠️ WARNING: HF_TOKEN not found in Secrets. Private models load nahi honge.")

# --- SETUP & CONFIG ---
client = Groq(api_key=GROQ_API_KEY)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
AST_MODEL_PATH = "aneela-pervez/FAKE-AUDIO"

# --- 🏷️ CATEGORIZED FORENSIC LABELS ---
FORENSIC_CATEGORIES = {
    "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_LABELS = [item for sublist in FORENSIC_CATEGORIES.values() for item in sublist]

# --- MODEL LOADING ---
print("🔄 Initializing Fuzzy Forensic Engine v3.9...")
whisper_model = whisper.load_model("base", device=DEVICE)
id_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base-960h")
id_model = AutoModelForAudioFrameClassification.from_pretrained("facebook/wav2vec2-base-960h").to(DEVICE).eval()

try:
    ast_extractor = AutoFeatureExtractor.from_pretrained(AST_MODEL_PATH)
    ast_model = AutoModelForAudioClassification.from_pretrained(AST_MODEL_PATH).to(DEVICE).eval()
except Exception as e:
    print(f"⚠️ AST Model issue (using default logic if fails): {e}")

# Note: laion_clap loads weights during initialization. 
clap_model = laion_clap.CLAP_Module(enable_fusion=False, amodel='HTSAT-tiny').to(DEVICE)
clap_model.load_ckpt(model_id=1)

def extract_dna_embeddings(y):
    # Ensure standard numpy array shape before processing
    y = np.ascontiguousarray(y)
    inputs = id_extractor(y, sampling_rate=16000, return_tensors="pt")
    
    # 🛡️ EXPLICIT TENSOR CONVERSION FIX:
    inputs = {k: (v.clone().detach() if isinstance(v, torch.Tensor) else torch.tensor(v)).to(DEVICE) for k, v in inputs.items()}
    
    with torch.no_grad():
        logits = id_model(**inputs).logits
        emb = torch.mean(logits, dim=1).cpu().numpy()
    return normalize(emb)

def analyze_voice_forensics(audio_path):
    try:
        y, sr = librosa.load(audio_path, sr=16000)
        y_48, _ = librosa.load(audio_path, sr=48000)
        y = y / (np.max(np.abs(y)) + 1e-9)
        y = np.ascontiguousarray(y) # Data alignment safety

        # 1. 🛡️ SEGMENT-WISE CROSS VALIDATION
        segments = np.array_split(y, 3)
        seg_embs = [extract_dna_embeddings(s) for s in segments]
        cos_sim = cosine_similarity(seg_embs[0], seg_embs[-1])[0][0]
        euc_dist = euclidean_distances(seg_embs[0], seg_embs[-1])[0][0]

        # 2. FEATURE EXTRACTION
        transcription = whisper_model.transcribe(audio_path)["text"]
        inputs = ast_extractor(y, sampling_rate=16000, return_tensors="pt")
        
        # 🛡️ EXPLICIT TENSOR CONVERSION FIX:
        inputs = {k: (v.clone().detach() if isinstance(v, torch.Tensor) else torch.tensor(v)).to(DEVICE) for k, v in inputs.items()}
        
        with torch.no_grad():
            ast_prob = torch.nn.functional.softmax(ast_model(**inputs).logits, dim=-1).cpu().numpy()[0][1]
            ast_score = ast_prob * 100

        # 🛡️ CLAP TENSOR FIX: Reverted to your strict `torch.from_numpy` logic to prevent errors
        audio_emb_np = clap_model.get_audio_embedding_from_data(x=[y_48])
        text_emb_np = clap_model.get_text_embedding(ALL_LABELS)
        
        # Explicitly converting NumPy arrays to PyTorch Tensors before calculating similarity
        audio_emb = torch.from_numpy(audio_emb_np).to(DEVICE)
        text_emb = torch.from_numpy(text_emb_np).to(DEVICE)
        c_sims = torch.nn.functional.cosine_similarity(audio_emb, text_emb)
        
        top_idx = torch.argmax(c_sims).item()
        clap_label = ALL_LABELS[top_idx]

        category = "UNKNOWN"
        for cat, labels in FORENSIC_CATEGORIES.items():
            if clap_label in labels:
                category = cat
                break

        num_edits = len(np.where(np.diff(librosa.onset.onset_strength(y=y, sr=16000)) > 6.5)[0])

        # --- 🧠 STEP 3: FUZZY LOGIC MEMBERSHIP SCORING ---
        mu_ai = min(1.0, max(0.0, (ast_score - 40) / 45))
        mu_spliced = min(1.0, max(0.0, (num_edits - 8) / 10))
        dna_conf = (cos_sim - 0.85) / 0.10
        cat_conf = 1.0 if category == "AUTHENTIC" else 0.0
        mu_auth = min(1.0, max(0.0, (dna_conf + cat_conf) / 2))

        # --- 🛡️ STEP 4: BALANCED FUZZY DECISION ENGINE (V4.0) ---
        if mu_auth > 0.80 and mu_spliced < 0.75:
            verdict = "✅ AUTHENTIC HUMAN VOICE"
            category_override = "AUTHENTIC"
        elif mu_ai > 0.85 and mu_auth < 0.80:
            verdict = "🚨 AI CLONE DETECTED"
            category_override = "AI_SYNTHETIC"
        else:
            verdict = "⚠️ MANIPULATED REAL VOICE (Spliced)"
            category_override = "MANIPULATED"

        # --- 🔥 ALIGNED PROFESSIONAL REASONING ---
        audit_context = (f"Verdict: {verdict}. Category: {category_override}. Metrics: ID {cos_sim:.2f}, "
                         f"Texture {ast_score:.1f}%, Cuts {num_edits}. "
                         f"Fuzzy Scores: AI={mu_ai:.2f}, Spliced={mu_spliced:.2f}, Auth={mu_auth:.2f}")

        prompt = (f"Act as a Senior Forensic Auditor. {audit_context} "
                  "Instruction: Justify the verdict. If identity stability is high (Auth > 0.8), "
                  "emphasize that DNA integrity confirms human origin despite high spectral texture. "
                  "Give a 2-line direct forensic conclusion.")

        res = client.chat.completions.create(model="llama-3.3-70b-versatile", messages=[{"role": "user", "content": prompt}])
        reasoning = res.choices[0].message.content

        evidence = f"Category: {category} | ID: {cos_sim:.2f} | Texture: {ast_score:.1f}% | Cuts: {num_edits}"
        return verdict, evidence, reasoning

    except Exception as e:
        return "Error", str(e), "N/A"