import sys
from unittest import mock

import pytest
import torch

# mock detection module
sys.modules["torchvision._C"] = mock.Mock()
import torchseg  # noqa


def get_encoders():
    exclude_encoders = [
        "senet154",
        "resnext101_32x16d",
        "resnext101_32x32d",
        "resnext101_32x48d",
    ]
    encoders = torchseg.encoders.get_encoder_names()
    encoders = [e for e in encoders if e not in exclude_encoders]
    encoders.append("tu-resnet34")  # for timm universal encoder
    return encoders


ENCODERS = get_encoders()
DEFAULT_ENCODER = "resnet18"


def get_sample(model_class):
    if model_class in [
        torchseg.Unet,
        torchseg.Linknet,
        torchseg.FPN,
        torchseg.PSPNet,
        torchseg.UnetPlusPlus,
        torchseg.MAnet,
    ]:
        sample = torch.ones([1, 3, 64, 64])
    elif model_class == torchseg.PAN:
        sample = torch.ones([2, 3, 256, 256])
    elif model_class == torchseg.DeepLabV3:
        sample = torch.ones([2, 3, 128, 128])
    else:
        raise ValueError(f"Not supported model class {model_class}")
    return sample


def _test_forward(model, sample, test_shape=False):
    with torch.no_grad():
        out = model(sample)
    if test_shape:
        assert out.shape[2:] == sample.shape[2:]


def _test_forward_backward(model, sample, test_shape=False):
    out = model(sample)
    out.mean().backward()
    if test_shape:
        assert out.shape[2:] == sample.shape[2:]


@pytest.mark.parametrize("encoder_name", ENCODERS)
@pytest.mark.parametrize("encoder_depth", [3, 5])
@pytest.mark.parametrize(
    "model_class",
    [
        torchseg.FPN,
        torchseg.PSPNet,
        torchseg.Linknet,
        torchseg.Unet,
        torchseg.UnetPlusPlus,
    ],
)
def test_forward(model_class, encoder_name, encoder_depth, **kwargs):
    if (
        model_class is torchseg.Unet
        or model_class is torchseg.UnetPlusPlus
        or model_class is torchseg.MAnet
    ):
        kwargs["decoder_channels"] = (16, 16, 16, 16, 16)[-encoder_depth:]
    if model_class in [
        torchseg.UnetPlusPlus,
        torchseg.Linknet,
    ] and encoder_name.startswith("mit_b"):
        return  # skip mit_b*
    if (
        model_class is torchseg.FPN
        and encoder_name.startswith("mit_b")
        and encoder_depth != 5
    ):
        return  # skip mit_b*
    model = model_class(
        encoder_name, encoder_depth=encoder_depth, encoder_weights=None, **kwargs
    )
    sample = get_sample(model_class)
    model.eval()
    if encoder_depth == 5 and model_class != torchseg.PSPNet:
        test_shape = True
    else:
        test_shape = False

    _test_forward(model, sample, test_shape)


@pytest.mark.parametrize(
    "model_class",
    [
        torchseg.PAN,
        torchseg.FPN,
        torchseg.PSPNet,
        torchseg.Linknet,
        torchseg.Unet,
        torchseg.UnetPlusPlus,
        torchseg.MAnet,
        torchseg.DeepLabV3,
    ],
)
def test_forward_backward(model_class):
    sample = get_sample(model_class)
    model = model_class(DEFAULT_ENCODER, encoder_weights=None)
    _test_forward_backward(model, sample)


@pytest.mark.parametrize(
    "model_class",
    [
        torchseg.PAN,
        torchseg.FPN,
        torchseg.PSPNet,
        torchseg.Linknet,
        torchseg.Unet,
        torchseg.UnetPlusPlus,
        torchseg.MAnet,
    ],
)
def test_aux_output(model_class):
    model = model_class(
        DEFAULT_ENCODER, encoder_weights=None, aux_params=dict(classes=2)
    )
    sample = get_sample(model_class)
    label_size = (sample.shape[0], 2)
    mask, label = model(sample)
    assert label.size() == label_size


@pytest.mark.parametrize("upsampling", [2, 4, 8])
@pytest.mark.parametrize("model_class", [torchseg.FPN, torchseg.PSPNet])
def test_upsample(model_class, upsampling):
    default_upsampling = 4 if model_class is torchseg.FPN else 8
    model = model_class(DEFAULT_ENCODER, encoder_weights=None, upsampling=upsampling)
    sample = get_sample(model_class)
    mask = model(sample)
    assert mask.size()[-1] / 64 == upsampling / default_upsampling


@pytest.mark.parametrize("model_class", [torchseg.FPN])
@pytest.mark.parametrize("in_channels", [1, 2, 4])
def test_in_channels(model_class, in_channels):
    sample = torch.ones([1, in_channels, 64, 64])
    model = model_class(DEFAULT_ENCODER, encoder_weights=None, in_channels=in_channels)
    model.eval()
    with torch.no_grad():
        model(sample)

    assert model.encoder._in_channels == in_channels


@pytest.mark.parametrize("encoder_name", ENCODERS)
def test_dilation(encoder_name):
    if (
        encoder_name in ["inceptionresnetv2", "xception", "inceptionv4"]
        or encoder_name.startswith("vgg")
        or encoder_name.startswith("densenet")
        or encoder_name.startswith("timm-res")
        or encoder_name.startswith("mit_b")
    ):
        return

    encoder = torchseg.encoders.get_encoder(encoder_name, output_stride=16)

    encoder.eval()
    with torch.no_grad():
        sample = torch.ones([1, 3, 64, 64])
        output = encoder(sample)

    shapes = [out.shape[-1] for out in output]
    assert shapes == [64, 32, 16, 8, 4, 4]  # last downsampling replaced with dilation


if __name__ == "__main__":
    pytest.main([__file__])
