"""Compose independent token augmentation, transformation and MIL pooling stages."""
from __future__ import annotations

import inspect
from typing import Any, Dict, Optional, Set, Tuple, Union

import torch
import torch.nn as nn

from .cross_shuffle import CrossTokenShuffle
from .mil_head_beta import MILHead


def _forward_param_names(module: nn.Module) -> Set[str]:
    try:
        return set(inspect.signature(module.forward).parameters.keys())
    except (TypeError, ValueError):
        return set()


class CrossShuffleStage(nn.Module):
    """Optional cross-batch token replacement; disabled stages act as the identity."""

    def __init__(self, use: bool = True) -> None:
        super().__init__()
        self._use = use
        self.op: Optional[CrossTokenShuffle]
        if use:
            self.op = CrossTokenShuffle()
        else:
            self.op = None

    def forward(
        self,
        x: torch.Tensor,
        y: Optional[torch.Tensor],
        *,
        apply: bool,
        replace_ratio: float = 0.5,
        with_replacement: bool = False,
        detach_inserted: bool = True,
        min_keep_original: int = 0,
    ) -> torch.Tensor:
        if not self._use or not apply or self.op is None or y is None:
            return x
        if y.dim() != 1 or y.shape[0] != x.shape[0]:
            raise ValueError(
                f"y must be 1D [B], got {tuple(y.shape)} for B={x.shape[0]}"
            )
        return self.op(
            x, y,
            replace_ratio=replace_ratio,
            with_replacement=with_replacement,
            detach_inserted=detach_inserted,
            min_keep_original=min_keep_original,
        )


class TransStage(nn.Module):
    """Apply an optional token transformation with supported keyword arguments.
    
    A missing core acts as the identity. Fixed-length masking can be passed
    through key_padding_mask without dropping tokens.
    """

    def __init__(self, core: Optional[nn.Module] = None) -> None:
        super().__init__()
        self.core = core
        self._params = _forward_param_names(core) if core is not None else set()

    def forward(
        self,
        x: torch.Tensor,
        *,
        p_shuffle: float = 0.0,
        p_drop: float = 0.0,
        key_padding_mask: Optional[torch.Tensor] = None,
        **kwargs: Any,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        if self.core is None:
            return x
        if "p_shuffle" in self._params:
            call_kw: Dict[str, Any] = {}
            for name, val in (
                ("p_shuffle", p_shuffle),
                ("p_drop", p_drop),
                ("key_padding_mask", key_padding_mask),
            ):
                if name in self._params:
                    call_kw[name] = val
            for k, v in kwargs.items():
                if k in self._params:
                    call_kw[k] = v
            return self.core(x, **call_kw)
        kw2: Dict[str, Any] = {}
        if key_padding_mask is not None and "src_key_padding_mask" in self._params:
            kw2["src_key_padding_mask"] = key_padding_mask
        for k, v in kwargs.items():
            if k in self._params:
                kw2[k] = v
        return self.core(x, **kw2) if kw2 else self.core(x)


class MILHeadStage(nn.Module):
    """Classify instances and pool tokens [B, N, D] using the wrapped MIL head."""

    def __init__(self, **mil_kw: Any) -> None:
        super().__init__()
        self.head = MILHead(**mil_kw)

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> Dict[str, Any]:
        return self.head(x, mask=mask)


class MILCore(nn.Module):
    """Compose Cross-Shuffle, token transformation and MIL head stages.
    
    Fixed-length bags can use attention masks instead of token dropping.
    Optional masks and transformation arguments are forwarded to supported stages.
    """

    def __init__(
        self,
        d_in: int,
        num_classes: int = 2,
        mil_head_pooling: str = "attn",
        mil_head_kwargs: Optional[Dict[str, Any]] = None,
        *,
        use_cross_shuffle: bool = True,
        trans: Optional[nn.Module] = None,
        allow_trans_token_drop: bool = False,
    ) -> None:
        super().__init__()
        self.allow_trans_token_drop = bool(allow_trans_token_drop)
        self.shuffle = CrossShuffleStage(use=use_cross_shuffle)
        self.trans = TransStage(core=trans)
        mil_kw: Dict[str, Any] = {"d_in": d_in, "num_classes": num_classes, "pooling": mil_head_pooling}
        if mil_head_kwargs:
            mil_kw.update(mil_head_kwargs)
        self.head = MILHeadStage(**mil_kw)

    def forward(
        self,
        tokens: torch.Tensor,
        y: Optional[torch.Tensor] = None,
        mask: Optional[torch.Tensor] = None,
        *,
        do_cross_shuffle: bool = False,
        replace_ratio: float = 0.5,
        with_replacement: bool = False,
        detach_inserted: bool = True,
        cross_shuffle_min_keep: int = 0,
        p_shuffle: float = 0.0,
        p_drop: float = 0.0,
        trans_key_padding_mask: Optional[torch.Tensor] = None,
        trans_extra: Optional[Dict[str, Any]] = None,
        return_dict: bool = False,
    ):
        x = self.shuffle(
            tokens, y,
            apply=do_cross_shuffle,
            replace_ratio=replace_ratio,
            with_replacement=with_replacement,
            detach_inserted=detach_inserted,
            min_keep_original=cross_shuffle_min_keep,
        )
        trans_p_drop = float(p_drop) if self.allow_trans_token_drop else 0.0
        tex = dict(trans_extra or {})
        trans_out = self.trans(
            x,
            p_shuffle=p_shuffle,
            p_drop=trans_p_drop,
            key_padding_mask=trans_key_padding_mask,
            **tex,
        )
        if isinstance(trans_out, tuple) and len(trans_out) == 2:
            x, trans_mil_mask = trans_out
            if mask is not None:
                raise ValueError(
                    "MILCore: `trans` returned (x, mil_valid); omit `mask` and use the trans-stage mask."
                )
            mask = trans_mil_mask
        else:
            x = trans_out  # type: ignore[assignment]
        if mask is not None and mask.shape != x.shape[:2]:
            raise ValueError(
                f"mask must be [B,N], got {tuple(mask.shape)} vs x {tuple(x.shape[:2])}"
            )
        mil_out = self.head(x, mask=mask)
        if not return_dict:
            return mil_out["bag_logits"]
        return {"bag_logits": mil_out["bag_logits"], "mil_head": mil_out, "features": x}

    def forward_logits(
        self,
        tokens: torch.Tensor,
        y: Optional[torch.Tensor] = None,
        mask: Optional[torch.Tensor] = None,
        **kwargs: Any,
    ) -> torch.Tensor:
        extra = dict(kwargs)
        extra.pop("return_dict", None)
        return self.forward(tokens, y, mask, return_dict=False, **extra)


CrossShuffleTransMIL = MILCore