"""目标公式的不动点判决（解析高斯场，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

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'])


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