import cv2
import torch
import torch.nn as nn
import numpy as np
import torch
import torch.nn as nn

class CustomScatter(nn.Module):
    def __init__(self, dim):
        super(CustomScatter, self).__init__()
        self.dim = dim

    def forward(self, src, indices, output):
        return torch.scatter(output, self.dim, indices, src)

if __name__ == '__main__':
    name = "CustomScatter"

    # Indices tensor and source tensor
    src = torch.tensor([3])  # Source values to scatter
    indices = torch.tensor([0])  # Indices where src will be scattered
    output = torch.zeros((1), dtype=src.dtype)  # Example size

    dim = 0 # Dimension along which to scatter

    model = CustomScatter(dim=dim)

    # Note: For scatter, inputs are indices and src
    inputs = [src, indices, output]
    outputs = [model(src, indices, output)]

    print("src: ", src, "\t src shape: ", src.shape)
    print("indices: ", indices, "\t indices shape: ", indices.shape)
    print("output_size: ", output.shape , "\t dim: ", dim)
    print("output: ", outputs[0], "\t output shape: ", outputs[0].shape)

    print("\n\t\t **** TORCH INFERENCE END ****\n\n")
    # Export the model to ONNX
    torch.onnx.export(
        model,
        (src, indices, output),  # Pass both indices and src as inputs to the model
        f"{name}.onnx",
        verbose=True,
        input_names=[f"input_{i}" for i in range(len(inputs))],
        output_names=["output"],
        # dynamic_axes={'indices': {0: 'batch_size'}, 'src': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
    )

    print("\n\t\t **** EXPORTED ONNX ****\n\n")


    net = cv2.dnn.readNetFromONNX(f"{name}.onnx");
    print("\n\t\t **** CV2 MODEL PARSED ****\n\n")

    for i in range(len(inputs)):
        inputs[i] = inputs[i].detach().numpy()
        inputs[i] = cv2.Mat(inputs[i])

    ## set inputs
    for i in range(len(inputs)):
        net.setInput(inputs[i], f"input_{i}")

    print("\n\t\t **** CV2 INPUTS SET ****\n\n")

    ## forward
    output_cv = net.forward()
    print("\n\t\t **** CV2 FORWARD END ****\n\n")

    # Adjust for output shape differences
    # output_cv = output_cv.reshape(outputs[0].shape)
    print("output_cv: ", output_cv, "\t output_cv shape: ", output_cv.shape)
    print("torch output: ", outputs[0], "\t torch output shape: ", outputs[0].shape)

    if np.allclose(output_cv, outputs[0], atol=1e-6):
        print("\033[92m\t\t **** TEST PASSED ****\n\n")
    else:
        print("\033[91m\t\t **** TEST FAILED ****\n\n")