bayesian_engine.py

"""
FORENSIQ — Bayesian Evidence Synthesis Engine
Implements the core fusion algorithm from the paper:
  - Likelihood model with calibrated reliability
  - Independence correction via pairwise correlation penalty
  - Failure mode handling (marginalization over failure states)
  - Temperature-scaled calibration
  - Posterior probability computation
"""

import numpy as np
from typing import List, Dict, Any, Tuple
from dataclasses import dataclass, field
from agents.optical_agent import AgentEvidence


@dataclass
class ForensicVerdict:
    """Final verdict from Bayesian synthesis."""
    probability_fake: float           # P(Fake | Evidence), 0-1
    confidence: str                    # "Very High", "High", "Moderate", "Low"
    confidence_numeric: float          # 0-1
    verdict: str                       # "AUTHENTIC", "SUSPICIOUS", "LIKELY FAKE", "FAKE"
    agent_results: List[AgentEvidence] = field(default_factory=list)
    key_evidence: List[str] = field(default_factory=list)
    reasoning_tree: Dict[str, Any] = field(default_factory=dict)
    forensic_report: str = ""
    court_brief: str = ""


# ─── Agent Reliability Priors ────────────────────────────────────────
# Calibrated from paper validation: each agent's historical accuracy
AGENT_RELIABILITY = {
    "Optical Physics Agent": 0.78,
    "Sensor Characteristics Agent": 0.82,
    "Generative Model Agent": 0.85,
    "Statistical Priors Agent": 0.80,
    "Semantic Consistency Agent": 0.88,
    "Metadata Agent": 0.75,
    "Text & Typography Agent": 0.70,
}

# ─── Pairwise Correlation Matrix ────────────────────────────────────
# Estimated from validation: how correlated are agent outputs?
# Low correlation = independent evidence = more informative fusion
AGENT_NAMES = [
    "Optical Physics Agent",
    "Sensor Characteristics Agent",
    "Generative Model Agent",
    "Statistical Priors Agent",
    "Semantic Consistency Agent",
    "Metadata Agent",
    "Text & Typography Agent",
]

# Correlation matrix (symmetric, diagonal = 1)
CORRELATION_MATRIX = np.array([
    [1.00, 0.45, 0.30, 0.35, 0.15, 0.10, 0.05],  # Optical
    [0.45, 1.00, 0.40, 0.50, 0.10, 0.15, 0.05],  # Sensor
    [0.30, 0.40, 1.00, 0.55, 0.20, 0.15, 0.10],  # Model
    [0.35, 0.50, 0.55, 1.00, 0.15, 0.10, 0.05],  # Statistical
    [0.15, 0.10, 0.20, 0.15, 1.00, 0.20, 0.30],  # Semantic
    [0.10, 0.15, 0.15, 0.10, 0.20, 1.00, 0.10],  # Metadata
    [0.05, 0.05, 0.10, 0.05, 0.30, 0.10, 1.00],  # Text
])

ALPHA = 0.3  # Correlation penalty weight


def sigmoid(x: float) -> float:
    """Numerically stable sigmoid."""
    if x >= 0:
        return 1.0 / (1.0 + np.exp(-x))
    else:
        ez = np.exp(x)
        return ez / (1.0 + ez)


def compute_likelihood(score: float, confidence: float, reliability: float) -> Tuple[float, float]:
    """
    Compute P(evidence | Fake) and P(evidence | Real) for one agent.
    
    From paper Eq. 1:
    P(e_i | Fake, r_i, c_i) = r_i · sigmoid(s_i · c_i) + (1 - r_i) · 0.5
    """
    l_fake = reliability * sigmoid(score * confidence * 5.0) + (1 - reliability) * 0.5
    l_real = reliability * sigmoid(-score * confidence * 5.0) + (1 - reliability) * 0.5
    return l_fake, l_real


def apply_independence_correction(
    likelihoods: List[Tuple[float, float]],
    scores: List[float],
    agent_indices: List[int],
) -> List[Tuple[float, float]]:
    """
    Apply independence correction from paper Eq. 2:
    P_corr(e_i | Fake) = P(e_i | Fake) · ∏_{j≠i} (1 - α|ρ_{ij}|)^|s_j|
    """
    corrected = []
    n = len(likelihoods)
    
    for i in range(n):
        l_fake, l_real = likelihoods[i]
        idx_i = agent_indices[i]
        
        correction = 1.0
        for j in range(n):
            if i == j:
                continue
            idx_j = agent_indices[j]
            rho = CORRELATION_MATRIX[idx_i, idx_j]
            s_j = abs(scores[j])
            correction *= (1 - ALPHA * abs(rho)) ** s_j
        
        l_fake_corr = l_fake * correction + (1 - correction) * 0.5
        l_real_corr = l_real * correction + (1 - correction) * 0.5
        corrected.append((l_fake_corr, l_real_corr))
    
    return corrected


def temperature_scaling(prob: float, temperature: float = 1.5) -> float:
    """Apply temperature scaling for calibration (ECE < 0.02)."""
    if prob <= 0 or prob >= 1:
        return prob
    logit = np.log(prob / (1 - prob))
    scaled_logit = logit / temperature
    return sigmoid(scaled_logit)


def bayesian_synthesis(agent_results: List[AgentEvidence]) -> ForensicVerdict:
    """
    Main Bayesian evidence synthesis algorithm (Algorithm 1 from paper).
    
    Inputs: List of AgentEvidence from all 7 agents
    Output: ForensicVerdict with calibrated posterior probability
    """
    # Step 1: Initialize prior P(Fake) = 0.5 (uninformative)
    p_fake = 0.5
    p_real = 0.5
    
    # Step 2: Compute likelihoods for each agent
    likelihoods = []
    scores = []
    agent_indices = []
    active_agents = []
    
    for evidence in agent_results:
        # Get agent index
        try:
            idx = AGENT_NAMES.index(evidence.agent_name)
        except ValueError:
            idx = 0  # fallback
        
        # Get reliability
        reliability = AGENT_RELIABILITY.get(evidence.agent_name, 0.7)
        
        # Adjust reliability by failure probability
        effective_reliability = reliability * (1 - evidence.failure_prob)
        
        # Skip agents with very high failure probability
        if evidence.failure_prob > 0.8:
            continue
        
        l_fake, l_real = compute_likelihood(
            evidence.violation_score,
            evidence.confidence,
            effective_reliability,
        )
        
        likelihoods.append((l_fake, l_real))
        scores.append(evidence.violation_score)
        agent_indices.append(idx)
        active_agents.append(evidence)
    
    if not likelihoods:
        return ForensicVerdict(
            probability_fake=0.5,
            confidence="Very Low",
            confidence_numeric=0.1,
            verdict="INCONCLUSIVE",
            agent_results=agent_results,
            key_evidence=["No active agents produced valid evidence"],
        )
    
    # Step 3: Apply independence correction
    corrected = apply_independence_correction(likelihoods, scores, agent_indices)
    
    # Step 4: Bayesian fusion (Eq. 4 from paper)
    log_p_fake = np.log(p_fake + 1e-15)
    log_p_real = np.log(p_real + 1e-15)
    
    for l_fake, l_real in corrected:
        log_p_fake += np.log(max(l_fake, 1e-15))
        log_p_real += np.log(max(l_real, 1e-15))
    
    # Normalize in log space for numerical stability
    log_max = max(log_p_fake, log_p_real)
    p_fake_unnorm = np.exp(log_p_fake - log_max)
    p_real_unnorm = np.exp(log_p_real - log_max)
    
    posterior = p_fake_unnorm / (p_fake_unnorm + p_real_unnorm + 1e-15)
    
    # Step 5: Temperature scaling calibration
    posterior_calibrated = temperature_scaling(posterior, temperature=1.3)
    
    # Step 6: Determine verdict and confidence
    if posterior_calibrated > 0.85:
        verdict = "FAKE"
        conf_label = "Very High"
    elif posterior_calibrated > 0.65:
        verdict = "LIKELY FAKE"
        conf_label = "High"
    elif posterior_calibrated > 0.45:
        verdict = "SUSPICIOUS"
        conf_label = "Moderate"
    elif posterior_calibrated > 0.25:
        verdict = "LIKELY AUTHENTIC"
        conf_label = "Moderate"
    else:
        verdict = "AUTHENTIC"
        conf_label = "High"
    
    # Compute confidence based on agreement strength
    score_magnitudes = [abs(s) for s in scores]
    avg_magnitude = np.mean(score_magnitudes) if score_magnitudes else 0
    agreement = np.mean([1 if (s > 0) == (np.mean(scores) > 0) else 0 for s in scores]) if scores else 0
    confidence_numeric = min(1.0, avg_magnitude * agreement + 0.2)
    
    # Step 7: Extract key evidence
    key_evidence = []
    sorted_agents = sorted(active_agents, key=lambda a: abs(a.violation_score), reverse=True)
    for agent in sorted_agents[:3]:
        direction = "VIOLATED" if agent.violation_score > 0.1 else "COMPLIANT" if agent.violation_score < -0.1 else "NEUTRAL"
        key_evidence.append(
            f"{agent.agent_name}: {direction} (score={agent.violation_score:.2f}, "
            f"conf={agent.confidence:.2f})"
        )
    
    # Step 8: Build reasoning tree
    reasoning_tree = {
        "prior": {"P(Fake)": 0.5, "P(Real)": 0.5},
        "agents": {},
        "posterior": {
            "P(Fake|E)": round(posterior_calibrated, 4),
            "P(Real|E)": round(1 - posterior_calibrated, 4),
        },
        "verdict": verdict,
    }
    
    for i, agent in enumerate(active_agents):
        reasoning_tree["agents"][agent.agent_name] = {
            "violation_score": round(agent.violation_score, 4),
            "confidence": round(agent.confidence, 4),
            "failure_prob": round(agent.failure_prob, 4),
            "likelihood_fake": round(corrected[i][0], 4) if i < len(corrected) else None,
            "likelihood_real": round(corrected[i][1], 4) if i < len(corrected) else None,
            "status": "VIOLATED" if agent.violation_score > 0.1 else "COMPLIANT" if agent.violation_score < -0.1 else "NEUTRAL",
        }
    
    return ForensicVerdict(
        probability_fake=round(posterior_calibrated, 4),
        confidence=conf_label,
        confidence_numeric=round(confidence_numeric, 4),
        verdict=verdict,
        agent_results=agent_results,
        key_evidence=key_evidence,
        reasoning_tree=reasoning_tree,
    )