from typing import Optional

import torch
from transformers import PreTrainedModel
from transformers.modeling_outputs import SemanticSegmenterOutput

from .configuration_cond_unet import CondUNetConfig

try:
    from .unet_attn import UNet2DAttn
except ModuleNotFoundError:
    # The release exporter bundles these modules into the Hub repository root.
    from nets.unet_attn import UNet2DAttn


class CondUNetForSemanticSegmentation(PreTrainedModel):
    config_class = CondUNetConfig
    main_input_name = "pixel_values"

    def __init__(self, config: CondUNetConfig):
        super().__init__(config)
        self.unet = UNet2DAttn(
            in_channels=config.in_channels,
            num_classes=config.num_labels,
            n_organs=config.n_organs,
            size=config.size,
            depth=config.depth,
            attn_start=config.attn_start,
            use_attn=config.use_attn,
            img_size=config.image_size,
            patch_size=config.patch_size,
            emb_dim=config.emb_dim,
            n_heads=config.n_heads,
            distill=False,
            distill_unet=False,
            use_dwt=config.use_dwt,
            wavelet=config.wavelet,
            dwt_bands=config.dwt_bands,
            use_shape=config.use_shape,
            shape_res=config.shape_res,
        )
        self.post_init()

    def forward(
        self,
        pixel_values: torch.FloatTensor,
        organ_id: Optional[torch.LongTensor] = None,
        labels: Optional[torch.FloatTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ):
        if organ_id is None:
            organ_id = torch.full(
                (pixel_values.shape[0],),
                self.config.unknown_organ_id,
                device=pixel_values.device,
                dtype=torch.long,
            )
        else:
            organ_id = organ_id.to(device=pixel_values.device, dtype=torch.long)

        outputs = self.unet(
            pixel_values=pixel_values,
            organ_id=organ_id,
            masks=labels,
            **kwargs,
        )
        logits = outputs["logits"]
        if logits.ndim == 3:
            logits = logits.unsqueeze(1)
        loss = outputs["loss"] if labels is not None else None
        if return_dict is False:
            return (loss, logits) if loss is not None else (logits,)
        return SemanticSegmenterOutput(loss=loss, logits=logits)