"""无老师出题（PyTorch）：A 均匀 / B 直线被挡 / C 贴立柱 / D 工作空间外；pos/pose 各半。
2026-10-18：无老师题没有参考时长与行程，exec_s/travel 置 None（此前写死 3.0/0.0 造成假时长预算）；
闭环仿真对这类题按定长 6 s 口径判到达，time_ok_1p3 / travel_ratio 只在有老师参考的题上统计。"""
import math
import torch

from ml.common.features import wrap180


def auto_chunk(kin, mb_per_q, lo=256, hi=8192):
    """2026-10-18 第五轮（审查 N-6）：分块大小按设备空闲显存定，用掉其 1/4。
    两处每位形峰值不同不是矛盾：kin.gaps 按 (胶囊对下限 + cap)² 筛候选对，cap=0.5 mm 时幸存约 2800 对（≈1.5 MB/位形），
    column_gap 用 cap=40 mm，40 mm 内幸存的对多一个量级（≈20 MB/位形）。CPU 上取下限 256。"""
    if kin.lo.device.type == 'cuda':
        free = torch.cuda.mem_get_info(kin.lo.device)[0]
        return int(max(lo, min(hi, free * 0.25 / (mb_per_q * 1e6))))
    return lo


def min_gap_chunked(kin, q, cap=0.5, chunk=None):
    """分块算最小间隙：一次 4 万位形会爆内存/显存（2026-10-18 实测出题 OOM）；块大小见 auto_chunk（≈1.5 MB/位形）"""
    chunk = chunk or auto_chunk(kin, 1.5)
    return torch.cat([kin.gaps(q[i:i + chunk], cap)['gm'].min(-1).values for i in range(0, q.shape[0], chunk)])


def rand_q(kin, n, g):
    u = torch.rand(n, 6, device=kin.lo.device, generator=g)
    return kin.lo + u * (kin.hi - kin.lo)


def free_q(kin, n, g, chunk=1024):
    out = []
    while sum(x.shape[0] for x in out) < n:
        q = rand_q(kin, chunk, g)
        ok = min_gap_chunked(kin, q) > 0
        out.append(q[ok])
    return torch.cat(out)[:n]


def straight_blocked(kin, a, b, n=24):
    d = torch.where(kin.cont, wrap180(b - a), b - a)
    s = torch.linspace(0, 1, n, device=a.device)
    pts = (a[:, None] + s[None, :, None] * d[:, None]).reshape(-1, 6)
    return (min_gap_chunked(kin, pts) <= 0).view(a.shape[0], n).any(-1)


def column_gap(kin, q, chunk=None):
    chunk = chunk or auto_chunk(kin, 20.0)                                     # cap=40 → ≈20 MB/位形
    out = []
    for s in range(0, q.shape[0], chunk):                                       # 与 min_gap_chunked 一样分块防 OOM
        gm = kin.gaps(q[s:s + chunk], 40)['gm']
        out.append(torch.stack([gm[:, g] for g in (2, 3, 4, 5, 6)], -1).min(-1).values)   # 组 0（立柱/底座）× 组 2..6
    return torch.cat(out)


@torch.no_grad()
def sample_tasks(kin, n, seed, mix=(0.3, 0.3, 0.3, 0.1)):
    g = torch.Generator(device=kin.lo.device); g.manual_seed(seed)
    want = [int(round(n * m)) for m in mix]; want[0] += n - sum(want)
    tasks = []
    def emit(cat, q0, b, T=None):
        pose = torch.rand(q0.shape[0], device=q0.device, generator=g) < 0.5
        if cat == 'D': pose[:] = False
        xf = kin.cum_xf(b)
        Tb = kin.fk(b, xf) if T is None else T
        Rb = kin.orient(b, xf).reshape(-1, 9)
        for i in range(q0.shape[0]):
            tasks.append({'cat': cat, 'mode': 'pose' if pose[i] else 'pos', 'q0': q0[i].tolist(), 'T': Tb[i].tolist(),
                          'R': Rb[i].tolist() if pose[i] else None, 'exec_s': None, 'travel': None})
    if want[0]:
        emit('A', free_q(kin, want[0], g), free_q(kin, want[0], g))
    for cat, k in (('B', want[1]), ('C', want[2])):
        got_a, got_b, have = [], [], 0
        while have < k:
            a, b = free_q(kin, 2048, g), free_q(kin, 2048, g)
            m = straight_blocked(kin, a, b)
            if cat == 'C': m &= column_gap(kin, b) < 8
            got_a.append(a[m]); got_b.append(b[m]); have += int(m.sum())
        if k: emit(cat, torch.cat(got_a)[:k], torch.cat(got_b)[:k])
    if want[3]:
        a, b = free_q(kin, want[3], g), rand_q(kin, want[3], g)
        o = kin.fk(torch.zeros(1, 6, device=a.device))
        T = o + (kin.fk(b) - o) * 1.6
        emit('D', a, b, T)
    return tasks


def neighbours(t, k, g):
    """±5 mm / ±3° 邻居"""
    out = []
    for _ in range(k):
        u = torch.randn(3, generator=g); u = u / u.norm() * 5.0
        nt = dict(t); nt['T'] = [x + float(d) for x, d in zip(t['T'], u)]
        if t.get('R'):
            ax = torch.randn(3, generator=g); ax = ax / ax.norm(); ang = math.radians(3.0) * float(torch.rand(1, generator=g))
            K = torch.tensor([[0, -ax[2], ax[1]], [ax[2], 0, -ax[0]], [-ax[1], ax[0], 0]])
            Rr = torch.eye(3) + math.sin(ang) * K + (1 - math.cos(ang)) * K @ K
            nt['R'] = (Rr @ torch.tensor(t['R']).view(3, 3)).flatten().tolist()
        nt['hard'] = True
        out.append(nt)
    return out