"""
llm_guard/multiagent_simulator.py — Multi-Agent Error Propagation Simulator
============================================================================
Simulates a 2-tier agent pipeline where Agent 1's errors can propagate to Agent 2.

Tests the independence assumption P(pipeline correct) = P(A1 correct) x P(A2 correct).

Background (exp105):
  - Empirical inter-agent correlation rho = +0.108 (weakly correlated failures)
  - P(B fails | A fails) = 0.679 vs P(B fails | A correct) = 0.574
  - Best A2A strategy: Agent B scores itself (AUROC 0.808)

Usage
-----
    from llm_guard.infrastructure.multiagent_simulator import MultiAgentSimulator

    sim = MultiAgentSimulator(n_trials=1000)
    result = sim.simulate(p1_error=0.3, p2_base_error=0.4, correlation=0.108)
    print(result.joint_failure_rate, result.independence_p_value)
Purpose:    Simulates a 2-tier agent pipeline to measure failure correlation and test the independence assumption P(pipeline correct) = P(A1 correct) × P(A2 correct). Based on exp105 empirical correlation ρ=0.108.
Exports:    MultiAgentSimulator, SimulationResult
Added:      v0.1.0
Status:     Experimental
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional

import numpy as np
from scipy import stats


@dataclass
class SimulationResult:
    """Result of one simulation run."""
    p1_error: float             # Input: Agent 1 error rate
    p2_base_error: float        # Input: Agent 2 base error rate (when A1 correct)
    p2_given_correct: float     # Observed P(A2 fails | A1 correct)
    p2_given_wrong: float       # Observed P(A2 fails | A1 wrong)
    joint_failure_rate: float   # Observed P(both fail)
    independent_expected: float # Predicted P(both fail) under independence
    correlation_coefficient: float  # Phi coefficient of actual correlation
    independence_p_value: float    # Chi-squared test p-value for independence
    independence_rejected: bool     # True if p_value < 0.05


class MultiAgentSimulator:
    """
    Monte Carlo simulator for 2-tier agent pipeline error propagation.

    Parameters
    ----------
    n_trials : int   Number of simulation trials. Default: 1000.
    seed     : int   Random seed. Default: 42.
    """

    def __init__(self, n_trials: int = 1000, seed: int = 42) -> None:
        self.n_trials = n_trials
        self.seed = seed

    def simulate(
        self,
        p1_error: float,
        p2_base_error: float,
        correlation: float = 0.108,
    ) -> SimulationResult:
        """
        Simulate joint failure in a 2-tier agent pipeline.

        Model: Agent 1 fails at rate p1_error.
               When A1 fails, A2's error rate increases by correlation * (1 - p2_base_error).
               When A1 succeeds, A2's error rate = p2_base_error.

        Parameters
        ----------
        p1_error       : float  Agent 1 failure probability.
        p2_base_error  : float  Agent 2 failure probability when A1 is correct.
        correlation    : float  Error correlation (from exp105: 0.108). Default: 0.108.

        Returns
        -------
        SimulationResult
        """
        rng = np.random.RandomState(self.seed)

        # Simulate n_trials
        a1_fails = rng.random(self.n_trials) < p1_error

        # A2 error rate is higher when A1 fails (error propagation)
        p2_given_wrong = min(1.0, p2_base_error + correlation * (1.0 - p2_base_error))
        p2_given_correct = p2_base_error  # baseline when A1 is correct

        p2_rates = np.where(a1_fails, p2_given_wrong, p2_given_correct)
        a2_fails = rng.random(self.n_trials) < p2_rates

        # Observed joint failure
        both_fail = a1_fails & a2_fails
        joint_failure_rate = float(both_fail.mean())
        independent_expected = float(p1_error * p2_base_error)

        # Observed conditional rates
        obs_p2_given_correct = float(a2_fails[~a1_fails].mean()) if (~a1_fails).sum() > 0 else 0.0
        obs_p2_given_wrong   = float(a2_fails[a1_fails].mean()) if a1_fails.sum() > 0 else 0.0

        # Chi-squared test for independence
        n_ff = int(both_fail.sum())
        n_fc = int((a1_fails & ~a2_fails).sum())
        n_cf = int((~a1_fails & a2_fails).sum())
        n_cc = int((~a1_fails & ~a2_fails).sum())
        contingency = np.array([[n_ff, n_fc], [n_cf, n_cc]])

        try:
            chi2, p_val, dof, expected_ct = stats.chi2_contingency(contingency)
            # Phi coefficient (effect size for 2x2)
            phi = np.sqrt(chi2 / self.n_trials) * np.sign(n_ff * n_cc - n_fc * n_cf)
        except (ValueError, ZeroDivisionError):
            p_val = 1.0
            phi = 0.0

        return SimulationResult(
            p1_error=p1_error,
            p2_base_error=p2_base_error,
            p2_given_correct=obs_p2_given_correct,
            p2_given_wrong=obs_p2_given_wrong,
            joint_failure_rate=joint_failure_rate,
            independent_expected=independent_expected,
            correlation_coefficient=float(phi),
            independence_p_value=float(p_val),
            independence_rejected=float(p_val) < 0.05,
        )

    def simulate_batch(
        self,
        p1_errors: List[float],
        p2_base_error: float,
        correlation: float = 0.108,
    ) -> List[SimulationResult]:
        """Run simulate() for each p1_error value."""
        return [self.simulate(p1, p2_base_error, correlation) for p1 in p1_errors]
