混合 Logit：随机系数、仿真积分与 MSLE

"""
混合 Logit（Random Parameters Logit）完整实现
- Halton 准蒙特卡洛抽样
- 仿真选择概率 P_ij = (1/R) sum_r L_ij(beta^r)
- 极大模拟似然 MSLE
- 个体偏好后验期望 E[beta_i | y_i, X]
依赖：pip install numpy scipy pandas
"""
import numpy as np
from scipy.stats import norm
from scipy.optimize import minimize

# =========================================================
# 1. Halton 序列生成器（准蒙特卡洛低差异抽样）
# =========================================================
def halton(dim: int, n_sample: int, leap: int = 1):
    """
    生成 dim 维 Halton 序列，返回 shape=(n_sample, dim)，值域 [0,1)
    dim: 随机系数的维度（用前 dim 个素数作基）
    """
    # 前 20 个素数，足够覆盖 dim<=20
    primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29,
              31, 37, 41, 43, 47, 53, 59, 61, 67, 71]
    assert dim <= len(primes), "维度超出预设素数表"
    H = np.zeros((n_sample, dim))
    for d in range(dim):
        p = primes[d]
        # van der Corput 一维序列
        n, base_inv = 0, 1.0 / p
        while n_sample >= (p ** n):
            n += 1
        seq = np.zeros(n_sample)
        for i in range(1, n_sample + 1):
            u = i * leap
            f, denom = 0.0, 1.0
            while u > 0:
                f += (u % p) / denom
                u //= p
                denom *= p
            seq[i - 1] = f
        H[:, d] = seq
    return H


def draws_normal(mean, std, R, dim, seed=0):
    """把 [0,1) 上的 Halton 反演为 N(mean, std^2) draws"""
    H = halton(dim, R)
    Z = norm.ppf(np.clip(H, 1e-6, 1 - 1e-6))  # 反 CDF
    return mean + std * Z                     # (R, dim)


# =========================================================
# 2. 模拟数据：N 个个体，J=3 个备选，K=2 个随机系数
# =========================================================
np.random.seed(2026)
N, J, K = 2000, 3, 2
# 备选属性 X[i, j, k]：价格（越负越好）+ 品质
X = np.random.normal(0, 1, (N, J, K))
# 真实 beta_i 服从 N(mu, Sigma)
mu_true = np.array([-1.5, 1.0])              # 价格为负效用
sig_true = np.array([0.5, 0.7])              # 异质性标准差
beta_true = mu_true + sig_true * np.random.randn(N, K)

# 生成选择：logit 核
def choice_from_beta(X, beta):
    """X:(N,J,K), beta:(N,K) -> y:(N,)"""
    V = np.einsum('njk,nk->nj', X, beta)     # (N,J)
    V -= V.max(axis=1, keepdims=True)        # log-sum-exp 数值稳定
    P = np.exp(V)
    P /= P.sum(axis=1, keepdims=True)
    y = np.array([np.random.choice(J, p=P[i]) for i in range(N)])
    return y

y = choice_from_beta(X, beta_true)


# =========================================================
# 3. ★ 仿真选择概率（MSLE 的核心）
# =========================================================
R = 100                                       # Halton 抽样数
halton_base = halton(K, R)                    # 固定 draws，避免目标函数抖动

def simulated_choice_probs(params):
    """
    params: 长度 2K，前 K 是均值 mu，后 K 是 log(std)（保证 std>0）
    返回 P_hat: (N, J)
    """
    mu = params[:K]
    std = np.exp(params[K:])                  # 保证正定
    # 把 Halton [0,1) 反演为正态 draws: (R, K)
    Z = norm.ppf(np.clip(halton_base, 1e-6, 1 - 1e-6))
    draws = mu + std * Z                      # (R, K)
    # 对每个 draw r 算 MNL 概率，再对 r 平均
    P = np.zeros((N, J))
    for r in range(R):
        V = X @ draws[r]                      # (N, J)
        V -= V.max(axis=1, keepdims=True)
        p_r = np.exp(V)
        p_r /= p_r.sum(axis=1, keepdims=True)
        P += p_r
    P /= R
    return P


# =========================================================
# 4. ★ 极大模拟似然（MSLE）目标函数
# =========================================================
def neg_loglik(params):
    P = simulated_choice_probs(params)
    P = np.clip(P, 1e-12, 1 - 1e-12)
    # 选中那一行的概率
    chosen = P[np.arange(N), y]
    return -np.sum(np.log(chosen))


# =========================================================
# 5. 优化：从 MNL 系数起步
# =========================================================
# 起步：普通 Logit 的均值估计（用全部数据压扁成 (N*J, K)）
Xflat = X.reshape(N * J, K)
yflat = (np.tile(np.arange(J), (N, 1)).reshape(-1) == np.repeat(y, J)).astype(float)
# 简单 OLS 起步
beta_start = np.linalg.lstsq(Xflat, yflat - yflat.mean(), rcond=None)[0]
params0 = np.concatenate([beta_start, np.log(np.ones(K) * 0.5)])

res = minimize(neg_loglik, params0, method='BFGS',
               options={'maxiter': 200, 'gtol': 1e-5})
mu_hat = res.x[:K]
std_hat = np.exp(res.x[K:])
print("=== MSLE 估计 ===")
print(f"mu  真值={mu_true}, 估计={mu_hat.round(3)}")
print(f"std 真值={sig_true}, 估计={std_hat.round(3)}")
print(f"对数似然 = {-res.fun:.1f}")


# =========================================================
# 6. ★ 个体偏好后验期望（posterior expectation）
# =========================================================
def posterior_beta_i(i, params):
    """对个体 i，用其观测 y_i 做贝叶斯加权平均"""
    mu = params[:K]; std = np.exp(params[K:])
    Z = norm.ppf(np.clip(halton_base, 1e-6, 1 - 1e-6))
    draws = mu + std * Z                      # (R, K)
    weights = np.zeros(R)
    for r in range(R):
        V = X[i] @ draws[r]                   # (J,)
        V -= V.max()
        p_r = np.exp(V); p_r /= p_r.sum()
        weights[r] = p_r[y[i]]                # 选 y_i 那一行的似然
    weights /= weights.sum()
    # 加权平均：E[beta | y_i] = sum_r beta^r * w_r
    return weights @ draws

# 抽查前 5 个个体的后验 vs 真值
print("\n=== 个体偏好后验 vs 真值（前 5 人）===")
for i in range(5):
    pb = posterior_beta_i(i, res.x)
    print(f"个体 {i}: 后验={pb.round(2)}  真值={beta_true[i].round(2)}")