"""General equilibrium with heterogeneous households: relaxation on the interest-rate path.

Given a solved two-class scenario (paths P, model m, unknowns u), the HA block replaces the capitalists:
households of both skills (three discount types each) hold all claims. The algorithm:
  1. calibrate the discount types so that stationary HA savings at the reference prices equal the two-class
     wealth W_0 (so that the HA economy starts on the same balance sheet);
  2. iterate: (a) given r, solve the production side with r exogenous (4 capital Eulers, banded Newton) ->
     supply of claims W^s_t and wages/transfers/price index; (b) solve the HA block backward and simulate ->
     demand for claims W^d_t; (c) update r_{t+1} <- r_{t+1} - lam * log(W^d_t / W^s_t) (excess demand for
     assets lowers the equilibrium return), with a secant-type step adaptation;
  3. stop when max_t |log(W^d/W^s)| < tol.
Outputs: equilibrium paths, distributional statistics along the transition, and welfare by type.
"""
from __future__ import annotations

import time

import numpy as np
from scipy.optimize import brentq

from .households import HouseholdBlock
from .steady import bgp_interest_rate
from .transition import KINDS, TwoClassModel


def calibrate_households(p, W0, wL, wH, transfer=0.0, P=1.0, masses=(0.5, 0.3, 0.2), betas_low=(0.93, 0.97),
                         verbose=True, **hb_kw):
    """Find the highest discount factor so that stationary HA savings equal W0 at the reference prices."""
    r = bgp_interest_rate(p)
    bmax = (1 + p.g_Z) ** p.gamma_u / (1 + r) - 1e-4     # no-explosion bound

    def excess(bh):
        hb = HouseholdBlock(p, (*betas_low, bh), masses, **hb_kw)
        pols = hb.solve_stationary(r, (wL, wH), transfer, P)
        dists = hb.stationary_distribution(pols, r)
        A, X = hb.aggregate(pols, dists)
        return np.log(A / W0), hb, pols, dists

    lo, hi = betas_low[1] + 0.002, bmax
    f_lo, f_hi = excess(lo)[0], excess(hi)[0]
    if f_lo > 0:
        bh = lo
    elif f_hi < 0:
        bh = hi
    else:
        bh = brentq(lambda b: excess(b)[0], lo, hi, xtol=1e-6)
    ex, hb, pols, dists = excess(bh)
    st = hb.wealth_stats(dists, pols)
    if verbose:
        print(f"  HA calibration: beta_high = {bh:.5f}, savings/W0 = {np.exp(ex):.4f}, stats = {st}")
    return hb, pols, dists, dict(beta_high=bh, stats=st)


def ha_inputs_from_paths(P, p):
    """Detrended wage and transfer paths for the HA block (per efficiency unit), price index, r."""
    T = len(P["Y"]) - 1
    wL = P["wL"][:T + 1]
    wH = P["wH"][:T + 1]
    gf = P["gf"][:T + 1]
    mass_workers = float(np.sum(p.H))
    transfer = P["transfers"][:T + 1] / mass_workers / gf           # per worker, detrended
    Pidx = P["P_idx"][:T + 1] / P["P_idx"][0]
    r = P["r"][:T + 2]
    return wL, wH, transfer, Pidx, r, gf


def _ha_demand(hb, r_path, wL, wH, transfer, Pidx, gf, dists0, T, x_term=None, scales=None):
    if x_term is None:
        x_term = hb.solve_stationary(r_path[T], (wL[T], wH[T]), transfer[T], Pidx[T], scale=1.0 if scales is None else scales[T])
    pols = hb.solve_transition(r_path[:T + 1], wL, wH, transfer, Pidx, x_term, scales=scales)
    A, X, dist_hist = hb.simulate_transition(pols, dists0, r_path, scales=scales)
    return A * gf[:T], X, pols, dist_hist


def grid_scales(Ws, gf, W0, T):
    """Grid scale per period: detrended aggregate wealth relative to 2026, never below one, smoothed."""
    sc = np.ones(T + 2)
    w = np.maximum(Ws[:T] / gf[:T] / W0, 1.0)
    sc[:T] = w
    sc[T:] = w[-1]
    return sc


def solve_ha_equilibrium(p, m: TwoClassModel, u, P, lam=1.0, tol=1e-4, maxit=60, verbose=True, hb_kw=None,
                         masses=(0.5, 0.3, 0.2), betas_low=(0.93, 0.97), r_bounds=(-0.02, 0.8), jac_every=6,
                         step_max=0.2, checkpoint=None, r_init=None):
    """Quasi-Newton on the return path: the Jacobian of the household block's log asset demand with respect to
    the return path is computed by finite differences (T household solves, cheap); the supply side's response is
    not differentiated (it has the opposite sign, so using the demand Jacobian alone over-estimates the required
    move, which the damping lam corrects). Steps are bounded by r_bounds."""
    hb_kw = hb_kw or {}
    T = m.T
    wL, wH, transfer, Pidx, r, gf = ha_inputs_from_paths(P, p)
    W0 = (1.0 - p.theta_F) * P["W"][0]
    hb, pols0, dists0, hcal = calibrate_households(p, W0, wL[0], wH[0], transfer[0], Pidx[0], masses, betas_low, verbose, **hb_kw)
    U6 = u.reshape(T, 6)
    u4 = U6[:, :5].ravel().copy()          # fixed-r mode: 4 log stocks + log X_K
    r_path = r.copy() if r_init is None else np.asarray(r_init, float).copy()
    hist = []
    t0 = time.time()
    J = None
    for it in range(maxit):
        m.set_fixed_r(r_path)
        u4, ok, h = m.solve(u4, tol=1e-9, verbose=False, maxit=30)
        Pp = m.evaluate(u4, full=True)
        wL, wH, transfer, Pidx, _, gf = ha_inputs_from_paths(Pp, p)
        Ws = (1.0 - p.theta_F) * Pp["W"][:T]          # households hold the claims not held by the citizen fund
        scales = grid_scales((1.0 - p.theta_F) * Pp["W"], gf, W0, T)
        x_term = hb.solve_stationary(r_path[T], (wL[T], wH[T]), transfer[T], Pidx[T], scale=scales[T])
        Wd, X, pols, dist_hist = _ha_demand(hb, r_path, wL, wH, transfer, Pidx, gf, dists0, T, x_term, scales)
        gap = np.log(Wd / Ws)
        err = float(np.max(np.abs(gap)))
        hist.append(err)
        if verbose:
            print(f"  HA-GE it {it:2d}: max|log Wd/Ws| = {err:.3e}, production converged={ok}, ({time.time() - t0:.0f}s)", flush=True)
        if err < tol:
            break
        if J is None or it % jac_every == 0:
            # Jacobian of log Wd (T) w.r.t. r_{t+1} (t = 0..T-1) at the current prices
            h_fd = 1e-4
            J = np.zeros((T, T))
            for s in range(T):
                rp = r_path.copy()
                rp[s + 1] += h_fd
                Wd_p, _, _, _ = _ha_demand(hb, rp, wL, wH, transfer, Pidx, gf, dists0, T, x_term, scales)
                J[:, s] = (np.log(Wd_p) - np.log(Wd)) / h_fd
            if verbose:
                print(f"  HA-GE: household Jacobian computed ({time.time() - t0:.0f}s)", flush=True)
        # supply responds with the opposite sign; approximate its elasticity by a diagonal term
        Jtot = J + np.diag(np.full(T, 5.0))   # d log Ws / d r ~ -duration (~5 years): crude, corrected by damping
        try:
            step = np.linalg.solve(Jtot, gap)
        except np.linalg.LinAlgError:
            step = gap / 20.0
        step = np.clip(lam * step, -step_max, step_max)
        r_new = np.clip(r_path[1:T + 1] - step, r_bounds[0], r_bounds[1])
        r_path[1:T + 1] = r_new
        r_path[T + 1] = r_path[T]
        if checkpoint:
            np.save(checkpoint, r_path)
    m.set_free_r()
    stats = []
    for t in range(T):
        st = hb.wealth_stats(dist_hist[t], {k: pols[k][t] for k in pols})
        st["year"] = int(P["years"][t])
        stats.append(st)
    out = dict(r=r_path, Wd=Wd, Ws=Ws, gap=gap, hist=hist, stats=stats, hcal=hcal, X=X, paths=Pp, u4=u4,
               converged=err < tol)
    return out, hb, pols, dist_hist