"""Main UD Parser model combining all components."""

from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from transformers import AutoModel

from ..config import ModelConfig
from .lemma_head import LemmaHead
from .dependency_head import DependencyHead
from .morpho_head import MorphoHead


def extract_word_representations(
    hidden_states: Tensor,
    word_start_mask: Tensor,
    word_mask: Tensor,
) -> Tensor:
    """Extract word-level representations from subword representations.

    Uses first-subword pooling: for each word, take the representation
    of its first subword token.

    Args:
        hidden_states: Subword-level representations [batch, subword_seq_len, hidden_dim]
        word_start_mask: Boolean mask indicating first subword of each word [batch, subword_seq_len]
        word_mask: Boolean mask for valid words [batch, max_words]

    Returns:
        word_repr: Word-level representations [batch, max_words, hidden_dim]
    """
    batch_size, _, hidden_dim = hidden_states.shape
    max_words = word_mask.size(1)

    # Initialize output tensor
    word_repr = torch.zeros(batch_size, max_words, hidden_dim, device=hidden_states.device)

    # Extract word representations for each batch
    for b in range(batch_size):
        # Get indices where word_start_mask is True
        word_indices = word_start_mask[b].nonzero(as_tuple=True)[0]
        n_words = min(len(word_indices), max_words)

        if n_words > 0:
            word_repr[b, :n_words] = hidden_states[b, word_indices[:n_words]]

    return word_repr


class UDParser(nn.Module):
    """Multi-task UD parsing model with LoRA adapters.

    Combines:
    - XLM-RoBERTa encoder with optional LoRA
    - Decomposed morphological tagging head (separate UPOS, XPOS, per-feat-category)
    - Lemmatization head (character-level seq2seq)
    - Dependency parsing head (biaffine attention)
    """

    def __init__(self, config: ModelConfig):
        """Initialize UDParser.

        Args:
            config: Model configuration.
        """
        super().__init__()
        self.config = config

        # Load base encoder (suppress "unexpected keys" report for lm_head
        # when loading base model from a masked-LM checkpoint)
        import transformers.utils.logging as _hf_logging
        _prev = _hf_logging.get_verbosity()
        _hf_logging.set_verbosity_error()
        try:
            self.encoder = AutoModel.from_pretrained(config.base_model)
        finally:
            _hf_logging.set_verbosity(_prev)
        hidden_size = self.encoder.config.hidden_size

        # Apply LoRA if enabled
        if config.use_lora:
            from peft import LoraConfig, get_peft_model

            lora_config = LoraConfig(
                r=config.lora_r,
                lora_alpha=config.lora_alpha,
                target_modules=config.lora_target_modules,
                lora_dropout=config.lora_dropout,
                bias="none",
                task_type="FEATURE_EXTRACTION",
            )
            self.encoder = get_peft_model(self.encoder, lora_config)

        # Enable gradient checkpointing to save memory
        if getattr(config, "gradient_checkpointing", False):
            self.encoder.gradient_checkpointing_enable()

        # Root embedding for dependency parsing (position 0 = root)
        self.root_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
        nn.init.normal_(self.root_embedding, mean=0, std=0.02)

        # Task heads
        self.morpho_head = MorphoHead(
            hidden_size=hidden_size,
            num_upos=config.num_upos,
            num_xpos=config.num_xpos,
            feat_vocab_sizes=config.feat_vocab_sizes,
            dropout=config.dropout,
        )

        # Morpho-to-lemma conditioning: project morpho logits and concat with word_repr
        total_morpho_logits = config.num_upos + sum(config.feat_vocab_sizes.values())
        self.morpho_to_lemma_proj = nn.Sequential(
            nn.Linear(total_morpho_logits, config.morpho_to_lemma_proj_dim),
            nn.GELU(),
            nn.Dropout(config.dropout),
        )
        lemma_input_size = hidden_size + config.morpho_to_lemma_proj_dim

        self.lemma_head = LemmaHead(
            hidden_size=lemma_input_size,
            char_embed_dim=config.char_embed_dim,
            char_vocab_size=config.char_vocab_size,
            input_proj_dim=config.lemma_input_proj_dim,
            input_proj_dropout=config.lemma_input_proj_dropout,
            cnn_filters=config.lemma_cnn_filters,
            cnn_kernel_sizes=config.lemma_cnn_kernel_sizes,
            cnn_dilations=config.lemma_cnn_dilations,
            cnn_paddings=config.lemma_cnn_paddings,
            dropout=config.dropout,
            gru_hidden_dim=config.lemma_gru_hidden_dim,
        )

        self.dependency_head = DependencyHead(
            hidden_size=hidden_size,
            arc_hidden_dim=config.arc_hidden_dim,
            rel_hidden_dim=config.rel_hidden_dim,
            num_deprels=config.num_deprels,
            dropout=config.dropout,
        )

        # Uncertainty-based multi-task loss weighting (Kendall et al. 2018)
        if config.use_uncertainty_weighting:
            self.log_vars = nn.ParameterDict({
                "upos": nn.Parameter(torch.zeros(1)),
                "xpos": nn.Parameter(torch.zeros(1)),
                "feats": nn.Parameter(torch.zeros(1)),
                "arc": nn.Parameter(torch.zeros(1)),
                "rel": nn.Parameter(torch.zeros(1)),
                "lemma": nn.Parameter(torch.zeros(1)),
            })

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Tensor,
        word_start_mask: Tensor,
        word_mask: Tensor,
        form_chars: Optional[Tensor] = None,
        upos_labels: Optional[Tensor] = None,
        xpos_labels: Optional[Tensor] = None,
        feats_labels: Optional[Tensor] = None,
        head_labels: Optional[Tensor] = None,
        deprel_labels: Optional[Tensor] = None,
        lemma_chars: Optional[Tensor] = None,
        **kwargs,
    ) -> dict[str, Tensor]:
        """Forward pass through the model.

        Args:
            input_ids: Subword token IDs [batch, subword_seq_len]
            attention_mask: Attention mask [batch, subword_seq_len]
            word_start_mask: First subword mask [batch, subword_seq_len]
            word_mask: Valid word mask [batch, max_words]
            form_chars: Input word characters [batch, max_words, max_char_len]
            upos_labels: UPOS labels [batch, max_words]
            xpos_labels: XPOS labels [batch, max_words]
            feats_labels: Feature labels [batch, max_words, num_feat_categories]
            head_labels: Head indices [batch, max_words]
            deprel_labels: Dependency relation labels [batch, max_words]
            lemma_chars: Target lemma characters [batch, max_words, max_lemma_len]

        Returns:
            Dictionary containing logits and optionally losses.
        """
        # Encode with transformer
        encoder_output = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )
        hidden_states = encoder_output.last_hidden_state  # [batch, subword_len, hidden]

        # Extract word-level representations
        word_repr = extract_word_representations(
            hidden_states, word_start_mask, word_mask
        )  # [batch, max_words, hidden]

        # Morphological tagging (decomposed UPOS, XPOS, per-feat-category)
        morpho_out = self.morpho_head(word_repr)

        # Dependency parsing - add root embedding at position 0
        batch_size = word_repr.size(0)
        root_expanded = self.root_embedding.expand(batch_size, -1, -1)  # [batch, 1, hidden]
        word_repr_with_root = torch.cat([root_expanded, word_repr], dim=1)  # [batch, 1+max_words, hidden]

        arc_logits, rel_logits = self.dependency_head(word_repr_with_root)
        # arc_logits: [batch, 1+max_words, 1+max_words]
        # rel_logits: [batch, 1+max_words, 1+max_words, num_deprels]
        # Remove row 0 (root doesn't have a head), keep column 0 (words can point to root)
        arc_logits = arc_logits[:, 1:, :]  # [batch, max_words, 1+max_words]
        rel_logits = rel_logits[:, 1:, :, :]  # [batch, max_words, 1+max_words, num_deprels]

        # Enrich word_repr with morpho predictions for lemmatizer
        morpho_logits_list = [morpho_out["upos"]]
        for cat in self.morpho_head.feat_categories:
            morpho_logits_list.append(morpho_out["feats"][cat])
        morpho_concat = torch.cat(morpho_logits_list, dim=-1)
        morpho_proj = self.morpho_to_lemma_proj(morpho_concat)
        word_repr_for_lemma = torch.cat([word_repr, morpho_proj], dim=-1)

        # Lemmatization
        lemma_logits = None
        if form_chars is not None:
            lemma_logits = self.lemma_head(word_repr_for_lemma, form_chars, lemma_chars)
            # lemma_logits: [batch, max_words, max_lemma_len, char_vocab_size]

        outputs = {
            "upos_logits": morpho_out["upos"],
            "xpos_logits": morpho_out["xpos"],
            "feats_logits": morpho_out["feats"],
            "arc_logits": arc_logits,
            "rel_logits": rel_logits,
            "lemma_logits": lemma_logits,
            "word_repr": word_repr,
            "word_repr_for_lemma": word_repr_for_lemma,
        }

        # Compute losses if labels provided
        if upos_labels is not None:
            upos_loss = self._compute_ce_loss(morpho_out["upos"], upos_labels, word_mask)
            xpos_loss = self._compute_ce_loss(morpho_out["xpos"], xpos_labels, word_mask)

            feat_losses = []
            if feats_labels is not None and morpho_out["feats"]:
                for i, cat in enumerate(self.morpho_head.feat_categories):
                    cat_logits = morpho_out["feats"][cat]
                    cat_labels = feats_labels[:, :, i]
                    feat_losses.append(self._compute_ce_loss(cat_logits, cat_labels, word_mask))
            feats_loss = (
                torch.stack(feat_losses).mean() if feat_losses
                else torch.tensor(0.0, device=upos_labels.device)
            )

            outputs["upos_loss"] = upos_loss
            outputs["xpos_loss"] = xpos_loss
            outputs["feats_loss"] = feats_loss
            outputs["morpho_loss"] = upos_loss + xpos_loss + self.config.feats_weight * feats_loss

        if head_labels is not None:
            outputs["arc_loss"] = self._compute_arc_loss(
                arc_logits, head_labels, word_mask
            )

        if deprel_labels is not None and head_labels is not None:
            outputs["rel_loss"] = self._compute_rel_loss(
                rel_logits, head_labels, deprel_labels, word_mask
            )

        if lemma_chars is not None and lemma_logits is not None:
            outputs["lemma_loss"] = self._compute_lemma_loss(
                lemma_logits, lemma_chars, word_mask
            )

        # Total loss
        if "morpho_loss" in outputs:
            if self.config.use_uncertainty_weighting:
                # Kendall et al. 2018: L = sum( exp(-log_var_i) * L_i + log_var_i )
                _dev = outputs["upos_loss"].device
                _zero = torch.tensor(0.0, device=_dev)
                outputs["loss"] = (
                    torch.exp(-self.log_vars["upos"]) * outputs["upos_loss"] + self.log_vars["upos"]
                    + torch.exp(-self.log_vars["xpos"]) * outputs["xpos_loss"] + self.log_vars["xpos"]
                    + torch.exp(-self.log_vars["feats"]) * outputs["feats_loss"] + self.log_vars["feats"]
                    + torch.exp(-self.log_vars["arc"]) * outputs.get("arc_loss", _zero) + self.log_vars["arc"]
                    + torch.exp(-self.log_vars["rel"]) * outputs.get("rel_loss", _zero) + self.log_vars["rel"]
                    + torch.exp(-self.log_vars["lemma"]) * outputs.get("lemma_loss", _zero) + self.log_vars["lemma"]
                )
                outputs["uncertainty_weights"] = {
                    k: torch.exp(-v).item() for k, v in self.log_vars.items()
                }
            else:
                outputs["loss"] = (
                    self.config.morpho_weight * outputs["morpho_loss"]
                    + self.config.arc_weight * outputs["arc_loss"]
                    + self.config.rel_weight * outputs.get("rel_loss", 0)
                    + self.config.lemma_weight * outputs.get("lemma_loss", 0)
                )

        return outputs

    def _compute_ce_loss(
        self, logits: Tensor, labels: Tensor, mask: Tensor
    ) -> Tensor:
        """Compute masked cross-entropy loss."""
        logits_flat = logits.view(-1, logits.size(-1))
        labels_flat = labels.view(-1)
        mask_flat = mask.view(-1).float()
        num_valid = mask_flat.sum().clamp(min=1)

        loss = F.cross_entropy(logits_flat, labels_flat, reduction="none")
        return (loss * mask_flat).sum() / num_valid

    def _compute_arc_loss(
        self, logits: Tensor, labels: Tensor, mask: Tensor
    ) -> Tensor:
        """Compute arc prediction loss.

        logits: [batch, max_words, 1+max_words] - includes root at position 0
        labels: [batch, max_words] - head indices (0=root, 1..N=words)
        """
        num_heads = logits.size(-1)  # 1 + max_words

        # Check for out-of-range labels
        if (labels >= num_heads).any():
            labels = labels.clamp(max=num_heads - 1)

        # Flatten (use reshape as logits may not be contiguous after slicing)
        logits_flat = logits.reshape(-1, num_heads)
        labels_flat = labels.reshape(-1)
        mask_flat = mask.reshape(-1)

        # Compute loss only on valid positions
        loss = F.cross_entropy(logits_flat, labels_flat, reduction="none")
        loss = (loss * mask_flat.float()).sum() / mask_flat.float().sum().clamp(min=1)
        return loss

    def _compute_rel_loss(
        self,
        rel_logits: Tensor,
        head_labels: Tensor,
        deprel_labels: Tensor,
        mask: Tensor,
    ) -> Tensor:
        """Compute dependency relation loss.

        Uses gold heads to select the relevant relation scores.
        rel_logits: [batch, max_words, 1+max_words, num_rels] - includes root at position 0
        head_labels: [batch, max_words] - head indices (0=root, 1..N=words)
        """
        batch_size, seq_len, num_heads, num_rels = rel_logits.shape

        # Clamp head_labels to valid range
        head_labels = head_labels.clamp(max=num_heads - 1)

        # Gather relation logits for gold heads
        # rel_logits: [batch, dep, head, num_rels]
        # head_labels: [batch, dep]
        head_indices = head_labels.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, 1, num_rels)
        # Select: [batch, dep, 1, num_rels] -> [batch, dep, num_rels]
        selected_rel_logits = rel_logits.gather(2, head_indices).squeeze(2)

        # Flatten and compute loss (use reshape for non-contiguous tensors)
        logits_flat = selected_rel_logits.reshape(-1, num_rels)
        labels_flat = deprel_labels.reshape(-1)
        mask_flat = mask.reshape(-1)

        loss = F.cross_entropy(logits_flat, labels_flat, reduction="none")
        loss = (loss * mask_flat.float()).sum() / mask_flat.float().sum().clamp(min=1)
        return loss

    def _compute_lemma_loss(
        self, logits: Tensor, labels: Tensor, word_mask: Tensor
    ) -> Tensor:
        """Compute lemmatization loss with per-character masking.

        Loss is computed only up to and including EOS (id=3) in each lemma.
        Positions after EOS (PAD tokens) are excluded from the loss to focus
        the gradient signal on actual character predictions.
        """
        batch_size, num_words, max_len, vocab_size = logits.shape
        EOS_ID = 3

        # Create character-level mask: True for positions up to and including first EOS
        is_eos = (labels == EOS_ID)  # [B, W, max_len]
        eos_cumsum = is_eos.cumsum(dim=-1)  # [B, W, max_len]
        # Before EOS: cumsum==0, at EOS: is_eos==True, after EOS: cumsum>=1 & is_eos==False
        char_mask = (eos_cumsum == 0) | is_eos  # [B, W, max_len]

        # Also apply word-level mask (exclude padded words entirely)
        word_mask_expanded = word_mask.unsqueeze(-1).expand_as(char_mask)
        final_mask = (char_mask & word_mask_expanded).float().reshape(-1)

        logits_flat = logits.reshape(-1, vocab_size)
        labels_flat = labels.reshape(-1)

        loss = F.cross_entropy(logits_flat, labels_flat, reduction="none")
        num_valid = final_mask.sum().clamp(min=1)
        return (loss * final_mask).sum() / num_valid

    def get_trainable_params(self) -> int:
        """Get number of trainable parameters."""
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    def get_total_params(self) -> int:
        """Get total number of parameters."""
        return sum(p.numel() for p in self.parameters())
