from __future__ import annotations

import torch
import torch.nn.functional as F
from torch import nn

from .unet import ConvBlock, Down, Up


class UNetInstance(nn.Module):
    """U-Net with dual heads: semantic logits + per-pixel embedding."""

    def __init__(
        self,
        in_channels: int = 1,
        num_classes: int = 3,
        base_channels: int = 32,
        embedding_dim: int = 8,
        normalize_embeddings: bool = True,
    ):
        super().__init__()
        c1 = base_channels
        c2 = base_channels * 2
        c3 = base_channels * 4
        c4 = base_channels * 8

        self.inc = ConvBlock(in_channels, c1)
        self.down1 = Down(c1, c2)
        self.down2 = Down(c2, c3)
        self.down3 = Down(c3, c4)

        self.up1 = Up(c4, c3)
        self.up2 = Up(c3, c2)
        self.up3 = Up(c2, c1)

        self.seg_head = nn.Conv2d(c1, num_classes, kernel_size=1)
        self.emb_head = nn.Conv2d(c1, embedding_dim, kernel_size=1)
        self.normalize_embeddings = bool(normalize_embeddings)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x = self.up1(x4, x3)
        x = self.up2(x, x2)
        x = self.up3(x, x1)

        seg_logits = self.seg_head(x)
        embeddings = self.emb_head(x)
        if self.normalize_embeddings:
            embeddings = F.normalize(embeddings, dim=1)
        return seg_logits, embeddings
