src/models/fusion.py

from dataclasses import dataclass
from typing import Dict, Optional, Tuple

import torch
import torch.nn as nn


@dataclass
class StructuralFusionConfig:
    """
    Cấu hình cho StructuralVisualFusion.

    Args:
        d_model: Hidden dimension của visual tokens từ ViT/VLM.
        topo_feature_dim: Số chiều feature topology/morphology mỗi patch.
        hidden_dim: Hidden dimension của MLP chiếu topo_features -> d_model.
        prior_strength: Hệ số spatial modulation ban đầu cho lesion prior.
        topo_strength: Hệ số enrichment ban đầu cho topo embedding.
        dropout: Dropout trong topo MLP.
        use_layer_norm: Chuẩn hóa visual/topo branch để ổn định training.
        learnable_gates: Nếu True, alpha/beta là tham số học được.
        normalize_prior: Nếu True, normalize prior_mask về [0, 1] theo từng sample.
        residual_scale: Scale residual cuối để tránh làm lệch mạnh pretrained tokens.
    """

    d_model: int
    topo_feature_dim: int
    hidden_dim: Optional[int] = None
    prior_strength: float = 0.30
    topo_strength: float = 0.10
    dropout: float = 0.10
    use_layer_norm: bool = True
    learnable_gates: bool = True
    normalize_prior: bool = True
    residual_scale: float = 1.0

    def __post_init__(self):
        if self.d_model <= 0:
            raise ValueError("d_model phải > 0")
        if self.topo_feature_dim <= 0:
            raise ValueError("topo_feature_dim phải > 0")
        if self.hidden_dim is None:
            self.hidden_dim = max(self.d_model // 2, self.topo_feature_dim * 4)
        if self.hidden_dim <= 0:
            raise ValueError("hidden_dim phải > 0")


class StructuralVisualFusion(nn.Module):
    """
    Fusion module nối ViT visual tokens với lesion prior và topology features.

    Forward logic:
        V_fused = V * (1 + alpha * P) + beta * MLP(F_topo)

    Trong đó:
        V: visual_tokens, shape [B, N, D]
        P: prior_mask, shape [B, N] hoặc [B, H, W]
        F_topo: topo_features, shape [B, N, F] hoặc [B, H, W, F]

    Output chính:
        fused_tokens, shape [B, N, D]
    """

    def __init__(self, config: StructuralFusionConfig):
        super().__init__()
        self.config = config

        self.visual_norm = nn.LayerNorm(config.d_model) if config.use_layer_norm else nn.Identity()
        self.topo_norm = nn.LayerNorm(config.topo_feature_dim) if config.use_layer_norm else nn.Identity()

        self.topo_projector = nn.Sequential(
            nn.Linear(config.topo_feature_dim, config.hidden_dim),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.hidden_dim, config.d_model),
            nn.Dropout(config.dropout),
        )

        if config.learnable_gates:
            self.prior_gate = nn.Parameter(torch.tensor(float(config.prior_strength)))
            self.topo_gate = nn.Parameter(torch.tensor(float(config.topo_strength)))
        else:
            self.register_buffer("prior_gate", torch.tensor(float(config.prior_strength)))
            self.register_buffer("topo_gate", torch.tensor(float(config.topo_strength)))

        self.output_norm = nn.LayerNorm(config.d_model) if config.use_layer_norm else nn.Identity()

    def forward(
        self,
        visual_tokens: torch.Tensor,
        prior_mask: Optional[torch.Tensor] = None,
        topo_features: Optional[torch.Tensor] = None,
        return_diagnostics: bool = True,
    ):
        """
        Args:
            visual_tokens: Tensor [B, N, D].
            prior_mask: Optional tensor [B, N] hoặc [B, H, W].
            topo_features: Optional tensor [B, N, F] hoặc [B, H, W, F].
            return_diagnostics: Nếu True trả về dict debug; nếu False chỉ trả fused_tokens.

        Returns:
            dict hoặc Tensor:
                {
                    "fused_tokens": [B, N, D],
                    "prior_weights": [B, N, 1] | None,
                    "topo_embedding": [B, N, D] | None,
                    "prior_gate": scalar,
                    "topo_gate": scalar,
                }
        """
        self._validate_visual_tokens(visual_tokens)
        batch_size, num_tokens, d_model = visual_tokens.shape

        fused_tokens = self.visual_norm(visual_tokens)
        prior_weights = None
        topo_embedding = None

        if prior_mask is not None:
            prior_weights = self._prepare_prior_mask(
                prior_mask,
                batch_size=batch_size,
                num_tokens=num_tokens,
                device=visual_tokens.device,
                dtype=visual_tokens.dtype,
            )
            fused_tokens = fused_tokens * (1.0 + self.prior_gate * prior_weights)

        if topo_features is not None:
            topo_features = self._prepare_topo_features(
                topo_features,
                batch_size=batch_size,
                num_tokens=num_tokens,
                device=visual_tokens.device,
                dtype=visual_tokens.dtype,
            )
            topo_features = self.topo_norm(topo_features)
            topo_embedding = self.topo_projector(topo_features)
            fused_tokens = fused_tokens + self.topo_gate * topo_embedding

        if self.config.residual_scale != 1.0:
            fused_tokens = visual_tokens + self.config.residual_scale * (fused_tokens - visual_tokens)

        fused_tokens = self.output_norm(fused_tokens)

        if not return_diagnostics:
            return fused_tokens

        return {
            "fused_tokens": fused_tokens,
            "prior_weights": prior_weights,
            "topo_embedding": topo_embedding,
            "prior_gate": self.prior_gate.detach().clone(),
            "topo_gate": self.topo_gate.detach().clone(),
        }

    def _validate_visual_tokens(self, visual_tokens: torch.Tensor) -> None:
        if not torch.is_tensor(visual_tokens):
            raise TypeError("visual_tokens phải là torch.Tensor")
        if visual_tokens.ndim != 3:
            raise ValueError(
                f"visual_tokens phải có shape [B, N, D], nhận {tuple(visual_tokens.shape)}"
            )
        if visual_tokens.shape[-1] != self.config.d_model:
            raise ValueError(
                f"visual_tokens D={visual_tokens.shape[-1]} không khớp d_model={self.config.d_model}"
            )

    def _prepare_prior_mask(
        self,
        prior_mask: torch.Tensor,
        batch_size: int,
        num_tokens: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        if not torch.is_tensor(prior_mask):
            prior_mask = torch.as_tensor(prior_mask)

        prior_mask = prior_mask.to(device=device, dtype=dtype)

        if prior_mask.ndim == 3:
            prior_mask = prior_mask.flatten(start_dim=1)
        elif prior_mask.ndim == 2:
            pass
        else:
            raise ValueError(
                f"prior_mask phải có shape [B, N] hoặc [B, H, W], nhận {tuple(prior_mask.shape)}"
            )

        if prior_mask.shape[0] != batch_size:
            raise ValueError(
                f"prior_mask batch={prior_mask.shape[0]} không khớp visual batch={batch_size}"
            )
        if prior_mask.shape[1] != num_tokens:
            raise ValueError(
                f"prior_mask N={prior_mask.shape[1]} không khớp visual tokens N={num_tokens}"
            )

        if self.config.normalize_prior:
            prior_mask = self._minmax_normalize_per_sample(prior_mask)

        return prior_mask.unsqueeze(-1)

    def _prepare_topo_features(
        self,
        topo_features: torch.Tensor,
        batch_size: int,
        num_tokens: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        if not torch.is_tensor(topo_features):
            topo_features = torch.as_tensor(topo_features)

        topo_features = topo_features.to(device=device, dtype=dtype)

        if topo_features.ndim == 4:
            topo_features = topo_features.flatten(start_dim=1, end_dim=2)
        elif topo_features.ndim == 3:
            pass
        else:
            raise ValueError(
                "topo_features phải có shape [B, N, F] hoặc [B, H, W, F], "
                f"nhận {tuple(topo_features.shape)}"
            )

        if topo_features.shape[0] != batch_size:
            raise ValueError(
                f"topo_features batch={topo_features.shape[0]} không khớp visual batch={batch_size}"
            )
        if topo_features.shape[1] != num_tokens:
            raise ValueError(
                f"topo_features N={topo_features.shape[1]} không khớp visual tokens N={num_tokens}"
            )
        if topo_features.shape[2] != self.config.topo_feature_dim:
            raise ValueError(
                f"topo_features F={topo_features.shape[2]} không khớp "
                f"topo_feature_dim={self.config.topo_feature_dim}"
            )

        return topo_features

    @staticmethod
    def _minmax_normalize_per_sample(values: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
        min_values = values.amin(dim=1, keepdim=True)
        max_values = values.amax(dim=1, keepdim=True)
        norm = (values - min_values) / (max_values - min_values + eps)
        return torch.nan_to_num(norm, nan=0.0, posinf=0.0, neginf=0.0)


def build_structural_visual_fusion(
    d_model: int,
    topo_feature_dim: int,
    **kwargs,
) -> StructuralVisualFusion:
    """Factory tiện dụng cho training scripts."""
    config = StructuralFusionConfig(
        d_model=d_model,
        topo_feature_dim=topo_feature_dim,
        **kwargs,
    )
    return StructuralVisualFusion(config)


__all__ = [
    "StructuralFusionConfig",
    "StructuralVisualFusion",
    "build_structural_visual_fusion",
]