from typing import Any, Callable, Union
import pytest

import torch
import torch.nn as nn

try:
    from colossalai._analyzer._subclasses import MetaTensor
except:
    pass

aten = torch.ops.aten

registered_meta = {
    ('aten.convolution.default', True): [    # (aten ops, requires_backward)
        (nn.Conv1d(in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2), torch.rand(2, 3, 4)),
        (nn.Conv2d(in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2), torch.rand(2, 3, 4, 4)),
        (nn.Conv3d(in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2), torch.rand(2, 3, 4, 4, 4)),
        (nn.ConvTranspose1d(in_channels=3, out_channels=4, kernel_size=2, padding=1, dilation=2), torch.rand(2, 3, 4)),
        (nn.ConvTranspose2d(in_channels=3, out_channels=4, kernel_size=2, padding=1,
                            dilation=2), torch.rand(2, 3, 4, 4)),
        (nn.ConvTranspose3d(in_channels=3, out_channels=4, kernel_size=2, padding=1,
                            dilation=2), torch.rand(2, 3, 4, 4, 4)),
    ],
    ('aten.native_batch_norm.default', True): [
        (nn.BatchNorm1d(4), torch.rand(2, 4)),
        (nn.BatchNorm2d(4), torch.rand(1, 4, 4, 4)),
        (nn.BatchNorm3d(4), torch.rand(1, 4, 4, 4, 4)),
    ],
    ('aten.native_layer_norm.default', True): [(nn.LayerNorm(4), torch.rand(1, 2, 3, 4)),],
    ('aten.avg_pool1d.default', True): [
        (nn.MaxPool1d(3, stride=2), torch.rand(4, 5, 5)),
        (nn.AvgPool1d(3, stride=2), torch.rand(4, 5, 5)),
        (nn.AdaptiveMaxPool1d(3), torch.rand(4, 5, 5)),
        (nn.AdaptiveAvgPool1d(3), torch.rand(4, 5, 5)),
    ],
    ('aten.avg_pool2d.default', True): [
        (nn.MaxPool2d((3, 2), stride=(2, 1)), torch.rand(2, 4, 5, 5)),
        (nn.AvgPool2d((3, 2), stride=(2, 1)), torch.rand(2, 4, 5, 5)),
        (nn.AdaptiveMaxPool2d((3, 2)), torch.rand(2, 4, 5, 5)),
        (nn.AdaptiveAvgPool2d((3, 2)), torch.rand(2, 4, 5, 5)),
    ],
    ('aten.relu.default', True): [
        (nn.ReLU(), torch.rand(4, 3, 1, 2)),
        (nn.LeakyReLU(), torch.rand(4, 3, 1, 2)),
        (nn.SiLU(), torch.rand(4, 3, 1, 2)),
        (nn.GELU(), torch.rand(4, 3, 1, 2)),
        (nn.ELU(), torch.rand(4, 3, 1, 2)),
        (nn.Sigmoid(), torch.rand(4, 3, 1, 2)),
        (nn.Tanh(), torch.rand(4, 3, 1, 2)),
        (nn.Hardswish(), torch.rand(4, 3, 1, 2)),
    ]
}


def compare_all(tensor: torch.Tensor, meta_tensor: torch.Tensor) -> Any:
    assert tensor.shape == meta_tensor.shape, f'the shape of tensor ({tensor.shape}) and meta tensor ({meta_tensor.shape}) does not match.'
    assert tensor.dtype == meta_tensor.dtype, f'the dtype of tensor ({tensor.dtype}) and meta tensor ({meta_tensor.dtype}) does not match.'
    assert tensor.stride() == meta_tensor.stride(
    ), f'the stride of tensor ({tensor.stride()}) and meta tensor ({meta_tensor.stride()}) does not match.'


def run_and_compare(f: Union[nn.Module, Callable], x: torch.Tensor, requires_backward=False) -> Any:
    x.requires_grad = requires_backward
    meta_x = MetaTensor(x)
    x_out, meta_out = f(x), f(meta_x)
    compare_all(x_out, meta_out)
    if requires_backward:
        x_out.sum().backward()
        meta_out.sum().backward()
        compare_all(x.grad, meta_x.grad)


@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12')
def test_meta_aten():
    for (aten_op, requires_backward), v in registered_meta.items():
        for f, x in v:
            run_and_compare(f, x, requires_backward)


if __name__ == '__main__':
    test_meta_aten()