# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Information aggregation and deliberation: closed forms and Monte Carlo (numpy only).

Notation follows `paper/sections/theory.tex`:
  * K agents observe signals s_k = v + eta_k, eta ~ N(0, Sigma).
  * Linear pool w (sum to one): MSE(w) = w' Sigma w.
  * Exchangeable Sigma = sigma^2 [(1-rho) I + rho 11']: MSE_mean = sigma^2 [(1-rho)/K + rho].
  * Opinion-only deliberation: x^{r+1} = (1-lambda) x^r + lambda T(x^r) 1, T = mean or median.
  * Evidence model: J evidence items z_j = v + c + u_j (c common, u_j iid), agent k uses the uniform
    average of its view A_k plus a provider-common error g_{m(k)} and an idiosyncratic reasoning error e_k.
    Evidence exchange enlarges the views; the implied evidence weights omega = (1/K) sum_k a_k determine
    the pooled MSE through sigma_u^2 ||omega||^2.
"""
from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np

# --------------------------------------------------------------------------------------------------
# Static linear pools (Propositions 1–3)
# --------------------------------------------------------------------------------------------------


def exchangeable_cov(K: int, sigma: float, rho: float) -> np.ndarray:
    """Sigma = sigma^2 [(1-rho) I + rho 11']; requires rho >= -1/(K-1) for positive semi-definiteness."""
    if K > 1 and rho < -1.0 / (K - 1):
        raise ValueError("rho below -1/(K-1): not a valid correlation matrix")
    return sigma**2 * ((1.0 - rho) * np.eye(K) + rho * np.ones((K, K)))


def mse_mean_closed_form(K: int, sigma: float, rho: float) -> float:
    """Proposition 1: MSE of the equal-weight mean of K exchangeable signals."""
    return sigma**2 * ((1.0 - rho) / K + rho)


def mse_pool(w: np.ndarray, Sigma: np.ndarray, b: np.ndarray | None = None) -> float:
    """MSE of the linear pool w (weights summing to one) with error covariance Sigma and bias vector b."""
    w = np.asarray(w, dtype=float)
    val = float(w @ Sigma @ w)
    if b is not None:
        val += float(w @ np.asarray(b, dtype=float)) ** 2
    return val


def optimal_weights(Sigma: np.ndarray) -> np.ndarray:
    """Bates–Granger weights w* = Sigma^{-1} 1 / (1' Sigma^{-1} 1) (unbiased signals, known Sigma)."""
    K = Sigma.shape[0]
    x = np.linalg.solve(Sigma, np.ones(K))
    return x / x.sum()


def mse_optimal(Sigma: np.ndarray) -> float:
    """Proposition 2: 1 / (1' Sigma^{-1} 1)."""
    K = Sigma.shape[0]
    return 1.0 / float(np.ones(K) @ np.linalg.solve(Sigma, np.ones(K)))


def mse_equal(Sigma: np.ndarray) -> float:
    """1' Sigma 1 / K^2."""
    K = Sigma.shape[0]
    return float(Sigma.sum()) / K**2


def nested_cov(n_per_provider: int, M: int, sigma_c: float, sigma_g: float, sigma_e: float) -> np.ndarray:
    """Proposition 3: eta_k = c + g_{m(k)} + e_k with K = n*M agents in M provider blocks."""
    K = n_per_provider * M
    block = np.kron(np.eye(M), np.ones((n_per_provider, n_per_provider)))
    return sigma_c**2 * np.ones((K, K)) + sigma_g**2 * block + sigma_e**2 * np.eye(K)


def mse_nested_closed_form(n_per_provider: int, M: int, sigma_c: float, sigma_g: float, sigma_e: float) -> float:
    K = n_per_provider * M
    return sigma_c**2 + sigma_g**2 / M + sigma_e**2 / K


def simulate_mean_mse(K: int, sigma: float, rho: float, n_draws: int, rng: np.random.Generator) -> float:
    """Monte Carlo MSE of the equal-weight mean under the exchangeable model (v = 0 without loss)."""
    c = rng.standard_normal(n_draws) * (sigma * np.sqrt(max(rho, 0.0)))
    e = rng.standard_normal((n_draws, K)) * (sigma * np.sqrt(1.0 - rho))
    if rho < 0:  # negative rho: draw from the full covariance instead of the factor representation
        L = np.linalg.cholesky(exchangeable_cov(K, sigma, rho))
        eta = rng.standard_normal((n_draws, K)) @ L.T
    else:
        eta = c[:, None] + e
    return float(np.mean(eta.mean(axis=1) ** 2))


# --------------------------------------------------------------------------------------------------
# Opinion-only deliberation (Propositions 4–5, 7)
# --------------------------------------------------------------------------------------------------


def deliberate_opinion_only(x0: np.ndarray, rounds: int, lam: float, anchor: str = "median") -> np.ndarray:
    """Apply x^{r+1} = (1-lam) x^r + lam T(x^r) 1 for `rounds` rounds to each row of x0 (draws x K).
    Returns an array (rounds+1, draws, K) with the trajectory."""
    x = np.asarray(x0, dtype=float)
    traj = np.empty((rounds + 1,) + x.shape)
    traj[0] = x
    for r in range(rounds):
        stat = np.median(x, axis=1) if anchor == "median" else x.mean(axis=1)
        x = (1.0 - lam) * x + lam * stat[:, None]
        traj[r + 1] = x
    return traj


def degroot_consensus_weights(W: np.ndarray) -> np.ndarray:
    """Left Perron vector pi of a primitive row-stochastic W (pi' W = pi', pi' 1 = 1)."""
    vals, vecs = np.linalg.eig(W.T)
    k = int(np.argmin(np.abs(vals - 1.0)))
    pi = np.real(vecs[:, k])
    return pi / pi.sum()


def inverse_variance_weights(Sigma: np.ndarray) -> np.ndarray:
    """Weights proportional to 1/Sigma_kk (ignore correlations), as used by the empirical aggregator."""
    w = 1.0 / np.diag(Sigma)
    return w / w.sum()


def deliberation_cov(Sigma: np.ndarray, rounds: int, lam: float) -> np.ndarray:
    """Error covariance after `rounds` rounds of mean-anchored linear deliberation: W^r Sigma W^r."""
    K = Sigma.shape[0]
    W = (1.0 - lam) * np.eye(K) + lam * np.ones((K, K)) / K
    Wr = np.linalg.matrix_power(W, rounds)
    return Wr @ Sigma @ Wr.T


# --------------------------------------------------------------------------------------------------
# Evidence model with exchange (Proposition 6) and placebo roles (Proposition 8)
# --------------------------------------------------------------------------------------------------


@dataclass
class EvidenceCommittee:
    """K agents, J evidence items, views as index lists; provider blocks; reasoning-error variances."""

    J: int
    views: list[np.ndarray]
    providers: np.ndarray  # provider index m(k)
    tau2: np.ndarray  # reasoning-error variance of agent k (prompt-specific), may depend on the slot
    sigma_c: float = 0.0
    sigma_u: float = 1.0
    sigma_g: float = 0.0

    @property
    def K(self) -> int:
        return len(self.views)

    def view_matrix(self, views: list[np.ndarray] | None = None) -> np.ndarray:
        """A (K x J) with uniform weights on each view (the BLUE within a view under iid evidence errors)."""
        views = self.views if views is None else views
        A = np.zeros((self.K, self.J))
        for k, idx in enumerate(views):
            A[k, np.asarray(idx)] = 1.0 / len(idx)
        return A

    def implied_evidence_weights(self, views: list[np.ndarray] | None = None) -> np.ndarray:
        return self.view_matrix(views).mean(axis=0)

    def mse_equal_closed_form(self, views: list[np.ndarray] | None = None, tau2: np.ndarray | None = None) -> float:
        """sigma_c^2 + sigma_u^2 ||omega||^2 + sigma_g^2 sum_m (n_m/K)^2 + (1/K^2) sum_k tau2_k."""
        omega = self.implied_evidence_weights(views)
        tau2 = self.tau2 if tau2 is None else tau2
        _, counts = np.unique(self.providers, return_counts=True)
        prov = self.sigma_g**2 * float(np.sum((counts / self.K) ** 2))
        return self.sigma_c**2 + self.sigma_u**2 * float(omega @ omega) + prov + float(tau2.sum()) / self.K**2

    def draw(self, n_draws: int, rng: np.random.Generator, views: list[np.ndarray] | None = None,
             tau2: np.ndarray | None = None, shocks: dict | None = None) -> np.ndarray:
        """Round-0 signals (n_draws x K) with v = 0. `shocks` allows common random numbers across variants."""
        tau2 = self.tau2 if tau2 is None else tau2
        if shocks is None:
            shocks = self.draw_shocks(n_draws, rng)
        A = self.view_matrix(views)
        c, u, g, e = shocks["c"], shocks["u"], shocks["g"], shocks["e"]
        return c[:, None] + u @ A.T + g[:, self.providers] + e * np.sqrt(tau2)[None, :]

    def draw_shocks(self, n_draws: int, rng: np.random.Generator) -> dict:
        M = int(self.providers.max()) + 1
        return {
            "c": rng.standard_normal(n_draws) * self.sigma_c,
            "u": rng.standard_normal((n_draws, self.J)) * self.sigma_u,
            "g": rng.standard_normal((n_draws, M)) * self.sigma_g,
            "e": rng.standard_normal((n_draws, self.K)),
        }


def exchange_views(views: list[np.ndarray], cited: list[np.ndarray] | None = None, J: int | None = None) -> list[np.ndarray]:
    """Evidence exchange: each agent adds the ids cited by the others (default: the union of all views)."""
    if cited is None:
        union = np.unique(np.concatenate([np.asarray(v) for v in views]))
        return [union.copy() for _ in views]
    return [np.unique(np.concatenate([np.asarray(v), np.asarray(c)])) for v, c in zip(views, cited)]


def random_partition_views(K: int, J: int, view_size: int, rng: np.random.Generator) -> list[np.ndarray]:
    """K views of `view_size` evidence ids each drawn without replacement (partially overlapping)."""
    return [np.sort(rng.choice(J, size=view_size, replace=False)) for _ in range(K)]


def placebo_mse(committee: EvidenceCommittee, perm: np.ndarray) -> float:
    """Equal-weight pooled MSE when the prompt-specific variances are permuted across slots (views fixed)."""
    return committee.mse_equal_closed_form(tau2=committee.tau2[perm])


def simulate_deliberation(committee: EvidenceCommittee, rounds: int, lam: float, n_draws: int,
                          rng: np.random.Generator, anchor: str = "median", exchange_at: int | None = None,
                          persuasion_sd: float = 0.0) -> dict[str, np.ndarray]:
    """Trajectory of dispersion and pooled errors under opinion-only deliberation, optional evidence
    exchange at round `exchange_at` (views enlarged to the union once) and optional persuasion noise
    (zero-mean revisions independent of the errors). Returns per-round arrays."""
    shocks = committee.draw_shocks(n_draws, rng)
    x = committee.draw(n_draws, rng, shocks=shocks)
    disp, mse_mean, mse_median = [], [], []

    def record(x):
        disp.append(float(np.mean(x.std(axis=1, ddof=1))))
        mse_mean.append(float(np.mean(x.mean(axis=1) ** 2)))
        mse_median.append(float(np.mean(np.median(x, axis=1) ** 2)))

    record(x)
    for r in range(1, rounds + 1):
        stat = np.median(x, axis=1) if anchor == "median" else x.mean(axis=1)
        x = (1.0 - lam) * x + lam * stat[:, None]
        if exchange_at is not None and r == exchange_at:
            # The exchange replaces each agent's evidence component by the enlarged-view component,
            # keeping the same evidence draws (common random numbers) and the same reasoning errors.
            A_old = committee.view_matrix()
            A_new = committee.view_matrix(exchange_views(committee.views))
            x = x + shocks["u"] @ (A_new - A_old).T
        if persuasion_sd > 0:
            x = x + rng.standard_normal(x.shape) * persuasion_sd
        record(x)
    return {"dispersion": np.array(disp), "mse_mean": np.array(mse_mean), "mse_median": np.array(mse_median)}