"""Model B — heterogeneous-agent stationary recursive competitive equilibrium with AI delegation and compute ownership.

Households (continuum): CRRA utility, idiosyncratic earnings e (Rouwenhorst), borrowing limit a_min, cash-on-hand state
x = (R + nu) a + w e, where nu is the household's idiosyncratic return term reflecting *decision quality*:
    self-managed (d=0):  nu ~ N(-m_H, sigma_rH^2)     (mean cost of mistakes + dispersion)
    delegated   (d=1):  nu ~ N(0,    sigma_rA^2)     at fee  phi0 + phi1 * a'  (goods, paid when saving)
Mean mistakes are a resource cost; dispersion is zero-sum redistribution.  Delegation d is chosen each period together
with savings a' (value-function iteration with a discrete choice, numba kernels).
Firms: task-based technology of Model A (age.economy.tasks) with compute stock K_Q.  Compute supply: K_Q = KQ0 (pQ/uQ)^eps_Q
(eps_Q = 0: fixed stock).  Ownership of compute rents: broad (added to every household's return) or concentrated (paid
to the wealthiest mass n_K as a proportional-to-wealth bonus return).
Equilibrium: r such that household savings A(r) = K(r) + pK_Q K_Q (capital market), compute market (A.4), L from the
earnings distribution.  Solved by bisection on r with warm starts.
"""
from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
from numba import njit
from scipy.optimize import brentq

from age.economy.tasks import TaskTech


# --------------------------------------------------------------------------- parameters
@dataclass
class HAParams:
    beta: float = 0.96
    gamma: float = 2.0
    alpha: float = 0.33
    delta: float = 0.06
    rho_e: float = 0.96
    sigma_e: float = 0.20
    n_e: int = 7
    a_min: float = 0.0
    a_max: float = 400.0
    n_a: int = 200
    # decision-quality returns
    m_H: float = 0.01          # mean cost of self-managed mistakes (per unit of wealth per year)
    sigma_rH: float = 0.04     # dispersion of self-managed idiosyncratic return
    sigma_rA: float = 0.01     # dispersion under delegation
    phi0: float = 0.02         # fixed fee (goods per year)
    phi1: float = 0.002        # proportional fee (per unit saved)
    n_gh: int = 5              # Gauss-Hermite nodes for the return term
    delegation_on: bool = True
    # technology / compute
    tasks: TaskTech = field(default_factory=lambda: TaskTech(alpha=0.33, iota_bar=0.5, psi0=1.0, psi1=3.0))
    KQ0: float = 0.0           # compute stock at pQ = uQ
    eps_Q: float = 0.0         # compute supply elasticity
    pK_Q: float = 1.0
    delta_Q: float = 0.20
    qbar: float = 0.002        # compute per delegated household-year
    ownership: str = 'broad'   # 'broad' | 'concentrated'
    n_K: float = 0.10          # mass of owners if concentrated (top of the wealth distribution)
    # numerics
    tol_v: float = 1e-6
    max_iter: int = 2000
    howard: int = 20


# --------------------------------------------------------------------------- discretisations
def rouwenhorst(n, rho, sigma):
    """Rouwenhorst (1995) discretisation of an AR(1) with unconditional s.d. sigma/sqrt(1-rho^2)."""
    p = (1 + rho) / 2
    P = np.array([[p, 1 - p], [1 - p, p]])
    for i in range(2, n):
        Pn = np.zeros((i + 1, i + 1))
        Pn[:i, :i] += p * P
        Pn[:i, 1:] += (1 - p) * P
        Pn[1:, :i] += (1 - p) * P
        Pn[1:, 1:] += p * P
        Pn[1:i] /= 2
        P = Pn
    sd = sigma / np.sqrt(1 - rho ** 2)
    grid = np.linspace(-sd * np.sqrt(n - 1), sd * np.sqrt(n - 1), n)
    return grid, P


def gauss_hermite(n):
    x, w = np.polynomial.hermite_e.hermegauss(n)
    return x, w / w.sum()


def make_grid(a_min, a_max, n):
    g = np.linspace(0, 1, n) ** 3
    return a_min + (a_max - a_min) * g


# --------------------------------------------------------------------------- numba kernels
@njit(cache=True)
def _interp(xg, yg, x):
    n = xg.shape[0]
    if x <= xg[0]:
        return yg[0] + (yg[1] - yg[0]) * (x - xg[0]) / (xg[1] - xg[0])
    if x >= xg[n - 1]:
        return yg[n - 1] + (yg[n - 1] - yg[n - 2]) * (x - xg[n - 1]) / (xg[n - 1] - xg[n - 2])
    lo, hi = 0, n - 1
    while hi - lo > 1:
        mid = (lo + hi) // 2
        if xg[mid] <= x:
            lo = mid
        else:
            hi = mid
    t = (x - xg[lo]) / (xg[hi] - xg[lo])
    return yg[lo] + t * (yg[hi] - yg[lo])


@njit(cache=True)
def _expected_value(V, xgrid, egrid, P, agrid, w, R, nu_nodes, nu_w, ie):
    """EV[a'] for current e-index ie: sum over e', nu of V(x' = (R+nu) a' + w e', e')."""
    na = agrid.shape[0]
    ne = egrid.shape[0]
    EV = np.zeros(na)
    for ia in range(na):
        s = 0.0
        for je in range(ne):
            pe = P[ie, je]
            if pe == 0.0:
                continue
            for k in range(nu_nodes.shape[0]):
                xp = (R + nu_nodes[k]) * agrid[ia] + w * egrid[je]
                s += pe * nu_w[k] * _interp(xgrid, V[:, je], xp)
        EV[ia] = s
    return EV


@njit(cache=True)
def _bellman(V, xgrid, egrid, P, agrid, w, R, gamma, beta, nuH, wH, nuA, wA, phi0, phi1, deleg_on):
    nx = xgrid.shape[0]
    ne = egrid.shape[0]
    na = agrid.shape[0]
    Vn = np.empty((nx, ne))
    pol_a = np.empty((nx, ne))
    pol_d = np.zeros((nx, ne), dtype=np.int64)
    pol_c = np.empty((nx, ne))
    for ie in range(ne):
        EV0 = _expected_value(V, xgrid, egrid, P, agrid, w, R, nuH, wH, ie)
        if deleg_on:
            EV1 = _expected_value(V, xgrid, egrid, P, agrid, w, R, nuA, wA, ie)
        else:
            EV1 = EV0
        for ix in range(nx):
            x = xgrid[ix]
            best = -1e300
            ba = agrid[0]
            bd = 0
            bc = 0.0
            # d = 0
            for ia in range(na):
                c = x - agrid[ia]
                if c <= 1e-12:
                    break
                u = (c ** (1.0 - gamma) - 1.0) / (1.0 - gamma) if gamma != 1.0 else np.log(c)
                v = u + beta * EV0[ia]
                if v > best:
                    best = v
                    ba = agrid[ia]
                    bd = 0
                    bc = c
            if deleg_on:
                for ia in range(na):
                    c = x - agrid[ia] - phi0 - phi1 * agrid[ia]
                    if c <= 1e-12:
                        break
                    u = (c ** (1.0 - gamma) - 1.0) / (1.0 - gamma) if gamma != 1.0 else np.log(c)
                    v = u + beta * EV1[ia]
                    if v > best:
                        best = v
                        ba = agrid[ia]
                        bd = 1
                        bc = c
            Vn[ix, ie] = best
            pol_a[ix, ie] = ba
            pol_d[ix, ie] = bd
            pol_c[ix, ie] = bc
    return Vn, pol_a, pol_d, pol_c


@njit(cache=True)
def _howard(V, xgrid, egrid, P, agrid, w, R, gamma, beta, nuH, wH, nuA, wA, pol_a, pol_d, pol_c, n_steps):
    nx = xgrid.shape[0]
    ne = egrid.shape[0]
    for _ in range(n_steps):
        Vn = np.empty((nx, ne))
        for ie in range(ne):
            EV0 = _expected_value(V, xgrid, egrid, P, agrid, w, R, nuH, wH, ie)
            EV1 = _expected_value(V, xgrid, egrid, P, agrid, w, R, nuA, wA, ie)
            for ix in range(nx):
                c = pol_c[ix, ie]
                u = (c ** (1.0 - gamma) - 1.0) / (1.0 - gamma) if gamma != 1.0 else np.log(c)
                ev = _interp(agrid, EV1 if pol_d[ix, ie] == 1 else EV0, pol_a[ix, ie])
                Vn[ix, ie] = u + beta * ev
        V = Vn
    return V


@njit(cache=True)
def _stationary(xgrid, egrid, P, pol_a, pol_d, w, R, nuH, wH, nuA, wA, n_iter, tol):
    """Non-stochastic simulation (Young 2010) of the distribution over (x, e)."""
    nx = xgrid.shape[0]
    ne = egrid.shape[0]
    lam = np.ones((nx, ne)) / (nx * ne)
    for it in range(n_iter):
        new = np.zeros((nx, ne))
        for ie in range(ne):
            for ix in range(nx):
                m = lam[ix, ie]
                if m == 0.0:
                    continue
                ap = pol_a[ix, ie]
                if pol_d[ix, ie] == 1:
                    nodes = nuA
                    wts = wA
                else:
                    nodes = nuH
                    wts = wH
                for je in range(ne):
                    pe = P[ie, je]
                    if pe == 0.0:
                        continue
                    for k in range(nodes.shape[0]):
                        xp = (R + nodes[k]) * ap + w * egrid[je]
                        # lottery onto grid
                        if xp <= xgrid[0]:
                            new[0, je] += m * pe * wts[k]
                        elif xp >= xgrid[nx - 1]:
                            new[nx - 1, je] += m * pe * wts[k]
                        else:
                            lo, hi = 0, nx - 1
                            while hi - lo > 1:
                                mid = (lo + hi) // 2
                                if xgrid[mid] <= xp:
                                    lo = mid
                                else:
                                    hi = mid
                            t = (xp - xgrid[lo]) / (xgrid[hi] - xgrid[lo])
                            new[lo, je] += m * pe * wts[k] * (1.0 - t)
                            new[hi, je] += m * pe * wts[k] * t
        diff = np.max(np.abs(new - lam))
        lam = new
        if diff < tol:
            break
    return lam


# --------------------------------------------------------------------------- household block
class HouseholdBlock:
    def __init__(self, p: HAParams):
        self.p = p
        self.egrid_log, self.P = rouwenhorst(p.n_e, p.rho_e, p.sigma_e)
        self.egrid = np.exp(self.egrid_log)
        # normalise mean labour efficiency to 1
        pi = self.stationary_e()
        self.egrid = self.egrid / (pi @ self.egrid)
        self.agrid = make_grid(p.a_min, p.a_max, p.n_a)
        self.xgrid = make_grid(p.a_min + 1e-3, p.a_max * 1.2, p.n_a)
        gh_x, gh_w = gauss_hermite(p.n_gh)
        self.nuH = -p.m_H + p.sigma_rH * gh_x
        self.wH = gh_w
        self.nuA = p.sigma_rA * gh_x
        self.wA = gh_w
        self.V = None

    def stationary_e(self):
        w, v = np.linalg.eig(self.P.T)
        pi = np.real(v[:, np.argmin(np.abs(w - 1))])
        return pi / pi.sum()

    def solve(self, R: float, w: float, bonus: np.ndarray | None = None):
        """Value function iteration with Howard steps. `bonus` (per x-grid point) adds a wealth-proportional return
        bonus for concentrated compute ownership: R_eff(x) = R + bonus[x]; implemented by shifting nu nodes per x is
        expensive, so bonus is applied as an extra income r_bonus*a' through the return R for the top group only
        via a separate solve (see Economy)."""
        p = self.p
        if self.V is None:
            c0 = np.maximum(self.xgrid[:, None] * 0.05, 1e-3) * np.ones((1, p.n_e))
            self.V = ((c0 ** (1 - p.gamma) - 1) / (1 - p.gamma)) / (1 - p.beta)
        V = self.V
        for it in range(p.max_iter):
            Vn, pol_a, pol_d, pol_c = _bellman(V, self.xgrid, self.egrid, self.P, self.agrid, w, R, p.gamma, p.beta,
                                               self.nuH, self.wH, self.nuA, self.wA, p.phi0, p.phi1, p.delegation_on)
            diff = np.max(np.abs(Vn - V))
            V = _howard(Vn, self.xgrid, self.egrid, self.P, self.agrid, w, R, p.gamma, p.beta,
                        self.nuH, self.wH, self.nuA, self.wA, pol_a, pol_d, pol_c, p.howard)
            if diff < p.tol_v:
                break
        self.V = V
        self.pol_a, self.pol_d, self.pol_c = pol_a, pol_d, pol_c
        self.n_iter = it + 1
        return V

    def distribution(self, R, w):
        self.lam = _stationary(self.xgrid, self.egrid, self.P, self.pol_a, self.pol_d, w, R,
                               self.nuH, self.wH, self.nuA, self.wA, 20000, 1e-11)
        return self.lam

    def aggregates(self):
        lam = self.lam
        A = float(np.sum(lam * self.pol_a))
        D = float(np.sum(lam * self.pol_d))
        C = float(np.sum(lam * self.pol_c))
        L = float(np.sum(lam.sum(axis=0) * self.egrid))
        fees = float(np.sum(lam * self.pol_d * (self.p.phi0 + self.p.phi1 * self.pol_a)))
        mistakes = float(np.sum(lam * (1 - self.pol_d) * self.p.m_H * self.pol_a))
        return dict(A=A, D=D, C=C, L=L, fees=fees, mistakes=mistakes)

    # ----- distributional statistics on end-of-period wealth a'
    def wealth_stats(self, extra_mass: float = 0.0, extra_wealth: float = 0.0):
        lam = self.lam
        a = self.pol_a.ravel()
        m = lam.ravel()
        d_all = self.pol_d.ravel()
        if extra_mass > 0:
            a = np.append(a, extra_wealth)
            m = np.append(m * (1 - extra_mass), extra_mass)
            d_all = np.append(d_all, 0)
        order = np.argsort(a)
        a, m = a[order], m[order]
        cm = np.cumsum(m)
        cw = np.cumsum(a * m)
        W = cw[-1]
        lorenz_x = cm / cm[-1]
        lorenz_y = cw / W
        gini = 1 - 2 * np.trapezoid(lorenz_y, lorenz_x)
        lx = np.concatenate([[0.0], lorenz_x])
        ly = np.concatenate([[0.0], lorenz_y])
        def top(share):
            return 1 - float(np.interp(1 - share, lx, ly))
        # delegation by wealth decile, households only (the owner class is outside the household block)
        a_h = self.pol_a.ravel(); m_h = lam.ravel(); d_h = self.pol_d.ravel()
        oh = np.argsort(a_h); cm_h = np.cumsum(m_h[oh]) / m_h.sum()
        dec = np.minimum((cm_h * 10).astype(int), 9)
        deleg_by_decile = np.array([np.sum(m_h[oh][dec == k] * d_h[oh][dec == k]) / max(np.sum(m_h[oh][dec == k]), 1e-12) for k in range(10)])
        return dict(gini=float(gini), top1=float(top(0.01)), top10=float(top(0.10)), bottom50=float(np.interp(0.5, lx, ly)),
                    lorenz_x=lorenz_x, lorenz_y=lorenz_y, deleg_by_decile=deleg_by_decile, mean_wealth=float(W))


# --------------------------------------------------------------------------- general equilibrium
class EconomyB:
    def __init__(self, p: HAParams):
        self.p = p
        self.hh = HouseholdBlock(p)
        self.uQ = (1 / p.beta - 1 + p.delta_Q) * p.pK_Q      # user cost of compute at the discount rate

    def compute_supply(self, pQ):
        p = self.p
        if p.KQ0 <= 0:
            return 0.0
        if p.eps_Q == 0:
            return p.KQ0
        return p.KQ0 * max(pQ / self.uQ, 1e-9) ** p.eps_Q

    def firm_side(self, r, L, D):
        """Given r, L and delegation share, clear the compute market and return prices and quantities."""
        p = self.p
        RK = r + p.delta
        def solve_for_KQ(KQ):
            QP = max(KQ - p.qbar * D, 1e-12) if KQ > 0 else 0.0
            T = p.tasks.bundle(QP, L)
            K = T * (p.alpha / RK) ** (1 / (1 - p.alpha))
            pr = p.tasks.prices(0.0, K, QP, L)
            return K, QP, pr
        if p.KQ0 <= 0:
            K, QP, pr = solve_for_KQ(0.0)
            KQ = 0.0
        elif p.eps_Q == 0:
            KQ = p.KQ0
            K, QP, pr = solve_for_KQ(KQ)
        else:
            f = lambda KQ: self.compute_supply(solve_for_KQ(KQ)[2]['pQ']) - KQ
            KQ = brentq(f, p.qbar * D + 1e-6, 50.0)
            K, QP, pr = solve_for_KQ(KQ)
        return dict(K=K, KQ=KQ, QP=QP, **pr)

    def excess_assets(self, r, D_guess=0.0, L=1.0):
        """Household savings minus K + pK_Q*KQ at interest rate r (inner fixed point on D for the compute market)."""
        p = self.p
        D = D_guess
        for _ in range(5):
            fs = self.firm_side(r, L, D)
            pQ, KQ, K = fs['pQ'], fs['KQ'], fs['K']
            rents = (pQ - self.uQ) * KQ if KQ > 0 else 0.0     # scarcity rents over user cost
            assets_total = K + p.pK_Q * KQ
            # fund return: capital earns r (= RK - delta); compute earns uQ (user cost) + rents
            if p.ownership == 'broad' or rents <= 0:
                R = 1 + r + rents / max(assets_total, 1e-9)
            else:
                R = 1 + r
            self.hh.solve(R, fs['w'])
            self.hh.distribution(R, fs['w'])
            agg = self.hh.aggregates()
            if abs(agg['D'] - D) < 1e-3:
                D = agg['D']
                break
            D = agg['D']
        self.last = dict(r=r, R=R, D=D, rents=rents, agg=agg, fs=fs, assets_total=assets_total)
        return agg['A'] - assets_total

    def solve(self, r_lo=-0.02, r_hi=None, tol=1e-4, verbose=False):
        p = self.p
        L = 1.0   # efficiency units normalised
        r_hi = r_hi if r_hi is not None else 1 / p.beta - 1 - 2e-4
        f = lambda r: self.excess_assets(r, L=L)
        # coarse scan for a sign change (excess assets is increasing in r), then Brent
        grid = np.linspace(r_lo, r_hi, 8)
        vals = [f(r) for r in grid]
        idx = next((i for i in range(len(grid) - 1) if vals[i] < 0 <= vals[i + 1]), None)
        if idx is None:
            raise RuntimeError(f'no sign change in excess assets on [{r_lo}, {r_hi}]: {np.round(vals, 3)}')
        r = brentq(f, grid[idx], grid[idx + 1], xtol=tol)
        f(r)
        last = self.last
        agg, fs = last['agg'], last['fs']
        ws = self.hh.wealth_stats()
        if p.ownership == 'concentrated' and last['rents'] > 0:
            # owner class (mass n_K) holds the right to the scarcity rents, valued as a perpetuity at r; its members are
            # appended to the wealth distribution as a point mass (two-class economy of Proposition 6)
            ws = self.hh.wealth_stats(extra_mass=p.n_K, extra_wealth=last['rents'] / p.n_K / max(r, 1e-4))
        Y = fs['Y']
        shares = dict(labor=fs['w'] * L / Y, capital=fs['RK'] * fs['K'] / Y, compute=fs['pQ'] * fs['QP'] / Y)
        return dict(r=r, R=last['R'], w=fs['w'], K=fs['K'], KQ=fs['KQ'], QP=fs['QP'], pQ=fs['pQ'], iota=fs['iota'],
                    Y=Y, C=agg['C'], D=agg['D'], fees=agg['fees'], mistakes=agg['mistakes'], rents=last['rents'],
                    shares=shares, gini=ws['gini'], top1=ws['top1'], top10=ws['top10'], bottom50=ws['bottom50'],
                    deleg_by_decile=ws['deleg_by_decile'], lorenz=(ws['lorenz_x'], ws['lorenz_y']),
                    mean_wealth=ws['mean_wealth'], n_iter_v=self.hh.n_iter)