"""目标公式的不动点判决（解析高斯场，float64，CPU，与任何已训模型无关）。

1-D：数据 x~N(0,σ²)、噪声 ε~N(0,1)、路径 z_t=(1−t)x+tε ⇒
  s²(t)=(1−t)²σ²+t²，v(z,t)=E[ε−x|z]=z·(t−(1−t)σ²)/s²(t)。
  因为 a(t)=v/z=(s²)′/(2s²)，PF-ODE 的流图有闭式 z_r = z_t·sqrt(s²(r)/s²(t)) ⇒ 平均速度真值也是闭式。
把"精确场"代入四种目标写法，看谁能收敛回真值。判据用**收敛阶**而不是绝对值：
  有结构性偏差的写法，区间减半时误差只能减半（线性）；没有偏差的应当减到 1/4（二阶）。
"""
import unittest

import torch
from torch.func import jvp

LAM = 0.5


class Field:
    def __init__(self, sig):
        self.sig2 = sig ** 2

    def s2(self, t):
        return (1 - t) ** 2 * self.sig2 + t ** 2

    def v(self, z, t):
        return z * (t - (1 - t) * self.sig2) / self.s2(t)

    def u(self, z, r, t):
        """平均速度真值：状态在时刻 t 的 z，要走到时刻 r"""
        return (z - z * torch.sqrt(self.s2(r) / self.s2(t))) / (t - r)

    def total_du(self, z, r, t, tv):
        """du(z,r,t) 沿切向量 tv 的全导数 = ∂_z u·tv + ∂_t u（r 槽固定）"""
        zz = z.clone().requires_grad_(True)
        tt = t.clone().requires_grad_(True)
        rr = r.clone().requires_grad_(False)
        gz, gt = torch.autograd.grad(self.u(zz, rr, tt), (zz, tt), grad_outputs=torch.ones((), dtype=z.dtype))
        return gz * tv + gt

    def targets(self, z, r, t, lam=LAM):
        r = torch.as_tensor(r, dtype=torch.float64); t = torch.as_tensor(t, dtype=torch.float64)
        z = torch.as_tensor(z, dtype=torch.float64)
        v = self.v(z, t)
        c = (1 - lam) * t + lam * r
        xh = z * (1 - t) * self.sig2 / self.s2(t)      # E[x|z_t]，eh = E[ε|z_t]；耦合速度 eh−xh = v
        eh = z * t / self.s2(t)
        z_c = (1 - c) * xh + c * eh
        d_mf = self.total_du(z, r, t, v)
        d1 = self.total_du(z, c, t, v)                  # 方案里的第一项：r 槽挪到 c，对 t 求导
        d2 = self.total_du(z_c, r, c, v)                # 第二项：状态与第二个时刻槽都在中点
        span = t - r
        return {'mf': v - span * d_mf,
                'sq': v - span * ((1 - lam) ** 2 * d1 + lam ** 2 * d2),
                'lin': v - span * ((1 - lam) * d1 + lam * d2),
                'mid': v - span * d2}

    def abs_rel(self, z, r, t):
        """返回 (相对真值的偏差, 真值本身)。σ=1 且 (r,t)=(0,1) 时 s²(0)=s²(1) ⇒ 流图是恒等、真值为 0，
        那是个退化点，只能用绝对偏差，别拿它算相对值。"""
        z_ = torch.as_tensor(z, dtype=torch.float64); r_ = torch.as_tensor(r, dtype=torch.float64); t_ = torch.as_tensor(t, dtype=torch.float64)
        u = self.u(z_, r_, t_)
        den = u.abs() if float(u.abs()) > 1e-9 else torch.tensor(1.0, dtype=torch.float64)
        return {k: float((val - u).abs() / den) for k, val in self.targets(z, r, t).items()}, float(u)


    def abs_err(self, z, r, t):
        """绝对偏差（除以固定尺度 |v(z,t)|，不除以会趋 0 的 |u_true|）"""
        z_ = torch.as_tensor(z, dtype=torch.float64); r_ = torch.as_tensor(r, dtype=torch.float64); t_ = torch.as_tensor(t, dtype=torch.float64)
        u = self.u(z_, r_, t_); sc = self.v(z_, t_).abs()
        return {k: float((val - u).abs() / sc) for k, val in self.targets(z, r, t).items()}


class MeanFlowTargetFixedPointTest(unittest.TestCase):
    def test_mf_identity_is_exact(self):
        """先证仪表：把精确场代 Eq.4 必须逐位还原平均速度，否则后面的数都白读（用绝对偏差，退化点也能过）"""
        for sig in (0.5, 1.0, 2.0):
            for (r, t) in ((0.0, 1.0), (0.0, 0.3), (0.2, 0.9)):
                z_ = torch.tensor(0.7, dtype=torch.float64)
                f = Field(sig)
                tgt = f.targets(0.7, r, t)
                with self.subTest(sig=sig, r=r, t=t):
                    self.assertLess(float((tgt['mf'] - f.u(z_, torch.tensor(r, dtype=torch.float64),
                                                          torch.tensor(t, dtype=torch.float64))).abs()), 1e-12)

    def test_convergence_order_separates_biased_from_consistent(self):
        """区间减半：sq 只能线性收敛（系数和 0.5 ⇒ 结构性少修一半），lin/mid 应当二次收敛"""
        for sig in (1.0, 2.0):
            f = Field(sig)
            rel = [f.abs_rel(0.7, 1.0 - s, 1.0)[0] for s in (0.2, 0.1, 0.05)]
            print('σ=%.1f 相对偏差 span=0.2/0.1/0.05：' % sig,
                  [{k: round(v, 4) for k, v in d.items()} for d in rel])
            for name, want in (('sq', 0.5), ('lin', 0.25), ('mid', 0.25)):
                ratios = [rel[1][name] / rel[0][name], rel[2][name] / rel[1][name]]
                with self.subTest(sig=sig, name=name):
                    for rt in ratios:
                        self.assertLess(rt, want * 1.5)
                        self.assertGreater(rt, want * 0.5)

    def test_full_span_step_is_unreliable_for_every_variant(self):
        """跨整段 (0,1) 时除 Eq.4 外都偏得很厉害；跨半段就明显变小 ⇒ "一步跨完"本身是难的，
        不是训练出来的（σ=2，偏差按 |v| 归一，避免除以会趋 0 的 |u_true|）"""
        f = Field(2.0)
        full = f.abs_err(0.7, 0.0, 1.0)
        half = f.abs_err(0.7, 0.5, 1.0)
        print('σ=2 整段 (0,1) 偏差：', {k: round(v, 3) for k, v in full.items()})
        print('σ=2 半段 (0.5,1) 偏差：', {k: round(v, 3) for k, v in half.items()})
        self.assertLess(full['mf'], 1e-12)
        self.assertGreater(full['sq'], 0.5)
        self.assertLess(half['sq'], full['sq'])
        self.assertLess(half['mid'], full['mid'])


class AvgVelocityTargetWiringTest(unittest.TestCase):
    """把三种写法接进 ml/model/meanflow.py 之后，再用同一个解析场验一遍接线：
    浮点从 float64 换到 float32，量级判据放宽，但"mf 无偏 / sq 只线性收敛"必须照样成立，
    否则说明 z_c、切向量、r 槽换位这几处接错了。"""
    SIG = 2.0

    class Stub:
        def __init__(self, sig2):
            self.sig2 = sig2

        def s2(self, t):
            return (1 - t) ** 2 * self.sig2 + t ** 2

        def decode(self, mem, z, r, t):
            return z * (1 - torch.sqrt(self.s2(r) / self.s2(t)))[:, None, None] / (t - r)[:, None, None]

    def bias(self, target, r, t, lam=0.5):
        from ml.model.meanflow import avg_velocity_target
        f = Field(self.SIG)
        z = torch.tensor([[0.7] * 4, [0.7] * 4], dtype=torch.float32)          # [2,4,6]? 形状自定，逐元素即可
        z = torch.full((2, 4, 6), 0.7)
        rt = torch.tensor([r, r]); tt = torch.tensor([t, t])
        xh = (z * (1 - tt)[:, None, None] * self.SIG ** 2 / f.s2(tt)[:, None, None])
        eh = (z * tt[:, None, None] / f.s2(tt)[:, None, None])
        u, tgt, _ = avg_velocity_target(self.Stub(self.SIG ** 2), None, xh, eh, rt, tt, target, lam)
        truth = f.u(torch.tensor(0.7, dtype=torch.float32), torch.tensor(r), torch.tensor(t))
        return float(((tgt - truth).abs() / truth.abs()).mean()), float(truth)

    def test_wiring_keeps_mf_unbiased_and_sq_first_order(self):
        b_mf, _ = self.bias('mf', 0.8, 1.0)
        self.assertLess(b_mf, 1e-4)
        rel = []
        for s in (0.2, 0.1, 0.05):
            rel.append({k: self.bias(k, 1.0 - s, 1.0)[0] for k in ('sq', 'lin', 'mid', 'mf')})
        print('float32 接线后的偏差：', [{k: round(v, 4) for k, v in d.items()} for d in rel])
        r_sq = rel[1]['sq'] / rel[0]['sq']
        r_mid = rel[1]['mid'] / rel[0]['mid']
        self.assertGreater(r_sq, 0.3)          # 只线性收敛 ⇒ 结构性偏差
        self.assertLess(r_mid, 0.35)           # 二次收敛 ⇒ 只剩积分误差
        self.assertLess(rel[0]['lin'], rel[0]['sq'])

    def test_default_path_is_bitwise_identical(self):
        """target='mf' 必须与重构前的原地写法逐位相同（默认训练路径不许被这次改动碰到）"""
        from ml.model.meanflow import avg_velocity_target
        torch.manual_seed(0)
        m = self.Stub(2.0)
        x = torch.randn(3, 5, 6); e = torch.randn(3, 5, 6)
        r = torch.rand(3) * 0.4; t = 0.6 + torch.rand(3) * 0.4
        u_new, tgt_new, z_new = avg_velocity_target(m, None, x, e, r, t, 'mf', 0.5)
        fn = lambda a, b, c: m.decode(None, a, b, c)
        z = (1 - t)[:, None, None] * x + t[:, None, None] * e
        v = e - x
        u0, dudt = jvp(fn, (z, r, t), (v, torch.zeros_like(r), torch.ones_like(t)))
        tgt0 = (v - (t - r)[:, None, None] * dudt).detach()
        for name, a_, b_ in (('u', u_new, u0), ('u_tgt', tgt_new, tgt0), ('z', z_new, z)):
            with self.subTest(name=name):
                torch.testing.assert_close(a_, b_, rtol=0, atol=0)


if __name__ == '__main__':
    unittest.main()