动态结构模型：DDCM、Rust 与 Hotz-Miller

"""
Rust (1987) Bus Engine Replacement —— NFXP 简化实现
====================================================
状态: x in {0, 1, ..., NX-1} 表示里程分箱
选择: d in {0,1} = {继续使用, 更换发动机}
当期效用:
    u(x, 0) = - theta1 * log(1 + x / theta2)   # 运行成本（负号表示成本）
    u(x, 1) = - RC                               # 重置成本（里程归 0）
转移:
    d=0: x' ~ Poisson(1) 加在 x 上（截断在 NX-1）
    d=1: x' = 0 + Poisson(1)
"""
import numpy as np
from scipy.optimize import minimize_scalar, minimize
from scipy.stats import poisson

# ---------- 1. 模型设定 ----------
NX = 90              # 状态数（里程分箱）
BETA = 0.95          # 贴现因子
RC = 11.7255         # 重置成本（校准值，Rust 论文里的量级）

# 真实参数（用于模拟数据）
theta1_true, theta2_true = 2.0, 10.0

# 状态向量
x = np.arange(NX)

# 转移概率矩阵 P[d] 是 (NX, NX): P[d][x, x_next] = Pr(x_next | x, d)
def build_transition(NX):
    P = [np.zeros((NX, NX)), np.zeros((NX, NX))]
    # 里程增加服从 Poisson(lambda=1)
    lam = 1.0
    k_max = 5
    pk = poisson.pmf(np.arange(k_max+1), lam)
    for xi in range(NX):
        for k in range(k_max+1):
            # d=0: 里程从 xi 增加 k
            xn0 = min(xi + k, NX - 1)
            P[0][xi, xn0] += pk[k]
            # d=1: 里程从 0 增加 k
            xn1 = min(k, NX - 1)
            P[1][xi, xn1] += pk[k]
    # 归一化（截断尾部）
    P[0] /= P[0].sum(axis=1, keepdims=True)
    P[1] /= P[1].sum(axis=1, keepdims=True)
    return P

P = build_transition(NX)

# ---------- 2. 当期效用 ----------
def flow_utility(theta1, theta2):
    """返回 (NX, 2) 的 flow utility 矩阵"""
    u0 = - theta1 * np.log(1 + x / theta2)   # 继续使用
    u1 = - RC * np.ones(NX)                   # 更换发动机
    return np.column_stack([u0, u1])

# ---------- 3. Bellman 算子 + VFI ----------
def bellman_operator(V, theta1, theta2):
    """
    V: (NX, 2) 上一期的 choice-specific value
    返回新的 V
    """
    u = flow_utility(theta1, theta2)   # (NX, 2)
    # ex-ante value: EV(x) = logsumexp over d of V(x, d)
    V_max = V.max(axis=1, keepdims=True)         # (NX, 1)
    EV = (V_max.squeeze(1) + np.log(np.exp(V - V_max).sum(axis=1)))  # (NX,)
    # 对每个 d: V_new(x, d) = u(x,d) + beta * sum_x' P[d][x,x'] * EV(x')
    V_new = np.zeros_like(V)
    for d in [0, 1]:
        EV_next = P[d] @ EV     # (NX,)
        V_new[:, d] = u[:, d] + BETA * EV_next
    return V_new

def solve_bellman(theta1, theta2, tol=1e-9, max_iter=5000):
    V = np.zeros((NX, 2))
    for it in range(max_iter):
        V_new = bellman_operator(V, theta1, theta2)
        diff = np.max(np.abs(V_new - V))
        V = V_new
        if diff < tol:
            break
    # 计算 CCP
    V_max = V.max(axis=1, keepdims=True)
    P_choice = np.exp(V - V_max) / np.exp(V - V_max).sum(axis=1, keepdims=True)
    return V, P_choice

# ---------- 4. 模拟数据 ----------
def simulate_data(P_choice, T=50, N=200, seed=0):
    """模拟 N 辆 bus 各 T 期的 (x_t, d_t)"""
    rng = np.random.default_rng(seed)
    data = []
    for i in range(N):
        xi = 0
        for t in range(T):
            # 根据 P_choice[xi, :] 抽 d
            d = rng.choice(2, p=P_choice[xi])
            data.append((xi, d))
            # 状态转移
            xn = rng.choice(NX, p=P[d][xi])
            xi = xn
    return np.array(data)

V_true, P_true = solve_bellman(theta1_true, theta2_true)
data = simulate_data(P_true, T=50, N=200)
print(f"模拟数据形状: {data.shape}")
print(f"更换率 (d=1 的比例): {data[:,1].mean():.3f}")

# ---------- 5. 似然函数 + MLE ----------
def neg_loglik(params, data):
    theta1, theta2 = params
    if theta2 <= 0:
        return 1e10
    try:
        V, P_choice = solve_bellman(theta1, theta2)
    except Exception:
        return 1e10
    # 似然: log P(d | x)
    x_obs = data[:, 0].astype(int)
    d_obs = data[:, 1].astype(int)
    p_obs = P_choice[x_obs, d_obs]
    p_obs = np.clip(p_obs, 1e-300, 1.0)
    nll = -np.sum(np.log(p_obs))
    return nll

# 外层优化
print("\n开始 MLE 估计 ...")
res = minimize(neg_loglik,
               x0=[1.0, 5.0],
               args=(data,),
               method='Nelder-Mead',
               options={'xatol':1e-4, 'fatol':1e-4, 'maxiter':50})
print(f"收敛: {res.message}")
print(f"估计 theta1 = {res.x[0]:.3f}  (真实 {theta1_true})")
print(f"估计 theta2 = {res.x[1]:.3f}  (真实 {theta2_true})")

# ---------- 6. 用估计参数画 CCP ----------
V_hat, P_hat = solve_bellman(*res.x)
print("\n更换概率 P(d=1 | x) 在不同里程分箱:")
for xi in [0, 10, 30, 50, 70, 89]:
    print(f"  x={xi:3d}  P(replace)={P_hat[xi,1]:.4f}")