"""
robot_lib.py —— 第 8 讲的工具箱：一个二维"绕障碍"任务、人类示范，以及三种模仿学习策略
=====================================================================================

【任务】一个点状机器人从下方出发（x ~ U[−0.05, 0.05]，y = −3），要走到上方的目标 (0, 3)（距离 < 0.3 算成功）。
    中间有一个半径 1 的圆形障碍，圆心在原点。机器人每一步移动一个"动作"a ∈ R²（位移，长度约 0.15），
    最多走 80 步。进入障碍就算失败（碰撞）。

【示范】200 条"人类示范"：每条随机地从左边或右边绕过障碍（各一半），路径是一条三次 Bézier 曲线，
    绕行的宽度、速度略有随机性。于是在起点附近、中线上，**同一个观测对应两种截然不同的动作**：
    往左上走，或者往右上走——示范是**多峰**的。

【三种策略】（都是"给定观测 o = 当前位置，输出动作"，都用同样大小的 MLP）
    RegressionPolicy     ：均方误差回归 a = g(o)。最优解是条件均值 E[a | o]——两个峰的平均，"直着往上走"。
    DiscretePolicy       ：把动作的方向离散成 24 个格子（每格 15°），用分类（交叉熵）学 p(格子 | o)，按概率抽样。
                           这是 RT-1 / RT-2 这类"把动作变成 token"的做法的最简单版本。
    FlowPolicy           ：条件流匹配（第 2、4 讲）学 p(a | o)：网络 u_θ(x, t, o)，从 N(0, I) 积分到动作。
                           horizon = H 时一次生成未来 H 步的动作（"动作块"，Diffusion Policy 的做法）。

【流与扩散联系】FlowPolicy 就是第 4 讲的条件生成，只是"条件"从类别 y 换成了观测 o，"数据"从图像换成了动作。
    Diffusion Policy（Chi et al. 2023）、π0（Black et al. 2024）的动作专家都是这个结构。
"""

import math
import time

import torch
import torch.nn as nn
import torch.nn.functional as F

OBST_R = 1.0
GOAL = torch.tensor([0.0, 3.0])
GOAL_R = 0.3
STEP = 0.15             # 示范的平均步长（动作的典型长度）
MAX_STEPS = 80


# =====================================================================
# 1. 示范
# =====================================================================

def bezier(p0, p1, p2, p3, s):
    """三次 Bézier 曲线在参数 s ∈ [0, 1] 处的点。p0..p3：(2,)；s：(m,) → (m, 2)。"""
    s = s.unsqueeze(1)
    return ((1 - s) ** 3) * p0 + 3 * ((1 - s) ** 2) * s * p1 + 3 * (1 - s) * s ** 2 * p2 + s ** 3 * p3


def make_demo(side, g, start_spread=0.05):
    """
    一条示范：从随机起点出发，从 side（+1 右、−1 左）绕过障碍到目标。

    【参数】side：+1 / −1；g：随机数生成器；start_spread：起点 x 的范围
    【思路】Bézier 曲线的两个控制点放在 (side · w, −1.5) 和 (side · w, 1.5)，w ~ U[1.9, 2.6] 决定绕多宽；
        先在曲线上密集取点，再按弧长每隔约 STEP（乘一个 U[0.85, 1.15] 的速度因子）取一个点。
        动作 = 相邻两点之差，再加一点高斯噪声（标准差 0.01，模拟人手的抖动）。
    【返回】(positions: (N, 2), actions: (N, 2))——第 i 个动作是在 positions[i] 处做的
    """
    x0 = ((torch.rand(1, generator=g) * 2 - 1) * start_spread).item()
    w = 1.9 + 0.7 * torch.rand(1, generator=g).item()
    speed = STEP * (0.85 + 0.3 * torch.rand(1, generator=g).item())
    p0, p3 = torch.tensor([x0, -3.0]), GOAL.clone()
    p1, p2 = torch.tensor([side * w, -1.5]), torch.tensor([side * w, 1.5])
    dense = bezier(p0, p1, p2, p3, torch.linspace(0, 1, 2001))
    seg = (dense[1:] - dense[:-1]).norm(dim=1)
    arc = torch.cat([torch.zeros(1), seg.cumsum(0)])
    marks = torch.arange(0, arc[-1].item(), speed)
    idx = torch.searchsorted(arc, marks).clamp(max=len(dense) - 1)
    pos = torch.cat([dense[idx], p3.unsqueeze(0)])
    act = pos[1:] - pos[:-1] + 0.01 * torch.randn(len(pos) - 1, 2, generator=g)
    return pos[:-1], act


def make_dataset(n_demos=200, seed=0, horizon=1, start_spread=0.05):
    """
    生成 n_demos 条示范，整理成 (观测, 动作块) 的训练对。
    start_spread：起点 x 的范围 ±start_spread。默认 ±0.05：起点几乎都在中线上，"往左还是往右"
        完全由示范者随机决定，观测分辨不出来（范围放宽到 ±0.3 时情况不同，见 01 号脚本）。

    【参数】horizon：动作块的长度 H。第 i 个样本的目标是 a_i, a_{i+1}, ..., a_{i+H−1}；
        超出示范末尾的部分重复最后一个动作（"继续往前走"）。
        【边界】如果补零（"到了就停"），在观测有噪声时，机器人离目标还差一点就以为到了、停下不动，
        会大量超时——这是补零的副作用，不是我们想研究的现象。
    【返回】dict(obs: (M, 2), act: (M, 2H), demos: [(positions, actions), ...], sides: (n_demos,))
    """
    g = torch.Generator().manual_seed(seed)
    obs, act, demos = [], [], []
    sides = torch.where(torch.rand(n_demos, generator=g) < 0.5, -1.0, 1.0)
    for k in range(n_demos):
        pos, a = make_demo(sides[k].item(), g, start_spread)
        demos.append((pos, a))
        padded = torch.cat([a, a[-1:].expand(horizon, 2)])
        for i in range(len(a)):
            obs.append(pos[i])
            act.append(padded[i:i + horizon].reshape(-1))
    return {"obs": torch.stack(obs), "act": torch.stack(act), "demos": demos, "sides": sides}


# =====================================================================
# 2. 网络与三种策略
# =====================================================================

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


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


class RegressionPolicy:
    """
    均方误差回归：a = g_θ(o)。

    【边界】最优解是 E[a | o]。观测只对应一种动作时它很好；对应两种动作时，它输出两者的平均——
        在起点的中线上就是"直着往上走"，正对着障碍。
    """

    def __init__(self, horizon=1, seed=0):
        torch.manual_seed(seed)
        self.H = horizon
        self.net = mlp(2, 2 * horizon)

    def fit(self, data, steps=5000, batch=512, lr=1e-3, seed=0, obs_noise=0.0):
        obs, act = data["obs"], data["act"] / STEP

        def loss_fn(g):
            idx = torch.randint(0, len(obs), (batch,), generator=g)
            o = obs[idx] + obs_noise * torch.randn(batch, 2, generator=g)
            return ((self.net(o) - act[idx]) ** 2).sum(dim=1).mean()
        return _train(self.net.parameters(), loss_fn, steps, lr, seed)

    @torch.no_grad()
    def act(self, obs, generator=None):
        """【返回】(n, 2H) 动作块（单位：坐标）。"""
        return self.net(obs) * STEP


class DiscretePolicy:
    """
    离散化动作：把动作方向分成 n_bins 个格子，长度固定为 STEP；分类网络给出每个格子的概率，按概率抽样。

    【边界】只适用于 horizon = 1；量化误差最大是半个格子（7.5°）；速度信息被丢掉了。
    """

    def __init__(self, n_bins=24, seed=0):
        torch.manual_seed(seed)
        self.n_bins = n_bins
        self.net = mlp(2, n_bins)

    def _bin(self, a):
        ang = torch.atan2(a[:, 1], a[:, 0]) % (2 * math.pi)
        return (ang / (2 * math.pi / self.n_bins)).long() % self.n_bins

    def fit(self, data, steps=5000, batch=512, lr=1e-3, seed=0, obs_noise=0.0):
        obs, labels = data["obs"], self._bin(data["act"][:, :2])

        def loss_fn(g):
            idx = torch.randint(0, len(obs), (batch,), generator=g)
            o = obs[idx] + obs_noise * torch.randn(batch, 2, generator=g)
            return F.cross_entropy(self.net(o), labels[idx])
        return _train(self.net.parameters(), loss_fn, steps, lr, seed)

    @torch.no_grad()
    def act(self, obs, generator=None):
        k = torch.multinomial(F.softmax(self.net(obs), dim=1), 1, generator=generator).squeeze(1)
        ang = (k.float() + 0.5) * (2 * math.pi / self.n_bins)
        return STEP * torch.stack([torch.cos(ang), torch.sin(ang)], dim=1)


class FlowPolicy:
    """
    条件流匹配策略：学 p(动作块 | 观测)。网络 u_θ(x, t, o)，x 是 2H 维（动作块除以 STEP，大约是单位尺度）。

    训练（第 2 讲 Algorithm 3 + 第 4 讲的条件）：a ~ 示范，ε ~ N(0, I)，t ~ U[0, 1]，
        x_t = (1 − t) ε + t a，损失 ‖u_θ(x_t, t, o) − (a − ε)‖²。
    采样：从 ε 出发，欧拉法 n_steps 步积分到 t = 1。

    【参数】horizon：动作块长度 H；n_steps：采样时的欧拉步数
    【边界】n_steps = 1 时，第一步的速度是 E[a | o] − ε，一步走完就是 E[a | o]——
        **一步的流策略就是回归策略**（第 7 讲 §1.1 的"一步塌缩到均值"）。
    """

    def __init__(self, horizon=1, seed=0, n_steps=10):
        torch.manual_seed(seed)
        self.H = horizon
        self.n_steps = n_steps
        self.net = mlp(2 * horizon + 1 + 2, 2 * horizon)

    def fit(self, data, steps=5000, batch=512, lr=1e-3, seed=0, obs_noise=0.0):
        obs, act = data["obs"], data["act"] / STEP

        def loss_fn(g):
            idx = torch.randint(0, len(obs), (batch,), generator=g)
            a, o = act[idx], obs[idx] + obs_noise * torch.randn(batch, 2, generator=g)
            eps = torch.randn(a.shape, generator=g)
            t = torch.rand(batch, 1, generator=g)
            xt = (1 - t) * eps + t * a
            return ((self.net(torch.cat([xt, t, o], dim=1)) - (a - eps)) ** 2).sum(dim=1).mean()
        return _train(self.net.parameters(), loss_fn, steps, lr, seed)

    @torch.no_grad()
    def act(self, obs, generator=None, n_steps=None):
        n_steps = n_steps or self.n_steps
        x = torch.randn(obs.shape[0], 2 * self.H, generator=generator)
        for k in range(n_steps):
            t = torch.full((obs.shape[0], 1), k / n_steps)
            x = x + self.net(torch.cat([x, t, obs], dim=1)) / n_steps
        return x * STEP


# =====================================================================
# 3. 闭环执行（rollout）与评估
# =====================================================================

@torch.no_grad()
def rollout(policy, n=500, seed=0, execute=1, push_prob=0.0, push_std=0.3, obs_noise=0.0, start_spread=0.05,
            record=False, **act_kwargs):
    """
    让 n 个机器人同时执行策略，直到成功、碰撞或用完 MAX_STEPS 步。

    【参数】
        execute   : 每次调用策略后执行动作块的前几步再重新规划（"执行长度" k ≤ H）。
                    k = 1：每一步都重新看观测（完全闭环）；k = H：整个块开环执行完。
        push_prob : 每一步受到"外力推动"的概率；push_std：推动位移的标准差（模拟被人碰了一下）
        obs_noise : 观测噪声的标准差：策略看到的是"真实位置 + 噪声"（传感器不准）
        start_spread：起点 x 的范围 ±start_spread（默认 ±0.05，和示范一致）
        record    : 是否记录轨迹（画图用）
    【返回】dict：success / collision / timeout（比例）、side（成功者从哪边绕：右边的比例）、
        jerk（相邻两步动作之差的平均长度，衡量"抖不抖"）、calls（平均每条轨迹调用策略的次数）、
        flips（在中线附近、障碍下方时，横向动作改变方向的平均次数——"左右摇摆"）、
        steps（成功者的平均步数）、traj（若 record）
    【边界】碰撞判定在每一步之后做：位置进入半径 OBST_R 的圆就算碰撞（不检查一步之内穿过障碍的情况，
        步长 0.15 相对半径 1 很小，影响可以忽略）。
    """
    g = torch.Generator().manual_seed(seed)
    x = torch.stack([(torch.rand(n, generator=g) * 2 - 1) * start_spread, torch.full((n,), -3.0)], dim=1)
    status = torch.zeros(n, dtype=torch.long)            # 0 进行中，1 成功，2 碰撞
    steps_taken = torch.zeros(n)
    calls = torch.zeros(n)
    prev_a = None
    jerk_sum, jerk_cnt = torch.zeros(n), torch.zeros(n)
    flips = torch.zeros(n)                                # 中线附近"左右摇摆"的次数
    min_x = x[:, 0].clone()
    max_x = x[:, 0].clone()
    plan, plan_idx = None, execute
    traj = [x.clone()] if record else None
    for step in range(MAX_STEPS):
        active = status == 0
        if not active.any():
            break
        if plan_idx >= execute:
            seen = x + obs_noise * torch.randn(n, 2, generator=g)              # 带噪声的观测
            plan = policy.act(seen, generator=g, **act_kwargs).reshape(n, -1, 2)
            plan_idx = 0
            calls += active.float()
        a = plan[:, plan_idx]
        plan_idx += 1
        if prev_a is not None:
            jerk_sum += ((a - prev_a).norm(dim=1)) * active
            jerk_cnt += active.float()
            near_mid = active & (x[:, 0].abs() < 0.5) & (x[:, 1] < -1.0)
            flips += (near_mid & (torch.sign(a[:, 0]) * torch.sign(prev_a[:, 0]) < 0)).float()
        prev_a = a
        push = (torch.rand(n, generator=g) < push_prob).float().unsqueeze(1) * push_std * torch.randn(n, 2, generator=g)
        x = torch.where(active.unsqueeze(1), x + a + push, x)
        steps_taken += active.float()
        min_x = torch.minimum(min_x, x[:, 0])
        max_x = torch.maximum(max_x, x[:, 0])
        status = torch.where(active & (x.norm(dim=1) < OBST_R), torch.full_like(status, 2), status)
        status = torch.where((status == 0) & ((x - GOAL).norm(dim=1) < GOAL_R), torch.full_like(status, 1), status)
        if record:
            traj.append(x.clone())
    ok = status == 1
    right = (max_x.abs() > min_x.abs())
    return {"success": ok.float().mean().item(), "collision": (status == 2).float().mean().item(),
            "timeout": (status == 0).float().mean().item(),
            "side": right[ok].float().mean().item() if ok.any() else float("nan"),
            "jerk": (jerk_sum.sum() / jerk_cnt.sum().clamp(min=1)).item(),
            "calls": calls.mean().item(), "flips": flips.mean().item(), "steps": steps_taken[ok].mean().item() if ok.any() else float("nan"),
            "traj": torch.stack(traj) if record else None, "status": status}