"""策略网络（方案 §6.1）：条件特征 → 8 个条件 token → 4 层 encoder；16 个动作 token（噪声 z + 时间嵌入 (r,t)）
→ 4 层 decoder（自注意力 + 交叉注意力，非自回归一次前向）→ 每 token 6 维（平均速度场 u）。
价值头 / 子目标头只依赖观测（encoder 输出）。
注意力用显式 matmul+softmax 实现：torch.func.jvp（MeanFlow 训练）需要前向模式自动微分，SDPA 融合内核不支持。"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from ml.common.features import F_DIM, H


class Attn(nn.Module):
    def __init__(self, d, nh):
        super().__init__()
        self.nh, self.dh = nh, d // nh
        self.q, self.k, self.v, self.o = (nn.Linear(d, d) for _ in range(4))

    def forward(self, x, mem=None):
        mem = x if mem is None else mem
        B, Lq, D = x.shape; Lk = mem.shape[1]
        q = self.q(x).view(B, Lq, self.nh, self.dh).transpose(1, 2)
        k = self.k(mem).view(B, Lk, self.nh, self.dh).transpose(1, 2)
        v = self.v(mem).view(B, Lk, self.nh, self.dh).transpose(1, 2)
        a = torch.softmax(q @ k.transpose(-1, -2) / math.sqrt(self.dh), dim=-1)
        return self.o((a @ v).transpose(1, 2).reshape(B, Lq, D))


class FFN(nn.Module):
    def __init__(self, d, f):
        super().__init__(); self.a = nn.Linear(d, f); self.b = nn.Linear(f, d)

    def forward(self, x):
        return self.b(F.gelu(self.a(x)))


class EncBlock(nn.Module):
    def __init__(self, d, nh, f):
        super().__init__(); self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d); self.sa, self.ff = Attn(d, nh), FFN(d, f)

    def forward(self, x):
        x = x + self.sa(self.n1(x))
        return x + self.ff(self.n2(x))


class DecBlock(nn.Module):
    def __init__(self, d, nh, f):
        super().__init__()
        self.n1, self.n2, self.n3 = nn.LayerNorm(d), nn.LayerNorm(d), nn.LayerNorm(d)
        self.sa, self.ca, self.ff = Attn(d, nh), Attn(d, nh), FFN(d, f)

    def forward(self, x, mem):
        x = x + self.sa(self.n1(x))
        x = x + self.ca(self.n2(x), mem)
        return x + self.ff(self.n3(x))


TEMB_WMAX = 50.0     # 时间嵌入最高角频率（rad / 单位 t）


def time_emb(t, dim):
    """2026-10-18 首次正式预训练后改（step 21850 闭环到达 0%）：原实现照扩散模型写成 sin(1000·t·f)，f∈[1e-3, 1]，
    最高角频率 1000 —— t∈[0,1] 上网络对 t 极其敏感，实测 last.pt 的 ∂u/∂t 均方根 6–12，而 v 只有约 1.1，
    MeanFlow 目标 u_tgt = v − (t−r)·du/dt 被自激的 JVP 项主导：MF 残差 2 万步一直停在 7–8、17k–19k 尖峰到 96、梯度 100% 被裁剪。
    改为角频率在 [0.5, 50] 几何分布（最快一维在 [0,1] 上约转 8 圈，足够分辨 t，导数量级降 20 倍）。"""
    half = dim // 2
    w = torch.exp(torch.linspace(math.log(0.5), math.log(TEMB_WMAX), half, device=t.device, dtype=t.dtype))
    a = t[:, None] * w[None]
    return torch.cat([torch.sin(a), torch.cos(a)], -1)


class Policy(nn.Module):
    def __init__(self, d=192, nh=8, ff=768, n_enc=4, n_dec=4, n_cond=8, f_dim=F_DIM, horizon=H, head='meanflow', K=8):
        super().__init__()
        self.d, self.n_cond, self.h, self.head, self.K = d, n_cond, horizon, head, K
        self.cond = nn.Sequential(nn.Linear(f_dim, 2 * d), nn.GELU(), nn.Linear(2 * d, n_cond * d))
        self.cpos = nn.Parameter(torch.zeros(1, n_cond, d))
        self.enc = nn.ModuleList(EncBlock(d, nh, ff) for _ in range(n_enc))
        self.enc_norm = nn.LayerNorm(d)
        self.ain = nn.Linear(6, d)
        self.apos = nn.Parameter(torch.zeros(1, horizon, d))
        self.temb = nn.Sequential(nn.Linear(2 * 64, d), nn.GELU(), nn.Linear(d, d))
        self.dec = nn.ModuleList(DecBlock(d, nh, ff) for _ in range(n_dec))
        self.dec_norm = nn.LayerNorm(d)
        self.aout = nn.Linear(d, 6)
        self.value = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, 1))
        self.sub = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, 13))
        if head == 'wta':   # 预定回退：同骨干 WTA K 模式头（方案 §6.2）
            self.mode_emb = nn.Parameter(torch.randn(1, K, 1, d) * 0.02)
            self.mode_logit = nn.Linear(d, 1)
        nn.init.normal_(self.cpos, std=0.02); nn.init.normal_(self.apos, std=0.02)
        nn.init.zeros_(self.aout.weight); nn.init.zeros_(self.aout.bias)

    def encode(self, feat):
        x = self.cond(feat).view(-1, self.n_cond, self.d) + self.cpos
        for b in self.enc: x = b(x)
        return self.enc_norm(x)

    def heads(self, mem):
        p = mem.mean(1)
        return self.value(p).squeeze(-1), self.sub(p)

    def decode(self, mem, z, r, t):
        """平均速度场 u(z, r, t)；z [B,16,6]，r,t [B]"""
        te = self.temb(torch.cat([time_emb(t, 64), time_emb(t - r, 64)], -1))
        x = self.ain(z) + self.apos + te[:, None]
        for b in self.dec: x = b(x, mem)
        return self.aout(self.dec_norm(x))

    def wta_modes(self, mem):
        """WTA：K 个模式一次解码 → act [B,K,16,6], logits [B,K]"""
        B = mem.shape[0]
        x = (self.ain.bias + self.apos)[:, None] + self.mode_emb          # [1,K,16,d]
        x = x.expand(B, -1, -1, -1).reshape(B * self.K, self.h, self.d)
        m = mem.repeat_interleave(self.K, 0)
        for b in self.dec: x = b(x, m)
        x = self.dec_norm(x)
        act = self.aout(x).view(B, self.K, self.h, 6)
        return act, self.mode_logit(x.mean(1)).view(B, self.K)

    def sample(self, feat, z):
        """部署接口（ONNX 同签名）：返回 (act [B,16,6] 归一化增量, value [B], sub [B,13])
        meanflow：一步生成 x = z − u(z, 0, 1)；
        wta：取第 0 行观测的 K 个模式，第 b 行输出模式 b mod K（部署时 8 行观测相同；前瞻只用各行的 value）。"""
        mem = self.encode(feat)
        v, s = self.heads(mem)
        B = feat.shape[0]
        if self.head == 'wta':
            act, _ = self.wta_modes(mem[:1])
            idx = torch.arange(B, device=feat.device) % self.K
            return act[0, idx] + 0.0 * z, v, s     # z 不参与计算，保留为输入以统一 ONNX 签名
        one, zero = torch.ones(B, device=feat.device, dtype=feat.dtype), torch.zeros(B, device=feat.device, dtype=feat.dtype)
        x = z - self.decode(mem, z, zero, one)
        return x, v, s

    def sample_shared(self, feat1, z):
        """部署接口（ONNX，2026-10-18 第五轮，R4 实测后改）：一次观测、K 条候选。
        feat1 [1,F]，z [K,16,6] → act [K,16,6]，value [K]，sub [K,13]（value/sub 对 K 行相同，展开只为保持运行时读法）。
        此前部署把同一观测复制 K=8 行喂 sample()，encoder 被重复算 8 次（约占总计算 29%）；这里只编码一次再展开给解码器。
        2026-10-18 第六轮：feat 行数也可以等于 K（前瞻：8 个候选末态各一行特征），此时逐行编码，结果与 sample(feat, z) 一致；
        expand 对 1→K 与 K→K 都成立，所以 ONNX 里 feat 轴是动态的（meta.feat_dyn）。"""
        mem = self.encode(feat1)
        v, s = self.heads(mem)
        K = z.shape[0]
        if self.head == 'wta':
            act, _ = self.wta_modes(mem[:1])
            idx = torch.arange(K, device=z.device) % self.K
            return act[0, idx] + 0.0 * z, v.expand(K), s.expand(K, -1)
        m = mem.expand(K, -1, -1)
        one = torch.ones(K, device=z.device, dtype=z.dtype)
        x = z - self.decode(m, z, torch.zeros_like(one), one)
        return x, v.expand(K), s.expand(K, -1)

    def sample_k(self, feat, K, z=None, steps=1):
        """仿真 / APG 用：每个观测 K 条候选 → act [B,K,16,6], value [B], sub [B,13]
        steps>1（诊断用，2026-10-19）：把 [0,1] 均分成 steps 段，x ← x − (t−r)·u(x, r, t) 逐段走；steps=1 即一步生成。"""
        mem = self.encode(feat)
        v, s = self.heads(mem)
        B = feat.shape[0]
        if self.head == 'wta':
            act, logit = self.wta_modes(mem)
            if K < self.K:
                top = logit.topk(K, -1).indices
                act = act.gather(1, top[:, :, None, None].expand(-1, -1, self.h, 6))
            return act, v, s
        z = torch.randn(B * K, self.h, 6, device=feat.device, dtype=feat.dtype) if z is None else z.reshape(B * K, self.h, 6)
        m = mem.repeat_interleave(K, 0)
        one = torch.ones(B * K, device=feat.device, dtype=feat.dtype)
        x = z
        for i in range(steps):
            t, r = 1 - i / steps, 1 - (i + 1) / steps
            x = x - (t - r) * self.decode(m, x, one * r, one * t)
        return x.view(B, K, self.h, 6), v, s


def build(cfg):
    m = cfg.get('model', {})
    return Policy(d=m.get('d', 192), nh=m.get('heads', 8), ff=m.get('ff', 768), n_enc=m.get('enc', 4), n_dec=m.get('dec', 4),
                  head=m.get('head', 'meanflow'), K=m.get('K', 8))


def wta_loss(model, mem, y, w_all=0.05):
    """WTA-L1 + 模式 CE；返回 (loss, 最优模式输出 [B,16,6])"""
    act, logit = model.wta_modes(mem)
    l1 = (act - y[:, None]).abs().mean(dim=(2, 3))                       # [B,K]
    best = l1.argmin(-1)
    loss = l1.gather(1, best[:, None]).mean() + w_all * l1.mean() + 0.1 * F.cross_entropy(logit, best)
    return loss, act[torch.arange(y.shape[0]), best]