"""Linear-quadratic delegation game (Model A, §1.8 of docs/model_specification.md).

Closed forms derived in scripts/00_lq_derivations.py (sympy) and re-implemented here in numpy so that every formula is
cross-checked by Monte Carlo in tests/unit/test_lq_game.py.

Conventions
-----------
theta ~ N(0, 1/tau_theta); human signal s_i = theta + e_i (precision tau_H); model-m signal s_m = theta + eta_m
(precision tau_A), common to all users of model m; M models, delegation share D split evenly.
Best response x_i = E_i[kappa0 * theta + alpha * X]; alpha < 0 strategic substitutes.
Payoff normalised by |U_xx| = 1; the social curvature on the aggregate action is (1 - alpha) under the CRS restriction
U_xX + U_XX = 0 (Lemma 2); the general case U_XX = -upsilon * U_xX is kept for Proposition 4(a).
"""
from __future__ import annotations

from dataclasses import dataclass, replace

import numpy as np


@dataclass(frozen=True)
class LQParams:
    kappa0: float = 1.0
    alpha: float = -0.3
    tau_theta: float = 1.0
    tau_H: float = 1.0
    tau_A: float = 4.0
    M: int = 1
    upsilon: float = 1.0       # U_XX = -upsilon U_xX ; 1 = CRS/pecuniary restriction
    U_Xtheta: float = 0.0      # direct aggregate-state interaction (0 under CRS)

    @property
    def beta_H(self):
        return self.tau_H / (self.tau_theta + self.tau_H)

    @property
    def beta_A(self):
        return self.tau_A / (self.tau_theta + self.tau_A)


# --------------------------------------------------------------------------- equilibrium strategies (A.8)
def strategies(p: LQParams, D: float) -> dict:
    bH, bA = p.beta_H, p.beta_A
    xi = 1.0 / (1.0 - p.alpha * (D / p.M) * (1.0 - bA))
    c = p.kappa0 / (1.0 - p.alpha * ((1.0 - D) * bH + D * bA * xi))
    aH, aA = bH * c, bA * xi * c
    Abar = (1.0 - D) * aH + D * aA
    return dict(aH=aH, aA=aA, xi=xi, c=c, Abar=Abar, kappa_full=p.kappa0 / (1.0 - p.alpha))


# --------------------------------------------------------------------------- second moments (A.9)
def moments(p: LQParams, D: float) -> dict:
    s = strategies(p, D)
    st2, sH2, sA2 = 1.0 / p.tau_theta, 1.0 / p.tau_H, 1.0 / p.tau_A
    aH, aA, Abar, kf = s['aH'], s['aA'], s['Abar'], s['kappa_full']
    V_noise = D ** 2 * aA ** 2 * sA2 / p.M
    V_track = (Abar - kf) ** 2 * st2
    V_agg = V_track + V_noise
    # dispersion E (x_i - X)^2
    V_disp_H = aH ** 2 * sH2 + (aH - Abar) ** 2 * st2 + V_noise
    V_disp_A = (aA - Abar) ** 2 * st2 + aA ** 2 * sA2 - 2 * aA * (D * aA / p.M) * sA2 + V_noise
    V_disp = (1.0 - D) * V_disp_H + D * V_disp_A
    # raw moments for the general quadratic welfare
    Ex2 = (1.0 - D) * aH ** 2 * (st2 + sH2) + D * aA ** 2 * (st2 + sA2)
    EX2 = Abar ** 2 * st2 + V_noise
    ExX = (1.0 - D) * aH * Abar * st2 + D * aA * (Abar * st2 + D * aA / p.M * sA2)
    Ext = ((1.0 - D) * aH + D * aA) * st2
    EXt = Abar * st2
    return dict(V_noise=V_noise, V_track=V_track, V_agg=V_agg, V_disp=V_disp,
                Ex2=Ex2, EX2=EX2, ExX=ExX, Ext=Ext, EXt=EXt, **s)


# --------------------------------------------------------------------------- welfare (A.11) and values
def welfare_gross(p: LQParams, D: float) -> float:
    """Per-capita expected payoff E U(x_i, X, theta) (constant terms dropped), general quadratic form."""
    m = moments(p, D)
    a = p.alpha
    return (-m['Ex2'] / 2.0 + a * m['ExX'] - p.upsilon * a * m['EX2'] / 2.0
            + p.kappa0 * m['Ext'] + p.U_Xtheta * m['EXt'])


def welfare_decomposed(p: LQParams, D: float) -> dict:
    """Angeletos–Pavan decomposition under the CRS restriction: W = W* - 1/2[(1-alpha) V_agg + V_disp]."""
    m = moments(p, D)
    return dict(loss_agg=0.5 * (1.0 - p.alpha) * m['V_agg'], loss_disp=0.5 * m['V_disp'],
                W_rel=-0.5 * ((1.0 - p.alpha) * m['V_agg'] + m['V_disp']))


def private_value(p: LQParams, D: float) -> float:
    """Gain in expected payoff of a measure-zero agent switching from own signal to model signal (units |U_xx|)."""
    m = moments(p, D)
    st2, sH2, sA2 = 1.0 / p.tau_theta, 1.0 / p.tau_H, 1.0 / p.tau_A
    g = p.kappa0 + p.alpha * m['Abar']
    Vx = g ** 2 * st2 + p.alpha ** 2 * (D * m['aA'] / p.M) ** 2 * p.M * sA2
    loss_H = Vx - (g * st2) ** 2 / (st2 + sH2)
    loss_A = Vx - (g * st2 + p.alpha * (D * m['aA'] / p.M) * sA2) ** 2 / (st2 + sA2)
    return 0.5 * (loss_H - loss_A)


def social_marginal_value(p: LQParams, D: float, h: float = 1e-6) -> float:
    """d E U / dD (gross of fees), central finite difference of the exact quadratic welfare."""
    D0, D1 = max(0.0, D - h), min(1.0, D + h)
    return (welfare_gross(p, D1) - welfare_gross(p, D0)) / (D1 - D0)


def wedge(p: LQParams, D: float) -> float:
    """Private minus social marginal value; identically 0 under Lemma 2 (Proposition 3)."""
    return private_value(p, D) - social_marginal_value(p, D)


# --------------------------------------------------------------------------- delegation equilibrium (A.10)
def delegation_equilibrium(p: LQParams, cost_cdf, scale: float = 1.0, tol: float = 1e-12) -> float:
    """Fixed point D = F(scale * PV(D) - 0) where cost_cdf(c) = Pr(kappa_i + fee <= c).

    PV is evaluated in units |U_xx|; `scale` converts to the units of the cost distribution.
    Solved by bisection on D in [0,1]; the map D -> F(scale*PV(D)) is continuous, so a fixed point exists (Brouwer).
    Uniqueness is checked numerically by scanning for sign changes of D - F(.).
    """
    f = lambda D: D - cost_cdf(scale * private_value(p, D))
    grid = np.linspace(0.0, 1.0, 201)
    vals = np.array([f(D) for D in grid])
    if vals[0] >= 0:
        return 0.0
    if vals[-1] <= 0:
        return 1.0
    idx = np.where(np.diff(np.sign(vals)) != 0)[0]
    lo, hi = grid[idx[0]], grid[idx[0] + 1]
    for _ in range(200):
        mid = 0.5 * (lo + hi)
        if f(mid) > 0:
            hi = mid
        else:
            lo = mid
        if hi - lo < tol:
            break
    return 0.5 * (lo + hi)


def planner_delegation(p: LQParams, cost_cdf, cost_quantile, scale: float = 1.0) -> float:
    """Planner's D: maximises E U(D) - ∫_0^D cost_quantile(q) dq / scale over a grid + golden refinement."""
    grid = np.linspace(0.0, 1.0, 401)

    def obj(D):
        if D <= 0:
            return welfare_gross(p, 0.0)
        qs = np.linspace(0.0, D, 400)
        fees = np.trapezoid(cost_quantile(qs), qs) / scale
        return welfare_gross(p, D) - fees

    vals = np.array([obj(D) for D in grid])
    return float(grid[int(np.argmax(vals))])


# --------------------------------------------------------------------------- Monte Carlo verifier
def simulate(p: LQParams, D: float, n_agents: int = 20000, n_draws: int = 4000, seed: int = 0) -> dict:
    """Simulate the signal game with the closed-form strategies; returns sample moments to compare with moments()."""
    rng = np.random.default_rng(seed)
    s = strategies(p, D)
    nD = int(round(D * n_agents))
    nH = n_agents - nD
    model_of = np.arange(nD) % p.M if nD > 0 else np.zeros(0, int)
    theta = rng.normal(0.0, p.tau_theta ** -0.5, n_draws)
    X = np.empty(n_draws)
    Ex2 = 0.0
    disp = 0.0
    for t in range(n_draws):
        e = rng.normal(0.0, p.tau_H ** -0.5, nH)
        eta = rng.normal(0.0, p.tau_A ** -0.5, p.M)
        xH = s['aH'] * (theta[t] + e)
        xA = s['aA'] * (theta[t] + eta[model_of]) if nD > 0 else np.zeros(0)
        x = np.concatenate([xH, xA])
        X[t] = x.mean()
        Ex2 += (x ** 2).mean()
        disp += ((x - X[t]) ** 2).mean()
    return dict(EX2=float(np.mean(X ** 2)), Ex2=Ex2 / n_draws, V_disp=disp / n_draws,
                EXt=float(np.mean(X * theta)), Abar_hat=float(np.sum(X * theta) / np.sum(theta ** 2)))