"""Heterogeneous-agent household block (Bewley–Huggett–Aiyagari with skill and discount-factor types).

State: skill s in {L, H} (permanent), discount type j (permanent), productivity e (Markov, Rouwenhorst),
assets a >= 0. Each period the household receives labor income w_{s,t} gf_t e + transfer_t, chooses
expenditure x (goods + housing, Cobb-Douglas inside the period, so the period utility is
(x/P_t)^{1-gamma}/(1-gamma) with P_t the housing-inclusive price index) and end-of-period savings a'
which earn r_{t+1}. Budget: x + a' = a + y. Euler: x^{-gamma} P^{gamma-1} = beta (1+r') E[x'^{-gamma} P'^{gamma-1}].

All quantities are detrended by gf_t = (1+g_Z)^t (hatted variables) so that the grid is stationary in the
frozen-technology reference. Solution by the endogenous grid method (Carroll 2006); distribution by
non-stochastic simulation on the asset grid (Young 2010). numba-accelerated.
"""
from __future__ import annotations

import numpy as np
from numba import njit


# ------------------------------------------------------------------ income process
def rouwenhorst(n, rho, sigma):
    """Rouwenhorst discretization of an AR(1) in logs with persistence rho and innovation s.d. sigma."""
    q = (1 + rho) / 2
    P = np.array([[q, 1 - q], [1 - q, q]])
    for i in range(2, n):
        Pn = np.zeros((i + 1, i + 1))
        Pn[:i, :i] += q * P
        Pn[:i, 1:] += (1 - q) * P
        Pn[1:, :i] += (1 - q) * P
        Pn[1:, 1:] += q * P
        Pn[1:-1] /= 2
        P = Pn
    psi = sigma * np.sqrt((n - 1) / (1 - rho ** 2))
    z = np.linspace(-psi, psi, n)
    pi = np.linalg.matrix_power(P.T, 2000)[:, 0]
    e = np.exp(z)
    e = e / np.sum(pi * e)       # mean one
    return e, P, pi


def asset_grid(a_max, n=300, curv=4.0):
    u = np.linspace(0, 1, n)
    return a_max * u ** curv


# ------------------------------------------------------------------ EGM step (numba)
@njit(cache=True)
def _egm_step(agrid, e, Pe, y_next_e, rnext, beta_eff, gamma, Pnow, Pnext, gnext, c_next_pol, a_grid_next):
    """One backward EGM step for a given (skill, type). Inputs are detrended:
    c_next_pol[e_idx, a_idx] = next-period expenditure policy on a_grid_next (assets at start of next period,
    detrended by gf_{t+1}); returns this period's policy x(a) on agrid and savings a'(a).
    beta_eff = beta * (1+g)^{-gamma} handles detrending; gnext = 1+g."""
    ne, na = c_next_pol.shape
    x_pol = np.empty((ne, na))
    ap_pol = np.empty((ne, na))
    # marginal utility next period at each (e', a'): x'^{-gamma} P'^{gamma-1}
    for ie in range(ne):
        # for each savings level a' on agrid (detrended by gf_t), next-period assets = (1+r') a' / gnext
        rhs = np.empty(na)
        for ia in range(na):
            an = (1.0 + rnext) * agrid[ia] / gnext
            s = 0.0
            for je in range(ne):
                # interpolate next-period policy at an
                # find bracket
                if an <= a_grid_next[0]:
                    xn = c_next_pol[je, 0]
                elif an >= a_grid_next[na - 1]:
                    # extrapolate linearly
                    xn = c_next_pol[je, na - 1] + (c_next_pol[je, na - 1] - c_next_pol[je, na - 2]) / (a_grid_next[na - 1] - a_grid_next[na - 2]) * (an - a_grid_next[na - 1])
                else:
                    lo, hi = 0, na - 1
                    while hi - lo > 1:
                        mid = (lo + hi) // 2
                        if a_grid_next[mid] <= an:
                            lo = mid
                        else:
                            hi = mid
                    wgt = (an - a_grid_next[lo]) / (a_grid_next[hi] - a_grid_next[lo])
                    xn = (1 - wgt) * c_next_pol[je, lo] + wgt * c_next_pol[je, hi]
                if xn < 1e-12:
                    xn = 1e-12
                s += Pe[ie, je] * xn ** (-gamma) * Pnext ** (gamma - 1.0)
            rhs[ia] = beta_eff * (1.0 + rnext) * s
        # endogenous grid: x = (rhs / P^{gamma-1})^{-1/gamma}; cash on hand m = x + a'
        m_end = np.empty(na)
        x_end = np.empty(na)
        for ia in range(na):
            x_end[ia] = (rhs[ia] / Pnow ** (gamma - 1.0)) ** (-1.0 / gamma)
            m_end[ia] = x_end[ia] + agrid[ia]
        # map back to the exogenous asset grid: m = a + y_e
        for ia in range(na):
            m = agrid[ia] + y_next_e[ie]
            if m <= m_end[0]:
                # borrowing constraint binds: a' = 0, x = m
                x_pol[ie, ia] = m
                ap_pol[ie, ia] = 0.0
            else:
                lo, hi = 0, na - 1
                if m >= m_end[na - 1]:
                    slope = (x_end[na - 1] - x_end[na - 2]) / (m_end[na - 1] - m_end[na - 2])
                    xx = x_end[na - 1] + slope * (m - m_end[na - 1])
                else:
                    while hi - lo > 1:
                        mid = (lo + hi) // 2
                        if m_end[mid] <= m:
                            lo = mid
                        else:
                            hi = mid
                    wgt = (m - m_end[lo]) / (m_end[hi] - m_end[lo])
                    xx = (1 - wgt) * x_end[lo] + wgt * x_end[hi]
                if xx > m:
                    xx = m
                if xx < 1e-12:
                    xx = 1e-12
                x_pol[ie, ia] = xx
                ap_pol[ie, ia] = m - xx
    return x_pol, ap_pol


@njit(cache=True)
def _forward(dist, ap_pol, agrid, Pe, rnext, gnext):
    """Push the distribution over (e, a) forward one period. dist[e, a] at start of t -> at start of t+1.
    Savings a' earn r' and are deflated by gnext before being placed on the grid."""
    ne, na = dist.shape
    new = np.zeros((ne, na))
    for ie in range(ne):
        for ia in range(na):
            mass = dist[ie, ia]
            if mass == 0.0:
                continue
            an = (1.0 + rnext) * ap_pol[ie, ia] / gnext
            # lottery on the grid
            if an <= agrid[0]:
                lo, w = 0, 1.0
            elif an >= agrid[na - 1]:
                lo, w = na - 2, 0.0
            else:
                lo, hi = 0, na - 1
                while hi - lo > 1:
                    mid = (lo + hi) // 2
                    if agrid[mid] <= an:
                        lo = mid
                    else:
                        hi = mid
                w = 1.0 - (an - agrid[lo]) / (agrid[lo + 1] - agrid[lo])
            for je in range(ne):
                pm = mass * Pe[ie, je]
                new[je, lo] += pm * w
                new[je, lo + 1] += pm * (1.0 - w)
    return new


# ------------------------------------------------------------------ household block
class HouseholdBlock:
    def __init__(self, p, betas, beta_masses, rho_e=0.95, sigma_e=0.20, ne=5, na=300, a_max=400.0):
        self.p = p
        self.betas = np.asarray(betas, float)          # discount types
        self.beta_masses = np.asarray(beta_masses, float)   # masses within each skill (sum to one)
        self.e, self.Pe, self.pi_e = rouwenhorst(ne, rho_e, sigma_e)
        self.agrid = asset_grid(a_max, na)
        self.ne, self.na = ne, na
        self.H = np.asarray(p.H, float)
        self.gamma = p.gamma_u
        self.g = p.g_Z

    # ----- stationary policy for constant prices (reference BGP)
    def solve_stationary(self, r, w_by_skill, transfer, P, tol=1e-10, maxit=5000, scale=1.0):
        """Stationary policies on the grid scaled by `scale` (grid points = scale * agrid)."""
        pols = {}
        grid = self.agrid * scale
        for s in range(2):
            y = w_by_skill[s] * self.e + transfer
            for j, b in enumerate(self.betas):
                beta_eff = b * (1 + self.g) ** (-self.gamma)
                x = np.tile(y[:, None] + 0.05 * grid[None, :], (1, 1))   # initial policy: consume income + 5% assets
                for it in range(maxit):
                    x_new, ap = _egm_step(grid, self.e, self.Pe, y, r, beta_eff, self.gamma, P, P, 1 + self.g, x, grid)
                    err = np.max(np.abs(x_new - x))
                    x = x_new
                    if err < tol:
                        break
                pols[(s, j)] = (x, ap)
        return pols

    def stationary_distribution(self, pols, r, tol=1e-12, maxit=20000, scale=1.0):
        dists = {}
        grid = self.agrid * scale
        for key, (x, ap) in pols.items():
            d = np.zeros((self.ne, self.na))
            d[:, 0] = self.pi_e
            for it in range(maxit):
                dn = _forward(d, ap, grid, self.Pe, r, 1 + self.g)
                err = np.max(np.abs(dn - d))
                d = dn
                if err < tol:
                    break
            dists[key] = d
        return dists

    def aggregate(self, pols, dists):
        """Aggregate end-of-period savings and expenditure (detrended), weighting types by masses."""
        A, X = 0.0, 0.0
        for s in range(2):
            for j in range(len(self.betas)):
                m = self.H[s] * self.beta_masses[j]
                x, ap = pols[(s, j)]
                d = dists[(s, j)]
                A += m * np.sum(d * ap)
                X += m * np.sum(d * x)
        return A, X

    # ----- transition given price paths (arrays of length T+1 for t = 0..T; r[t+1] is the return on a'_t)
    def solve_transition(self, r, wL, wH, transfer, P, x_terminal, scales=None):
        """Backward induction from terminal policies x_terminal[(s,j)] (stationary at T on grid*scales[T]);
        returns policies for t = 0..T-1 on the grid scaled by scales[t] (grid follows aggregate wealth so that
        the upper bound never binds). Inputs are detrended wages per efficiency unit; r[t] is the return
        between t-1 and t."""
        T = len(r) - 1
        scales = np.ones(T + 2) if scales is None else np.asarray(scales, float)
        pols = {key: [None] * T for key in x_terminal}
        for s in range(2):
            w_path = wL if s == 0 else wH
            for j, b in enumerate(self.betas):
                beta_eff = b * (1 + self.g) ** (-self.gamma)
                x_next = x_terminal[(s, j)][0]
                for t in range(T - 1, -1, -1):
                    y = w_path[t] * self.e + transfer[t]
                    x, ap = _egm_step(self.agrid * scales[t], self.e, self.Pe, y, r[t + 1], beta_eff, self.gamma, P[t], P[t + 1], 1 + self.g, x_next, self.agrid * scales[t + 1])
                    pols[(s, j)][t] = (x, ap)
                    x_next = x
        return pols

    def simulate_transition(self, pols, dists0, r, scales=None):
        """Forward simulation; returns aggregate savings A_t (detrended, t=0..T-1) and the distributions."""
        T = len(next(iter(pols.values())))
        scales = np.ones(T + 2) if scales is None else np.asarray(scales, float)
        A = np.zeros(T)
        X = np.zeros(T)
        dists = {k: d.copy() for k, d in dists0.items()}
        hist = []
        for t in range(T):
            snap = {}
            for s in range(2):
                for j in range(len(self.betas)):
                    key = (s, j)
                    m = self.H[s] * self.beta_masses[j]
                    x, ap = pols[key][t]
                    d = dists[key]
                    A[t] += m * np.sum(d * ap)
                    X[t] += m * np.sum(d * x)
                    snap[key] = d
                    dists[key] = _forward(d, ap, self.agrid * scales[t + 1], self.Pe, r[t + 1], 1 + self.g)
            hist.append(snap)
        return A, X, hist

    # ----- distribution statistics
    def wealth_stats(self, dists, pols=None):
        """Gini and top shares of end-of-period savings (or start-of-period assets if pols is None)."""
        vals, wts = [], []
        for s in range(2):
            for j in range(len(self.betas)):
                key = (s, j)
                m = self.H[s] * self.beta_masses[j]
                d = dists[key]
                a = pols[key][1] if pols is not None else np.tile(self.agrid, (self.ne, 1))
                vals.append(a.ravel())
                wts.append((m * d).ravel())
        v = np.concatenate(vals)
        w = np.concatenate(wts)
        w = w / w.sum()
        o = np.argsort(v)
        v, w = v[o], w[o]
        cw = np.cumsum(w)
        cv = np.cumsum(v * w)
        tot = cv[-1]
        gini = 1 - 2 * np.sum(w * (cv - 0.5 * v * w)) / tot if tot > 0 else np.nan
        def top(share):
            idx = np.searchsorted(cw, 1 - share)
            return 1 - (cv[idx - 1] / tot if idx > 0 else 0.0)
        return dict(gini=float(gini), top1=float(top(0.01)), top10=float(top(0.10)), bottom50=float(cv[np.searchsorted(cw, 0.5)] / tot), mean=float(tot))