# Copyright contributors to the Terratorch project

import gc

import pytest
import torch

from terratorch.models import SMPModelFactory
from terratorch.models.backbones.prithvi_vit import PRETRAINED_BANDS

NUM_CHANNELS = 6
NUM_CLASSES = 2
EXPECTED_SEGMENTATION_OUTPUT_SHAPE = (1, NUM_CLASSES, 224, 224)
EXPECTED_REGRESSION_OUTPUT_SHAPE = (1, 224, 224)
EXPECTED_CLASSIFICATION_OUTPUT_SHAPE = (1, NUM_CLASSES)


@pytest.fixture(scope="session")
def model_factory() -> SMPModelFactory:
    return SMPModelFactory()


@pytest.fixture(scope="session")
def model_input() -> torch.Tensor:
    return torch.ones((1, NUM_CHANNELS, 224, 224))


@pytest.mark.parametrize("backbone", ["tu-convnext_atto.d2_in1k"])
@pytest.mark.parametrize("model", ["DeepLabV3"])
def test_create_segmentation_model(backbone, model, model_factory: SMPModelFactory, model_input):
    model = model_factory.build_model(
        "segmentation",
        backbone=backbone,
        model=model,
        in_channels=NUM_CHANNELS,
        bands=PRETRAINED_BANDS,
        pretrained=False,
        num_classes=NUM_CLASSES,
    )
    model.eval()

    with torch.no_grad():
        assert model(model_input).output.shape == EXPECTED_SEGMENTATION_OUTPUT_SHAPE

    gc.collect()


@pytest.mark.skip(
    reason="It looks these models are no more supported by `smp`, maybe it's better to remove them in the future."
)
@pytest.mark.parametrize("backbone", ["timm-regnetx_002"])
@pytest.mark.parametrize("model", ["Unet"])
def test_create_segmentation_model_no_in_channels(backbone, model, model_factory: SMPModelFactory, model_input):
    model = model_factory.build_model(
        "segmentation",
        backbone=backbone,
        model=model,
        backbone_out_channels=768,
        backbone_output_stride=1,
        bands=PRETRAINED_BANDS,
        pretrained=False,
        num_classes=NUM_CLASSES,
        smp_encoder_depth=5,
    )
    model.eval()

    with torch.no_grad():
        assert model(model_input).output.shape == EXPECTED_SEGMENTATION_OUTPUT_SHAPE

    gc.collect()


@pytest.mark.skip(
    reason="It looks these models are no more supported by `smp`, maybe it's better to remove them in the future."
)
@pytest.mark.parametrize("backbone", ["timm-regnetx_002"])
@pytest.mark.parametrize("model", ["Unet"])
def test_create_model_with_extra_bands(backbone, model, model_factory: SMPModelFactory):
    model = model_factory.build_model(
        "segmentation",
        backbone=backbone,
        model=model,
        in_channels=NUM_CHANNELS + 1,
        backbone_out_channels=768,
        backbone_output_stride=1,
        bands=[*PRETRAINED_BANDS, 7],  # add an extra band
        pretrained=False,
        num_classes=NUM_CLASSES,
        smp_encoder_depth=5,
    )
    model.eval()
    model_input = torch.ones((1, NUM_CHANNELS + 1, 224, 224))
    with torch.no_grad():
        assert model(model_input).output.shape == EXPECTED_SEGMENTATION_OUTPUT_SHAPE

    gc.collect()
