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

"""
Generate expected outputs for DeformConv tests using torchvision.ops.deform_conv2d.
Run with: .venv/bin/python onnxruntime/test/providers/cpu/nn/deform_conv_expected_gen.py
Outputs C++-friendly std::vector<float> initializer lists for pasting into deform_conv_op_test.cc

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

import torch
import torchvision.ops


def _pair(x: int | tuple[int, int]) -> tuple[int, int]:
    if isinstance(x, int):
        return (x, x)
    return x


def to_cpp_list(t: torch.Tensor, fmt="{:.6f}") -> str:
    """Flatten tensor in NCHW order and format as C++ initializer list."""
    t = t.detach().float().contiguous()
    return ", ".join(fmt.format(x) + "f" for x in t.flatten().tolist())


def run_case(
    name: str,
    batch_sz: int,
    n_in: int,
    n_out: int,
    n_weight_grps: int,
    n_offset_grps: int,
    kernel_h: int,
    kernel_w: int,
    stride: tuple[int, int] | int,
    pad: tuple[int, int] | int,
    dilation: tuple[int, int] | int,
    in_h: int,
    in_w: int,
    seed: int = 42,
):
    """Build inputs with seed, run deform_conv2d, print C++ snippets."""
    torch.manual_seed(seed)
    stride_h, stride_w = _pair(stride)
    pad_h, pad_w = _pair(pad)
    dil_h, dil_w = _pair(dilation)

    out_h = (in_h + 2 * pad_h - (dil_h * (kernel_h - 1) + 1)) // stride_h + 1
    out_w = (in_w + 2 * pad_w - (dil_w * (kernel_w - 1) + 1)) // stride_w + 1

    x = torch.rand(batch_sz, n_in, in_h, in_w, dtype=torch.float32)
    offset = torch.randn(batch_sz, n_offset_grps * 2 * kernel_h * kernel_w, out_h, out_w, dtype=torch.float32)
    mask = torch.randn(batch_sz, n_offset_grps * kernel_h * kernel_w, out_h, out_w, dtype=torch.float32)
    weight = torch.randn(n_out, n_in // n_weight_grps, kernel_h, kernel_w, dtype=torch.float32)
    bias = torch.randn(n_out, dtype=torch.float32)

    # Standard answer from torchvision
    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,
    )

    # ONNX pads = [top, left, bottom, right] (symmetric: single pad_h, pad_w expanded)
    pads_onnx = [pad_h, pad_w, pad_h, pad_w]

    print(f"// --- {name} (seed={seed}) ---")
    print(f"// Shapes: X({batch_sz},{n_in},{in_h},{in_w}) W({n_out},{n_in // n_weight_grps},{kernel_h},{kernel_w})")
    print(f"// stride=({stride_h},{stride_w}) pad=({pad_h},{pad_w}) dilation=({dil_h},{dil_w})")
    print(f"// out_h={out_h} out_w={out_w}")
    print()
    print("std::vector<float> X = {" + to_cpp_list(x) + "};")
    print("std::vector<float> W = {" + to_cpp_list(weight) + "};")
    print("std::vector<float> offset = {" + to_cpp_list(offset) + "};")
    print("std::vector<float> B = {" + to_cpp_list(bias) + "};")
    print("std::vector<float> mask = {" + to_cpp_list(mask) + "};")
    print("std::vector<float> expected_Y = {" + to_cpp_list(out) + "};")
    print()
    print(
        "// Params: kernel_shape={"
        f"{kernel_h}, {kernel_w}"
        "}, stride={"
        f"{stride_h}, {stride_w}"
        "}, pads={"
        + ", ".join(map(str, pads_onnx))
        + "}, dilations={"
        + f"{dil_h}, {dil_w}"
        + "}, group="
        + str(n_weight_grps)
        + ", offset_group="
        + str(n_offset_grps)
    )
    print()
    return out


def main():
    print("// Generated by deform_conv_expected_gen.py (torchvision.ops.deform_conv2d)")
    print()

    # Case 1: Same config as PyTorch TestDeformConv.get_fn_args (small batch for readability)
    run_case(
        "PyTorch get_fn_args style (batch=1)",
        batch_sz=1,
        n_in=6,
        n_out=2,
        n_weight_grps=2,
        n_offset_grps=3,
        kernel_h=3,
        kernel_w=2,
        stride=(2, 1),
        pad=(1, 0),
        dilation=(2, 1),
        in_h=5,
        in_w=4,
        seed=42,
    )

    # Case 2: No mask (mask optional) - same config, then expected with mask=None
    torch.manual_seed(42)
    n_in, n_out = 6, 2
    n_weight_grps, n_offset_grps = 2, 3
    kH, kW = 3, 2  # noqa: N806
    stride_h, stride_w = 2, 1
    pad_h, pad_w = 1, 0
    dil_h, dil_w = 2, 1
    in_h, in_w = 5, 4
    batch_sz = 1
    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

    x = torch.rand(batch_sz, n_in, in_h, in_w, dtype=torch.float32)
    offset = torch.randn(batch_sz, n_offset_grps * 2 * 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_no_mask = 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=None,
    )
    print("// --- Same inputs, no mask (expected_Y when mask is omitted) ---")
    print("std::vector<float> expected_Y_no_mask = {" + to_cpp_list(out_no_mask) + "};")
    print()

    # Case 3: groups=2, offset_group=2, non-zero offset (for GroupsWithNonZeroOffset test)
    run_case(
        "Groups with non-zero offset (batch=1, 2 groups)",
        batch_sz=1,
        n_in=4,
        n_out=2,
        n_weight_grps=2,
        n_offset_grps=2,
        kernel_h=2,
        kernel_w=2,
        stride=(1, 1),
        pad=(0, 0),
        dilation=(1, 1),
        in_h=3,
        in_w=3,
        seed=123,
    )


if __name__ == "__main__":
    main()
