"""A direct predictor with its correcting network as an ONNX file under the benchmark's contract:
python export_fb.py NAME OUT.onnx [--check N]

Reads ml/runs/NAME.pt (train_fb.py) and its base network, and writes a graph from boundary_r_cos, boundary_z_sin
(5, 9), requirement_metrics (4), n_field_periods and n_coils_per_half_period to coil_x_n, coil_y_n, coil_z_n (K, 15)
and currents (K) in amperes, with the metadata key supported_tracks set to fixed_shape. The graph runs the base
network, computes the feedback of train_fb.py for its coilset with tensor operations (feedback() below, the same
quantities as fast.feedback), and applies the corrector, once per round. Prints the largest difference between the
graph's outputs in onnxruntime and the networks' in PyTorch, and with --check N the largest difference between
feedback() and fast.feedback on N dev instances.
"""
from __future__ import annotations

import sys
import time
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
import torch
from torch import nn

from physics import Physics, coil_basis
from train import Net
from train_fb import Corrector

HERE = Path(__file__).resolve().parent


def cross(a, b):
    return torch.stack([a[..., 1] * b[..., 2] - a[..., 2] * b[..., 1], a[..., 2] * b[..., 0] - a[..., 0] * b[..., 2],
                        a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0]], -1)


class Feedback(nn.Module):
    """fast.feedback for one coilset, with the pair sums as tensors: constants are the grids' angles, the coil
    modes at 64 and 128 parameters and the symmetry maps."""

    def __init__(self, nfp, k):
        super().__init__()
        ph = Physics(nfp, k, "cpu")
        self.k, self.J = k, ph.J
        buf = lambda name, v: self.register_buffer(name, torch.as_tensor(np.asarray(v), dtype=torch.float32))      # noqa: E731
        for name, (th, p) in (("f", ph.grid(32, int(round(3.2 * k)), 0.5, 0.5)), ("s", ph.grid(32, 16 * ph.J, 0.5, 0.5, full=True)),
                              ("a", (ph.grid(64, 48)[0], ph.grid(64, 48)[1] * 2.0))):
            ang = th[:, None] * ph.m[None] - p[:, None] * ph.n[None]
            buf(name + "_cos", torch.cos(ang).T)
            buf(name + "_sin", torch.sin(ang).T)
            buf(name + "_cp", torch.cos(p))
            buf(name + "_sp", torch.sin(p))
        buf("m", ph.m)
        buf("n", ph.n)
        b64, b128 = coil_basis(64), coil_basis(128)
        buf("b0", b64[0])
        buf("b1", b64[1] * (2 * np.pi / 64))
        buf("c1", b128[1])
        buf("c2", b128[2])
        buf("maps", ph.maps)
        buf("sign", ph.sign)
        mine = torch.arange(k).repeat_interleave(64)                                             # the coil of each point of the unique coils
        theirs = torch.cat([mine, torch.full(((ph.J - 1) * k * 64,), -1)])                         # and of each point of the full coilset's first image
        self.register_buffer("mine", mine.float())
        self.register_buffer("theirs", theirs.float())

    def points(self, rc, zs, name):
        c, s, cp, sp = (getattr(self, f"{name}_{v}") for v in ("cos", "sin", "cp", "sp"))
        R, Z = rc @ c, zs @ s
        return R, Z, cp, sp

    def forward(self, coef, cur, rc, zs, target):
        """coef (K, 3, 15), cur (K,) in megaamperes, rc and zs (45,), target (4,): the field feedback (X,) and the
        slacks (3 K,)."""
        k = self.k
        # the minor radius
        R, Z, _, _ = self.points(rc, zs, "a")
        R, Z = R.reshape(48, 64), Z.reshape(48, 64)
        area = 0.5 * (R * torch.roll(Z, -1, 1) - torch.roll(R, -1, 1) * Z).sum(1).abs()
        a = (area.mean() / np.pi).sqrt()
        # the field at the fixed grid
        R, Z, cp, sp = self.points(rc, zs, "f")
        Rt, Zt = -(rc * self.m) @ self.f_sin, (zs * self.m) @ self.f_cos
        Rp, Zp = (rc * self.n) @ self.f_sin, -(zs * self.n) @ self.f_cos
        x = torch.stack([R * cp, R * sp, Z], -1)
        nrm = cross(torch.stack([Rt * cp, Rt * sp, Zt], -1), torch.stack([Rp * cp - R * sp, Rp * sp + R * cp, Zp], -1))
        nrm = nrm / (nrm * nrm).sum(-1, keepdim=True).sqrt()
        r, e = (coef @ self.b0).transpose(1, 2), (coef @ self.b1).transpose(1, 2)                  # (K, P, 3)
        rj = torch.einsum("jac,kpc->jkpa", self.maps, r)
        ej = torch.einsum("jac,kpc->jkpa", self.maps, e)
        w = (self.sign[:, None] * cur[None, :])[..., None, None]
        we, rq = (w * ej).reshape(-1, 3), rj.reshape(-1, 3)
        wc = cross(we, rq)
        d2 = (x * x).sum(-1)[:, None] + (rq * rq).sum(-1)[None, :] - 2.0 * x @ rq.T
        inv = d2.clamp(min=1e-10).rsqrt()
        inv3 = inv * inv * inv
        field = cross(inv3 @ we, x) - inv3 @ wc
        signed = (field * nrm).sum(-1) / (field * field).sum(-1).sqrt().clamp(min=1e-30) / target[3]
        # the slacks per coil
        u = r.reshape(-1, 3)
        d2 = (u * u).sum(-1)[:, None] + (rq * rq).sum(-1)[None, :] - 2.0 * u @ rq.T + 1e6 * (self.mine[:, None] == self.theirs[None, :]).float()
        cc = d2.reshape(k, -1).amin(1).clamp(min=0).sqrt() / (a * target[0])
        R, Z, cp, sp = self.points(rc, zs, "s")
        surf = torch.stack([R * cp, R * sp, Z], -1)
        d2 = (u * u).sum(-1)[:, None] + (surf * surf).sum(-1)[None, :] - 2.0 * u @ surf.T
        cp_ = d2.reshape(k, -1).amin(1).clamp(min=0).sqrt() / (a * target[1])
        r1, r2 = (coef @ self.c1).transpose(1, 2), (coef @ self.c2).transpose(1, 2)
        kap = (cross(r1, r2) ** 2).sum(-1).sqrt() / (r1 * r1).sum(-1).sqrt().clamp(min=1e-12) ** 3
        return signed, torch.cat([cc, cp_, kap.amax(1) * (a / target[2])]) - 1.0


class Wrapped(nn.Module):
    def __init__(self, ck, base):
        super().__init__()
        self.k, self.rounds = int(base["k"]), int(ck["rounds"])
        n_x, n_y = int(base["varies"].sum()), self.k * 46
        self.net = Net(n_x, n_y, int(base["width"]), int(base["depth"]))
        self.net.load_state_dict(base["net"])
        self.fix = Corrector(n_x, n_y, 32 * int(round(3.2 * self.k)), 3 * self.k, int(ck["width"]), int(ck["depth"]))
        self.fix.load_state_dict(ck["fix"])
        self.feedback = Feedback(int(base["nfp"]), self.k)
        self.register_buffer("pick", torch.nonzero(base["varies"]).reshape(-1))
        for name in ("x_mean", "x_std", "y_mean", "y_std"):
            self.register_buffer(name, base[name].float())

    def forward(self, boundary_r_cos, boundary_z_sin, requirement_metrics, n_field_periods, n_coils_per_half_period):
        rc, zs = boundary_r_cos.reshape(45), boundary_z_sin.reshape(45)
        x = torch.cat([rc, zs, requirement_metrics.log()])
        x = ((x.index_select(0, self.pick) - self.x_mean) / self.x_std).reshape(1, -1)
        y = self.net(x)
        for _ in range(self.rounds):
            full = y.reshape(-1) * self.y_std + self.y_mean
            field, slack = self.feedback(full[:self.k * 45].reshape(self.k, 3, 15), full[self.k * 45:], rc, zs, requirement_metrics)
            y = y + self.fix(x, y, field.reshape(1, -1), slack.reshape(1, -1))
        full = y.reshape(-1) * self.y_std + self.y_mean
        full = full + 0.0 * (n_field_periods + n_coils_per_half_period).to(full.dtype)          # the shape inputs are part of the contract
        coef = full[:self.k * 45].reshape(self.k, 3, 15)
        return coef[:, 0], coef[:, 1], coef[:, 2], full[self.k * 45:] * 1e6


def main(name, out, *flags):
    ck = torch.load(HERE / "runs" / f"{name}.pt", map_location="cpu")
    base = torch.load(HERE / "runs" / f"{ck['base']}.pt", map_location="cpu")
    model = Wrapped(ck, base).eval()
    if "--check" in flags:
        import fast
        from train import Data
        n = int(flags[flags.index("--check") + 1])
        data, ph = Data(int(base["nfp"]), model.k, "cuda"), Physics(int(base["nfp"]), model.k, "cuda")
        rows = data.dev_rows[torch.linspace(0, len(data.dev_rows) - 1, n, device="cuda").long()]
        keys, target = data.keys[data.boundary[rows]], data.target[rows]
        coef, cur = data.coef[rows].reshape(-1, model.k, 3, 15), data.cur[rows]
        f, s = fast.feedback(ph, coef, cur, keys[:, :45], keys[:, 45:], target, ph.minor_radius(keys[:, :45], keys[:, 45:]))
        worst = [0.0, 0.0]
        with torch.no_grad():
            for i in range(n):
                g = model.feedback(coef[i].cpu(), cur[i].cpu(), keys[i, :45].cpu(), keys[i, 45:].cpu(), target[i].cpu())
                worst = [max(worst[0], float((g[0] - f[i].cpu()).abs().max())), max(worst[1], float((g[1] - s[i].cpu()).abs().max()))]
        print(f"feedback against the kernels on {n} dev instances: largest difference in the field {worst[0]:.2e}, in the slacks {worst[1]:.2e}")
    feed = (torch.randn(5, 9) * 0.02, torch.randn(5, 9) * 0.02, torch.tensor([0.4, 0.8, 3.0, 0.05]), torch.tensor(3), torch.tensor(5))
    feed[0][0, 4], feed[0][1, 4], feed[1][1, 4] = 1.0, 0.15, 0.15
    names_in = ["boundary_r_cos", "boundary_z_sin", "requirement_metrics", "n_field_periods", "n_coils_per_half_period"]
    names_out = ["coil_x_n", "coil_y_n", "coil_z_n", "currents"]
    torch.onnx.export(model, feed, out, input_names=names_in, output_names=names_out, opset_version=17, dynamo=False, do_constant_folding=False)
    proto = onnx.load(out)
    proto.metadata_props.append(onnx.StringStringEntryProto(key="supported_tracks", value="fixed_shape"))
    onnx.save(proto, out)
    onnx.checker.check_model(out)
    opts = ort.SessionOptions()
    opts.intra_op_num_threads = 1
    sess = ort.InferenceSession(out, opts, providers=["CPUExecutionProvider"])
    inputs = {n: (v.numpy().astype(np.int64) if v.dtype == torch.int64 else v.numpy()) for n, v in zip(names_in, feed)}
    got = sess.run(None, inputs)
    t0 = time.perf_counter()
    for _ in range(20):
        sess.run(None, inputs)
    took = (time.perf_counter() - t0) / 20
    with torch.no_grad():
        want = model(*feed)
    worst = max(float(np.abs(g - w.numpy()).max() / max(1.0, float(w.abs().max()))) for g, w in zip(got, want))
    print(f"wrote {out}: {Path(out).stat().st_size / 2 ** 20:.1f} MiB, {model.rounds} round(s), dev feasible in training {ck.get('dev_feasible')}, "
          f"largest relative difference from PyTorch {worst:.2e}, {1e3 * took:.1f} ms a run on one thread")


if __name__ == "__main__":
    main(*sys.argv[1:])