反事实求解：Exact Hat Algebra 帽子代数

"""
Exact Hat Algebra 求解器（EK 模型）
- N=10 国例子
- 输入: 基准贸易份额矩阵 Lambda、基准支出 X、贸易成本冲击 tau_hat
- 输出: 工资变化 hat_w、福利变化 hat_W
- 不需要知道 T_i、w_i 的水平值！
"""
import numpy as np

def hat_ek(Lambda, X, tau_hat, theta, tol=1e-9, max_iter=50000, damp=0.5):
    """
    Lambda   : (N,N) 基准贸易份额，Lambda[n,i] = X_{ni}/X_n（观测）
    X        : (N,)   基准总支出（=w_i*L_i，观测）
    tau_hat  : (N,N)  贸易成本变化率，hat tau_{ni}
    theta    : scalar，Fréchet 形状参数
    返回:
      hat_w    : (N,) 工资变化
      hat_lam  : (N,N) 贸易份额变化
      hat_W    : (N,) 福利变化
    """
    N = len(X)
    hat_w = np.ones(N)

    for it in range(max_iter):
        # Step 1: 帽子份额
        # num[n,i] = (hat_w[i] * tau_hat[n,i])^(-theta)
        num = (hat_w[None, :] * tau_hat) ** (-theta)        # (N,N)
        # 分母: sum_k Lambda[n,k] * num[n,k]
        denom = (Lambda * num).sum(axis=1, keepdims=True)   # (N,1)
        hat_lam = num / denom                                # (N,N)

        # Step 2: 隐含销售 sales[i] = sum_n Lambda[n,i] * hat_lam[n,i] * hat_w[n] * X[n]
        # 注意 X[n] 是 n 国总支出
        weighted = Lambda * hat_lam * (hat_w[None, :] * X[None, :])  # (N,N)
        sales = weighted.sum(axis=0)   # 对 n 求和，得到 i 国销售  (N,)

        # Step 3: 过剩销售
        excess = sales / X

        # Step 4: 更新 hat_w（带阻尼）
        hat_w_new = hat_w * excess ** damp

        # 标准化：保持 sum(hat_w * X) = sum(X)（去掉工资的整体缩放）
        hat_w_new = hat_w_new * X.sum() / (hat_w_new * X).sum()

        # 收敛
        if np.max(np.abs(excess - 1.0)) < tol:
            hat_w = hat_w_new
            break
        hat_w = hat_w_new
    else:
        print("Warning: 未收敛")

    # 最终 hat_lam
    num = (hat_w[None, :] * tau_hat) ** (-theta)
    denom = (Lambda * num).sum(axis=1, keepdims=True)
    hat_lam = num / denom

    # 福利变化（ACR 公式）
    hat_W = np.diag(hat_lam) ** (-1.0 / theta)
    return hat_w, hat_lam, hat_W


# ============ 1. 构造 N=10 国例子 ============
np.random.seed(42)
N = 10
country_names = [f"国家{i}" for i in range(N)]

# 模拟基准贸易份额 Lambda（行=进口国，列=出口国，行和=1）
# 方法：先生成一个"吸引力"矩阵，再按行归一化
attract = np.random.uniform(0.1, 1.0, size=(N, N))
np.fill_diagonal(attract, np.random.uniform(2, 5, N))  # 国内偏好
Lambda = attract / attract.sum(axis=1, keepdims=True)

# 基准支出 X（GDP，对数正态）
X = np.exp(np.random.normal(10, 1.0, N))

# 贸易弹性
theta = 5.0

print("=" * 60)
print("基准贸易份额矩阵（前 5 行 × 5 列）：")
print(np.round(Lambda[:5, :5], 3))
print(f"\n基准支出 X（前 5 国）: {np.round(X[:5], 1)}")
print(f"国内份额 lambda_nn（前 5 国）: {np.round(np.diag(Lambda)[:5], 3)}")

# ============ 2. 反事实：所有跨国贸易成本下降 10% ============
tau_hat_free = np.ones((N, N))
tau_hat_free[~np.eye(N, dtype=bool)] = 0.9   # 跨国成本 *0.9（下降 10%）

hat_w, hat_lam, hat_W = hat_ek(Lambda, X, tau_hat_free, theta)
print("\n" + "=" * 60)
print("反事实：跨国贸易成本下降 10%")
print(f"  工资变化 hat_w（前 5 国）: {np.round(hat_w[:5], 4)}")
print(f"  福利变化 hat_W（前 5 国）: {np.round(hat_W[:5], 4)}")
print(f"  平均福利变化: {np.mean(hat_W):.4f}")

# ============ 3. 反事实：贸易战（大国加关税 25%） ============
# 假设 0 号国家对所有其他国家加 25% 关税
tau_hat_war = np.ones((N, N))
tau_hat_war[0, 1:] = 1.25   # 国家 0 从 j!=0 进口成本 *1.25
hat_w_war, hat_lam_war, hat_W_war = hat_ek(Lambda, X, tau_hat_war, theta)
print("\n" + "=" * 60)
print("反事实：国家 0 对所有他国加 25% 关税")
print(f"  国家 0 福利变化: {hat_W_war[0]:.4f} (即 {(hat_W_war[0]-1)*100:.2f}%)")
print(f"  其他国家平均福利变化: {np.mean(hat_W_war[1:]):.4f}")
print(f"  世界平均福利变化: {np.mean(hat_W_war):.4f}")
# 经典预测：加关税国自身福利下降（大国有利论在小模型里不成立），全球福利下降

# ============ 4. 反事实：完全自给自足 ============
tau_hat_autarky = np.ones((N, N))
tau_hat_autarky[~np.eye(N, dtype=bool)] = 1e6   # 跨国成本无穷大
hat_w_auto, hat_lam_auto, hat_W_auto = hat_ek(Lambda, X, tau_hat_autarky, theta)
print("\n" + "=" * 60)
print("反事实：完全自给自足（跨国成本 -> ∞）")
print(f"  福利变化 hat_W（应 < 1）: {np.round(hat_W_auto[:5], 4)}")
# 注意：hat_W_auto = hat_lam_nn^(-1/theta)，自给自足下 hat_lam_nn -> 1/lam_nn，所以 hat_W_auto = lam_nn^(1/theta) < 1
# 与 04 页 ACR 公式一致：W_trade/W_autarky = lam_nn^(-1/theta)

# 用 scipy.optimize.fixed_point + Anderson 加速
from scipy.optimize import fixed_point

def hat_ek_fast(log_hat_w, Lambda, X, tau_hat, theta):
    hat_w = np.exp(log_hat_w)
    num = (hat_w[None, :] * tau_hat) ** (-theta)
    denom = (Lambda * num).sum(axis=1, keepdims=True)
    hat_lam = num / denom
    weighted = Lambda * hat_lam * (hat_w[None, :] * X[None, :])
    sales = weighted.sum(axis=0)
    excess = sales / X
    # 不动点：log(hat_w_new) = log(hat_w) + log(excess)
    return np.log(hat_w) + np.log(excess)

log_hat_w_star = fixed_point(
    hat_ek_fast, np.zeros(N),
    args=(Lambda, X, tau_hat_free, theta),
    method='anderson', xtol=1e-10, maxiter=1000
)
hat_w_star = np.exp(log_hat_w_star)
print("Anderson 加速后的 hat_w:", np.round(hat_w_star[:5], 4))