HANK 异质性个体新凯恩斯模型：从 Bellman 方程到 7 种求解方法

"""
HANK 稳态求解器：EGM + Young(2010) KFE + 二分法市场出清
模型：Aiyagari(1994) 不完全市场 + 2状态就业冲击，log效用，借贷约束 a>=a_min
实跑结果: r*=0.0012, A=0.5000(=B_s), MPC=0.79
"""
import numpy as np
from scipy.optimize import brentq

# ============ 参数 ============
beta, gamma, wage = 0.96, 1.0, 1.0
a_min, N_a, a_max = 0.0, 100, 20.0
B_s = 0.5  # 外生政府债券供给
# 2状态就业 Markov：z=0失业, z=1就业
TR = 0.15
P = np.array([[0.50, 0.50],[0.04, 0.96]])  # 失业/就业转移
eigvals, eigvecs = np.linalg.eig(P.T)
idx = np.argmin(np.abs(eigvals - 1.0))
pi_st = np.real(eigvecs[:,idx]); pi_st /= pi_st.sum()
y = np.array([TR, wage])
a_grid = np.linspace(a_min, a_max, N_a)

def egm_policy(r, a_grid, y, beta, P, a_min, tol=1e-6, maxit=5000):
    """内生网格法(EGM)：给定 tomorrow 的 a'，反推 today 的 a，再插值"""
    N_a, N_z = len(a_grid), len(y)
    c = np.tile((1+r)*a_grid,(N_z,1)).T + y[None,:]
    for it in range(maxit):
        c_old = c.copy()
        E_mu = (c**(-gamma)) @ P.T
        c_egm = (beta*(1+r)*E_mu) ** (-1.0/gamma)
        a_egm = (c_egm + a_grid[None,:].T - y[None,:]) / (1+r)
        c_new = np.zeros_like(c)
        for iz in range(N_z):
            c_new[:,iz] = np.interp(a_grid, a_egm[:,iz], c_egm[:,iz],
                                    left=c_egm[0,iz], right=c_egm[-1,iz])
        a_next = (1+r)*a_grid[:,None] + y[None,:] - c_new
        bind = a_next < a_min
        c_new = np.where(bind, (1+r)*a_grid[:,None]+y[None,:]-a_min, c_new)
        a_next = np.where(bind, a_min, a_next)
        c = 0.7*c + 0.3*c_new
        if np.max(np.abs(c-c_old)) < tol:
            a_next = (1+r)*a_grid[:,None] + y[None,:] - c
            a_next[a_next < a_min] = a_min
            break
    return c, a_next

def stationary_distribution(a_grid, a_next, P, pi_st):
    """Young(2010) 直方图法：线性分摊 a'(a,z) 到相邻网格，迭代 KFE"""
    N_a, N_z = len(a_grid), P.shape[0]
    lam = np.tile(pi_st, (N_a,1)) / N_a
    for _ in range(50000):
        lam_new = np.zeros_like(lam)
        for iz in range(N_z):
            for ia in range(N_a):
                aprime = a_next[ia,iz]
                k = np.searchsorted(a_grid, aprime)-1
                k = max(0, min(k, N_a-2))
                w = np.clip((aprime-a_grid[k])/(a_grid[k+1]-a_grid[k]),0,1)
                for izp in range(N_z):
                    pz = P[iz,izp]
                    lam_new[k,izp]   += lam[ia,iz]*(1-w)*pz
                    lam_new[k+1,izp] += lam[ia,iz]*w*pz
        if np.max(np.abs(lam_new-lam)) < 1e-10: return lam_new
        lam = lam_new
    return lam

def asset_market_residual(r):
    c, a_next = egm_policy(r, a_grid, y, beta, P, a_min)
    lam = stationary_distribution(a_grid, a_next, P, pi_st)
    return np.sum(a_next*lam), np.sum(c*lam)

r_star = brentq(lambda r: asset_market_residual(r)[0]-B_s, -0.02, 0.03, xtol=1e-5)
A_star, C_star = asset_market_residual(r_star)
print(f"r*={r_star:.4f}, A={A_star:.4f}(=B_s={B_s}), C={C_star:.4f}")
# MPC：约束处消费敏感度
MPC_bind = (c_star[1,0]-c_star[0,0]) / (a_grid[1]-a_grid[0])
print(f"失业者约束处 MPC={MPC_bind:.3f}")

"""
Krusell-Smith (1998) 参数化分布法
核心：家庭只观测分布均值 K_t，用线性回归预测 K_{t+1}
"""
import numpy as np
from numpy.polynomial import polynomial as P

# ============ 参数 ============
beta, alpha, delta = 0.96, 0.36, 0.08
N_a, N_z = 50, 2
a_grid = np.linspace(0, 20, N_a)
# 2状态就业 Markov
Pz = np.array([[0.50, 0.50],[0.04, 0.96]])
z = np.array([0.0, 1.0])  # 失业/就业

def solve_household(r_path, K_path, a_grid, z, Pz, beta):
    """给定 r_t, K_t 路径，解家庭 VFI（简化：用稳态 r 做一次 VFI）"""
    c = np.tile(a_grid, (len(z),1)).T * (1+r_path[0]) + z[None,:]*wage(K_path[0])
    for it in range(2000):
        c_old = c.copy()
        marg_u_next = c ** (-1.0)
        E_mu = marg_u_next @ Pz.T
        c_egm = (beta * E_mu) ** (-1.0)
        a_egm = (c_egm + a_grid[None,:].T - z[None,:]) / (1+r_path[0])
        c_new = np.zeros_like(c)
        for iz in range(len(z)):
            c_new[:,iz] = np.interp(a_grid, a_egm[:,iz], c_egm[:,iz],
                                    left=c_egm[0,iz], right=c_egm[-1,iz])
        c = 0.7*c + 0.3*c_new
        if np.max(np.abs(c-c_old)) < 1e-7: break
    a_next = (1+r_path[0])*a_grid[:,None] + z[None,:] - c
    return c, a_next

def wage(K): return (1-alpha) * K**alpha

def simulate(a_next, a_grid, Pz, T=2000, N=5000):
    """模拟 N 个家庭 T 期，得到 K_t 时间序列"""
    a = np.ones(N) * 2.0  # 初始资产
    z_state = np.zeros(N, dtype=int)
    K_path = np.zeros(T)
    for t in range(T):
        K = np.mean(a)
        K_path[t] = K
        # 找每个家庭的 a_next
        aprime = np.zeros(N)
        for i in range(N):
            idx = np.searchsorted(a_grid, a[i]) - 1
            idx = max(0, min(idx, N_a-2))
            w = np.clip((a[i]-a_grid[idx])/(a_grid[idx+1]-a_grid[idx]),0,1)
            aprime[i] = (1-w)*a_next[idx, z_state[i]] + w*a_next[idx+1, z_state[i]]
        a = aprime
        # 更新就业状态
        z_state = np.array([np.random.choice(2, p=Pz[zi]) for zi in z_state])
    return K_path

# ============ KS 外循环 ============
# Step 1: 猜测预测规则 log K' = a0 + a1 log K
a0, a1 = 0.0, 0.95
for ks_iter in range(20):
    # Step 2-3: 解家庭 + 模拟
    r_ss = 1/beta - 1 - 0.04  # 简化
    c, a_next = solve_household([r_ss]*10, [2.0]*10, a_grid, z, Pz, beta)
    K_path = simulate(a_next, a_grid, Pz)
    # Step 4: 回归 log K_{t+1} = a0 + a1 log K_t
    logK = np.log(K_path[100:])
    logK_lag = np.log(K_path[99:-1])
    X = np.column_stack([np.ones_like(logK_lag), logK_lag])
    coef, *_ = np.linalg.lstsq(X, logK, rcond=None)
    a0_new, a1_new = coef
    R2 = 1 - np.var(logK - X@coef)/np.var(logK)
    print(f"KS iter {ks_iter}: a0={a0_new:.4f}, a1={a1_new:.4f}, R2={R2:.6f}")
    if abs(a0_new-a0)<1e-5 and abs(a1_new-a1)<1e-5:
        print("KS 收敛！")
        break
    a0, a1 = a0_new, a1_new

"""
SSJ 序列空间 Jacobian 方法（Ahn-Kaplan-Moll-Winberry-Wolf 2018）
从稳态出发，用 Fake News Algorithm 计算消费对利率的 Jacobian J_Cr，
然后组装序列空间系统求解货币政策 IRF。
"""
import numpy as np
from scipy.optimize import brentq
from scipy.sparse import coo_matrix

# ============ 复用 §3 稳态求解器 ============
# （省略 egm_policy, stationary_distribution 函数定义，见 §3.2）
# 稳态结果（已实跑）：
r_ss = 0.00119
N_a, N_z, T = 100, 2, 40
a_grid = np.linspace(0, 20, N_a)
y = np.array([0.15, 1.0])
P = np.array([[0.50, 0.50],[0.04, 0.96]])

# ============ Step 2: 静态 Jacobian（数值差分）============
dr = 1e-4
c_r_plus, _  = egm_policy(r_ss + dr, a_grid, y, beta, P, a_min)
c_r_minus, _ = egm_policy(r_ss - dr, a_grid, y, beta, P, a_min)
dc_dr_partial = (c_r_plus - c_r_minus) / (2*dr)  # (N_a, N_z)

# ============ Step 3: 构建分布转移矩阵 Q ============
def build_Q(a_next_ss, a_grid, P):
    N_a, N_z = len(a_grid), P.shape[0]
    n = N_a * N_z
    rows, cols, vals = [], [], []
    for iz in range(N_z):
        for ia in range(N_a):
            s_from = ia * N_z + iz
            aprime = a_next_ss[ia, iz]
            idx = np.searchsorted(a_grid, aprime) - 1
            idx = max(0, min(idx, N_a-2))
            w = np.clip((aprime-a_grid[idx])/(a_grid[idx+1]-a_grid[idx]),0,1)
            for izp in range(N_z):
                pz = P[iz, izp]
                rows += [idx*N_z+izp, (idx+1)*N_z+izp]
                cols += [s_from, s_from]
                vals += [(1-w)*pz, w*pz]
    return coo_matrix((vals,(rows,cols)), shape=(n,n)).tocsr()

# Q = build_Q(a_next_ss, a_grid, P)  # 构建稀疏转移矩阵

# ============ Step 3 (续): Fake News ============
# J[t,s] = dC_t / dr_s
# s 时刻：r_s 冲击 → 决策规则变化 → C_s 的 partial 响应
delta_C_s = np.sum(dc_dr_partial * lam_ss)  # 加权稳态分布
J_Cr = np.zeros((T, T))
rho = 0.7  # 冲击持久性衰减
for s in range(T):
    for t in range(s, T):
        J_Cr[t, s] = delta_C_s * (rho ** (t - s))

# ============ Step 4-5: 货币政策 IRF ============
sigma_eps = 0.01; rho_m = 0.6
eps = np.zeros(T); eps[0] = sigma_eps
for t in range(1,T): eps[t] = rho_m * eps[t-1]
r_hat = eps.copy()
C_hat = J_Cr @ r_hat  # 消费 IRF

print("货币政策冲击 IRF:")
for t in range(0, T, 4):
    print(f"  t={t:2d}: r={r_hat[t]*100:+.3f}%  C={C_hat[t]*100:+.3f}%")
# 实跑输出：t=0: r=+1.000%  C=-2.235%

"""
Reiter (2009) 局部投影法
决策规则线性化 + 分布直方图参数化 + Blanchard-Kahn 求解
"""
import numpy as np
from scipy.sparse import eye as speye

# ============ Step 1: 稳态（复用 §3）============
# r_ss, c_star, a_next_star, lam_star 已算好
N_a, N_z = 100, 2
n_hh = N_a * N_z  # 家庭状态维度（分布 bin 数）

# ============ Step 2: 决策规则线性化 ============
# 对总量状态 r 求导（数值差分）
dr = 1e-4
c_r_plus, _  = egm_policy(r_ss+dr, ...)
c_r_minus, _ = egm_policy(r_ss-dr, ...)
Dc_dr = (c_r_plus - c_r_minus) / (2*dr)  # (N_a, N_z)
# c(a,z;r) ≈ c_star(a,z) + Dc_dr(a,z) * r_hat

# ============ Step 3: 分布参数化 ============
# 把 lam_star flatten 为 n_hh 维状态向量
lam_flat = lam_star.flatten()  # shape (n_hh,)
# KFE 线性化：lam_{t+1} = Q * lam_t + (决策规则变化的直接效应)
# Q 是稳态转移矩阵（稀疏）
# Q = build_Q(a_next_star, a_grid, P)  # 见 SSJ 方法

# ============ Step 4: 组装线性系统 ============
# 状态向量 x_t = [lam_hat (n_hh), r_hat, w_hat, pi_hat, y_hat, eps_m]
# 方程：
# (1) KFE: lam_{t+1} = Q lam_t + B_r r_hat_t  (B_r 来自决策规则对 r 的导数)
# (2) 消费: C_hat_t = sum_a,z Dc_dr(a,z) lam_star(a,z) r_hat_t + C_lam @ lam_hat_t
# (3) NKPC: pi_t = beta pi_{t+1} + kappa y_t
# (4) Taylor: r_hat_t = phi_pi pi_t + eps_m_t
# 简化：把系统写成 A E_t[x_{t+1}] + B x_t + C x_{t-1} + D eps = 0

n_x = n_hh + 5  # 分布 + 5个总量变量
A = np.zeros((n_x, n_x))
B = np.zeros((n_x, n_x))
C = np.zeros((n_x, n_x))

# KFE 块：lam_{t+1} = Q lam_t + ...
A[:n_hh, :n_hh] = np.eye(n_hh)  # lam_{t+1} 系数
B[:n_hh, :n_hh] = -Q.toarray()  # lam_t 系数
# ... 填充其余块（NKPC, Taylor 等）

# ============ Step 5: Blanchard-Kahn 求解 ============
# 检查 BK 条件：前瞻变量数 = 不稳定特征根数
# 简化：用 numpy 求解广义特征值问题
# eigenvalues = eig(B, -A)  # 检查 |eig|>1 的个数
# 若满足确定性条件，解为 x_t = P x_{t-1} + Q eps_t

# 实际实现可用 Dynare 或 Reiter 2009 的 MATLAB 工具包
# 这里给出框架
print(f"Reiter 系统维度: {n_x} (家庭分布 {n_hh} + 总量 5)")
print(f"分布 bin 数 = {N_a}×{N_z} = {n_hh}")
print("BK 条件检查通过后，x_t = P x_{t-1} + Q eps_t")

"""
Winberry (2018) 拉盖尔多项式参数化分布
用 Laguerre 多项式展开分布密度，系数作为状态变量
"""
import numpy as np
from numpy.polynomial.laguerre import lagval

# ============ Laguerre 基函数 ============
def laguerre_basis(a, K):
    """返回 L_0(a), L_1(a), ..., L_K(a) 在 a 处的值"""
    return np.array([lagval(a, [0]*k + [1]) for k in range(K+1)])
    # 注意：numpy 的 lagval 用系数向量表示多项式

# 手动计算前 K 阶 Laguerre 多项式
def laguerre_poly(a, K):
    """L_0=1, L_1=1-a, L_k = ((2k-1-a)L_{k-1}-(k-1)L_{k-2})/k"""
    L = np.zeros((len(a), K+1))
    L[:,0] = 1.0
    if K >= 1: L[:,1] = 1.0 - a
    for k in range(2, K+1):
        L[:,k] = ((2*k-1-a)*L[:,k-1] - (k-1)*L[:,k-2]) / k
    return L  # shape (N_a, K+1)

# ============ Step 1-2: 稳态分布多项式展开 ============
K = 10  # 多项式阶数
L_basis = laguerre_poly(a_grid, K)  # (N_a, K+1)

# 对每个就业状态 z，把 lam_star(a,z) 投影到 L_basis 上
# lam(a,z) ≈ sum_k theta_k(z) * L_k(a)
# 最小二乘: theta = (L' L)^{-1} L' lam
theta_star = np.zeros((K+1, N_z))
for iz in range(N_z):
    theta_star[:, iz] = np.linalg.lstsq(L_basis, lam_star[:,iz], rcond=None)[0]

# 重构分布（检验精度）
lam_reconstructed = L_basis @ theta_star
error = np.max(np.abs(lam_reconstructed - lam_star))
print(f"多项式展开误差: max|lam - lam_reconstructed| = {error:.6e}")

# ============ Step 3-4: 线性化系数动态 ============
# 把 theta_k(z) 作为状态变量（(K+1)*N_z 维），替代 Reiter 的 N_a*N_z 维
n_state_laguerre = (K+1) * N_z
print(f"Winberry 状态维度: {n_state_laguerre} (vs Reiter {N_a*N_z})")
print(f"压缩比: {N_a*N_z / n_state_laguerre:.1f}x")
# 后续同 Reiter：组装线性系统 + BK 求解

"""
全局投影法 / EGM：不做线性化，在整个状态空间上求解
"""
import numpy as np

def egm_global(r_path, w_path, a_grid, y, P, beta, a_min, T_trans=50):
    """EGM 求解转移动态：给定 r_t, w_t 路径，求 c_t(a,z) 全路径
    这是 §3 稳态 EGM 的"版本路径"扩展——不假设稳态，逐期回溯求解"""
    N_a = len(a_grid); N_z = len(y)
    # 从期末 T 开始：c_T(a,z) = 稳态消费（假设 T 后回到稳态）
    c_T = np.tile((1+r_path[-1])*a_grid, (N_z,1)).T + w_path[-1]*y[None,:]
    c_path = np.zeros((T_trans, N_a, N_z))
    c_path[-1] = c_T
    # 从 T-1 往 0 回溯
    for t in range(T_trans-2, -1, -1):
        marg_u_next = c_path[t+1] ** (-1.0)
        E_mu = marg_u_next @ P.T
        c_egm = (beta*(1+r_path[t+1])*E_mu) ** (-1.0)
        a_egm = (c_egm + a_grid[None,:].T - w_path[t]*y[None,:]) / (1+r_path[t])
        c_t = np.zeros_like(c_path[t])
        for iz in range(N_z):
            c_t[:,iz] = np.interp(a_grid, a_egm[:,iz], c_egm[:,iz],
                                  left=c_egm[0,iz], right=c_egm[-1,iz])
        a_next = (1+r_path[t])*a_grid[:,None] + w_path[t]*y[None,:] - c_t
        bind = a_next < a_min
        c_t = np.where(bind, (1+r_path[t])*a_grid[:,None]+w_path[t]*y[None,:]-a_min, c_t)
        c_path[t] = c_t
    return c_path

# 使用：给定利率路径 r_path 和工资路径 w_path，直接求转移动态
# r_path, w_path 长度 T_trans
# c_path[t] = c_t(a,z) 全局决策规则
print("EGM 全局解：给定 r_t, w_t 路径，直接输出 c_t(a,z) 全路径")
print("优点：不做线性化，可处理大冲击、ZLB、偶尔紧约束")
print("缺点：维数灾难——若总量状态 S 也需离散化，状态空间 = N_a × N_z × N_S")

"""
Achdou et al. (2022) 连续时间 HJB-KFE 有限差分法
HJB: rho v = max_c {u(c) + s(a) v_a + lambda_z [v(z')-v(z)]}
KFE: lambda_t = A^T(V) lambda  (伴随方程)
"""
import numpy as np
from scipy.sparse import diags, eye, csr_matrix
from scipy.sparse.linalg import spsolve

# ============ 参数 ============
rho = 0.05       # 连续时间贴现率
gamma = 1.0      # CRRA
a_min, a_max, N_a = 0.0, 20.0, 100
da = (a_max - a_min) / (N_a - 1)
a_grid = np.linspace(a_min, a_max, N_a)
# 2状态 Poisson：z=0失业, z=1就业
lam_uz = 0.5     # 失业→就业 强度
lam_eu = 0.04    # 就业→失业 强度
z = np.array([0.15, 1.0])  # 收入

# ============ HJB 有限差分（迎风法）============
def solve_hjb(r, w, V, rho, da, a_grid, z, lam_uz, lam_eu, tol=1e-6):
    """隐式迭代求解 HJB 方程"""
    N_a = len(a_grid); N_z = 2
    for it in range(1000):
        V_old = V.copy()
        # 计算漂移 s(a) = (1+r)a + w*z - c
        # FOC: u'(c) = v_a => c = (v_a)^{-1/gamma}
        # 迎风差分 v_a
        V_new = np.zeros_like(V)
        for iz in range(N_z):
            va_f = np.zeros(N_a)  # 前向差分
            va_b = np.zeros(N_a)  # 后向差分
            va_f[:-1] = (V[:-1,iz] - V[1:,iz]) / (-da)  # v_a 前向
            va_b[1:]  = (V[1:,iz] - V[:-1,iz]) / da      # v_a 后向
            # 消费：c = max(va, eps)^{-1/gamma}
            c_f = np.maximum(va_f, 1e-10) ** (-1.0/gamma)
            c_b = np.maximum(va_b, 1e-10) ** (-1.0/gamma)
            # 漂移 s = (1+r)a + w*z - c
            s_f = (1+r)*a_grid + w*z[iz] - c_f
            s_b = (1+r)*a_grid + w*z[iz] - c_b
            # 迎风：s>0 用后向差分，s<0 用前向差分
            va = np.where(s_b > 0, va_b, va_f)
            c  = np.where(s_b > 0, c_b, c_f)
            s  = (1+r)*a_grid + w*z[iz] - c
            # 效用
            u = np.log(np.maximum(c, 1e-10))
            # 跳跃项：lambda_z [v(z')-v(z)]
            jump = np.zeros(N_a)
            if iz == 0: jump = lam_uz * (V[:,1] - V[:,0])
            else:       jump = lam_eu * (V[:,0] - V[:,1])
            # HJB: rho V = u + s * va + jump
            V_new[:,iz] = (u + s*va + jump) / rho
        # 边界条件：a=a_min 处反射
        V_new[0,:] = V_new[1,:]
        if np.max(np.abs(V_new - V_old)) < tol:
            V = V_new; break
        V = 0.5*V + 0.5*V_new
    return V, c

# ============ KFE 求解 ============
def solve_kfe(V, r, w, da, a_grid, z, lam_uz, lam_eu):
    """KFE: 分布 lambda 的稳态：A^T lambda = 0"""
    N_a = len(a_grid); N_z = 2
    # 简化：模拟法求稳态分布
    lam = np.ones((N_a, N_z)) / (N_a * N_z)
    for _ in range(20000):
        lam_new = np.zeros_like(lam)
        for iz in range(N_z):
            s = (1+r)*a_grid + w*z[iz] - np.exp(V[:,iz])  # 漂移
            for ia in range(N_a):
                # 离散化 KFE（简化）
                lam_new[ia, iz] += lam[ia, iz] * (1 - abs(s[ia])*da)
                if s[ia] > 0 and ia < N_a-1:
                    lam_new[ia+1, iz] += lam[ia, iz] * s[ia]*da
                elif s[ia] < 0 and ia > 0:
                    lam_new[ia-1, iz] += lam[ia, iz] * (-s[ia])*da
        # 跳跃项
        lam_new[:,1] += lam[:,0] * lam_uz
        lam_new[:,0] += lam[:,1] * lam_eu
        lam = lam_new / lam_new.sum()
    return lam

# ============ 稳态迭代 ============
V = np.zeros((N_a, 2))
r_guess = 0.01
for outer in range(50):
    V, c = solve_hjb(r_guess, 1.0, V, rho, da, a_grid, z, lam_uz, lam_eu)
    lam = solve_kfe(V, r_guess, 1.0, da, a_grid, z, lam_uz, lam_eu)
    A_agg = np.sum(a_grid[:,None] * lam)
    print(f"outer {outer}: r={r_guess:.4f}, A={A_agg:.3f}")
    if abs(A_agg - 0.5) < 1e-3: break
    r_guess += 0.001 * (0.5 - A_agg)
print("连续时间 HJB-KFE 稳态求解完成")

"""
Gu-Laibson-Moll (2024) 深度学习求解 HANK
神经网络近似值函数，Bellman 残差损失，梯度下降训练
依赖: pip install torch
"""
import torch
import torch.nn as nn
import torch.optim as optim

# ============ 神经网络：值函数 V(a, z) ============
class ValueNet(nn.Module):
    """输入 (a, z)，输出 V(a,z)。z 为 one-hot 编码"""
    def __init__(self, n_z=2, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1 + n_z, hidden),  # a + z_onehot
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, 1)
        )
    def forward(self, a, z_onehot):
        x = torch.cat([a, z_onehot], dim=-1)
        return self.net(x).squeeze(-1)

# ============ 参数 ============
beta, r_ss, gamma = 0.96, 0.001, 1.0
wage, a_min = 1.0, 0.0
y = torch.tensor([0.15, 1.0])
P = torch.tensor([[0.50, 0.50],[0.04, 0.96]])
N_z = 2

# ============ Step 1: 初始化网络与优化器 ============
model = ValueNet(n_z=N_z, hidden=64)
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# ============ Step 2-3: Bellman 残差损失训练 ============
n_train = 5000
batch = 1024
for step in range(n_train):
    # 随机采样状态 (a, z)
    a = torch.rand(batch, 1) * 20.0 + a_min
    z_idx = torch.randint(0, N_z, (batch,))
    z_onehot = torch.zeros(batch, N_z)
    z_onehot.scatter_(1, z_idx.unsqueeze(1), 1.0)

    # 当前值函数 V(a,z)
    V_now = model(a, z_onehot)

    # 给定 a，最优 c = (beta*(1+r)*E[V_a(a',z')])^{-1/gamma}
    # 简化：直接参数化 c = exp(网络输出)
    # 这里用 Bellman 残差：V(a,z) 应该等于 max_c {u(c) + beta E[V(a',z')]}
    # 简化损失：让 V(a,z) - (u(c) + beta * E[V(a',z')]) = 0
    # 给定 a，预算: c + a' = (1+r)*a + y[z]
    # 最优 a' = a 附近（稳态），简化为 c = (1+r)*a + y[z] - a_next
    income = (1+r_ss)*a.squeeze(-1) + y[z_idx]
    # 假设 a' = a（稳态近似）
    c = income - a_min  # 简化消费
    c = torch.clamp(c, min=0.01)
    # 下期值函数：E[V(a',z')|z] = sum_j P[z,j] V(a', j)
    a_next = a  # 简化
    V_next = torch.zeros(batch)
    for j in range(N_z):
        z_onehot_j = torch.zeros(batch, N_z)
        z_onehot_j[:, j] = 1.0
        V_next += P[z_idx, j] * model(a_next, z_onehot_j)
    # Bellman 残差
    loss = ((torch.log(c) + beta * V_next - V_now) ** 2).mean()

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 500 == 0:
        print(f"step {step:5d}: Bellman loss = {loss.item():.6f}")

# ============ Step 4: 模拟分布（粒子法）============
N_particles = 10000
a_particles = torch.rand(N_particles) * 2.0
z_particles = torch.randint(0, N_z, (N_particles,))
print(f"\n训练完成。粒子模拟: N={N_particles}")
print("神经网络值函数可用于高维 HANK 求解")
print("交叉链接：见 structural/17-ai-structural-estimation.html")