core/analyzer.py

"""DHVANI production audio authenticity analyzer."""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path

import librosa
import numpy as np
import torch
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification

from core.config import (
    CUSTOM_MODEL_WEIGHT,
    FAKE_THRESHOLD,
    HEURISTIC_WEIGHT,
    INDIA_E2E_WEIGHT,
    INDIA_MODEL_WEIGHT,
    MAX_DURATION_SEC,
    MIN_DURATION_SEC,
    SAMPLE_RATE,
    SEGMENT_HOP_SEC,
    SEGMENT_SEC,
    SOVEREIGN_HEURISTIC_AGREE_MIN,
    SOVEREIGN_MIXED_MEDIA_CAP,
    SOVEREIGN_MIXED_MEDIA_MIN,
    SOVEREIGN_MIXED_MEDIA_MIN_SEC,
    SOVEREIGN_PREMIUM_TTS_BOOST,
    SOVEREIGN_PREMIUM_TTS_MIN,
    SOVEREIGN_PREMIUM_TTS_ML_FLOOR,
    SOVEREIGN_ML_HIGH_CONFIDENCE,
    SOVEREIGN_ML_ONLY_CAP,
    SOVEREIGN_ML_ONLY_MIN_SEC,
    SOVEREIGN_ML_TRUST_CEILING,
    SPECTRA_WEIGHT,
    SUSPICIOUS_THRESHOLD,
    custom_model_path,
    india_e2e_model_path,
    india_model_path,
    secondary_model_ids,
    sovereign_mode,
)
from core.custom_runner import CustomDhvaniDetector
from core.heuristics import (
    acoustic_synthetic_breakdown,
    acoustic_synthetic_score,
    mixed_media_score,
    premium_tts_likelihood,
)
from core.reasoning import build_sovereign_reasoning, reasoning_to_dicts
from core.india_e2e_runner import IndiaE2EDetector
from core.spectra_runner import SpectraDetector
from core.verdict import build_verdict

logger = logging.getLogger("dhvani.analyzer")


@dataclass
class AnalysisResult:
    synthetic_probability: float
    authentic_probability: float
    verdict: str
    verdict_en: str
    title: str
    summary: str
    action: str
    use_case: str
    model_id: str
    duration_sec: float
    model_scores: dict[str, float] = field(default_factory=dict)
    heuristic_score: float = 0.0
    segment_peak: float = 0.0
    detector: str = "spectra-aasist3+ensemble"
    custom_score: float | None = None
    india_score: float | None = None
    india_e2e_score: float | None = None
    benchmark: str | None = None
    sovereign: bool = False
    mixed_media_score: float | None = None
    media_note: str | None = None
    reasoning: list[dict] = field(default_factory=list)

    def to_api_dict(self) -> dict:
        return {
            "synthetic_probability": self.synthetic_probability,
            "authentic_probability": self.authentic_probability,
            "verdict": self.verdict,
            "verdict_en": self.verdict_en,
            "title": self.title,
            "summary": self.summary,
            "action": self.action,
            "use_case": self.use_case,
            "model_id": self.model_id,
            "duration_sec": self.duration_sec,
            "model_scores": self.model_scores,
            "heuristic_score": self.heuristic_score,
            "segment_peak": self.segment_peak,
            "detector": self.detector,
            "custom_score": self.custom_score,
            "india_score": self.india_score,
            "india_e2e_score": self.india_e2e_score,
            "benchmark": self.benchmark,
            "sovereign": self.sovereign,
            "mixed_media_score": self.mixed_media_score,
            "media_note": self.media_note,
            "reasoning": self.reasoning,
        }


class _SecondaryRunner:
    def __init__(self, model_id: str, device: str) -> None:
        self.model_id = model_id
        self.device = device
        self._extractor = None
        self._model = None

    def load(self) -> None:
        if self._model is not None:
            return
        logger.info("Loading secondary model %s", self.model_id)
        self._extractor = AutoFeatureExtractor.from_pretrained(self.model_id)
        self._model = AutoModelForAudioClassification.from_pretrained(self.model_id)
        self._model.to(self.device)
        self._model.eval()

    def fake_probability(self, waveform: np.ndarray) -> float:
        self.load()
        inputs = self._extractor(
            waveform,
            sampling_rate=SAMPLE_RATE,
            return_tensors="pt",
            padding=True,
        )
        inputs = {key: value.to(self.device) for key, value in inputs.items()}
        with torch.no_grad():
            logits = self._model(**inputs).logits
            probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
        id2label = self._model.config.id2label
        label_probs = {id2label[i].lower(): float(probs[i]) for i in range(len(probs))}
        return _pick_fake_probability(label_probs, probs)

    @property
    def loaded(self) -> bool:
        return self._model is not None


def _pick_fake_probability(label_probs: dict[str, float], probs: np.ndarray) -> float:
    fake_keys = ("fake", "spoof", "synthetic", "deepfake", "bonafide_neg", "label_1", "1")
    real_keys = ("real", "bonafide", "genuine", "authentic", "label_0", "0")
    fake_prob = _pick_from_labels(label_probs, fake_keys)
    real_prob = _pick_from_labels(label_probs, real_keys)
    if fake_prob is None and real_prob is None:
        return float(probs[1]) if len(probs) == 2 else float(np.max(probs))
    if fake_prob is None:
        return max(0.0, 1.0 - real_prob)
    return fake_prob


def _pick_from_labels(label_probs: dict[str, float], keys: tuple[str, ...]) -> float | None:
    for key in keys:
        if key in label_probs:
            return label_probs[key]
    for label, prob in label_probs.items():
        if any(key in label for key in keys):
            return prob
    return None


def _segment_waveforms(waveform: np.ndarray, sr: int) -> list[np.ndarray]:
    seg_len = int(SEGMENT_SEC * sr)
    hop = int(SEGMENT_HOP_SEC * sr)
    if waveform.size <= seg_len:
        return [waveform]
    segments = [waveform[start : start + seg_len] for start in range(0, waveform.size - seg_len + 1, hop)]
    return segments or [waveform[:seg_len]]


class VoiceAnalyzer:
    """
    Production stack:
    - Primary: Spectra-AASIST3 (SOTA open-source, InTheWild EER 1.2%)
    - Secondary: wav2vec2 anti-spoofing pair
    - Tertiary: acoustic heuristics
    """

    def __init__(self) -> None:
        self._device = "cuda" if torch.cuda.is_available() else "cpu"
        self._sovereign = sovereign_mode()
        self._e2e = IndiaE2EDetector(checkpoint_path=india_e2e_model_path(), device=self._device)
        self._spectra = SpectraDetector(device=self._device)
        self._india = CustomDhvaniDetector(checkpoint_path=india_model_path(), device=self._device)
        self._custom = CustomDhvaniDetector(checkpoint_path=custom_model_path(), device=self._device)
        self._secondary = [] if self._sovereign else [_SecondaryRunner(mid, self._device) for mid in secondary_model_ids()]

    def warmup(self) -> None:
        if self._sovereign and self._e2e.enabled:
            self._e2e.load()
            return
        if self._e2e.enabled:
            self._e2e.load()
        if not self._sovereign:
            self._spectra.load()
        if self._india.enabled:
            self._india.load()
        if self._custom.enabled:
            self._custom.load()
        for runner in self._secondary:
            runner.load()

    @property
    def models_loaded(self) -> list[str]:
        loaded = []
        if self._sovereign and self._e2e.loaded:
            return [self._e2e.model_id]
        if self._e2e.loaded:
            loaded.append(self._e2e.model_id)
        if self._spectra.loaded:
            loaded.append(self._spectra.model_id)
        if self._india.loaded:
            loaded.append(self._india.model_id)
        if self._custom.loaded:
            loaded.append(self._custom.model_id)
        loaded.extend(r.model_id for r in self._secondary if r.loaded)
        return loaded

    def analyze_file(self, audio_path: str | Path) -> AnalysisResult:
        waveform, _ = librosa.load(str(audio_path), sr=SAMPLE_RATE, mono=True)
        duration_sec = len(waveform) / SAMPLE_RATE

        if duration_sec < MIN_DURATION_SEC:
            raise ValueError(
                f"Audio too short ({duration_sec:.1f}s). Minimum {MIN_DURATION_SEC}s required."
            )

        max_samples = int(MAX_DURATION_SEC * SAMPLE_RATE)
        if len(waveform) > max_samples:
            waveform = waveform[:max_samples]

        if self._sovereign and self._e2e.enabled:
            return self._analyze_sovereign(waveform, duration_sec)

        model_scores: dict[str, float] = {}

        e2e_prob = 0.0
        if self._e2e.enabled:
            try:
                e2e_prob = self._e2e.fake_probability(waveform)
                model_scores[self._e2e.model_id] = e2e_prob
            except Exception as exc:
                logger.error("DHVANI-INDIA-E2E failed: %s", exc)

        try:
            spectra_prob = self._spectra.fake_probability(waveform)
            model_scores[self._spectra.model_id] = spectra_prob
        except Exception as exc:
            logger.error("Spectra-AASIST3 failed: %s", exc)
            spectra_prob = 0.0

        india_prob = 0.0
        if self._india.enabled:
            try:
                india_prob = self._india.fake_probability(waveform)
                model_scores[self._india.model_id] = india_prob
            except Exception as exc:
                logger.error("DHVANI-INDIA model failed: %s", exc)

        custom_prob = 0.0
        if self._custom.enabled:
            try:
                custom_prob = self._custom.fake_probability(waveform)
                model_scores[self._custom.model_id] = custom_prob
            except Exception as exc:
                logger.error("DHVANI custom model failed: %s", exc)

        segments = _segment_waveforms(waveform, SAMPLE_RATE)
        secondary_peak = 0.0
        for runner in self._secondary:
            peak = 0.0
            for segment in segments:
                try:
                    peak = max(peak, runner.fake_probability(segment))
                except Exception as exc:
                    logger.warning("Secondary segment failed for %s: %s", runner.model_id, exc)
            model_scores[runner.model_id] = peak
            secondary_peak = max(secondary_peak, peak)

        heuristic_score = acoustic_synthetic_score(waveform)

        indigenous_peak = max(
            e2e_prob * INDIA_E2E_WEIGHT if self._e2e.enabled else 0.0,
            india_prob * INDIA_MODEL_WEIGHT if self._india.enabled else 0.0,
            custom_prob * CUSTOM_MODEL_WEIGHT if self._custom.enabled else 0.0,
        )
        ml_peak = max(spectra_prob * SPECTRA_WEIGHT, secondary_peak, indigenous_peak)
        segment_peak = ml_peak
        boosted_heuristic = heuristic_score * HEURISTIC_WEIGHT

        final_fake = max(ml_peak, boosted_heuristic)
        if e2e_prob >= 0.50:
            final_fake = max(final_fake, e2e_prob * 0.98)
        if spectra_prob >= 0.50 and india_prob >= 0.50:
            final_fake = max(final_fake, spectra_prob, india_prob)
        elif spectra_prob >= 0.50 and custom_prob >= 0.50:
            final_fake = max(final_fake, spectra_prob, custom_prob)
        elif india_prob >= 0.75:
            final_fake = max(final_fake, india_prob * 0.98)
        elif spectra_prob >= 0.75:
            final_fake = max(final_fake, spectra_prob * 0.98)
        elif custom_prob >= 0.75:
            final_fake = max(final_fake, custom_prob * 0.98)
        if custom_prob >= 0.40:
            final_fake = max(final_fake, custom_prob * 0.95)
        if spectra_prob >= 0.30 and secondary_peak >= 0.30:
            final_fake = max(final_fake, (spectra_prob + secondary_peak) / 2 + 0.10)
        if heuristic_score >= 0.45:
            final_fake = max(final_fake, heuristic_score * 0.95, SUSPICIOUS_THRESHOLD + 0.20)
        if ml_peak >= 0.35 and heuristic_score >= 0.35:
            final_fake = max(final_fake, 0.52)

        final_fake = min(1.0, final_fake)
        final_real = max(0.0, 1.0 - final_fake)

        detail = build_verdict(final_fake)
        parts = []
        if self._e2e.enabled:
            parts.append("DHVANI-INDIA-E2E")
        if not self._sovereign:
            parts.append("Spectra-AASIST3")
        if self._india.enabled:
            parts.append("DHVANI-INDIA")
        elif self._custom.enabled:
            parts.append("DHVANI-custom")
        if self._secondary:
            parts.append("wav2vec2-backup")
        model_label = " + ".join(parts) or "DHVANI"
        if self._e2e.enabled and self._sovereign:
            detector_tag = "dhvani-india-e2e-sovereign"
        elif self._e2e.enabled:
            detector_tag = "dhvani-india-e2e+ensemble"
        elif self._india.enabled:
            detector_tag = "spectra-aasist3+india"
        elif self._custom.enabled:
            detector_tag = "spectra-aasist3+custom"
        else:
            detector_tag = "spectra-aasist3+ensemble"
        bench_tag = "DHVANI-INDIA-Bench-v1" if (self._e2e.enabled or self._india.enabled) else None

        return AnalysisResult(
            synthetic_probability=round(final_fake * 100, 1),
            authentic_probability=round(final_real * 100, 1),
            verdict=detail.verdict,
            verdict_en=detail.verdict_en,
            title=detail.title,
            summary=detail.summary,
            action=detail.action,
            use_case=detail.use_case,
            model_id=model_label,
            duration_sec=round(duration_sec, 2),
            model_scores={k: round(v * 100, 1) for k, v in model_scores.items()},
            heuristic_score=round(heuristic_score * 100, 1),
            segment_peak=round(segment_peak * 100, 1),
            detector=detector_tag,
            custom_score=round(custom_prob * 100, 1) if self._custom.enabled else None,
            india_score=round(india_prob * 100, 1) if self._india.enabled else None,
            india_e2e_score=round(e2e_prob * 100, 1) if self._e2e.enabled else None,
            benchmark=bench_tag,
            sovereign=self._sovereign and self._e2e.enabled,
        )

    def _analyze_sovereign(self, waveform: np.ndarray, duration_sec: float) -> AnalysisResult:
        """100% DHVANI-owned inference — no foreign models."""
        e2e_prob = self._e2e.fake_probability(waveform)
        heuristic_breakdown = acoustic_synthetic_breakdown(waveform)
        heuristic_score = acoustic_synthetic_score(waveform)
        media_score = mixed_media_score(waveform)
        premium_tts_score = premium_tts_likelihood(waveform)
        media_note = None
        ml_only_cap_applied = False
        mixed_media_cap_applied = False
        premium_tts_boost_applied = False
        high_confidence_trust = e2e_prob >= SOVEREIGN_ML_HIGH_CONFIDENCE

        if e2e_prob < SOVEREIGN_ML_TRUST_CEILING:
            final_fake = e2e_prob
        elif e2e_prob < FAKE_THRESHOLD:
            heuristic_nudge = heuristic_score * HEURISTIC_WEIGHT * 0.15
            final_fake = min(e2e_prob + heuristic_nudge, SUSPICIOUS_THRESHOLD - 0.01)
        else:
            final_fake = e2e_prob * 0.95
            if heuristic_score >= 0.45:
                final_fake = max(final_fake, (e2e_prob + heuristic_score) / 2)

        # Long edited video with music: cap even when ML is confident (real speech over beds).
        if (
            duration_sec >= SOVEREIGN_MIXED_MEDIA_MIN_SEC
            and media_score >= SOVEREIGN_MIXED_MEDIA_MIN
            and e2e_prob >= FAKE_THRESHOLD
        ):
            final_fake = min(final_fake, SOVEREIGN_MIXED_MEDIA_CAP)
            mixed_media_cap_applied = True
            media_note = (
                "Edited video or background music detected. "
                "Mixed media often triggers false AI alerts — use a clean voice-only clip for orders."
            )
        elif e2e_prob < SOVEREIGN_ML_HIGH_CONFIDENCE:
            if (
                duration_sec >= SOVEREIGN_ML_ONLY_MIN_SEC
                and e2e_prob >= FAKE_THRESHOLD
                and heuristic_score < SOVEREIGN_HEURISTIC_AGREE_MIN
                and e2e_prob < 0.95
            ):
                final_fake = min(final_fake, SOVEREIGN_ML_ONLY_CAP)
                ml_only_cap_applied = True

        # Unseen premium TTS (e.g. new ElevenLabs voice): ML uncertain, physics silent, studio-clean.
        long_mixed_media = (
            duration_sec >= SOVEREIGN_MIXED_MEDIA_MIN_SEC
            and media_score >= SOVEREIGN_MIXED_MEDIA_MIN
        )
        if (
            not mixed_media_cap_applied
            and not ml_only_cap_applied
            and not long_mixed_media
            and e2e_prob < FAKE_THRESHOLD
            and e2e_prob >= SOVEREIGN_PREMIUM_TTS_ML_FLOOR
            and premium_tts_score >= SOVEREIGN_PREMIUM_TTS_MIN
            and heuristic_score < 0.12
        ):
            boosted = max(
                final_fake,
                min(
                    SOVEREIGN_PREMIUM_TTS_BOOST,
                    e2e_prob * 1.15 + premium_tts_score * 0.48,
                ),
            )
            if boosted > final_fake + 0.01:
                final_fake = boosted
                premium_tts_boost_applied = True

        final_fake = min(1.0, final_fake)
        detail = build_verdict(final_fake)
        summary = detail.summary
        if media_note:
            summary = f"{summary} {media_note}"

        reasoning = reasoning_to_dicts(
            build_sovereign_reasoning(
                e2e_prob=e2e_prob,
                heuristic_score=heuristic_score,
                media_score=media_score,
                duration_sec=duration_sec,
                final_fake=final_fake,
                model_id=self._e2e.model_id,
                heuristic_breakdown=heuristic_breakdown,
                ml_only_cap_applied=ml_only_cap_applied,
                mixed_media_cap_applied=mixed_media_cap_applied,
                premium_tts_boost_applied=premium_tts_boost_applied,
                premium_tts_score=premium_tts_score,
                high_confidence_trust=high_confidence_trust
                and not ml_only_cap_applied
                and not mixed_media_cap_applied
                and not premium_tts_boost_applied,
            )
        )

        return AnalysisResult(
            synthetic_probability=round(final_fake * 100, 1),
            authentic_probability=round((1.0 - final_fake) * 100, 1),
            verdict=detail.verdict,
            verdict_en=detail.verdict_en,
            title=detail.title,
            summary=summary,
            action=detail.action,
            use_case=detail.use_case,
            model_id="DHVANI-INDIA-E2E (sovereign)",
            duration_sec=round(duration_sec, 2),
            model_scores={self._e2e.model_id: round(e2e_prob * 100, 1)},
            heuristic_score=round(heuristic_score * 100, 1),
            segment_peak=round(e2e_prob * 100, 1),
            detector="dhvani-india-e2e-sovereign",
            india_e2e_score=round(e2e_prob * 100, 1),
            benchmark="DHVANI-INDIA-Bench-v1",
            sovereign=True,
            mixed_media_score=round(media_score * 100, 1),
            media_note=media_note,
            reasoning=reasoning,
        )