# **IMPORTANT:** Based on your fusion layer expecting 29 channels,
        # but the concatenated tensor now has 18 channels, there's still a mismatch.

from functools import partial
from typing import List, Dict, Any
from omegaconf import DictConfig
from pytorch_toolbelt.modules import (
    ACT_RELU,
    decoders as D,
    EncoderModule,
    ResidualDeconvolutionUpsample2d,
)
from torch import nn, Tensor
import torch
import timm
from hydra.utils import instantiate
import random
import torch.nn.functional as F
from xview3.centernet.models.custom_encoder import CustomTimmEncoder, DebugUNetDecoder
from .unet_blocks import IRBlock, get_unet_block, SqueezeExcitation
from .necks import FPNNeck
from xview3.nas import ProgressiveNAS
from .heads.yolo_head import YOLOHead
from .transformer_classifier import ConvolutionalClassifier
from xview3.common import FilterResponseNormalization, ResidualBlock, FeaturePyramidNetwork
from ..constants import (
    CENTERNET_OUTPUT_CLASS_MAP,
    CENTERNET_OUTPUT_OBJECTNESS_MAP,
    CENTERNET_OUTPUT_OFFSET,
    CENTERNET_OUTPUT_SIZE,
    CENTERNET_OUTPUT_VESSEL_MAP,
    CENTERNET_OUTPUT_FISHING_MAP,
    CENTERNET_OUTPUT_YOLO,
    CENTERNET_TARGET_YOLO
)
from xview3.centernet.models.heads.decoupled_head import DecoupledHeadGroupNormLateShuffle
from xview3.centernet.models.heads.yolo_head import YOLOHead
from ..losses.circlenet_loss import CircleNetLoss, YOLOLoss
from xview3.centernet.bboxer.multilabel_circle_coder import MultilabelCircleNetCoder

__all__ = ["CenterNetUNetModel"]
from xview3.utils import TorchDevice
device = TorchDevice.choose_torch_device()

class CenterNetUNetModel(nn.Module):
    def __init__(
        self,
        encoder: EncoderModule,
        decoder_features: List[int],
        heads: List[nn.Module],
        box_coder,
        neck_out_channels: int = 256,
        input_upsample_factor: int = 1,
        activation=ACT_RELU,
        num_extra_blocks=0,
        block_type="UnetBlock",
        upsample_type="bilinear",
        extra_block=partial(IRBlock, stride=2),
        num_classes=2,
        anchors=[(10, 13), (16, 30), (33, 23)],
        nas=None,
    ):
        super().__init__()
        
        # Initial Setup
        self.output_stride = getattr(encoder, 'strides', [32])[0]
        self.initial_upsample = (
            nn.UpsamplingBilinear2d(scale_factor=input_upsample_factor)
            if input_upsample_factor != 1 else nn.Identity()
        )
        unet_block = get_unet_block(block_type, activation=activation)
        self.encoder = encoder
        self.extra_encoder_blocks = nn.ModuleList(
            [extra_block(encoder.channels[-1], encoder.channels[-1]) for _ in range(num_extra_blocks)]
        )

        # Upsample Block Selection
        if upsample_type in ['nearest', 'bilinear']:
            upsample_block_cls = partial(
                nn.Upsample,
                scale_factor=2,
                mode=upsample_type,
                align_corners=False
            )
        elif upsample_type == 'rdtsc':
            upsample_block_cls = partial(ResidualDeconvolutionUpsample2d, scale_factor=2)
        elif upsample_type == 'shuffle':
            upsample_block_cls = partial(nn.PixelShuffle, upscale_factor=2)
        else:
            raise ValueError(f"Unknown upsample_type {upsample_type}")

        # Encoder Channels
        if hasattr(encoder, 'channels'):
            encoder_channels = list(encoder.channels)
        else:
            encoder_channels = [32, 24, 40, 112, 320]  # Default channels if not provided
        encoder_channels += [encoder_channels[-1]] * num_extra_blocks

        # Head, NAS, YOLO Head, Classifier Initialization
        self.heads = nn.ModuleList(heads)  # Initialize heads
        self.head_output_channels = getattr(self.heads[0], 'out_channels', 256)
        self.output_channels = self.head_output_channels  # **Added**

        self.nas = nas  # Initialize nas

        self.yolo_head = YOLOHead(num_classes=num_classes, anchors=anchors)
        self.yolo_output_channels = self.yolo_head.num_anchors * (5 + num_classes)  # Adjust as per YOLO implementation

        self.classifier = ConvolutionalClassifier(in_channels=neck_out_channels, num_classes=num_classes)
        self.classifier_output_channels = getattr(self.classifier, 'out_channels', num_classes)

        self.box_coder = box_coder

        # Set neck_output_channels directly without multiplication
        self.neck_output_channels = neck_out_channels

        # Decoder Channels Assignment
        self.decoder_channels = decoder_features

        num_encoder_blocks = len(encoder_channels) - 1 

        if len(self.decoder_channels) < num_encoder_blocks:
            self.decoder_channels.extend(
                [self.decoder_channels[-1]] * (num_encoder_blocks - len(self.decoder_channels))
            )
        elif len(self.decoder_channels) > num_encoder_blocks:
            self.decoder_channels = self.decoder_channels[:num_encoder_blocks]

        self.num_decoded_features = len(encoder_channels)

        # Initialize Decoder
        self.decoder = DebugUNetDecoder(
            encoder_channels=encoder_channels,
            decoder_channels=self.decoder_channels,
            unet_block=unet_block,
            upsample_block=upsample_block_cls,
            center=False,
        )

        # Initialize Neck with correct decoder_channels
        self.neck = FPNNeck(channels=self.decoder_channels, out_channels=neck_out_channels)
        # self.neck_output_channels is already set to neck_out_channels

        # Initialize neck_conv with the correct in_channels
        self.neck_conv = nn.Conv2d(
            in_channels=self.neck_output_channels,  # 256
            out_channels=neck_out_channels,         # 256
            kernel_size=1
        )
        self.batch_norm = FilterResponseNormalization(self.neck_output_channels)

        # Calculate total_input_channels based on the corrected neck_output_channels
        # Calculate total_input_channels based on the actual number of channels
        self.total_input_channels = (
            self.head_output_channels +
            self.yolo_output_channels +
            self.classifier_output_channels
        )
        print(f"Total Input Channels for Fusion: {self.total_input_channels}")
        print(f"head_output_channels: {self.head_output_channels}")
        print(f"yolo_output_channels: {self.yolo_output_channels}")
        print(f"classifier_output_channels: {self.classifier_output_channels}")
        

        # Adjust channels if necessary
        self.yolo_adjust = nn.Conv2d(self.yolo_output_channels, self.head_output_channels, kernel_size=1)
        self.classifier_adjust = nn.Conv2d(self.classifier_output_channels, self.head_output_channels, kernel_size=1)

        #self.fusion = nn.Conv2d(self.head_output_channels * 3, self.head_output_channels, kernel_size=1)
        # Assuming total_input_channels should be 13 (based on your concatenation)
        #self.fusion = nn.Conv2d(13, self.head_output_channels, kernel_size=1)
        #self.fusion = nn.Conv2d(279, self.head_output_channels, kernel_size=1)
        #self.fusion = nn.Conv2d(29, 6, kernel_size=1)
        self.fusion = nn.Conv2d(29, self.head_output_channels, kernel_size=1)
        self.dropout = nn.Dropout(p=0.5)

        # Initialize Loss Functions
        self.circle_net_loss = CircleNetLoss()
        self.yolo_loss = YOLOLoss(num_classes=num_classes, anchors=anchors)
        self.criterion = YOLOLoss(num_classes, anchors)

    def forward(self, x):
        # Initial Upsampling
        x = self.initial_upsample(x)
        
        # Encoder and Decoder
        feature_maps = self.encoder(x)
        decoded_features = self.decoder(feature_maps)

        # Neck Processing
        neck_features = self.neck(decoded_features)
        neck_features = self.batch_norm(neck_features)
        neck_features = self.dropout(neck_features)
        neck_features = self.neck_conv(neck_features)

        # Heads
        head_output = self.heads[0](neck_features)
        yolo_output = self.yolo_head(neck_features)
        classification_output = self.classifier(neck_features)

        # Extracting Head Outputs
        objectness_map = head_output[CENTERNET_OUTPUT_OBJECTNESS_MAP]  # [B, 1, H, W]
        vessel_map = head_output[CENTERNET_OUTPUT_VESSEL_MAP]          # [B, 1, H, W]
        fishing_map = head_output[CENTERNET_OUTPUT_FISHING_MAP]        # [B, 1, H, W]
        size_map = head_output[CENTERNET_OUTPUT_SIZE]                  # [B, 1, H, W]
        offset_map = head_output[CENTERNET_OUTPUT_OFFSET]              # [B, 2, H, W]

        # Adjusting Yolo Output Shape if Necessary
        if len(yolo_output.shape) == 4:
            yolo_output = yolo_output.squeeze(1)  # Remove the extra dimension if it's of size 1

        # Upsample YOLO and Classification Outputs to Match Target Size
        target_size = objectness_map.size()[-2:]  # (H, W)

        if yolo_output.size()[-2:] != target_size:
            yolo_output = F.interpolate(yolo_output, size=target_size, mode='bilinear', align_corners=False)
        
        if classification_output.size()[-2:] != target_size:
            classification_output = F.interpolate(classification_output, size=target_size, mode='bilinear', align_corners=False)

        # Adjust Channels if Necessary
        yolo_output = self.yolo_adjust(yolo_output)               # [B, 6, H, W]
        classification_output = self.classifier_adjust(classification_output)  # [B, 6, H, W]

        # Concatenate All Relevant Feature Maps
        concatenated = torch.cat([
            objectness_map,         # [B, 1, H, W]
            vessel_map,             # [B, 1, H, W]
            fishing_map,            # [B, 1, H, W]
            size_map,               # [B, 1, H, W]
            offset_map,             # [B, 2, H, W]
            yolo_output,            # [B, 6, H, W]
            classification_output   # [B, 6, H, W]
        ], dim=1)  # Total Channels: 1 + 1 + 1 + 1 + 2 + 6 + 6 = 18

        # **IMPORTANT:** Based on your fusion layer expecting 29 channels,
        # but the concatenated tensor now has 18 channels, there's still a mismatch.
        # This suggests that yolo_output and classification_output adjustments may be incorrect.
        # Let's verify and correct the number of channels.

        # **Verification Prints:**
        print(f"objectness_map shape: {objectness_map.shape}")          # [B, 1, H, W]
        print(f"vessel_map shape: {vessel_map.shape}")                  # [B, 1, H, W]
        print(f"fishing_map shape: {fishing_map.shape}")                # [B, 1, H, W]
        print(f"size_map shape: {size_map.shape}")                      # [B, 1, H, W]
        print(f"offset_map shape: {offset_map.shape}")                  # [B, 2, H, W]
        print(f"yolo_output shape: {yolo_output.shape}")                # [B, 6, H, W]
        print(f"classification_output shape: {classification_output.shape}")  # [B, 6, H, W]
        print(f"Total concatenated channels: {concatenated.shape[1]}")  # Should be 1+1+1+1+2+6+6 = 18

        # **Assertion for Debugging**
        assert concatenated.shape[1] == self.total_input_channels, \
            f"Expected {self.total_input_channels} channels, but got {concatenated.shape[1]}"

        # Fusion Layer
        fused_output = self.fusion(concatenated)  # [B, 6, H, W]

        # Prepare Outputs
        outputs = {
            CENTERNET_OUTPUT_OBJECTNESS_MAP: objectness_map,
            CENTERNET_OUTPUT_SIZE: size_map,
            CENTERNET_OUTPUT_OFFSET: offset_map,
            CENTERNET_OUTPUT_VESSEL_MAP: vessel_map,
            CENTERNET_OUTPUT_FISHING_MAP: fishing_map,
            CENTERNET_OUTPUT_YOLO: yolo_output,  # Include YOLO output
            # Add `fused_output` if necessary
        }

        return outputs

    def calculate_loss(self, outputs, targets):
        # Ensure yolo_targets are included in the targets dictionary
        if CENTERNET_TARGET_YOLO not in targets:
            raise KeyError(f"Expected key '{CENTERNET_TARGET_YOLO}' not found in targets.")

        outputs[CENTERNET_TARGET_YOLO] = targets[CENTERNET_TARGET_YOLO]
        circle_net_loss = self.circle_net_loss(**outputs, **targets)

        # Calculate YOLO loss
        yolo_loss = self.yolo_loss(outputs[CENTERNET_OUTPUT_YOLO], outputs[CENTERNET_TARGET_YOLO])

        # Combine losses
        total_loss = circle_net_loss['objectness'] + circle_net_loss['vessel'] + circle_net_loss['fishing'] + circle_net_loss['size'] + yolo_loss

        return total_loss

    @classmethod
    def from_config(cls, **config):
        # Extract and remove 'encoder' from config
        encoder_config = config.pop('encoder')
        num_channels = config.pop('num_channels', 3)
        
        # Extract 'model_name' from encoder_config
        model_name = encoder_config.pop('model_name', 'tf_efficientnet_b5_ns')
        pretrained = encoder_config.pop('pretrained', False)
        
        # Initialize the encoder module
        encoder_module = CustomTimmEncoder(
            model_name=model_name,
            num_channels=num_channels,
            features_only=True,
            pretrained=pretrained
        )
        
        # Adjust encoder input channels
        encoder_module = adjust_encoder_input_channels(encoder_module, num_channels)
        
        # Extract other parameters with pop to remove them from config
        input_upsample_factor = config.pop("input_upsample_factor", 1)
        activation = config.pop('activation', ACT_RELU)
        num_extra_blocks = config.pop('num_extra_blocks', 0)
        num_classes = config.pop("num_classes", 2)
        anchors = config.pop("anchors", [(10, 13), (16, 30), (33, 23)])
        
        # Extract decoder configuration
        decoder_config = config.pop('decoder')
        decoder_features = decoder_config.pop('channels')
        block_type = decoder_config.pop("block_type", "UnetBlock")
        upsample_type = decoder_config.pop("upsample_type", "bilinear")
        
        # Extract neck_out_channels
        neck_out_channels = config.pop('neck_out_channels', 256)
        
        # Extract heads configuration and initialize
        heads_config = config.pop('heads')
        heads = []
        for head_config in heads_config:
            if head_config._target_ == 'xview3.centernet.models.heads.yolo_head.YOLOHead':
                head = instantiate(head_config, num_classes=num_classes, anchors=anchors)
            else:
                head = instantiate(head_config, channels=neck_out_channels)
            heads.append(head)
        
        # Extract box_coder configuration and initialize
        box_coder_config = config.pop('box_coder')
        box_coder = instantiate(
            box_coder_config,
            output_stride=encoder_module.strides[0] // (heads[0].upsample_factor * input_upsample_factor),
        )
        
        # **Define the evaluation function**
        def evaluate_architecture(model, train_loader, val_loader):
            return model.evaluate_architecture(model, train_loader, val_loader)

        # **Initialize ProgressiveNAS with the evaluation function**
        nas_config = config.pop('nas', {})  # Extract nas configuration
        nas = ProgressiveNAS(
            base_model=cls(
                encoder=encoder_module,
                decoder_features=decoder_features,
                heads=heads,
                box_coder=box_coder,
                neck_out_channels=neck_out_channels,
                input_upsample_factor=input_upsample_factor,
                activation=activation,
                block_type=block_type,
                upsample_type=upsample_type,
                num_extra_blocks=num_extra_blocks,
                num_classes=num_classes,  # Pass num_classes
                anchors=anchors,
                **config  # Remaining config without 'neck_out_channels'
            ),
            search_space=nas_config.get('search_space', {}),  # Access search_space from nas_config
            evaluation_function=evaluate_architecture  # Pass the evaluation function
        )

        return cls(
            encoder=encoder_module,
            decoder_features=decoder_features,
            heads=heads,
            box_coder=box_coder,
            nas=nas,  # Pass nas to the class
            neck_out_channels=neck_out_channels,
            input_upsample_factor=input_upsample_factor,
            activation=activation,
            block_type=block_type,
            upsample_type=upsample_type,
            num_extra_blocks=num_extra_blocks,
            num_classes=num_classes,  # Pass num_classes
            anchors=anchors,
            **config  # Remaining config without 'neck_out_channels'
        )

def adjust_encoder_input_channels(encoder_module, in_chans):
    # [Existing code...]
    if hasattr(encoder_module, 'encoder'):
        base_encoder = encoder_module.encoder
    else:
        base_encoder = encoder_module

    if hasattr(base_encoder, 'conv_stem'):
        old_conv = base_encoder.conv_stem
     
        new_conv = nn.Conv2d(
            in_channels=in_chans,
            out_channels=old_conv.out_channels,
            kernel_size=old_conv.kernel_size,
            stride=old_conv.stride,
            padding=old_conv.padding,
            bias=old_conv.bias is not None
        )
        
        with torch.no_grad():
            num_existing_in_channels = min(old_conv.in_channels, in_chans)
            new_conv.weight[:, :num_existing_in_channels, ...] = old_conv.weight[:, :num_existing_in_channels, ...]
            if in_chans > old_conv.in_channels:
                torch.nn.init.kaiming_normal_(new_conv.weight[:, old_conv.in_channels:, ...])

        base_encoder.conv_stem = new_conv
        
        print(f"[DEBUG] After adjustment, encoder.conv_stem.in_channels: {base_encoder.conv_stem.in_channels}")
    else:
        raise NotImplementedError("Adjusting input channels for this encoder is not implemented.")
    return encoder_module

class EnhancedCenterNetUNet(nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.backbone = CenterNetUNetModel(in_channels, num_classes)
        self.fpn = FeaturePyramidNetwork([64, 128, 256, 512], 256)
        self.se_blocks = nn.ModuleList([SqueezeExcitation(c) for c in [64, 128, 256, 512]])

    def forward(self, x):
        features = self.backbone(x)
        for i, feature in enumerate(features):
            features[i] = self.se_blocks[i](feature)
        return self.fpn(features)