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

import pytest
import torch
import torch.nn as nn

from torchgeo.models import ChangeMixin, ChangeStar, ChangeStarFarSeg

BACKBONE = ["resnet18", "resnet34", "resnet50", "resnet101"]
IN_CHANNELS = [64, 128]
INNNR_CHANNELS = [16, 32, 64]
NC = [1, 2, 4]
SF = [4, 8, 1]


class TestChangeStar:
    @torch.no_grad()
    def test_changestar_farseg_classes(self) -> None:
        model = ChangeStarFarSeg(
            classes=4, backbone="resnet50", backbone_pretrained=False
        )
        x = torch.randn(2, 2, 3, 128, 128)
        y = model(x)

        assert y["bi_seg_logit"].shape[2] == 4

    @torch.no_grad()
    def test_changestar_farseg_output_size(self) -> None:
        model = ChangeStarFarSeg(
            classes=4, backbone="resnet50", backbone_pretrained=False
        )
        model.eval()
        x = torch.randn(2, 2, 3, 128, 128)
        y = model(x)

        assert y["bi_seg_logit"].shape[3] == 128 and y["bi_seg_logit"].shape[4] == 128
        assert y["change_prob"].shape[2] == 128 and y["change_prob"].shape[3] == 128

        model.train()
        y = model(x)

        assert y["bi_seg_logit"].shape[3] == 128 and y["bi_seg_logit"].shape[4] == 128
        assert y["bi_change_logit"].shape[3] == 128
        assert y["bi_change_logit"].shape[4] == 128

    @pytest.mark.parametrize("backbone", BACKBONE)
    def test_valid_changestar_farseg_backbone(self, backbone: str) -> None:
        ChangeStarFarSeg(classes=4, backbone=backbone, backbone_pretrained=False)

    def test_invalid_changestar_farseg_backbone(self) -> None:
        match = "unknown backbone: anynet."
        with pytest.raises(ValueError, match=match):
            ChangeStarFarSeg(classes=4, backbone="anynet", backbone_pretrained=False)

    @torch.no_grad()
    @pytest.mark.parametrize("inc", IN_CHANNELS)
    @pytest.mark.parametrize("innerc", INNNR_CHANNELS)
    @pytest.mark.parametrize("nc", NC)
    @pytest.mark.parametrize("sf", SF)
    def test_changemixin_output_size(
        self, inc: int, innerc: int, nc: int, sf: int
    ) -> None:
        m = ChangeMixin(
            in_channels=inc, inner_channels=innerc, num_convs=nc, scale_factor=sf
        )

        y = m(torch.rand(3, 2, inc // 2, 32, 32))
        assert y[0].shape == y[1].shape
        assert y[0].shape == (3, 1, 32 * sf, 32 * sf)

    @torch.no_grad()
    def test_changestar(self) -> None:
        dense_feature_extractor = nn.modules.Sequential(
            nn.modules.Conv2d(3, 32, 3, 1, 1),
            nn.modules.BatchNorm2d(32),
            nn.modules.ReLU(),
            nn.modules.MaxPool2d(3, 2, 1),
        )

        seg_classifier = nn.modules.Sequential(
            nn.modules.Conv2d(32, 2, 3, 1, 1),
            nn.modules.UpsamplingBilinear2d(scale_factor=2.0),
        )

        m = ChangeStar(
            dense_feature_extractor,
            seg_classifier,
            ChangeMixin(
                in_channels=32 * 2, inner_channels=16, num_convs=4, scale_factor=2.0
            ),
        )
        m.eval()

        y = m(torch.rand(3, 2, 3, 64, 64))
        assert y["bi_seg_logit"].shape == (3, 2, 2, 64, 64)
        assert y["change_prob"].shape == (3, 1, 64, 64)

    @torch.no_grad()
    def test_changestar_invalid_inference_mode(self) -> None:
        dense_feature_extractor = nn.modules.Sequential(
            nn.modules.Conv2d(3, 32, 3, 1, 1),
            nn.modules.BatchNorm2d(32),
            nn.modules.ReLU(),
            nn.modules.MaxPool2d(3, 2, 1),
        )

        seg_classifier = nn.modules.Sequential(
            nn.modules.Conv2d(32, 2, 3, 1, 1),
            nn.modules.UpsamplingBilinear2d(scale_factor=2.0),
        )

        match = "Unknown inference_mode: random"
        with pytest.raises(ValueError, match=match):
            ChangeStar(
                dense_feature_extractor,
                seg_classifier,
                ChangeMixin(
                    in_channels=32 * 2, inner_channels=16, num_convs=4, scale_factor=2.0
                ),
                inference_mode="random",
            )

    @torch.no_grad()
    @pytest.mark.parametrize("inference_mode", ["t1t2", "t2t1", "mean"])
    def test_changestar_inference_output_size(self, inference_mode: str) -> None:
        dense_feature_extractor = nn.modules.Sequential(
            nn.modules.Conv2d(3, 32, 3, 1, 1),
            nn.modules.BatchNorm2d(32),
            nn.modules.ReLU(),
            nn.modules.MaxPool2d(3, 2, 1),
        )
        CLASSES = 2
        seg_classifier = nn.modules.Sequential(
            nn.modules.Conv2d(32, CLASSES, 3, 1, 1),
            nn.modules.UpsamplingBilinear2d(scale_factor=2.0),
        )

        m = ChangeStar(
            dense_feature_extractor,
            seg_classifier,
            ChangeMixin(
                in_channels=32 * 2, inner_channels=16, num_convs=4, scale_factor=2.0
            ),
            inference_mode=inference_mode,
        )
        m.eval()

        x = torch.randn(2, 2, 3, 128, 128)
        y = m(x)

        assert y["bi_seg_logit"].shape == (2, 2, CLASSES, 128, 128)
        assert y["change_prob"].shape == (2, 1, 128, 128)
