# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

"""
Generate DeformConv ONNX model and test data for cross-platform validation.

Based on ONNX DeformConv spec (opset 19+): https://onnx.ai/onnx/operators/onnx__DeformConv.html
Uses a moderately complex config: groups=2, offset_group=2, 2x2 kernel, non-zero offsets.
Reference output from torchvision.ops.deform_conv2d.

Limitation: Uses symmetric padding only. PyTorch padding=(pad_h, pad_w) and ONNX pads
[pad_top, pad_left, pad_bottom, pad_right] = [pad_h, pad_w, pad_h, pad_w]. Asymmetric
pads (e.g. pad_top != pad_bottom) would require PyTorch API support and are not generated.

Run from repo root:
  python onnxruntime/test/testdata/nn/deform_conv_test_gen.py

Outputs:
  - deform_conv_test.onnx
  - deform_conv_test_data.npz (X, W, offset, B, mask, expected_Y)
  - deform_conv_test_data.inc (C++ arrays for op test)
"""

from pathlib import Path

import numpy as np
from onnx import TensorProto, checker, helper, save

try:
    import onnxruntime as ort
except ImportError:
    ort = None
import torch
import torchvision.ops

# Config: groups=2, offset_group=2, 2x2 kernel (from deform_conv_expected_gen Case 3)
BATCH = 1
N_IN = 4
N_OUT = 2
N_WEIGHT_GRPS = 2
N_OFFSET_GRPS = 2
KH, KW = 2, 2
STRIDE_H, STRIDE_W = 1, 1
PAD_H, PAD_W = 0, 0
DIL_H, DIL_W = 1, 1
IN_H, IN_W = 3, 3
SEED = 123

OUT_H = (IN_H + 2 * PAD_H - (DIL_H * (KH - 1) + 1)) // STRIDE_H + 1
OUT_W = (IN_W + 2 * PAD_W - (DIL_W * (KW - 1) + 1)) // STRIDE_W + 1


def _generate_reference():
    """Generate inputs and expected output via torchvision.ops.deform_conv2d."""
    torch.manual_seed(SEED)
    x = torch.rand(BATCH, N_IN, IN_H, IN_W, dtype=torch.float32)
    offset = torch.randn(BATCH, N_OFFSET_GRPS * 2 * KH * KW, OUT_H, OUT_W, dtype=torch.float32)
    mask = torch.randn(BATCH, N_OFFSET_GRPS * KH * KW, OUT_H, OUT_W, dtype=torch.float32)
    weight = torch.randn(N_OUT, N_IN // N_WEIGHT_GRPS, KH, KW, dtype=torch.float32)
    bias = torch.randn(N_OUT, dtype=torch.float32)

    out = torchvision.ops.deform_conv2d(
        x,
        offset,
        weight,
        bias=bias,
        stride=(STRIDE_H, STRIDE_W),
        padding=(PAD_H, PAD_W),
        dilation=(DIL_H, DIL_W),
        mask=mask,
    )

    return {
        "X": x.numpy(),
        "W": weight.numpy(),
        "offset": offset.numpy(),
        "B": bias.numpy(),
        "mask": mask.numpy(),
        "expected_Y": out.numpy(),
    }


def _build_onnx_model():
    """Build DeformConv ONNX model. ONNX pads = [pad_top, pad_left, pad_bottom, pad_right]."""
    # Symmetric padding only: (pad_h, pad_w) -> [pad_h, pad_w, pad_h, pad_w]
    pads = [PAD_H, PAD_W, PAD_H, PAD_W]

    node = helper.make_node(
        "DeformConv",
        inputs=["X", "W", "offset", "B", "mask"],
        outputs=["Y"],
        kernel_shape=[KH, KW],
        strides=[STRIDE_H, STRIDE_W],
        pads=pads,
        dilations=[DIL_H, DIL_W],
        group=N_WEIGHT_GRPS,
        offset_group=N_OFFSET_GRPS,
    )

    graph = helper.make_graph(
        [node],
        "DeformConvTest",
        [
            helper.make_tensor_value_info("X", TensorProto.FLOAT, [BATCH, N_IN, IN_H, IN_W]),
            helper.make_tensor_value_info("W", TensorProto.FLOAT, [N_OUT, N_IN // N_WEIGHT_GRPS, KH, KW]),
            helper.make_tensor_value_info(
                "offset", TensorProto.FLOAT, [BATCH, N_OFFSET_GRPS * 2 * KH * KW, OUT_H, OUT_W]
            ),
            helper.make_tensor_value_info("B", TensorProto.FLOAT, [N_OUT]),
            helper.make_tensor_value_info("mask", TensorProto.FLOAT, [BATCH, N_OFFSET_GRPS * KH * KW, OUT_H, OUT_W]),
        ],
        [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [BATCH, N_OUT, OUT_H, OUT_W])],
    )

    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 19)])
    checker.check_model(model)
    return model


def _to_cpp_array(name: str, arr: np.ndarray) -> str:
    """Format numpy array as C++ initializer list."""
    flat = arr.flatten().tolist()
    vals = ", ".join(f"{x:.9g}f" for x in flat)
    return f"static const std::vector<float> {name} = {{{vals}}};"


def _write_cpp_inc(data: dict, inc_path: Path) -> None:
    """Write C++ include file with test data."""
    lines = [
        "// Auto-generated by deform_conv_test_gen.py - do not edit",
        "",
        "#include <vector>",
        "",
        _to_cpp_array("kDeformConvOnnxTest_X", data["X"]),
        _to_cpp_array("kDeformConvOnnxTest_W", data["W"]),
        _to_cpp_array("kDeformConvOnnxTest_offset", data["offset"]),
        _to_cpp_array("kDeformConvOnnxTest_B", data["B"]),
        _to_cpp_array("kDeformConvOnnxTest_mask", data["mask"]),
        _to_cpp_array("kDeformConvOnnxTest_expected_Y", data["expected_Y"]),
        "",
    ]
    inc_path.write_text("\n".join(lines), encoding="utf-8")


def main():
    # Output to testdata/ root (same as layernorm.onnx, attention_past_state.onnx, etc.)
    script_dir = Path(__file__).resolve().parent
    assert script_dir.name == "nn", "Script must live in testdata/nn/"
    testdata_root = script_dir.parent
    model_path = testdata_root / "deform_conv_test.onnx"
    data_path = testdata_root / "deform_conv_test_data.npz"
    inc_path = testdata_root / "deform_conv_test_data.inc"

    print("Generating reference via torchvision.ops.deform_conv2d...")
    data = _generate_reference()

    print("Building ONNX model...")
    model = _build_onnx_model()
    save(model, str(model_path))
    print(f"  Saved {model_path}")

    np.savez(str(data_path), **data)
    print(f"  Saved {data_path}")

    _write_cpp_inc(data, inc_path)
    print(f"  Saved {inc_path}")

    # Validate with onnxruntime if available
    if ort is not None:
        print("Validating with ONNX Runtime...")
        sess = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"])
        ort_out = sess.run(
            ["Y"],
            {
                "X": data["X"],
                "W": data["W"],
                "offset": data["offset"],
                "B": data["B"],
                "mask": data["mask"],
            },
        )[0]

        rtol, atol = 1e-4, 1e-4
        if np.allclose(ort_out, data["expected_Y"], rtol=rtol, atol=atol):
            print("  PASS: ORT output matches reference.")
        else:
            diff = np.abs(ort_out.astype(np.float64) - data["expected_Y"].astype(np.float64))
            print(f"  FAIL: max |diff|={diff.max()}, mean={diff.mean()}")
    else:
        print("  (onnxruntime not installed; skip validation)")


if __name__ == "__main__":
    main()
