"""Symbolic derivation of the linear-quadratic delegation game (Model A core).

Agents i in [0,1] choose an action x_i (log-deviation of next-period capital holdings
from the full-information steady state).  Quadratic payoff

    U(x_i, X, theta) = U_xx x_i^2/2 + U_xX x_i X + U_XX X^2/2 + U_xt x_i theta + U_Xt X theta (+ terms without x)

normalized by |U_xx|:  best response  x_i = E_i[ k0 * theta + alpha * X ],
    k0 = -U_xt/U_xx,  alpha = -U_xX/U_xx  (alpha<0: strategic substitutes, as in capital accumulation).

Information:  theta ~ N(0, 1/tau_t).  Non-delegators (mass 1-D) see s_i = theta + e_i, e_i ~ N(0, 1/tau_H), iid.
Delegators (mass D, split evenly across M foundation models) see s_m = theta + eta_m, eta_m ~ N(0, 1/tau_A) iid across m,
*common* to all users of model m.

This script derives (and prints as LaTeX-ready expressions):
  1. equilibrium linear strategies a_H, a_A and aggregate X = Abar*theta + (D a_A / M) sum_m eta_m;
  2. aggregate "model-noise" variance and the aggregate tracking error;
  3. welfare decomposition (Angeletos-Pavan 2007):  W = W* - (1/2)|U_xx|[ (1-alpha_soc) E(X-X*)^2 + E(x_i-X)^2 ] - fees
     where the social curvature is (U_xx+2U_xX+U_XX)/U_xx; under the CRS/pecuniary restriction U_xX+U_XX=0 it equals (1-alpha);
  4. private value of delegation for a measure-zero agent vs. the social marginal value dW/dD -> the delegation wedge.
All results are numerically cross-checked in tests/unit/test_lq_game.py.
"""
import sympy as sp

D, M, alpha, k0 = sp.symbols('D M alpha kappa_0', real=True)
tau_t, tau_H, tau_A = sp.symbols('tau_theta tau_H tau_A', positive=True)
bH = tau_H/(tau_t+tau_H)          # Bayesian weight, human signal
bA = tau_A/(tau_t+tau_A)          # Bayesian weight, AI signal
sig_t2, sig_H2, sig_A2 = 1/tau_t, 1/tau_H, 1/tau_A

# ---- 1. Equilibrium strategies -------------------------------------------------
# Non-delegator: x_i = k0 E_i theta + alpha E_i X ;  E_i eta_m = 0  ->  a_H = (k0 + alpha*Abar) bH
# Delegator m:   E_m X = Abar E_m theta + (D a_A/M) E_m eta_m ,  E_m theta = bA s_m , E_m eta_m = (1-bA) s_m
#                -> a_A = bA (k0+alpha Abar) + alpha (D/M)(1-bA) a_A
c = sp.symbols('c')                        # c = k0 + alpha*Abar
xi = 1/(1 - alpha*(D/M)*(1-bA))            # delegator damping factor (xi<1 if alpha<0)
aH = bH*c
aA = bA*c*xi
Abar = (1-D)*aH + D*aA
c_sol = sp.solve(sp.Eq(c, k0 + alpha*Abar), c)[0]
aH, aA, Abar = [sp.simplify(e.subs(c, c_sol)) for e in (aH, aA, Abar)]
kappa_full = k0/(1-alpha)                  # full-information equilibrium response X* = kappa_full*theta

# ---- 2. Aggregate law of motion ----------------------------------------------------
# X = Abar*theta + (D*aA/M) * sum_m eta_m   ->  Var(model noise) = D^2 aA^2 sig_A2 / M
V_noise = D**2*aA**2*sig_A2/M
V_track = (Abar-kappa_full)**2*sig_t2
V_agg = V_track + V_noise                  # E (X - X*)^2

# ---- 3. Dispersion  E (x_i - X)^2 --------------------------------------------------
# non-delegator: x_i - X = aH e_i + (aH-Abar) theta - (D aA/M) sum eta
# delegator m  : x_i - X = aA eta_m - (D aA/M) sum eta + (aA-Abar) theta
V_disp_H = aH**2*sig_H2 + (aH-Abar)**2*sig_t2 + V_noise
V_disp_A = (aA-Abar)**2*sig_t2 + sp.expand((aA*sig_A2)*(aA) - 2*aA*(D*aA/M)*sig_A2 + V_noise)  # Var(aA eta_m - (D aA/M) sum eta)
V_disp = (1-D)*V_disp_H + D*V_disp_A

# ---- 4. Welfare and the delegation wedge --------------------------------------------
# Social curvature ratio (U_xx+2U_xX+U_XX)/U_xx = 1 - 2*alpha + U_XX/U_xx ; with U_XX = -U_xX (CRS) -> 1 - alpha.
phi = sp.symbols('phi', positive=True)     # fee + disutility, in units of |U_xx|/2
alpha_soc = 1-alpha
W = -(sp.Rational(1,2))*(alpha_soc*V_agg + V_disp) - D*phi   # per-capita welfare relative to full information (|U_xx| normalised to 1)

# Private value of delegation for a measure-zero agent: target  x^ = k0 theta + alpha X.
# Loss = (1/2) Var(x^ | own signal).  For a non-delegator: Var(k0 theta + alpha X | s_i).
# For a delegator of model m: Var(k0 theta + alpha X | s_m).
# X = Abar theta + (D aA/M) sum_m eta_m.  Let g = k0 + alpha*Abar.  x^ = g theta + alpha (D aA/M) sum eta.
g = k0 + alpha*Abar
# Var(x^) = g^2 sig_t2 + alpha^2 (D aA/M)^2 M sig_A2
Vx = g**2*sig_t2 + alpha**2*(D*aA/M)**2*M*sig_A2
# human: Cov(x^, s_i) = g sig_t2 ; Var(s_i) = sig_t2 + sig_H2
PV_H_loss = sp.simplify(Vx - (g*sig_t2)**2/(sig_t2+sig_H2))
# AI model m: Cov(x^, s_m) = g sig_t2 + alpha (D aA/M) sig_A2 ; Var(s_m) = sig_t2 + sig_A2
PV_A_loss = sp.simplify(Vx - (g*sig_t2 + alpha*(D*aA/M)*sig_A2)**2/(sig_t2+sig_A2))
private_value = sp.simplify(sp.Rational(1,2)*(PV_H_loss - PV_A_loss))     # gain in units |U_xx|
social_mv = sp.simplify(sp.diff(W, D) + phi)                               # d(W + D phi)/dD : gross social marginal value
wedge = sp.simplify(private_value - social_mv)                              # >0 : competitive over-delegation

if __name__ == '__main__':
    sp.init_printing()
    print('a_H =', sp.simplify(aH)); print('a_A =', sp.simplify(aA)); print('Abar =', sp.simplify(Abar))
    print('V_noise =', V_noise)
    print('private value =', private_value)
    print('social marginal value (gross) =', social_mv)
    print('wedge =', wedge)
    # No-strategic-interaction limit alpha -> 0
    print('wedge | alpha=0 :', sp.simplify(wedge.subs(alpha, 0)))
    print('private | alpha=0 :', sp.simplify(private_value.subs(alpha, 0)))
    print('social  | alpha=0 :', sp.simplify(social_mv.subs(alpha, 0)))
    # Series of wedge in D around 0
    print('wedge series in D:', sp.series(wedge, D, 0, 2))