"""Encoder + BIO head + emoji whole-token decoder."""

import math
import torch
import torch.nn as nn
from transformers import AutoModel
from torchcrf import CRF
from src import emoji_vocab


class BIOHead(nn.Module):
    def __init__(self, hidden_size: int, num_labels: int = 3, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.linear = nn.Linear(hidden_size, num_labels)
        self.crf = CRF(num_labels, batch_first=True)

    def forward(self, hidden_states):
        """Return emission scores (batch, seq_len, num_labels)."""
        return self.linear(self.dropout(hidden_states))

    def loss(self, emissions, labels, attention_mask=None):
        """CRF negative log-likelihood loss.

        Uses attention_mask as the CRF mask. Positions labelled -100
        (CLS, SEP, non-first subwords) are replaced with O (0) — correct target.
        """
        safe_labels = labels.clone()
        safe_labels[safe_labels == -100] = 0
        if attention_mask is not None:
            mask = attention_mask.bool()
        else:
            mask = torch.ones(emissions.shape[:2], dtype=torch.bool, device=emissions.device)
        return -self.crf(emissions, safe_labels, mask=mask, reduction="mean")

    def decode(self, emissions, attention_mask=None):
        """Viterbi decoding. Returns list[list[int]] aligned to valid tokens."""
        if attention_mask is not None:
            mask = attention_mask.bool()
        else:
            mask = torch.ones(emissions.shape[:2], dtype=torch.bool, device=emissions.device)
        return self.crf.decode(emissions, mask=mask)


class _SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 32):
        super().__init__()
        pe = torch.zeros(1, max_len, d_model)
        pos = torch.arange(max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[0, :, 0::2] = torch.sin(pos * div)
        pe[0, :, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe)

    def forward(self, x):
        return x + self.pe[:, : x.shape[1], :]


class EmojiDecoder(nn.Module):
    """Whole-emoji Transformer decoder.

    Each output token is a complete emoji (one of ~5,225 in the Unicode standard),
    not a raw character/byte.  Embeddings are initialised from XLM-R descriptions
    of each emoji so the decoder starts with semantic priors.

    Training  : teacher-forced, all spans in one batched forward call.
    Inference : all K spans decoded together per step (vectorised over K).
    """

    def __init__(
        self,
        d_model: int = 512,
        nhead: int = 8,
        num_layers: int = 4,
        dim_feedforward: int = 2048,
        vocab_size: int = emoji_vocab.VOCAB_SIZE,
        max_length: int = 16,
        encoder_hidden_size: int = 768,
    ):
        super().__init__()
        self.d_model = d_model
        self.vocab_size = vocab_size
        self.max_length = max_length

        self.input_embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = _SinusoidalPositionalEncoding(d_model, max_len=max_length + 1)
        # Project encoder memory from hidden_size → d_model for cross-attention
        self.memory_proj = nn.Linear(encoder_hidden_size, d_model)
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            batch_first=True,
        )
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.output_projection = nn.Linear(d_model, vocab_size)

    def _build_memory(self, span_embeddings, encoder_hidden_states, encoder_attention_mask):
        """Prepend span embedding as position-0 of cross-attention memory.

        Each span gets a unique anchor token at position 0 that the decoder can
        attend to specifically, preventing all spans from collapsing to the same
        sentence-level representation.
        """
        K = span_embeddings.shape[0]
        device = span_embeddings.device
        projected = self.memory_proj(encoder_hidden_states)          # (K, src_len, d_model)
        span_token = span_embeddings.unsqueeze(1)                    # (K, 1, d_model)
        memory = torch.cat([span_token, projected], dim=1)           # (K, 1+src_len, d_model)

        mem_pad_mask = None
        if encoder_attention_mask is not None:
            span_valid = torch.ones(K, 1, dtype=encoder_attention_mask.dtype, device=device)
            extended = torch.cat([span_valid, encoder_attention_mask], dim=1)
            mem_pad_mask = ~extended.bool()
        return memory, mem_pad_mask

    def forward(
        self,
        span_embeddings,              # (K, d_model)
        encoder_hidden_states,        # (K, src_len, hidden_size)
        encoder_attention_mask=None,  # (K, src_len)
        target_ids=None,              # (K, tgt_len) whole-emoji token IDs
    ):
        K = span_embeddings.shape[0]
        device = span_embeddings.device
        memory, mem_pad_mask = self._build_memory(
            span_embeddings, encoder_hidden_states, encoder_attention_mask
        )

        if target_ids is not None:
            # ── Teacher-forced training ──────────────────────────────────────
            tgt_input = target_ids[:, :-1]
            tgt_len = tgt_input.shape[1]
            embedded = self.input_embedding(tgt_input)
            embedded = self.pos_encoding(embedded)
            causal_mask = nn.Transformer.generate_square_subsequent_mask(tgt_len, device=device)
            decoder_out = self.decoder(
                embedded, memory,
                tgt_mask=causal_mask,
                memory_key_padding_mask=mem_pad_mask,
            )
            return self.output_projection(decoder_out), None

        else:
            # ── Vectorised autoregressive inference ──────────────────────────
            sequences = torch.full((K, 1), emoji_vocab.BOS_ID, dtype=torch.long, device=device)
            finished = torch.zeros(K, dtype=torch.bool, device=device)
            for step in range(self.max_length):
                embedded = self.input_embedding(sequences)
                embedded = self.pos_encoding(embedded)
                decoder_out = self.decoder(
                    embedded, memory,
                    memory_key_padding_mask=mem_pad_mask,
                )
                next_ids = self.output_projection(decoder_out[:, -1, :]).argmax(dim=-1)
                next_ids = next_ids.masked_fill(finished, emoji_vocab.PAD_ID)
                sequences = torch.cat([sequences, next_ids.unsqueeze(1)], dim=1)
                finished = finished | (next_ids == emoji_vocab.EOS_ID)
                if finished.all():
                    break
            return None, [seq.tolist() for seq in sequences]


class EmojinizeEncoderModel(nn.Module):
    """Full model: XLM-R encoder (optionally LoRA) + BIO-CRF head + emoji decoder."""

    def __init__(
        self,
        encoder_model_id: str = "FacebookAI/xlm-roberta-base",
        d_model: int = 256,
        nhead: int = 4,
        decoder_layers: int = 2,
        dim_feedforward: int = 512,
        lora_rank: int = 16,
    ):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(encoder_model_id)
        self.encoder_hidden_size = self.encoder.config.hidden_size
        self.lora_rank = lora_rank

        if lora_rank > 0:
            from peft import get_peft_model, LoraConfig
            lora_cfg = LoraConfig(
                r=lora_rank,
                lora_alpha=lora_rank * 2,
                target_modules=["query", "key", "value", "dense"],
                lora_dropout=0.05,
                bias="none",
            )
            self.encoder = get_peft_model(self.encoder, lora_cfg)

        self.bio_head = BIOHead(self.encoder_hidden_size, num_labels=3)

        self.emoji_decoder = EmojiDecoder(
            d_model=d_model,
            nhead=nhead,
            num_layers=decoder_layers,
            dim_feedforward=dim_feedforward,
            encoder_hidden_size=self.encoder_hidden_size,
        )
        self.span_proj = nn.Linear(self.encoder_hidden_size * 3, d_model)

        self.log_sigma_bio = nn.Parameter(torch.zeros(1))
        self.log_sigma_emoji = nn.Parameter(torch.zeros(1))

    @torch.no_grad()
    def init_emoji_embeddings(self, tokenizer, device):
        """Initialise emoji decoder input embeddings from XLM-R descriptions.

        Called once after model creation. For each of the ~5,225 emoji in the
        vocabulary, encodes its Unicode description (e.g. '🥤' → 'cup with straw')
        with the frozen encoder and projects to d_model via memory_proj.  This
        gives the decoder a semantic prior for every emoji — including ones rare
        or absent in training data.
        """
        enc_was_training = self.encoder.training
        self.encoder.eval()

        descriptions = [emoji_vocab.get_description(em) or "emoji" for em in emoji_vocab.ALL_EMOJI]
        d_model = self.emoji_decoder.d_model
        batch_size = 256
        all_hidden = []

        for i in range(0, len(descriptions), batch_size):
            batch = descriptions[i : i + batch_size]
            tok = tokenizer(batch, return_tensors="pt", padding=True,
                            truncation=True, max_length=16).to(device)
            out = self.encoder(**tok)
            mask = tok["attention_mask"].unsqueeze(-1).float()
            emb = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1)
            all_hidden.append(emb.cpu())

        emoji_hidden = torch.cat(all_hidden, dim=0)  # (num_emoji, hidden_size)
        emoji_embs = self.emoji_decoder.memory_proj(emoji_hidden.to(device))  # → (num_emoji, d_model)

        emb_weight = self.emoji_decoder.input_embedding.weight
        emb_weight.requires_grad_(True)   # temporarily unfreeze to allow copy_
        special_embs = emb_weight[:emoji_vocab._SPECIAL].clone()
        full_embs = torch.cat([special_embs, emoji_embs], dim=0)
        emb_weight.copy_(full_embs)
        # Freeze emoji embeddings to preserve XLM-R semantic initialisation.
        emb_weight.requires_grad_(False)

        if enc_was_training:
            self.encoder.train()
        print(f"  ✓ emoji embeddings initialised from XLM-R ({len(emoji_vocab.ALL_EMOJI)} emoji, frozen)")

    def _pool_spans(self, hidden_states, attention_mask, span_info):
        embs, enc_hiddens, enc_masks = [], [], []
        for batch_idx, spans in enumerate(span_info):
            h = hidden_states[batch_idx]
            m = attention_mask[batch_idx]
            for start, end in spans:
                span_h = h[start:end]
                pool = torch.cat([span_h.mean(0), h[start], h[end - 1]])
                embs.append(self.span_proj(pool))
                enc_hiddens.append(h)
                enc_masks.append(m)
        if not embs:
            return None, None, None
        return (
            torch.stack(embs),
            torch.stack(enc_hiddens),
            torch.stack(enc_masks),
        )

    def forward(self, input_ids, attention_mask, span_info=None, target_bytes=None):
        encoder_out = self.encoder(input_ids, attention_mask)
        hidden_states = encoder_out.last_hidden_state
        result = {"bio_logits": self.bio_head(hidden_states)}

        if span_info is not None:
            span_embs, enc_hidden, enc_mask = self._pool_spans(
                hidden_states, attention_mask, span_info
            )
            if span_embs is not None:
                if target_bytes is not None:
                    logits, _ = self.emoji_decoder(
                        span_embs, enc_hidden, enc_mask, target_ids=target_bytes
                    )
                    result["emoji_logits"] = logits
                else:
                    _, decoded = self.emoji_decoder(span_embs, enc_hidden, enc_mask)
                    result["decoded_seqs"] = decoded
        return result