"""
=============================================================================
CS525 · 第 13 讲模块：数据归因的实验台（attrib_lab.py）
=============================================================================

对应讲义：lectures/lecture13_data_attribution.md §0.4

-----------------------------------------------------------------------------
【这个模块提供什么】
-----------------------------------------------------------------------------
    任务        linear_task（两个高斯类、少量翻转的标签）、curved_task（弯曲的决策边界，需要非线性模型）
    逻辑回归    fit_logreg（牛顿法，训练到收敛——凸问题，有唯一解）、lr_margin、lr_loss_grads、lr_hessian
    小网络      train_mlps（一次并行训练很多个单隐层网络：不同的训练子集、不同的初始化）、mlp_margins、
                mlp_output_grads（每个样本的输出对参数的梯度）、pick
    归因方法    lr_influence（逻辑回归的影响函数）、ggn_influence（网络的影响函数，用 Gauss–Newton 近似 Hessian）、
                trak_scores（TRAK：线性化 + 随机投影 + 残差加权）、fit_datamodels（子集指示向量 → 输出的岭回归）
    评价        spearman、lds（线性数据建模分数，TRAK 论文的指标）

-----------------------------------------------------------------------------
【记号】
-----------------------------------------------------------------------------
    "输出" f(x)：模型对 x 的 logit（第 1 类的对数几率）。归因要预测的就是：训练集变了，f(x) 怎样变。
    "影响"：把训练样本 i 的权重增加一点点（ε），f(x_test) 的变化率 df/dε。去掉样本 i ≈ 权重减少 1/n，
            所以"去掉 i 后的变化" ≈ −影响 / n。
"""

import numpy as np


def sigmoid(z):
    """数值稳定的 sigmoid（截断到 ±50，避免 exp 溢出的警告）。"""
    return 1.0 / (1.0 + np.exp(-np.clip(z, -50, 50)))


def with_bias(X):
    """在特征后面加一列 1（偏置项）。"""
    return np.c_[X, np.ones(len(X))]


# ============================================================================= 任务


def linear_task(n, rng, d=10, flip=0.05, sep=1.0):
    """
    两个高斯类（沿第 0 维分开），以 flip 的概率翻转标签。

    【返回】(X, y, flipped)：flipped 标记哪些样本的标签被翻转了（"脏数据"的真相）
    """
    y = rng.integers(0, 2, n)
    mu = np.zeros(d)
    mu[0] = sep
    X = rng.normal(size=(n, d)) + np.where(y[:, None] == 1, mu, -mu)
    flipped = rng.random(n) < flip
    return X, np.where(flipped, 1 - y, y), flipped


def curved_task(n, rng, d=5, flip=0.05):
    """弯曲的决策边界：y = [x0 + 0.8·x1² − 0.8 > 0]，再以 flip 的概率翻转。线性模型学不好，需要网络。"""
    X = rng.normal(size=(n, d))
    y = (X[:, 0] + 0.8 * X[:, 1] ** 2 - 0.8 > 0).astype(int)
    flipped = rng.random(n) < flip
    return X, np.where(flipped, 1 - y, y), flipped


# ============================================================================= 逻辑回归


def fit_logreg(X, y, weights=None, lam=1e-2, theta0=None, iters=50):
    """
    带 L2 正则的逻辑回归，用牛顿法训练到收敛。

    【参数】weights：每个样本的权重（None = 全 1）；lam：正则化系数；theta0：初始值（做留一法时从完整模型出发，收敛更快）
    【返回】参数向量（最后一维是偏置）
    【思路】目标 = (1/n) Σ w_i ℓ_i(θ) + (λ/2)||θ||²，凸函数；牛顿步 θ ← θ − H⁻¹ g，通常十几步就收敛到机器精度。
    【边界】n = 0 时返回全 0（只有正则项）。
    """
    n = len(y)
    th = np.zeros(X.shape[1] + 1) if theta0 is None else theta0.copy()
    if n == 0:
        return th
    w = np.ones(n) if weights is None else weights
    Xb = with_bias(X)
    for _ in range(iters):
        p = sigmoid(Xb @ th)
        g = Xb.T @ (w * (p - y)) / n + lam * th
        H = (Xb * (w * p * (1 - p))[:, None]).T @ Xb / n + lam * np.eye(len(th))
        step = np.linalg.solve(H, g)
        th -= step
        if np.abs(step).max() < 1e-11:
            break
    return th


def lr_margin(theta, X):
    return with_bias(X) @ theta


def lr_loss_grads(theta, X, y):
    """每个样本的损失对参数的梯度 ∇ℓ_i = (p_i − y_i)·[x_i, 1]，形状 (n, d+1)。"""
    Xb = with_bias(X)
    return Xb * (sigmoid(Xb @ theta) - y)[:, None]


def lr_hessian(theta, X, y, lam):
    """训练目标的 Hessian：(1/n) Σ p_i(1 − p_i) x_i x_iᵀ + λI。"""
    Xb = with_bias(X)
    p = sigmoid(Xb @ theta)
    return (Xb * (p * (1 - p))[:, None]).T @ Xb / len(y) + lam * np.eye(Xb.shape[1])


def lr_influence(theta, X, y, Xq, lam):
    """
    逻辑回归的影响函数（Koh & Liang 2017）：把训练样本 i 的权重增加 ε，查询点的输出 f(x_q) 的变化率。

    【返回】(查询数, n)：I[q, i] = −∇f(x_q)ᵀ H⁻¹ ∇ℓ_i
    【思路】最优参数满足 ∇[目标 + ε·ℓ_i] = 0。对 ε 求导（隐函数定理）：dθ/dε = −H⁻¹ ∇ℓ_i。
            再用链式法则：df(x_q)/dε = ∇f(x_q)ᵀ dθ/dε。f 是 logit，∇f(x_q) = [x_q, 1]。
    【边界】只在目标是强凸、θ 是精确最优解时成立；去掉一个样本 ≈ ε = −1/n，所以"去掉 i 的效果" ≈ −I / n。
    """
    Hinv = np.linalg.inv(lr_hessian(theta, X, y, lam))
    return -(with_bias(Xq) @ Hinv @ lr_loss_grads(theta, X, y).T)


# ============================================================================= 小网络（并行训练很多个）


def train_mlps(X, y, masks, seeds, hidden=16, steps=300, lr=0.5, wd=1e-3):
    """
    并行训练 M 个单隐层网络 f(x) = vᵀ tanh(Wᵀx + b) + c，全批量梯度下降。

    【参数】masks：(M, n) 的 0/1 矩阵，第 m 个网络只在 masks[m] = 1 的样本上训练（训练子集）；
            seeds：M 个整数，决定每个网络的初始化（同一个 seed = 同一个初始化）；hidden、steps、lr、wd：结构和优化设置
    【返回】参数 (W (M, d, h), b (M, h), v (M, h), c (M,))
    【思路】把 M 个网络的参数叠成一个批次，用 einsum 一次算完所有网络的前向和反向——训练 1000 个网络大约 5 秒。
            每个网络的损失 = 它的训练子集上的平均对数损失 + (wd/2)·||参数||²。
    【数据联系】datamodels 需要在成千上万个随机子集上各训练一个模型；这是让它在笔记本上可行的唯一办法。
    【numpy 小知识】np.einsum("nd,mdh->mnh", X, W)：对每个网络 m，算 X @ W[m]。
    """
    masks = np.asarray(masks, float)
    M, n = masks.shape
    d = X.shape[1]
    W = np.stack([np.random.default_rng(s).normal(size=(d, hidden)) / np.sqrt(d) for s in seeds])
    v = np.stack([np.random.default_rng(s + 10 ** 6).normal(size=hidden) / np.sqrt(hidden) for s in seeds])
    b = np.zeros((M, hidden))
    c = np.zeros(M)
    cnt = np.maximum(masks.sum(1), 1)
    for _ in range(steps):
        A = np.tanh(np.einsum("nd,mdh->mnh", X, W) + b[:, None, :])
        f = np.einsum("mnh,mh->mn", A, v) + c[:, None]
        r = (sigmoid(f) - y[None, :]) * masks / cnt[:, None]          # dLoss/df（只在训练子集上）
        gv = np.einsum("mn,mnh->mh", r, A) + wd * v
        gc = r.sum(1)
        dZ = r[:, :, None] * v[:, None, :] * (1 - A ** 2)
        gW = np.einsum("nd,mnh->mdh", X, dZ) + wd * W
        gb = dZ.sum(1)
        W -= lr * gW
        b -= lr * gb
        v -= lr * gv
        c -= lr * gc
    return W, b, v, c


def train_mlps_batched(X, y, masks, seeds, batch=1000, **kw):
    """分批调用 train_mlps（避免一次分配太大的数组）。返回参数元组。"""
    parts = [train_mlps(X, y, masks[k:k + batch], seeds[k:k + batch], **kw) for k in range(0, len(masks), batch)]
    return tuple(np.concatenate([p[j] for p in parts]) for j in range(4))


def mlp_margins(params, X):
    """所有网络对 X 的输出（logit），形状 (M, 查询数)。"""
    W, b, v, c = params
    A = np.tanh(np.einsum("nd,mdh->mnh", X, W) + b[:, None, :])
    return np.einsum("mnh,mh->mn", A, v) + c[:, None]


def pick(params, m):
    """取出第 m 个网络（仍然是"批次"的形状，批次大小为 1）。"""
    return tuple(p[m:m + 1] for p in params)


def mlp_output_grads(params, m, Xq):
    """
    第 m 个网络的输出 f(x) 对它全部参数的梯度，每个查询点一行，形状 (查询数, 参数数)。

    【思路】f = vᵀ tanh(Wᵀx + b) + c：∂f/∂v = a，∂f/∂c = 1，∂f/∂b = v ⊙ (1 − a²)，∂f/∂W = x ⊗ (v ⊙ (1 − a²))。
    """
    W, b, v, _ = params
    A = np.tanh(Xq @ W[m] + b[m])
    dz = v[m] * (1 - A ** 2)
    gW = np.einsum("qd,qh->qdh", Xq, dz).reshape(len(Xq), -1)
    return np.c_[gW, dz, A, np.ones(len(Xq))]


# ============================================================================= 归因方法


def ggn_influence(params, m, X, y, Xq, damp=1e-3):
    """
    网络的影响函数：在训练好的参数处，用 Gauss–Newton 近似代替真正的 Hessian（加阻尼 damp）。

    【返回】(查询数, n)：−∇f(x_q)ᵀ (G + damp·I)⁻¹ ∇ℓ_i，G = (1/n) Σ p_i(1 − p_i) ∇f_i ∇f_iᵀ
    【思路】网络的损失不是凸的，真正的 Hessian 可能不正定；Gauss–Newton 近似（把网络在当前参数处线性化）总是半正定，
            阻尼让它可逆。这也是 TRAK 和 EK-FAC（Grosse et al. 2023）的出发点。
    【边界】假设"训练到了一个最优点、去掉一个样本后还停在附近"——对网络来说两条都不太成立（Basu et al. 2021 说影响函数"脆弱"）。
    """
    Gtr = mlp_output_grads(params, m, X)
    Gq = mlp_output_grads(params, m, Xq)
    p = sigmoid(mlp_margins(pick(params, m), X)[0])
    H = (Gtr * (p * (1 - p))[:, None]).T @ Gtr / len(y) + damp * np.eye(Gtr.shape[1])
    return -(Gq @ np.linalg.solve(H, (Gtr * (p - y)[:, None]).T))


def trak_scores(params, m, X, y, Xq, k, rng, damp=1e-3):
    """
    TRAK（Park et al. 2023）的一个模型的分数。

    【参数】k：随机投影的维数；rng：投影矩阵的随机数
    【返回】(查询数, n)
    【思路】五步里的前三步和第五步（第四步"集成"由调用者对多个模型取平均）：
            1. 线性化：把网络在训练好的参数处看成"特征 = ∇f(x)"的线性模型；
            2. 随机投影：φ(x) = Pᵀ ∇f(x)，P 是 (参数数 × k) 的高斯矩阵（Johnson–Lindenstrauss：内积大致保持），
               把动辄几百万维的梯度压到几千维；
            3. 在线性化的模型上用影响函数：τ(x_q, i) = φ(x_q)ᵀ (ΦᵀΦ)⁻¹ φ(x_i) · (y_i − p_i)。
               最后的 (y_i − p_i) 是残差：模型已经很有把握的样本，影响小（论文里的 Q 矩阵）。
    【边界】TRAK 论文还对分数做软阈值稀疏化（第 5 步）；这里省略。
    """
    Gtr = mlp_output_grads(params, m, X)
    Gq = mlp_output_grads(params, m, Xq)
    P = rng.normal(size=(Gtr.shape[1], k)) / np.sqrt(k)
    Phi, Phq = Gtr @ P, Gq @ P
    p = sigmoid(mlp_margins(pick(params, m), X)[0])
    K = Phi.T @ Phi + damp * len(y) * np.eye(k)
    return (Phq @ np.linalg.solve(K, Phi.T)) * (y - p)[None, :]


def fit_datamodels(masks, outputs, ridge=0.01):
    """
    Datamodels（Ilyas et al. 2022）：对每个查询点，用岭回归拟合 输出 ≈ θᵀ 1_S + 常数。

    【参数】masks：(M, n) 训练子集的指示向量；outputs：(M, 查询数) 在这些子集上训练出的模型的输出；
            ridge：正则化系数（乘以 M）
    【返回】(查询数, n)：每个训练样本的"权重"——包含它会让查询点的输出平均增加多少
    【思路】论文用 L1 正则（稀疏），这里用岭回归（有闭式解，所有查询点一次解完）。
    """
    M, n = masks.shape
    A = np.c_[masks, np.ones(M)]
    coef = np.linalg.solve(A.T @ A + ridge * M * np.eye(n + 1), A.T @ outputs)
    return coef[:n].T


def spearman(a, b):
    """Spearman 秩相关（没有并列时）。"""
    ra, rb = np.argsort(np.argsort(a)), np.argsort(np.argsort(b))
    return float(np.corrcoef(ra, rb)[0, 1])


def lds(scores, eval_masks, eval_outputs):
    """
    线性数据建模分数（linear datamodeling score，TRAK 论文的指标）。

    【参数】scores：(查询数, n) 归因分数；eval_masks：(M_e, n) 新抽的训练子集；eval_outputs：(M_e, 查询数) 在这些子集上
            真正训练出的模型的输出（通常对几个随机种子取平均）
    【返回】对每个查询点：预测的输出（子集里样本的分数之和）和真实输出在 M_e 个子集上的 Spearman 相关，再对查询点取平均
    【思路】一个好的归因方法应该能回答反事实问题："换一个训练子集，输出会怎样变？"——这正是 LDS 衡量的。
    """
    pred = eval_masks @ scores.T
    return float(np.mean([spearman(pred[:, j], eval_outputs[:, j]) for j in range(eval_outputs.shape[1])]))


if __name__ == "__main__":
    rng = np.random.default_rng(0)
    X, y, _ = curved_task(100, rng)
    Xt, yt, _ = curved_task(200, np.random.default_rng(1), flip=0.0)
    P = train_mlps(X, y, np.ones((200, 100)), list(range(200)))
    acc = np.mean((mlp_margins(P, Xt) > 0) == yt[None, :])
    print(f"并行训练了 200 个网络（不同的初始化）；它们在弯曲边界任务上的平均测试准确率 {acc:.3f}")
    th = fit_logreg(X, y)
    print(f"逻辑回归在同一个任务上的测试准确率 {np.mean((lr_margin(th, Xt) > 0) == yt):.3f}（线性模型画不出弯曲的边界）")