from __future__ import annotations

import argparse
from pathlib import Path

import onnx
from onnx import TensorProto, helper

from custom_ops import (
    mesh_topology_extract,
    model,
    scalar,
    sparse_conv3d,
    sparse_downsample,
    sparse_subdivide,
    sparse_upsample,
    sparse_window_attention,
    vi,
)


def save(m, path: Path):
    path.parent.mkdir(parents=True, exist_ok=True)
    onnx.checker.check_model(m, full_check=False)
    onnx.save(m, path)
    print(path)


def build_all(out: Path):
    g = helper.make_graph(
        [sparse_subdivide(["feats", "coords"], ["out_feats", "out_coords"])],
        "SparseSubdivide",
        [vi("feats", TensorProto.FLOAT16, ["N", "C"]), vi("coords", TensorProto.INT32, ["N", 4])],
        [vi("out_feats", TensorProto.FLOAT16, ["N8", "C"]), vi("out_coords", TensorProto.INT32, ["N8", 4])],
    )
    save(model(g), out / "SparseSubdivide.onnx")

    g = helper.make_graph(
        [sparse_downsample(["feats", "coords"], ["out_feats", "out_coords", "inverse", "count"], factor_x=2, factor_y=2, factor_z=2)],
        "SparseDownsample",
        [vi("feats", TensorProto.FLOAT16, ["N", "C"]), vi("coords", TensorProto.INT32, ["N", 4])],
        [vi("out_feats", TensorProto.FLOAT16, ["M", "C"]), vi("out_coords", TensorProto.INT32, ["M", 4]), vi("inverse", TensorProto.INT32, ["N"]), scalar("count")],
    )
    save(model(g), out / "SparseDownsample.onnx")

    g = helper.make_graph(
        [sparse_upsample(["feats", "target_coords", "inverse"], ["out_feats", "out_coords"])],
        "SparseUpsample",
        [vi("feats", TensorProto.FLOAT16, ["M", "C"]), vi("target_coords", TensorProto.INT32, ["N", 4]), vi("inverse", TensorProto.INT32, ["N"])],
        [vi("out_feats", TensorProto.FLOAT16, ["N", "C"]), vi("out_coords", TensorProto.INT32, ["N", 4])],
    )
    save(model(g), out / "SparseUpsample.onnx")

    g = helper.make_graph(
        [sparse_window_attention(["qkv", "coords"], ["out"], window_size=8, shift_x=0, shift_y=0, shift_z=0)],
        "SparseWindowAttention",
        [vi("qkv", TensorProto.FLOAT16, ["N", 3, "H", "D"]), vi("coords", TensorProto.INT32, ["N", 4])],
        [vi("out", TensorProto.FLOAT16, ["N", "H", "D"])],
    )
    save(model(g), out / "SparseWindowAttention.onnx")

    g = helper.make_graph(
        [sparse_conv3d(["feats", "coords", "weight", "bias"], ["out_feats", "out_coords", "count"], out_channels=16, kernel_size=3, stride=1, dilation=1, padding=0, subm=True, spatial_x=64, spatial_y=64, spatial_z=64, batch_size=1)],
        "SparseConv3D",
        [vi("feats", TensorProto.FLOAT16, ["N", 8]), vi("coords", TensorProto.INT32, ["N", 4]), vi("weight", TensorProto.FLOAT16, [16, 3, 3, 3, 8]), vi("bias", TensorProto.FLOAT16, [16])],
        [vi("out_feats", TensorProto.FLOAT16, ["M", 16]), vi("out_coords", TensorProto.INT32, ["M", 4]), scalar("count")],
    )
    save(model(g), out / "SparseConv3D.onnx")

    g = helper.make_graph(
        [mesh_topology_extract(["verts_grid", "sdf", "cube_idx", "beta", "alpha", "gamma", "colors_grid"], ["vertices", "faces", "colors", "vertex_count", "face_count"], resolution=64, no_sigmoid=True)],
        "MeshTopologyExtract",
        [
            vi("verts_grid", TensorProto.FLOAT, ["Vg", 3]),
            vi("sdf", TensorProto.FLOAT, ["Vg"]),
            vi("cube_idx", TensorProto.INT64, ["Nc", 8]),
            vi("beta", TensorProto.FLOAT, ["Nc", 12]),
            vi("alpha", TensorProto.FLOAT, ["Nc", 8]),
            vi("gamma", TensorProto.FLOAT, ["Nc"]),
            vi("colors_grid", TensorProto.FLOAT, ["Vg", 6]),
        ],
        [vi("vertices", TensorProto.FLOAT, ["Nv", 3]), vi("faces", TensorProto.INT32, ["Nf", 3]), vi("colors", TensorProto.FLOAT, ["Nv", 6]), scalar("vertex_count"), scalar("face_count")],
    )
    save(model(g), out / "MeshTopologyExtract.onnx")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="custom-onnx")
    args = ap.parse_args()
    build_all(Path(args.out))