"""A trained direct predictor as an ONNX file under the benchmark's contract: python export.py NAME OUT.onnx

Reads ml/runs/NAME.pt 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 scalings of the inputs and outputs are constants of
the graph. Prints the largest difference between the graph's outputs in onnxruntime and the network's in PyTorch.
"""
from __future__ import annotations

import sys
from pathlib import Path

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

from train import Net

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


class Wrapped(nn.Module):
    def __init__(self, ck):
        super().__init__()
        self.k = int(ck["k"])
        self.net = Net(int(ck["varies"].sum()), self.k * 45 + self.k, int(ck["width"]), int(ck["depth"]))
        self.net.load_state_dict(ck["net"])
        self.register_buffer("pick", torch.nonzero(ck["varies"]).reshape(-1))
        for name in ("x_mean", "x_std", "y_mean", "y_std"):
            self.register_buffer(name, ck[name].float())

    def forward(self, boundary_r_cos, boundary_z_sin, requirement_metrics, n_field_periods, n_coils_per_half_period):
        x = torch.cat([boundary_r_cos.reshape(45), boundary_z_sin.reshape(45), requirement_metrics.log()])
        x = (x.index_select(0, self.pick) - self.x_mean) / self.x_std
        y = self.net(x.reshape(1, -1)).reshape(-1) * self.y_std + self.y_mean
        # the two shape inputs are part of the contract; the graph is for one shape, so they add nothing
        y = y + 0.0 * (n_field_periods + n_coils_per_half_period).to(y.dtype)
        coef = y[:self.k * 45].reshape(self.k, 3, 15)
        return coef[:, 0], coef[:, 1], coef[:, 2], y[self.k * 45:] * 1e6


def main(name, out):
    ck = torch.load(HERE / "runs" / f"{name}.pt", map_location="cpu")
    model = Wrapped(ck).eval()
    feed = (torch.randn(5, 9) * 0.05, torch.randn(5, 9) * 0.05, torch.tensor([0.4, 0.8, 3.0, 0.05]), torch.tensor(3), torch.tensor(5))
    feed[0][0, 4] = 1.0
    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)
    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)
    sess = ort.InferenceSession(out, providers=["CPUExecutionProvider"])
    got = sess.run(None, {n: (v.numpy().astype(np.int64) if v.dtype == torch.int64 else v.numpy()) for n, v in zip(names_in, feed)})
    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, dev feasible in training {ck.get('dev_feasible')}, "
          f"largest relative difference from PyTorch {worst:.2e}, inputs {[i.name for i in sess.get_inputs()]}")


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