"""Morphological tagging head with per-category classifiers for UPOS, XPOS, and morphological features."""

import torch
import torch.nn as nn
from torch import Tensor


class MorphoHead(nn.Module):
    """Classification head with separate outputs for UPOS, XPOS, and each morphological feature category.

    Includes a feature co-occurrence refinement layer that takes the initial
    per-category logits and refines them jointly, allowing the model to learn
    cross-category dependencies (e.g., Case+Number patterns in Polish).
    """

    def __init__(
        self,
        hidden_size: int,
        num_upos: int,
        num_xpos: int,
        feat_vocab_sizes: dict[str, int],
        dropout: float = 0.1,
    ):
        super().__init__()
        self.num_upos = num_upos
        self.num_xpos = num_xpos
        self.feat_categories = list(feat_vocab_sizes.keys())

        mid_size = hidden_size // 2

        self.backbone = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden_size, mid_size),
            nn.LayerNorm(mid_size),
            nn.GELU(),
            nn.Dropout(dropout),
        )

        self.upos_head = nn.Linear(mid_size, num_upos)
        self.xpos_head = nn.Linear(mid_size, num_xpos)
        self.feat_heads = nn.ModuleDict({
            cat: nn.Linear(mid_size, vocab_size)
            for cat, vocab_size in feat_vocab_sizes.items()
        })

        # Feature co-occurrence refinement:
        # Concatenate all initial feat logits + UPOS logits, refine through MLP,
        # then produce residual corrections for each category.
        total_feat_logits = sum(feat_vocab_sizes.values()) + num_upos
        refine_hidden = 256
        self.feat_refine = nn.Sequential(
            nn.Linear(total_feat_logits, refine_hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(refine_hidden, refine_hidden),
            nn.GELU(),
            nn.Dropout(dropout),
        )
        self.feat_refine_heads = nn.ModuleDict({
            cat: nn.Linear(refine_hidden, vocab_size)
            for cat, vocab_size in feat_vocab_sizes.items()
        })

    def forward(self, hidden_states: Tensor) -> dict[str, Tensor | dict[str, Tensor]]:
        """Compute morphological label logits.

        Args:
            hidden_states: Word-level representations [batch, seq_len, hidden_size]

        Returns:
            Dict with keys:
                "upos": [batch, seq_len, num_upos]
                "xpos": [batch, seq_len, num_xpos]
                "feats": dict mapping category name to [batch, seq_len, vocab_size]
        """
        h = self.backbone(hidden_states)

        upos_logits = self.upos_head(h)
        xpos_logits = self.xpos_head(h)

        # Initial per-category feat predictions
        initial_feats = {cat: self.feat_heads[cat](h) for cat in self.feat_categories}

        # Co-occurrence refinement: concatenate all feat logits + UPOS, refine jointly
        all_logits = [upos_logits] + [initial_feats[cat] for cat in self.feat_categories]
        combined = torch.cat(all_logits, dim=-1)  # [batch, seq, total_feat_logits]
        refined_h = self.feat_refine(combined)

        # Add residual corrections to initial feat predictions
        refined_feats = {
            cat: initial_feats[cat] + self.feat_refine_heads[cat](refined_h)
            for cat in self.feat_categories
        }

        return {
            "upos": upos_logits,
            "xpos": xpos_logits,
            "feats": refined_feats,
        }
