#centernet_unet.py
from functools import partial
from typing import List, Dict
from omegaconf import DictConfig
from pytorch_toolbelt.modules import (
    ACT_RELU,
    decoders as D,
    EncoderModule,
    ResidualDeconvolutionUpsample2d,
)
from torch import nn
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 NAS as NeuralArchitectureSearch
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,
)
__all__ = ["CenterNetUNetModel"]

class CenterNetUNetModel(nn.Module):
    def __init__(
        self,
        encoder: EncoderModule,
        decoder_features: List[int],
        head: nn.Module,
        box_coder,
        neck_out_channels: int = 256,  # Add this parameter
        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)],
        **kwargs
    ):
        super().__init__()
        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 = {
            "nearest": nn.UpsamplingNearest2d,
            "bilinear": nn.UpsamplingBilinear2d,
            "rdtsc": ResidualDeconvolutionUpsample2d,
            "shuffle": nn.PixelShuffle,
        }[upsample_type]

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

        self.head = head
        self.head_output_channels = head.out_channels  # Ensure your head module defines 'out_channels'

        # Ensure self.nas is a module that can be called during forward
        self.nas = NeuralArchitectureSearch(
            search_space={
                'num_layers': (2, 5),
                'num_filters': (32, 128),
                'kernel_size': (3, 5),
                'activation': ['relu', 'leaky_relu']
            },
            evaluation_function=self.evaluate_architecture
        )
        self.nas_output_channels = self.nas.output_channels  # Define this in your NAS module

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

        self.classifier = ConvolutionalClassifier(in_channels=neck_out_channels, num_classes=num_classes)
        self.classifier_output_channels = self.classifier.out_channels
        self.box_coder = box_coder
        self.head_output_channels = getattr(self.head, 'out_channels', 256)
        self.nas_output_channels = getattr(self.nas, 'output_channels', 256)
        self.yolo_output_channels = self.yolo_head.num_anchors * (num_classes + 5)
        self.classifier_output_channels = getattr(self.classifier, 'out_channels', num_classes)

        # Ensure all channel values are integers
        total_input_channels = (
            int(self.head_output_channels) +
            int(self.nas_output_channels) +
            int(self.yolo_output_channels) +
            int(self.classifier_output_channels)
        )
        self.fusion = nn.Conv2d(total_input_channels, self.head_output_channels, kernel_size=1)

        self.dropout = nn.Dropout(p=0.5)

        self.decoder_channels = decoder_features

        encoder_channels = list(encoder.channels)
        num_blocks = len(encoder_channels) - 1  # UNetDecoder expects decoder_features of length len(encoder_channels) - 1

        # Adjust decoder_channels to match num_blocks
        if len(self.decoder_channels) < num_blocks:
            # Extend decoder_channels by repeating the last value
            self.decoder_channels.extend(
                [self.decoder_channels[-1]] * (num_blocks - len(self.decoder_channels))
            )
        elif len(self.decoder_channels) > num_blocks:
            # Trim decoder_channels to the required length
            self.decoder_channels = self.decoder_channels[:num_blocks]

        self.num_decoded_features = len(encoder_channels)
        num_decoder_stages = self.num_decoded_features

        self.neck_output_channels = neck_out_channels * self.num_decoded_features
   
        self.decoder = DebugUNetDecoder(
            encoder_channels,
            self.decoder_channels,
            unet_block=unet_block,
            upsample_block=upsample_type,
            center=False,  # Ensure upsample_block is correctly passed
        )
        self.neck = FPNNeck(self.decoder_channels, neck_out_channels)
        self.neck_conv = nn.Conv2d(
            in_channels=self.neck_output_channels,
            out_channels=neck_out_channels,
            kernel_size=1
        )
        self.batch_norm = FilterResponseNormalization(self.neck_output_channels)

    def evaluate_architecture(self, architecture):
        # Implement a simple evaluation function
        return random.random()  # Replace with actual evaluation logic

    def forward(self, x):
        x = self.initial_upsample(x)
        print(f"[DEBUG] x shape after initial_upsample: {x.shape}")
        feature_maps = self.encoder(x)
        print(f"[DEBUG] feature_maps type: {type(feature_maps)}")
        if isinstance(feature_maps, list):
            for idx, fmap in enumerate(feature_maps):
                print(f"[DEBUG] feature_maps[{idx}] shape: {fmap.shape}")
        else:
            feature_maps = [feature_maps]
        for extra_block in self.extra_encoder_blocks:
            feature_maps.append(extra_block(feature_maps[-1]))
        decoded_features = self.decoder(feature_maps)
        print(f"[DEBUG] decoded_features type: {type(decoded_features)}")
        if isinstance(decoded_features, list):
            for idx, fmap in enumerate(decoded_features):
                print(f"[DEBUG] decoded_features[{idx}] shape: {fmap.shape}")
        else:
            decoded_features = [decoded_features]
        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)

        #self.head.channels = self.nas_output_channels  # Assuming self.nas_output_channels is the correct channel size
        head_output = self.head(neck_features)
        nas_output = self.nas(neck_features)
        yolo_output = self.yolo_head(neck_features)
        classification_output = self.classifier(neck_features)

        # Fusion layer
        fused_output = self.fusion(torch.cat([head_output, nas_output, yolo_output, classification_output], dim=1))

        return {
            CENTERNET_OUTPUT_OBJECTNESS_MAP: fused_output[:, :1],
            CENTERNET_OUTPUT_VESSEL_MAP: fused_output[:, 1:2],
            CENTERNET_OUTPUT_FISHING_MAP: fused_output[:, 2:3],
            CENTERNET_OUTPUT_SIZE: fused_output[:, 3:4],
            CENTERNET_OUTPUT_OFFSET: fused_output[:, 4:6],
        }

    @classmethod
    def from_config(cls, **config):
        # Extract and pop the encoder configuration
        encoder_config = config.pop('encoder')
        num_channels = config.pop('num_channels', 3)
        model_name = encoder_config.get('model_name', 'tf_efficientnet_b5_ns')
        pretrained = encoder_config.get('pretrained', False)

        # Instantiate the custom encoder
        encoder_module = CustomTimmEncoder(
            model_name=model_name,
            num_channels=num_channels,
            features_only=True,
            pretrained=pretrained
        )

        # Adjust input channels if necessary
        encoder_module = adjust_encoder_input_channels(encoder_module, num_channels)

        # Pop other configurations
        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)])

        # Pop decoder configuration
        decoder_config = config.pop('decoder')
        encoder_channels = list(encoder_module.channels)  
        decoder_features = decoder_config['channels']
        block_type = decoder_config.get("block_type", "UnetBlock")
        upsample_type = decoder_config.get("upsample_type", "bilinear")

        neck_out_channels = config.get('neck_out_channels', 256)
        neck_output_channels = neck_out_channels * len(decoder_features)
        num_blocks = len(encoder_channels) - 1

        # Adjust decoder_channels to match num_blocks
        if len(decoder_features) < num_blocks:
            decoder_features.extend([decoder_features[-1]] * (num_blocks - len(decoder_features)))
        elif len(decoder_features) > num_blocks:
            decoder_features = decoder_features[:num_blocks]

        # Instantiate head and box_coder
        head_config = config.pop('head')
        head = instantiate(head_config, channels=neck_output_channels)

        box_coder_config = config.pop('box_coder')
        box_coder = instantiate(
            box_coder_config,
            output_stride=encoder_module.strides[0] // (head.upsample_factor * input_upsample_factor),
        )

        # Now initialize the model instance with the remaining config
        return cls(
            encoder=encoder_module,
            decoder_features=decoder_features,
            head=head,
            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,
            anchors=anchors,
            **config
        )

def adjust_encoder_input_channels(encoder_module, in_chans):
    # Access the underlying encoder model
    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
        )
        # Initialize new_conv weights
        with torch.no_grad():
            # Copy weights for existing channels
            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:
                # Initialize added 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)


#custom_encoder.py

import timm
import torch.nn as nn
import torch
import torch.nn.functional as F
from pytorch_toolbelt.modules.decoders import UNetDecoder

class CustomTimmEncoder(nn.Module):
    def __init__(self, model_name, num_channels=2, features_only=True, pretrained=False):
        super(CustomTimmEncoder, self).__init__()
        self.encoder = timm.create_model(
            model_name,
            features_only=features_only,
            in_chans=num_channels,
            pretrained=pretrained
        )
        # Extract channels and strides
        self.channels = [f['num_chs'] for f in self.encoder.feature_info]
        self.strides = [f['reduction'] for f in self.encoder.feature_info]
        self.feature_info = self.encoder.feature_info
        
    def forward(self, x):
        features = self.encoder(x)
        for idx, feature in enumerate(features):
            print(f"[DEBUG] Feature {idx} shape: {feature.shape}")
        return features
    
class DebugUNetDecoder(UNetDecoder):
    def __init__(self, encoder_channels, decoder_channels, unet_block, upsample_block, center=False):
        super(DebugUNetDecoder, self).__init__(
            feature_maps=encoder_channels,
            decoder_features=decoder_channels,
            unet_block=unet_block,
            upsample_block=upsample_block
            # Remove 'center=center' if the base class does not accept 'center'
        )
        if not hasattr(self, 'center'):
            self.center = None
        if not hasattr(self, 'up_blocks'):
            self.up_blocks = nn.ModuleList()  # Initialize an empty ModuleList or handle as needed # Initialize an empty ModuleList or handle as needed

    def forward(self, features):
        x = features[-1]
        x = self.center(x) if self.center is not None else x  # Apply center block if it exists
        for i in reversed(range(len(features) - 1)):
            encoder_input = features[i]
            decoder_block = self.up_blocks[i]
            x = decoder_block(x, encoder_input)
            print(f"[DEBUG] Decoder stage {i}:")
            print(f"    x shape after up_block: {x.shape}")
            print(f"    encoder_input shape:       {encoder_input.shape}")
        return x


#the error
python train_multilabel.py +config=configs/multilabel_circle_net.yaml
DEBUG:hydra.core.utils:Setting JobRuntime:name=UNKNOWN_NAME
DEBUG:hydra.core.utils:Setting JobRuntime:name=train_multilabel
Dataset valid_sar_cv
        ignore_low_confidence_detections False
        ignore_low_confidence_labels     False
        filter_low_confidence_objects    False
Train dataset
        Unknown vessels 90
        Unknown fishing 11242
Valid dataset
        Unknown vessels 20
        Unknown fishing 2778
Preparing training dataset: 100%|███| 35/35 [00:00<00:00, 694.31it/s]
[DEBUG] self.box_coder in dataset.__init__: MultilabelCircleNetCoder(image_size=(512, 512), stride=4, max_objects=512, labels_encoding=circle, heatmap_encoding=umich, fixed_radius=None, ignore_index=-1, ignore_low_confidence_detections=False)
Preparing validation dataset: 100%|███| 7/7 [00:00<00:00, 568.54it/s]
Train instances 13933 dataset 35 sampler 8192
Valid instances 3655 dataset 1039
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 528, in <module>
    main()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 119, in run
    ret = run_job(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 524, in main
    CenternetPipeline(cfg).train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 415, in train
    model = self.build_model(config, self.loaders)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 202, in build_model
    model = super().build_model(config)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/pipeline.py", line 395, in build_model
    model = self.get_model(config["model"])
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 379, in get_model
    model: nn.Module = get_detection_model(model_config)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/factory.py", line 27, in get_detection_model
    return instantiate(model_config, _recursive_=False)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 226, in instantiate
    return instantiate_node(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 347, in instantiate_node
    return _call_target(_target_, partial, args, kwargs, full_key)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 92, in _call_target
    return _target_(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 202, in from_config
    encoder_module = CustomTimmEncoder(
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/custom_encoder.py", line 10, in __init__
    self.encoder = timm.create_model(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/timm/models/_factory.py", line 117, in create_model
    model = create_fn(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/timm/models/_registry.py", line 143, in _fn
    warnings.warn(f'Mapping deprecated model name {deprecated_name} to current {current_name}.', stacklevel=2)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/warnings.py", line 109, in _showwarnmsg
    sw(msg.message, msg.category, msg.filename, msg.lineno,
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 53, in warn_with_traceback
    traceback.print_stack(file=log)
/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/timm/models/_factory.py:117: UserWarning:

Mapping deprecated model name tf_efficientnet_b5_ns to current tf_efficientnet_b5.ns_jft_in1k.

[DEBUG] After adjustment, encoder.conv_stem.in_channels: 2
Error executing job with overrides: ['+config=configs/multilabel_circle_net.yaml']
Traceback (most recent call last):
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 92, in _call_target
    return _target_(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 247, in from_config
    return cls(
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 132, in __init__
    self.decoder = DebugUNetDecoder(
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/custom_encoder.py", line 29, in __init__
    super(DebugUNetDecoder, self).__init__(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/pytorch_toolbelt/modules/decoders/unet.py", line 51, in __init__
    elif issubclass(upsample_block, nn.Upsample):
TypeError: issubclass() arg 1 must be a class

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 528, in <module>
    main()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 524, in main
    CenternetPipeline(cfg).train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 415, in train
    model = self.build_model(config, self.loaders)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 202, in build_model
    model = super().build_model(config)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/pipeline.py", line 395, in build_model
    model = self.get_model(config["model"])
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 379, in get_model
    model: nn.Module = get_detection_model(model_config)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/factory.py", line 27, in get_detection_model
    return instantiate(model_config, _recursive_=False)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 226, in instantiate
    return instantiate_node(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 347, in instantiate_node
    return _call_target(_target_, partial, args, kwargs, full_key)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 97, in _call_target
    raise InstantiationException(msg) from e
hydra.errors.InstantiationException: Error in call to target 'xview3.centernet.models.centernet_unet.CenterNetUNetModel.from_config':
TypeError('issubclass() arg 1 must be a class')
full_key: model
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers %