"""Inference pipeline for COMBO."""

from pathlib import Path
from typing import Optional

import numpy as np
import torch
from torch import Tensor
from transformers import AutoTokenizer, PreTrainedTokenizer

from ..config import ModelConfig
from ..data.conllu_reader import Document, Sentence, Token
from ..data.conllu_writer import CoNLLUWriter
from ..data.label_encoder import MorphoEncoder, CharVocab, LabelEncoder
from ..data.tokenization import tokenize_and_align
from ..model import UDParser
from .mst import decode_mst


class Predictor:
    """Inference pipeline for new text."""

    def __init__(
        self,
        model: UDParser,
        tokenizer: PreTrainedTokenizer,
        morpho_encoder: MorphoEncoder,
        deprel_encoder: LabelEncoder,
        char_vocab: CharVocab,
        device: torch.device,
        max_length: int = 512,
        use_lemma_beam_search: bool = False,
        lemma_beam_width: int = 5,
        use_morpho_constraints: bool = False,
    ):
        self.model = model.to(device)
        self.model.eval()
        self.tokenizer = tokenizer
        self.morpho_encoder = morpho_encoder
        self.deprel_encoder = deprel_encoder
        self.char_vocab = char_vocab
        self.device = device
        self.max_length = max_length
        self.use_lemma_beam_search = use_lemma_beam_search
        self.lemma_beam_width = lemma_beam_width
        self.use_morpho_constraints = use_morpho_constraints

    @classmethod
    def from_checkpoint(
        cls,
        checkpoint_path: Path,
        encoders_dir: Path,
        config: ModelConfig,
        device: Optional[torch.device] = None,
        max_length: int = 512,
    ) -> "Predictor":
        """Load predictor from checkpoint."""
        if device is None:
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Load encoders
        morpho_encoder = MorphoEncoder.load(encoders_dir / "morpho_encoder.json")
        deprel_encoder = LabelEncoder.load(encoders_dir / "deprel_encoder.json")
        char_vocab = CharVocab.load(encoders_dir / "char_vocab.json")

        # Update config with encoder sizes
        config.num_upos = morpho_encoder.num_upos
        config.num_xpos = morpho_encoder.num_xpos
        config.feat_vocab_sizes = morpho_encoder.feat_vocab_sizes
        config.feat_categories = morpho_encoder.feat_categories
        config.num_deprels = len(deprel_encoder)
        config.char_vocab_size = len(char_vocab)

        # Load model
        model = UDParser(config)
        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
        model.load_state_dict(checkpoint["model_state_dict"], strict=False)

        # Load tokenizer
        tokenizer = AutoTokenizer.from_pretrained(config.base_model)

        return cls(
            model=model,
            tokenizer=tokenizer,
            morpho_encoder=morpho_encoder,
            deprel_encoder=deprel_encoder,
            char_vocab=char_vocab,
            device=device,
            max_length=max_length,
            use_lemma_beam_search=config.use_lemma_beam_search,
            lemma_beam_width=config.lemma_beam_width,
            use_morpho_constraints=config.use_morpho_constraints,
        )

    def predict(
        self,
        words: list[str],
        sent_id: str = "pred_0",
    ) -> Sentence:
        """Predict for a single sentence (pre-tokenized words).

        Args:
            words: List of words (pre-tokenized).
            sent_id: Sentence identifier.

        Returns:
            Sentence with all predictions.
        """
        # Tokenize and align
        aligned = tokenize_and_align(self.tokenizer, words, self.max_length)

        # Prepare tensors
        input_ids = torch.tensor([aligned.input_ids], dtype=torch.long, device=self.device)
        attention_mask = torch.tensor(
            [aligned.attention_mask], dtype=torch.long, device=self.device
        )
        word_start_mask = torch.tensor(
            [aligned.word_start_mask], dtype=torch.bool, device=self.device
        )
        word_mask = torch.ones(1, len(words), dtype=torch.bool, device=self.device)

        # Prepare character input for lemmatization
        form_chars = self._prepare_char_input(words)

        with torch.no_grad():
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                word_start_mask=word_start_mask,
                word_mask=word_mask,
                form_chars=form_chars,
            )

        # Decode outputs
        return self._decode_outputs(outputs, words, sent_id)

    def predict_batch(
        self,
        sentences: list[list[str]],
        sent_ids: Optional[list[str]] = None,
    ) -> list[Sentence]:
        """Predict for multiple sentences.

        Args:
            sentences: List of sentences (each is a list of words).
            sent_ids: Optional sentence identifiers.

        Returns:
            List of Sentence objects.
        """
        if sent_ids is None:
            sent_ids = [f"pred_{i}" for i in range(len(sentences))]

        predictions = []
        for words, sent_id in zip(sentences, sent_ids):
            pred = self.predict(words, sent_id)
            predictions.append(pred)
        return predictions

    def predict_text(
        self,
        text: str,
        sent_id: str = "pred_0",
    ) -> Sentence:
        """Predict for raw text (simple whitespace tokenization).

        Args:
            text: Raw text to parse.
            sent_id: Sentence identifier.

        Returns:
            Sentence with all predictions.
        """
        words = text.split()
        pred = self.predict(words, sent_id)
        pred.text = text
        return pred

    def _prepare_char_input(self, words: list[str]) -> torch.Tensor:
        """Prepare character input tensor for lemmatization."""
        max_len = max(len(w) + 1 for w in words) if words else 1  # +1 for EOS
        max_len = min(max_len, self.char_vocab.max_len)

        char_tensor = torch.full(
            (1, len(words), max_len),
            self.char_vocab.pad_id,
            dtype=torch.long,
            device=self.device,
        )

        for i, word in enumerate(words):
            char_ids = self.char_vocab.encode(word) + [self.char_vocab.eos_id]
            char_ids = char_ids[:max_len]  # truncate if needed
            char_tensor[0, i, : len(char_ids)] = torch.tensor(
                char_ids, dtype=torch.long, device=self.device
            )

        return char_tensor

    def _decode_outputs(
        self,
        outputs: dict[str, torch.Tensor],
        words: list[str],
        sent_id: str,
    ) -> Sentence:
        """Decode model outputs to a Sentence."""
        upos_preds = outputs["upos_logits"][0].argmax(dim=-1)  # [n_words]
        xpos_preds = outputs["xpos_logits"][0].argmax(dim=-1)  # [n_words]
        feats_logits = outputs["feats_logits"]  # dict[str, [1, n_words, vocab]]
        feats_preds = {
            cat: logits[0].argmax(dim=-1)  # [n_words]
            for cat, logits in feats_logits.items()
        }

        arc_logits = outputs["arc_logits"][0]  # [n_words, n_words+1] (col 0=ROOT)
        rel_logits = outputs["rel_logits"][0]  # [n_words, n_words+1, n_rels]

        # Decode heads with MST (Chu-Liu-Edmonds) to guarantee a valid tree
        n_words = len(words)
        arc_scores = arc_logits[:n_words, :n_words+1].cpu().numpy()  # [n_words, n_words+1]
        energy = np.full((n_words + 1, n_words + 1), -1e9)
        energy[:, 1:n_words+1] = arc_scores.T  # energy[parent, child] = score
        heads, _ = decode_mst(energy, n_words + 1, has_labels=False)
        arc_preds = heads[1:n_words+1].tolist()

        # Enforce valid tree: exactly one root, no self-loops
        for i in range(n_words):
            if arc_preds[i] == i + 1:  # self-loop
                arc_preds[i] = 0

        root_indices = [i for i, h in enumerate(arc_preds) if h == 0]
        if len(root_indices) == 0:
            best_root = max(range(n_words), key=lambda i: arc_scores[i, 0])
            arc_preds[best_root] = 0
        elif len(root_indices) > 1:
            best_root = max(root_indices, key=lambda i: arc_scores[i, 0])
            best_root_1indexed = best_root + 1
            for i in root_indices:
                if i != best_root:
                    arc_preds[i] = best_root_1indexed

        # Decode lemmas
        lemma_preds = None
        if outputs["lemma_logits"] is not None:
            if self.use_lemma_beam_search:
                word_repr_for_lemma = outputs["word_repr_for_lemma"]
                form_chars = self._prepare_char_input(words)
                lemma_preds = self.model.lemma_head.beam_decode(
                    word_repr_for_lemma, form_chars,
                    beam_width=self.lemma_beam_width,
                )[0]  # remove batch dim
            else:
                lemma_preds = outputs["lemma_logits"].argmax(dim=-1)[0]  # [seq, max_len]

        tokens = []
        for i, word in enumerate(words):
            # Decode morphological labels
            upos = self.morpho_encoder.decode_upos(upos_preds[i].item())
            xpos = self.morpho_encoder.decode_xpos(xpos_preds[i].item())
            feat_ids = [feats_preds[cat][i].item() for cat in self.morpho_encoder.feat_categories]

            # Apply morphological constraint filtering
            if self.use_morpho_constraints:
                feat_ids = self.morpho_encoder.constrain_feats(upos_preds[i].item(), feat_ids)

            feats = self.morpho_encoder.decode_feats(feat_ids)

            # Decode dependency
            head = arc_preds[i]
            rel_id = rel_logits[i, head].argmax().item()
            deprel = self.deprel_encoder.decode(rel_id)

            # Decode lemma
            if lemma_preds is not None:
                lemma_ids = lemma_preds[i].tolist()
                lemma = self.char_vocab.decode(lemma_ids)
                if not lemma:
                    lemma = word
            else:
                lemma = word

            tokens.append(
                Token(
                    id=str(i + 1),  # 1-indexed
                    form=word,
                    lemma=lemma,
                    upos=upos,
                    xpos=xpos,
                    feats=feats,
                    head=head,
                    deprel=deprel,
                )
            )

        return Sentence(
            tokens=tokens,
            sent_id=sent_id,
            text=" ".join(words),
        )

    def save_predictions(
        self,
        predictions: list[Sentence],
        output_path: Path,
    ) -> None:
        """Save predictions to CoNLL-U file."""
        CoNLLUWriter.write(predictions, output_path)
