"""
manifold_lib.py —— 第 9 讲的工具箱：圆周 S¹ 与旋转群 SO(3) 上的流匹配
=====================================================================================

【这个模块提供什么】
    1. 圆周 S¹：角度的"折返"（wrap）、圆周上的距离、包裹高斯混合（数据）、打分；
       两种流匹配：测地线插值（正确）与"把角度当实数"的朴素插值（错误的对照）。
    2. 旋转群 SO(3)：hat / vee、指数映射（Rodrigues 公式）、对数映射（含转角接近 π 的稳健处理）、
       测地距离、均匀抽样、集中在某个旋转附近的分布（切空间里的高斯）、投影回 SO(3)（SVD）；
       黎曼流匹配：R_t = R_0 exp(t log(R_0ᵀ R_1))，体坐标系下的速度是常数 log(R_0ᵀ R_1)。
    3. 网络与训练：TimeMLP（输入任意特征 + 时间），train_loop（通用的 Adam 循环）。

【约定】t = 0 是噪声（S¹ 上的均匀分布 / SO(3) 上的均匀分布），t = 1 是数据，和全课一致。
    SO(3) 上的速度用"体坐标系"：dR/dt = R · hat(ω)，ω ∈ R³。精确计算用 float64。
"""

import math
import time

import torch
import torch.nn as nn

DT = torch.float64


# =====================================================================
# 1. 圆周 S¹
# =====================================================================

def wrap(x):
    """
    把角度折返到 (−π, π]。

    【Python 小知识】torch.remainder(a, b) 的结果和 b 同号（总在 [0, b) 里），
        所以 remainder(x + π, 2π) − π 落在 [−π, π)；这里不区分 −π 和 π（它们是同一个点）。
    """
    return torch.remainder(x + math.pi, 2 * math.pi) - math.pi


def circ_dist(a, b):
    """圆周上两点之间的（测地）距离：|wrap(a − b)| ∈ [0, π]。"""
    return wrap(a - b).abs()


class WrappedGMM:
    """
    圆周上的"包裹高斯混合"：先从 Σ_k w_k N(μ_k, σ_k²) 抽一个实数，再折返到 (−π, π]。

    【参数】means、stds、weights：各 K 个
    【边界】σ 小（≤ 0.5）时包裹几乎不起作用，"离峰几个标准差"的判断仍然有效。
    """

    def __init__(self, means, stds, weights):
        self.means = torch.tensor(means, dtype=DT)
        self.stds = torch.tensor(stds, dtype=DT)
        self.weights = torch.tensor(weights, dtype=DT)

    def sample(self, n, generator=None):
        k = torch.multinomial(self.weights, n, replacement=True, generator=generator)
        return wrap(self.means[k] + self.stds[k] * torch.randn(n, generator=generator, dtype=DT))

    def report(self, theta, far=4.0):
        """
        给一批角度打分：各峰的比例与权重的 TV、峰内散度（理想值 1）、离所有峰都超过 far 个标准差的比例。

        【返回】dict(fractions, tv, spread, far)
        """
        d = circ_dist(theta.unsqueeze(1), self.means.unsqueeze(0)) / self.stds          # (n, K)
        k = d.argmin(dim=1)
        frac = torch.bincount(k, minlength=len(self.means)).double() / len(theta)
        dn = d.gather(1, k.unsqueeze(1)).squeeze(1)
        return {"fractions": frac.tolist(), "tv": 0.5 * (frac - self.weights).abs().sum().item(),
                "spread": torch.sqrt((dn ** 2).mean()).item(), "far": (dn > far).double().mean().item()}


def circle_energy_distance(a, b):
    """
    圆周上两组样本的能量距离：把角度放到单位圆上 (cos, sin)，用弦长当距离（第 1 讲 energy_distance 的圆周版）。
    """
    xa = torch.stack([torch.cos(a), torch.sin(a)], dim=1)
    xb = torch.stack([torch.cos(b), torch.sin(b)], dim=1)
    n, m = len(a), len(b)
    return (2 * torch.cdist(xa, xb).mean() - torch.cdist(xa, xa).sum() / (n * (n - 1))
            - torch.cdist(xb, xb).sum() / (m * (m - 1))).item()


# =====================================================================
# 2. 旋转群 SO(3)
# =====================================================================

def hat(w):
    """(…, 3) → (…, 3, 3) 反对称矩阵，满足 hat(w) v = w × v。"""
    z = torch.zeros_like(w[..., 0])
    return torch.stack([torch.stack([z, -w[..., 2], w[..., 1]], -1),
                        torch.stack([w[..., 2], z, -w[..., 0]], -1),
                        torch.stack([-w[..., 1], w[..., 0], z], -1)], -2)


def vee(M):
    """hat 的逆：(…, 3, 3) 反对称矩阵 → (…, 3)。"""
    return torch.stack([M[..., 2, 1], M[..., 0, 2], M[..., 1, 0]], -1)


def so3_exp(w):
    """
    指数映射 exp: so(3) → SO(3)（Rodrigues 公式）：绕单位轴 n = w/|w| 转 θ = |w|，
        R = I + sin θ K + (1 − cos θ) K²，K = hat(n)。

    【边界】θ → 0 时 sin θ/θ、(1 − cos θ)/θ² 用泰勒展开，避免 0/0。
    """
    th = w.norm(dim=-1, keepdim=True).unsqueeze(-1)                          # (…, 1, 1)
    K = hat(w)
    small = th < 1e-6
    a = torch.where(small, 1 - th ** 2 / 6, torch.sin(th) / th.clamp(min=1e-12))
    b = torch.where(small, 0.5 - th ** 2 / 24, (1 - torch.cos(th)) / (th ** 2).clamp(min=1e-24))
    eye = torch.eye(3, dtype=w.dtype).expand(K.shape)
    return eye + a * K + b * (K @ K)


def rot_angle(R):
    """旋转矩阵的转角 θ ∈ [0, π]：cos θ = (tr R − 1)/2（clamp 防止舍入误差超出 [−1, 1]）。"""
    c = ((R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2]) - 1) / 2
    return torch.arccos(c.clamp(-1.0, 1.0))


def so3_log(R):
    """
    对数映射 log: SO(3) → so(3)，返回旋转向量 w（|w| = θ ∈ [0, π]）。

    【思路】一般情况：w = θ / (2 sin θ) · vee(R − Rᵀ)。
        θ → 0：用 vee(R − Rᵀ)/2（一阶近似）。
        θ → π：sin θ → 0，上式不稳定。改用对称部分 B = (R + Rᵀ)/2 = cos θ I + (1 − cos θ) n nᵀ
            求出 n nᵀ，取对角元最大的那一列归一化得到 ±n，符号由 vee(R − Rᵀ) = 2 sin θ n 决定。
    【边界】θ 恰好等于 π 时 n 和 −n 表示同一个旋转，返回其中之一。
    """
    th = rot_angle(R)
    skew = vee(R - R.transpose(-1, -2))                                     # = 2 sin θ · n
    s = torch.sin(th)
    w_main = (th / (2 * s.clamp(min=1e-12))).unsqueeze(-1) * skew
    w_small = 0.5 * skew
    B = 0.5 * (R + R.transpose(-1, -2))
    nn_t = (B - torch.cos(th)[..., None, None] * torch.eye(3, dtype=R.dtype)) / (1 - torch.cos(th)).clamp(min=1e-12)[..., None, None]
    diag = torch.diagonal(nn_t, dim1=-2, dim2=-1)
    j = diag.argmax(dim=-1)
    col = torch.gather(nn_t, -1, j[..., None, None].expand(*nn_t.shape[:-1], 1)).squeeze(-1)
    n = col / col.norm(dim=-1, keepdim=True).clamp(min=1e-12)
    sign = torch.where((n * skew).sum(-1, keepdim=True) < 0, -1.0, 1.0).to(R.dtype)
    w_pi = sign * n * th.unsqueeze(-1)
    w = torch.where((th < 1e-6).unsqueeze(-1), w_small, w_main)
    return torch.where((th > 3.0).unsqueeze(-1), w_pi, w)


def geodesic_dist(R1, R2):
    """SO(3) 上的测地距离：R1ᵀ R2 的转角，∈ [0, π]。"""
    return rot_angle(R1.transpose(-1, -2) @ R2)


def random_rotations(n, generator=None):
    """
    SO(3) 上的均匀分布（Haar 测度）：4 维标准高斯归一化得到单位四元数，再转成旋转矩阵。

    【为什么这样是均匀的】单位四元数在 S³ 上均匀，而 S³ → SO(3) 是 2 对 1 的覆盖且保持测度。
        转角 θ 的密度是 (1 − cos θ)/π——集中在 π 附近，**不是**均匀的（02 号脚本检验）。
    """
    q = torch.randn(n, 4, generator=generator, dtype=DT)
    q = q / q.norm(dim=1, keepdim=True)
    w, x, y, z = q.unbind(1)
    return torch.stack([
        torch.stack([1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], -1),
        torch.stack([2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)], -1),
        torch.stack([2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], -1)], -2)


def project_to_so3(M):
    """
    把任意 3×3 矩阵投影到最近的旋转矩阵（Frobenius 范数下）：M = U S Vᵀ → U diag(1, 1, det(UVᵀ)) Vᵀ。
    """
    U, _, Vh = torch.linalg.svd(M)
    d = torch.det(U @ Vh)
    D = torch.diag_embed(torch.stack([torch.ones_like(d), torch.ones_like(d), d], -1))
    return U @ D @ Vh


def orth_error(R):
    """离开 SO(3) 有多远：‖RᵀR − I‖_F 和 |det R − 1|。"""
    eye = torch.eye(3, dtype=R.dtype)
    return (R.transpose(-1, -2) @ R - eye).flatten(-2).norm(dim=-1), (torch.det(R) - 1).abs()


class SO3Mixture:
    """
    SO(3) 上的"集中分布的混合"：第 k 个峰 R = M_k exp(hat(σ ε))，ε ~ N(0, I₃)（切空间里的高斯）。

    【参数】centers：(K, 3, 3)；sigma：每个峰的尺度（弧度）；weights：K 个权重
    """

    def __init__(self, centers, sigma, weights):
        self.centers = centers.to(DT)
        self.sigma = sigma
        self.weights = torch.tensor(weights, dtype=DT)

    def sample(self, n, generator=None):
        k = torch.multinomial(self.weights, n, replacement=True, generator=generator)
        eps = self.sigma * torch.randn(n, 3, generator=generator, dtype=DT)
        return self.centers[k] @ so3_exp(eps)

    def report(self, R):
        """
        打分：各峰比例的 TV、峰内散度（到最近峰的均方根距离 / (σ√3)，理想值约 1）、
        离所有峰都超过 4σ√3 的比例（"落在峰之间"）。
        """
        d = geodesic_dist(R.unsqueeze(1), self.centers.unsqueeze(0))                    # (n, K)
        k = d.argmin(dim=1)
        frac = torch.bincount(k, minlength=len(self.weights)).double() / len(R)
        dn = d.gather(1, k.unsqueeze(1)).squeeze(1)
        scale = self.sigma * math.sqrt(3)
        return {"fractions": frac.tolist(), "tv": 0.5 * (frac - self.weights).abs().sum().item(),
                "spread": (torch.sqrt((dn ** 2).mean()) / scale).item(), "far": (dn > 4 * scale).double().mean().item()}


# =====================================================================
# 3. 网络与训练
# =====================================================================

class TimeMLP(nn.Module):
    """输入 (特征, t)，输出 n_out 维。【参数】n_in：特征维数（不含 t）；n_out；hidden；depth"""

    def __init__(self, n_in, n_out, hidden=128, depth=4):
        super().__init__()
        layers, width = [], n_in + 1
        for _ in range(depth):
            layers += [nn.Linear(width, hidden), nn.SiLU()]
            width = hidden
        layers.append(nn.Linear(width, n_out))
        self.net = nn.Sequential(*layers)

    def forward(self, feat, t):
        if not torch.is_tensor(t):
            t = torch.full((feat.shape[0], 1), float(t), dtype=feat.dtype)
        return self.net(torch.cat([feat, t], dim=1))


def train_loop(model, loss_fn, steps=6000, lr=1e-3, seed=0):
    """通用训练循环：Adam，最后 20% 的步数学习率线性降到 0。loss_fn(g) → 标量。【返回】训练用时（秒）"""
    g = torch.Generator().manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda k: min(1.0, (steps - k) / (0.2 * steps)))
    t0 = time.perf_counter()
    model.train()
    for _ in range(steps):
        loss = loss_fn(g)
        opt.zero_grad()
        loss.backward()
        opt.step()
        sched.step()
    model.eval()
    return time.perf_counter() - t0


def angle_features(theta):
    """把角度（任意形状 (n, m)）变成网络输入 (cos θ, sin θ)：在 ±π 处连续，和圆周的拓扑一致。"""
    return torch.cat([torch.cos(theta), torch.sin(theta)], dim=1)