"""Lemmatization head using dilated CNN with GRU autoregressive refinement."""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from typing import Optional


class LemmaHead(nn.Module):
    """Dilated CNN lemmatizer with GRU autoregressive refinement.

    First predicts lemma characters in parallel using a stack of dilated Conv1d
    layers, then refines predictions left-to-right using a GRUCell. During
    training, the GRU uses teacher forcing with gold targets. During inference,
    the GRU feeds back its own predictions autoregressively.

    The GRU receives the form character embedding at each position as a copy
    shortcut, making it easy to learn prefix-copying behavior common in
    morphologically rich languages.
    """

    SOS_ID = 2  # Start-of-sequence token ID

    def __init__(
        self,
        hidden_size: int = 1024,
        char_embed_dim: int = 300,
        char_vocab_size: int = 512,
        input_proj_dim: int = 32,
        input_proj_dropout: float = 0.25,
        cnn_filters: Optional[list[int]] = None,
        cnn_kernel_sizes: Optional[list[int]] = None,
        cnn_dilations: Optional[list[int]] = None,
        cnn_paddings: Optional[list[int]] = None,
        dropout: float = 0.1,
        gru_hidden_dim: int = 512,
    ):
        super().__init__()
        self.char_vocab_size = char_vocab_size
        self.char_embed_dim = char_embed_dim
        self.gru_hidden_dim = gru_hidden_dim

        if cnn_filters is None:
            cnn_filters = [256, 256, 256]
        if cnn_kernel_sizes is None:
            cnn_kernel_sizes = [3, 3, 3, 1]
        if cnn_dilations is None:
            cnn_dilations = [1, 2, 4, 1]
        if cnn_paddings is None:
            cnn_paddings = [1, 2, 4, 0]

        # Character embedding (shared between CNN input and GRU input)
        self.char_embedding = nn.Embedding(char_vocab_size, char_embed_dim, padding_idx=0)

        # Project encoder word representation to small dim
        self.input_projection = nn.Sequential(
            nn.Linear(hidden_size, input_proj_dim),
            nn.Tanh(),
            nn.Dropout(input_proj_dropout),
        )

        # Build dilated CNN stack
        # Final layer outputs char_vocab_size
        all_filters = list(cnn_filters) + [char_vocab_size]
        cnn_input_dim = char_embed_dim + input_proj_dim
        input_dims = [cnn_input_dim] + all_filters[:-1]

        self.conv_layers = nn.ModuleList()
        self.activations = nn.ModuleList()
        for i in range(len(all_filters)):
            self.conv_layers.append(nn.Conv1d(
                in_channels=input_dims[i],
                out_channels=all_filters[i],
                kernel_size=cnn_kernel_sizes[i],
                stride=1,
                padding=cnn_paddings[i],
                dilation=cnn_dilations[i],
            ))
            # GELU for hidden layers, identity for output layer
            if i < len(all_filters) - 1:
                self.activations.append(nn.GELU())
            else:
                self.activations.append(nn.Identity())

        self.dropout = nn.Dropout(dropout)

        # GRU autoregressive refinement
        # Input: CNN logits + previous char embedding + form char embedding (copy shortcut)
        self.gru_refine = nn.GRUCell(
            input_size=char_vocab_size + char_embed_dim + char_embed_dim,
            hidden_size=gru_hidden_dim,
        )

        # Output projection: GRU hidden state -> vocab logits
        self.output_proj = nn.Linear(gru_hidden_dim, char_vocab_size)

        # Learned projection for GRU hidden state initialization
        self.h_init_proj = nn.Sequential(
            nn.Linear(char_vocab_size, gru_hidden_dim),
            nn.Tanh(),
        )

        # Scheduled sampling ratio: probability of using model predictions
        # instead of gold targets during training (set externally by trainer)
        self.scheduled_sampling_ratio = 0.0

    def _cnn_forward(self, word_repr: Tensor, input_chars: Tensor) -> Tensor:
        """Run CNN forward pass only (without GRU refinement).

        Args:
            word_repr: Word-level representations [batch, num_words, hidden_size]
            input_chars: Input word characters [batch, num_words, max_char_len]

        Returns:
            logits: Character logits [batch, num_words, max_char_len, vocab_size]
        """
        batch_size, num_words, max_char_len = input_chars.shape

        # 1. Embed characters: [B, W, max_char, char_embed_dim]
        char_embeds = self.dropout(self.char_embedding(input_chars))

        # 2. Project encoder representations: [B, W, input_proj_dim]
        proj = self.input_projection(word_repr)

        # 3. Broadcast to each character position: [B, W, max_char, input_proj_dim]
        proj = proj.unsqueeze(2).expand(-1, -1, max_char_len, -1)

        # 4. Concatenate: [B, W, max_char, char_embed_dim + input_proj_dim]
        x = torch.cat([char_embeds, proj], dim=-1)

        # 5. Reshape for Conv1d: [B*W, features, max_char]
        x = x.transpose(2, 3).reshape(batch_size * num_words, -1, max_char_len)

        # 6. Apply dilated CNN stack
        for conv, act in zip(self.conv_layers, self.activations):
            x = act(conv(x))

        # 7. Reshape back: [B, W, char_vocab_size, max_char] -> [B, W, max_char, char_vocab_size]
        x = x.reshape(batch_size, num_words, self.char_vocab_size, max_char_len)
        return x.transpose(2, 3)

    def _gru_refine(
        self,
        cnn_logits: Tensor,
        target_chars: Optional[Tensor] = None,
        form_chars: Optional[Tensor] = None,
    ) -> Tensor:
        """Apply GRU autoregressive refinement to CNN logits.

        Uses scheduled sampling during training: with probability
        self.scheduled_sampling_ratio, feeds back the model's own predictions
        instead of gold targets. This bridges the train/eval gap that causes
        cascading errors during greedy inference.

        Args:
            cnn_logits: CNN output logits [batch, num_words, max_char_len, vocab_size]
            target_chars: Gold target chars for teacher forcing [batch, num_words, max_char_len].
                          If None or not training, uses greedy autoregressive decoding.
            form_chars: Input form characters [batch, num_words, max_char_len] for copy shortcut.

        Returns:
            Refined logits [batch, num_words, max_char_len, vocab_size]
        """
        batch_size, num_words, max_char_len, vocab_size = cnn_logits.shape
        bw = batch_size * num_words
        device = cnn_logits.device
        H = self.gru_hidden_dim

        # Flatten batch and words: [B*W, max_char_len, vocab_size]
        cnn_flat = cnn_logits.reshape(bw, max_char_len, vocab_size)

        # Precompute form character embeddings for copy shortcut
        if form_chars is not None:
            form_flat = form_chars.reshape(bw, max_char_len)
            form_embs = self.char_embedding(form_flat)  # [B*W, max_char_len, char_embed_dim]
        else:
            form_embs = torch.zeros(bw, max_char_len, self.char_embed_dim, device=device)

        # Initialize GRU hidden state via learned projection
        h = self.h_init_proj(cnn_flat[:, 0, :])  # [B*W, H]

        # SOS embedding as initial "previous char"
        sos_ids = torch.full((bw,), self.SOS_ID, dtype=torch.long, device=device)

        # Prepare teacher forcing embeddings if available
        has_targets = target_chars is not None and self.training
        if has_targets:
            target_flat = target_chars.reshape(bw, max_char_len)
            # Shift right: at position t, feed target[t-1]; at position 0, feed SOS
            shifted = torch.cat([sos_ids.unsqueeze(1), target_flat[:, :-1]], dim=1)
            shifted_embs = self.char_embedding(shifted)  # [B*W, max_char_len, char_embed_dim]

        prev_char_emb = self.char_embedding(sos_ids)  # [B*W, char_embed_dim]
        refined = []

        for t in range(max_char_len):
            cnn_t = cnn_flat[:, t, :]  # [B*W, vocab_size]
            form_t = form_embs[:, t, :]  # [B*W, char_embed_dim]

            if has_targets:
                if self.scheduled_sampling_ratio > 0.0 and t > 0:
                    # Scheduled sampling: mix gold and predicted previous chars
                    use_pred = torch.rand(bw, device=device) < self.scheduled_sampling_ratio
                    pred_ids = logits_t.argmax(dim=-1)
                    pred_emb = self.char_embedding(pred_ids)
                    gold_emb = shifted_embs[:, t, :]
                    prev_char_emb = torch.where(
                        use_pred.unsqueeze(-1), pred_emb, gold_emb
                    )
                else:
                    prev_char_emb = shifted_embs[:, t, :]
            else:
                if t > 0:
                    # Greedy: feed predicted char for next step
                    pred_ids = logits_t.argmax(dim=-1)
                    prev_char_emb = self.char_embedding(pred_ids)

            gru_input = torch.cat([cnn_t, prev_char_emb, form_t], dim=-1)
            h = self.gru_refine(gru_input, h)
            logits_t = self.output_proj(h)  # [B*W, vocab_size]
            refined.append(logits_t)

        refined_logits = torch.stack(refined, dim=1)  # [B*W, max_char_len, vocab_size]
        return refined_logits.reshape(batch_size, num_words, max_char_len, vocab_size)

    def forward(
        self,
        word_repr: Tensor,
        input_chars: Tensor,
        target_chars: Optional[Tensor] = None,
    ) -> Tensor:
        """Predict lemma characters with CNN + GRU refinement.

        Args:
            word_repr: Word-level representations [batch, num_words, hidden_size]
            input_chars: Input word characters [batch, num_words, max_char_len]
            target_chars: Gold lemma characters for teacher forcing [batch, num_words, max_char_len].

        Returns:
            logits: Character logits [batch, num_words, max_char_len, vocab_size]
        """
        cnn_logits = self._cnn_forward(word_repr, input_chars)
        return self._gru_refine(cnn_logits, target_chars, form_chars=input_chars)

    def decode(
        self,
        word_repr: Tensor,
        input_chars: Tensor,
        eos_id: int = 3,
    ) -> Tensor:
        """Decode lemmas (greedy autoregressive).

        Args:
            word_repr: Word-level representations [batch, num_words, hidden_size]
            input_chars: Input word characters [batch, num_words, max_char_len]
            eos_id: End-of-sequence token ID (unused, kept for compatibility).

        Returns:
            predictions: Predicted character IDs [batch, num_words, max_char_len]
        """
        with torch.no_grad():
            logits = self.forward(word_repr, input_chars, target_chars=None)
            return logits.argmax(dim=-1)

    def beam_decode(
        self,
        word_repr: Tensor,
        input_chars: Tensor,
        beam_width: int = 5,
        eos_id: int = 3,
    ) -> Tensor:
        """Decode lemmas using beam search over GRU refinement.

        Args:
            word_repr: [batch, num_words, hidden_size]
            input_chars: [batch, num_words, max_char_len]
            beam_width: Number of beams to keep.
            eos_id: EOS token ID.

        Returns:
            predictions: [batch, num_words, max_char_len]
        """
        with torch.no_grad():
            batch_size, num_words, max_char_len = input_chars.shape
            device = input_chars.device
            bw = batch_size * num_words
            V = self.char_vocab_size
            H = self.gru_hidden_dim

            # Run CNN forward
            cnn_logits = self._cnn_forward(word_repr, input_chars)
            cnn_flat = cnn_logits.reshape(bw, max_char_len, V)

            # Precompute form character embeddings for copy shortcut
            form_flat = input_chars.reshape(bw, max_char_len)
            form_embs = self.char_embedding(form_flat)  # [bw, max_char_len, char_embed_dim]

            # Beam state: [bw, beam_width]
            # Start with 1 beam, expand to beam_width after first step
            beam_scores = torch.zeros(bw, 1, device=device)  # log-probs
            beam_h = self.h_init_proj(cnn_flat[:, 0, :]).unsqueeze(1)  # [bw, 1, H]
            beam_prev = torch.full((bw, 1), self.SOS_ID, dtype=torch.long, device=device)
            beam_seqs = torch.zeros(bw, 1, max_char_len, dtype=torch.long, device=device)
            beam_ended = torch.zeros(bw, 1, dtype=torch.bool, device=device)

            for t in range(max_char_len):
                cur_k = beam_scores.size(1)
                # CNN logits at this position, expanded for beams
                cnn_t = cnn_flat[:, t, :].unsqueeze(1).expand(-1, cur_k, -1)  # [bw, k, V]

                # Form char embedding at this position, expanded for beams
                form_t = form_embs[:, t, :].unsqueeze(1).expand(-1, cur_k, -1)  # [bw, k, char_embed_dim]

                # Previous char embeddings
                prev_emb = self.char_embedding(beam_prev)  # [bw, k, char_embed_dim]

                # GRU step for all beams
                gru_in = torch.cat([
                    cnn_t.reshape(bw * cur_k, V),
                    prev_emb.reshape(bw * cur_k, self.char_embed_dim),
                    form_t.reshape(bw * cur_k, self.char_embed_dim),
                ], dim=-1)
                h_flat = beam_h.reshape(bw * cur_k, H)
                h_new = self.gru_refine(gru_in, h_flat)  # [bw*k, H]

                # Project to logits and compute log-softmax
                logits_new = self.output_proj(h_new)  # [bw*k, V]
                log_probs = F.log_softmax(logits_new.reshape(bw, cur_k, V), dim=-1)

                # Mask ended beams: only allow PAD (id=0) continuation
                if beam_ended.any():
                    ended_mask = beam_ended.unsqueeze(-1).expand_as(log_probs)
                    pad_scores = torch.full_like(log_probs, -1e9)
                    pad_scores[:, :, 0] = 0.0  # PAD gets score 0 for ended beams
                    log_probs = torch.where(ended_mask, pad_scores, log_probs)

                # Combine with existing beam scores
                total = beam_scores.unsqueeze(-1) + log_probs  # [bw, k, V]
                total = total.reshape(bw, cur_k * V)

                # Top-k
                k = min(beam_width, total.size(-1))
                topk_scores, topk_idx = total.topk(k, dim=-1)  # [bw, k]

                beam_idx = topk_idx // V
                char_idx = topk_idx % V

                # Gather old state
                batch_arange = torch.arange(bw, device=device).unsqueeze(1).expand(-1, k)
                beam_scores = topk_scores
                beam_h = h_new.reshape(bw, cur_k, H)[batch_arange, beam_idx]  # [bw, k, H]

                # Update sequences
                beam_seqs = beam_seqs[batch_arange, beam_idx]  # [bw, k, max_char_len]
                beam_seqs[:, :, t] = char_idx

                # Track ended beams
                beam_ended = beam_ended[batch_arange, beam_idx] | (char_idx == eos_id)
                beam_prev = char_idx

                if beam_ended.all():
                    break

            # Select best beam per word
            best_idx = beam_scores.argmax(dim=-1)  # [bw]
            best_seqs = beam_seqs[torch.arange(bw, device=device), best_idx]  # [bw, max_char_len]

            return best_seqs.reshape(batch_size, num_words, max_char_len)
