#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_block,
        )
        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
        # In from_config method
        # 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)
        
        # 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 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 after applying your suggestions


DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x308b596a0>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x308b7ad00>
DEBUG:rasterio._io:Output nodata value read from file: -32768.0
DEBUG:rasterio._io:Output nodata values: [-32768.0]
DEBUG:rasterio._io:all_valid: False
DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=18600, row_off=6922, width=512, height=512)
DEBUG:rasterio._io:IO window xoff=18600.0 yoff=6922.0 width=512.0 height=512.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x308b7ad00>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x308b596a0> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x308b596a0>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x308b7ad00>
[DEBUG] Channel 0 shape: (512, 512)
[DEBUG] Channel 1 shape: (512, 512)
(192.0, 0.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(201.0, 78.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(356.0, 181.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(316.0, 275.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(262.0, 245.0, 2.856621265411377) [<class 'float'>, <class 'float'>, <class 'float'>]
(33.0, 416.0, 14.250475883483887) [<class 'float'>, <class 'float'>, <class 'float'>]
Error executing job with overrides: ['+config=configs/multilabel_circle_net.yaml']
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 206, in build_model
    model = nas.search(config.nas.num_iterations, loaders['train'], loaders['valid'])
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 36, in search
    scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 36, in <listcomp>
    scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 112, in evaluate_model
    outputs = model(inputs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 162, in forward
    decoded_features = self.decoder(feature_maps)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/custom_encoder.py", line 30, in forward
    x = self.center(x) if self.center is not None else x  # Apply center block if it exists
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1729, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'DebugUNetDecoder' object has no attribute 'center'
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers %