# Copyright (c) TorchGeo Contributors. All rights reserved.
# Licensed under the MIT License.

"""ChangeStar implementations."""

import torch
import torch.nn as nn
from einops import rearrange
from torch import Tensor
from torch.nn.modules import Module
from torchvision.models._api import WeightsEnum

from .farseg import FarSeg


class ChangeMixin(Module):
    """This module enables any segmentation model to detect binary change.

    The common usage is to attach this module on a segmentation model without the
    classification head.

    If you use this model in your research, please cite the following paper:

    * https://arxiv.org/abs/2108.07002
    """

    def __init__(
        self,
        in_channels: int = 128 * 2,
        inner_channels: int = 16,
        num_convs: int = 4,
        scale_factor: float = 4.0,
    ) -> None:
        """Initializes a new ChangeMixin module.

        Args:
            in_channels: sum of channels of bitemporal feature maps
            inner_channels: number of channels of inner feature maps
            num_convs: number of convolution blocks
            scale_factor: number of upsampling factor
        """
        super().__init__()
        layers: list[Module] = [
            nn.modules.Sequential(
                nn.modules.Conv2d(in_channels, inner_channels, 3, 1, 1),
                nn.modules.BatchNorm2d(inner_channels),
                nn.modules.ReLU(True),
            )
        ]
        layers += [
            nn.modules.Sequential(
                nn.modules.Conv2d(inner_channels, inner_channels, 3, 1, 1),
                nn.modules.BatchNorm2d(inner_channels),
                nn.modules.ReLU(True),
            )
            for _ in range(num_convs - 1)
        ]

        cls_layer = nn.modules.Conv2d(inner_channels, 1, 3, 1, 1)

        layers.append(cls_layer)
        layers.append(nn.modules.UpsamplingBilinear2d(scale_factor=scale_factor))

        self.convs = nn.modules.Sequential(*layers)

    def forward(self, bi_feature: Tensor) -> list[Tensor]:
        """Forward pass of the model.

        Args:
            bi_feature: input bitemporal feature maps of shape [b, t, c, h, w]

        Returns:
            a list of bidirected output predictions
        """
        batch_size = bi_feature.size(0)
        t1t2 = torch.cat([bi_feature[:, 0, :, :, :], bi_feature[:, 1, :, :, :]], dim=1)
        t2t1 = torch.cat([bi_feature[:, 1, :, :, :], bi_feature[:, 0, :, :, :]], dim=1)

        c1221 = self.convs(torch.cat([t1t2, t2t1], dim=0))
        c12, c21 = torch.split(c1221, batch_size, dim=0)
        return [c12, c21]


class ChangeStar(Module):
    """The base class of the network architecture of ChangeStar.

    ChangeStar is composed of an any segmentation model and a ChangeMixin module.
    This model is mainly used for binary/multi-class change detection under bitemporal
    supervision and single-temporal supervision. It features the property of
    segmentation architecture reusing, which is helpful to integrate advanced dense
    prediction (e.g., semantic segmentation) network architecture into change detection.

    For multi-class change detection, semantic change prediction can be inferred by a
    binary change prediction from the ChangeMixin module and two semantic predictions
    from the Segmentation model.

    If you use this model in your research, please cite the following paper:

    * https://arxiv.org/abs/2108.07002
    """

    def __init__(
        self,
        dense_feature_extractor: Module,
        seg_classifier: Module,
        changemixin: ChangeMixin,
        inference_mode: str = 't1t2',
    ) -> None:
        """Initializes a new ChangeStar model.

        Args:
            dense_feature_extractor: module for dense feature extraction, typically a
                semantic segmentation model without semantic segmentation head.
            seg_classifier: semantic segmentation head, typically a convolutional layer
                followed by an upsampling layer.
            changemixin: :class:`torchgeo.models.ChangeMixin` module
            inference_mode: name of inference mode ``'t1t2'`` | ``'t2t1'`` | ``'mean'``.
                ``'t1t2'``: concatenate bitemporal features in the order of t1->t2;
                ``'t2t1'``: concatenate bitemporal features in the order of t2->t1;
                ``'mean'``: the weighted mean of the output of ``'t1t2'`` and ``'t1t2'``
        """
        super().__init__()
        self.dense_feature_extractor = dense_feature_extractor
        self.seg_classifier = seg_classifier
        self.changemixin = changemixin

        if inference_mode not in ['t1t2', 't2t1', 'mean']:
            raise ValueError(f'Unknown inference_mode: {inference_mode}')
        self.inference_mode = inference_mode

    def forward(self, x: Tensor) -> dict[str, Tensor]:
        """Forward pass of the model.

        Args:
            x: a bitemporal input tensor of shape [B, T, C, H, W]

        Returns:
            a dictionary containing bitemporal semantic segmentation logit and binary
            change detection logit/probability
        """
        _, t, _, _, _ = x.shape
        x = rearrange(x, 'b t c h w -> (b t) c h w')
        # feature extraction
        bi_feature = self.dense_feature_extractor(x)
        # semantic segmentation
        bi_seg_logit = self.seg_classifier(bi_feature)
        bi_seg_logit = rearrange(bi_seg_logit, '(b t) c h w -> b t c h w', t=t)

        bi_feature = rearrange(bi_feature, '(b t) c h w -> b t c h w', t=t)
        # change detection
        c12, c21 = self.changemixin(bi_feature)

        results: dict[str, Tensor] = {}
        if not self.training:
            results.update({'bi_seg_logit': bi_seg_logit})
            if self.inference_mode == 't1t2':
                results.update({'change_prob': c12.sigmoid()})
            elif self.inference_mode == 't2t1':
                results.update({'change_prob': c21.sigmoid()})
            elif self.inference_mode == 'mean':
                results.update(
                    {
                        'change_prob': torch.stack([c12, c21], dim=0)
                        .sigmoid_()
                        .mean(dim=0)
                    }
                )
        else:
            results.update(
                {
                    'bi_seg_logit': bi_seg_logit,
                    'bi_change_logit': torch.stack([c12, c21], dim=1),
                }
            )
        return results


class ChangeStarFarSeg(ChangeStar):
    """The network architecture of ChangeStar(FarSeg).

    ChangeStar(FarSeg) is composed of a FarSeg model and a ChangeMixin module.

    If you use this model in your research, please cite the following paper:

    * https://arxiv.org/abs/2108.07002
    """

    def __init__(
        self,
        backbone: str = 'resnet50',
        classes: int = 1,
        backbone_weights: WeightsEnum | None = None,
    ) -> None:
        """Initializes a new ChangeStarFarSeg model.

        Args:
            backbone: name of ResNet backbone
            classes: number of output segmentation classes
            backbone_weights: Pre-trained model weights to use.

        .. versionadded:: 0.9
           The *backbone_weights* parameter.
        """
        model = FarSeg(
            backbone=backbone, classes=classes, backbone_weights=backbone_weights
        )
        seg_classifier: Module = model.decoder.classifier
        model.decoder.classifier = nn.modules.Identity()  # type: ignore[assignment]

        super().__init__(
            dense_feature_extractor=model,
            seg_classifier=seg_classifier,
            changemixin=ChangeMixin(
                in_channels=128 * 2, inner_channels=16, num_convs=4, scale_factor=4.0
            ),
            inference_mode='t1t2',
        )
