"""
Source url: https://github.com/Karel911/TRACER
Author: Min Seok Lee and Wooseok Shin
Modified by Nikita Selin (OPHoperHPO)[https://github.com/OPHoperHPO].
License: Apache License 2.0
Changes:
    - Refactored code
    - Removed unused code
    - Added comments
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Optional, Tuple

from torch import Tensor

from carvekit.ml.arch.tracerb7.efficientnet import EfficientEncoderB7
from carvekit.ml.arch.tracerb7.att_modules import (
    RFB_Block,
    aggregation,
    ObjectAttention,
)


class TracerDecoder(nn.Module):
    """Tracer Decoder"""

    def __init__(
        self,
        encoder: EfficientEncoderB7,
        features_channels: Optional[List[int]] = None,
        rfb_channel: Optional[List[int]] = None,
    ):
        """
        Initialize the tracer decoder.

        Args:
            encoder: The encoder to use.
            features_channels: The channels of the backbone features at different stages. default: [48, 80, 224, 640]
            rfb_channel: The channels of the RFB features. default: [32, 64, 128]
        """
        super().__init__()
        if rfb_channel is None:
            rfb_channel = [32, 64, 128]
        if features_channels is None:
            features_channels = [48, 80, 224, 640]
        self.encoder = encoder
        self.features_channels = features_channels

        # Receptive Field Blocks
        features_channels = rfb_channel
        self.rfb2 = RFB_Block(self.features_channels[1], features_channels[0])
        self.rfb3 = RFB_Block(self.features_channels[2], features_channels[1])
        self.rfb4 = RFB_Block(self.features_channels[3], features_channels[2])

        # Multi-level aggregation
        self.agg = aggregation(features_channels)

        # Object Attention
        self.ObjectAttention2 = ObjectAttention(
            channel=self.features_channels[1], kernel_size=3
        )
        self.ObjectAttention1 = ObjectAttention(
            channel=self.features_channels[0], kernel_size=3
        )

    def forward(self, inputs: torch.Tensor) -> Tensor:
        """
        Forward pass of the tracer decoder.

        Args:
            inputs: Preprocessed images.

        Returns:
            Tensors of segmentation masks and mask of object edges.
        """
        features = self.encoder(inputs)
        x3_rfb = self.rfb2(features[1])
        x4_rfb = self.rfb3(features[2])
        x5_rfb = self.rfb4(features[3])

        D_0 = self.agg(x5_rfb, x4_rfb, x3_rfb)

        ds_map0 = F.interpolate(D_0, scale_factor=8, mode="bilinear")

        D_1 = self.ObjectAttention2(D_0, features[1])
        ds_map1 = F.interpolate(D_1, scale_factor=8, mode="bilinear")

        ds_map = F.interpolate(D_1, scale_factor=2, mode="bilinear")
        D_2 = self.ObjectAttention1(ds_map, features[0])
        ds_map2 = F.interpolate(D_2, scale_factor=4, mode="bilinear")

        final_map = (ds_map2 + ds_map1 + ds_map0) / 3

        return torch.sigmoid(final_map)