"""
cnn_branch.py — ResNet-50 spatial feature extractor for satellite imagery.

Backbone: ResNet-50 pretrained on ImageNet (torchvision).
Modified to accept 4-channel input (RGB + NIR) by extending the first
convolutional layer. The final FC classification head is removed,
producing a 2048-dimensional spatial feature vector.
"""

import logging

import torch
import torch.nn as nn
from torchvision.models import resnet50, ResNet50_Weights

from src.training.config import IMAGE_CHANNELS, CNN_FEATURE_DIM

logger = logging.getLogger(__name__)


class CNNBranch(nn.Module):
    """
    ResNet-50 backbone modified for 4-channel satellite image input.

    Architecture:
        Input:  (batch, 4, 128, 128)
        Output: (batch, 2048) spatial feature vector

    The first conv layer is extended from 3→4 input channels by
    averaging the pretrained weights and appending them as the NIR
    channel weights. All subsequent layers use pretrained ImageNet weights.
    """

    def __init__(
        self,
        in_channels: int = IMAGE_CHANNELS,
        out_features: int = CNN_FEATURE_DIM,
        pretrained: bool = True,
        freeze_early: bool = False,
    ):
        """
        Args:
            in_channels: Number of input channels (4 for RGB+NIR).
            out_features: Output feature dimensionality (2048).
            pretrained: Whether to use ImageNet pretrained weights.
            freeze_early: Whether to freeze early layers (layer1, layer2).
        """
        super().__init__()

        # Load pretrained ResNet-50
        if pretrained:
            logger.info("Loading pretrained ResNet-50 weights (ImageNet)...")
            weights = ResNet50_Weights.IMAGENET1K_V2
            backbone = resnet50(weights=weights)
        else:
            backbone = resnet50(weights=None)

        # Modify first conv layer to accept 4 channels
        original_conv1 = backbone.conv1
        self.conv1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=64,
            kernel_size=7,
            stride=2,
            padding=3,
            bias=False,
        )

        # Initialize: copy pretrained RGB weights, average for NIR channel
        if pretrained and in_channels == 4:
            with torch.no_grad():
                # Copy RGB weights
                self.conv1.weight[:, :3, :, :] = original_conv1.weight
                # NIR channel: average of RGB weights
                self.conv1.weight[:, 3:, :, :] = original_conv1.weight.mean(dim=1, keepdim=True)
        elif pretrained and in_channels == 3:
            self.conv1.weight = original_conv1.weight

        # Copy remaining layers
        self.bn1 = backbone.bn1
        self.relu = backbone.relu
        self.maxpool = backbone.maxpool
        self.layer1 = backbone.layer1
        self.layer2 = backbone.layer2
        self.layer3 = backbone.layer3
        self.layer4 = backbone.layer4
        self.avgpool = backbone.avgpool

        # Optionally freeze early layers for transfer learning
        if freeze_early:
            for param in [self.conv1, self.bn1, self.layer1, self.layer2]:
                if isinstance(param, nn.Module):
                    for p in param.parameters():
                        p.requires_grad = False
            logger.info("Froze early layers (conv1, bn1, layer1, layer2).")

        self.out_features = out_features
        logger.info(f"CNNBranch initialized: {in_channels}ch → {out_features}d features")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass through ResNet-50 backbone.

        Args:
            x: Input tensor of shape (batch, 4, 128, 128).

        Returns:
            Feature vector of shape (batch, 2048).
        """
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)  # (batch, 2048, H', W')

        x = self.avgpool(x)  # (batch, 2048, 1, 1)
        x = torch.flatten(x, 1)  # (batch, 2048)

        return x

    def get_feature_maps(self, x: torch.Tensor) -> torch.Tensor:
        """
        Get intermediate feature maps before global average pooling.
        Used by Grad-CAM for visualization.

        Returns:
            Feature maps of shape (batch, 2048, H', W').
        """
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        return x  # (batch, 2048, 4, 4) for 128×128 input


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    model = CNNBranch(in_channels=4, pretrained=True)

    # Test forward pass
    dummy = torch.randn(2, 4, 128, 128)
    features = model(dummy)
    print(f"Input: {dummy.shape} → Output: {features.shape}")

    fmaps = model.get_feature_maps(dummy)
    print(f"Feature maps: {fmaps.shape}")

    # Parameter count
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Parameters: {total:,} total, {trainable:,} trainable")