"""Inference pipeline for COMBO."""

import logging
from pathlib import Path
from typing import Optional

import numpy as np
import torch
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence
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, AlignedTokens
from ..model import UDParser
from .mst import decode_mst

logger = logging.getLogger(__name__)

TOKEN_EMBEDDINGS = False


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",
        input_tokens: list[Token] | None = None,
    ) -> 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)

        use_amp = self.device.type == "cuda"
        with torch.no_grad():
            with torch.autocast(
                device_type=self.device.type,
                dtype=torch.float16,
                enabled=use_amp,
            ):
                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, input_tokens=input_tokens)

    def predict_batch(
        self,
        sentences: list[list[str]],
        sent_ids: Optional[list[str]] = None,
        batch_size: int = 32,
        input_tokens_list: Optional[list[list[Token]]] = None,
    ) -> list[Sentence]:
        """Predict for multiple sentences with batched GPU inference.

        Sentences are sorted by length to minimize padding overhead, then
        processed in batches with mixed precision on CUDA devices.

        Args:
            sentences: List of sentences (each is a list of words).
            sent_ids: Optional sentence identifiers.
            batch_size: Number of sentences per forward pass.
            input_tokens_list: Optional list of input Token lists (one per
                sentence) to preserve multiword info in output.

        Returns:
            List of Sentence objects in original input order.
        """
        if not sentences:
            return []

        if sent_ids is None:
            sent_ids = [f"pred_{i}" for i in range(len(sentences))]

        # Fall back to single-sentence predict for batch_size=1
        if batch_size <= 1:
            return [
                self.predict(
                    words, sent_id,
                    input_tokens=(input_tokens_list[i] if input_tokens_list else None),
                )
                for i, (words, sent_id) in enumerate(zip(sentences, sent_ids))
            ]

        # Sort by subword token count (descending) to minimize padding
        indexed = list(enumerate(zip(sentences, sent_ids)))
        aligned_cache: list[AlignedTokens] = [
            tokenize_and_align(self.tokenizer, words, self.max_length)
            for words, _ in zip(sentences, sent_ids)
        ]
        indexed.sort(key=lambda x: len(aligned_cache[x[0]].input_ids), reverse=True)

        use_amp = self.device.type == "cuda"
        results: list[Optional[Sentence]] = [None] * len(sentences)

        for batch_start in range(0, len(indexed), batch_size):
            batch_items = indexed[batch_start : batch_start + batch_size]
            batch_indices = [idx for idx, _ in batch_items]
            batch_words = [sentences[idx] for idx in batch_indices]
            batch_sent_ids = [sent_ids[idx] for idx in batch_indices]
            batch_aligned = [aligned_cache[idx] for idx in batch_indices]

            # Collate into padded tensors
            tensors = self._collate_batch(batch_words, batch_aligned)

            with torch.no_grad():
                with torch.autocast(
                    device_type=self.device.type,
                    dtype=torch.float16,
                    enabled=use_amp,
                ):
                    outputs = self.model(**tensors)

            # Decode each sentence in the batch
            for i, (orig_idx, (words, sid)) in enumerate(
                zip(batch_indices, zip(batch_words, batch_sent_ids))
            ):
                in_toks = (input_tokens_list[orig_idx]
                           if input_tokens_list else None)
                results[orig_idx] = self._decode_outputs_at(
                    outputs, i, words, sid, input_tokens=in_toks,
                )

        return results

    def _collate_batch(
        self,
        batch_words: list[list[str]],
        batch_aligned: list[AlignedTokens],
    ) -> dict[str, Tensor]:
        """Collate multiple sentences into padded batch tensors."""
        pad_token_id = self.tokenizer.pad_token_id or 1

        input_ids_list = [
            torch.tensor(a.input_ids, dtype=torch.long) for a in batch_aligned
        ]
        attention_mask_list = [
            torch.tensor(a.attention_mask, dtype=torch.long) for a in batch_aligned
        ]
        word_start_mask_list = [
            torch.tensor(a.word_start_mask, dtype=torch.bool) for a in batch_aligned
        ]

        input_ids = pad_sequence(
            input_ids_list, batch_first=True, padding_value=pad_token_id
        ).to(self.device)
        attention_mask = pad_sequence(
            attention_mask_list, batch_first=True, padding_value=0
        ).to(self.device)
        word_start_mask = pad_sequence(
            word_start_mask_list, batch_first=True, padding_value=False
        ).to(self.device)

        word_counts = [len(w) for w in batch_words]
        max_words = max(word_counts)

        word_mask = torch.zeros(
            len(batch_words), max_words, dtype=torch.bool, device=self.device
        )
        for i, count in enumerate(word_counts):
            word_mask[i, :count] = True

        # Character input for lemmatization
        max_char_len = max(
            max((len(w) + 1 for w in words), default=1) for words in batch_words
        )
        max_char_len = min(max_char_len, self.char_vocab.max_len)

        form_chars = torch.full(
            (len(batch_words), max_words, max_char_len),
            self.char_vocab.pad_id,
            dtype=torch.long,
            device=self.device,
        )
        for b, words in enumerate(batch_words):
            for w, word in enumerate(words):
                char_ids = self.char_vocab.encode(word) + [self.char_vocab.eos_id]
                char_ids = char_ids[:max_char_len]
                form_chars[b, w, : len(char_ids)] = torch.tensor(
                    char_ids, dtype=torch.long, device=self.device
                )

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "word_start_mask": word_start_mask,
            "word_mask": word_mask,
            "form_chars": form_chars,
        }

    def _decode_outputs_at(
        self,
        outputs: dict[str, Tensor],
        batch_idx: int,
        words: list[str],
        sent_id: str,
        input_tokens: list[Token] | None = None,
    ) -> Sentence:
        """Decode model outputs for a single sentence within a batch."""
        upos_preds = outputs["upos_logits"][batch_idx].argmax(dim=-1)
        xpos_preds = outputs["xpos_logits"][batch_idx].argmax(dim=-1)
        feats_logits = outputs["feats_logits"]
        feats_preds = {
            cat: logits[batch_idx].argmax(dim=-1)
            for cat, logits in feats_logits.items()
        }

        arc_logits = outputs["arc_logits"][batch_idx]
        rel_logits = outputs["rel_logits"][batch_idx]

        n_words = len(words)
        arc_scores = arc_logits[:n_words, :n_words + 1].cpu().numpy()
        energy = np.full((n_words + 1, n_words + 1), -1e9)
        energy[:, 1:n_words + 1] = arc_scores.T
        heads, _ = decode_mst(energy, n_words + 1, has_labels=False)
        arc_preds = heads[1:n_words + 1].tolist()

        for i in range(n_words):
            if arc_preds[i] == i + 1:
                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

        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"][batch_idx:batch_idx + 1]
                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]
            else:
                lemma_preds = outputs["lemma_logits"].argmax(dim=-1)[batch_idx]

        tokens = []
        for i, word in enumerate(words):
            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]

            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)

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

            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

            mwt = None
            misc = "_"
            if input_tokens is not None and i < len(input_tokens):
                mwt = input_tokens[i].multiword
                misc = input_tokens[i].misc

            embeddings = {}
            if TOKEN_EMBEDDINGS:
                if "upos_xpos_feats_embedding" in outputs:
                    embeddings["upos_xpos_feats"] = outputs["upos_xpos_feats_embedding"][batch_idx][i].cpu().tolist()
                if "deprel_embedding" in outputs:
                    embeddings["deprel"] = outputs["deprel_embedding"][batch_idx][i].cpu().tolist()

            tokens.append(
                Token(
                    id=str(i + 1),
                    form=word,
                    lemma=lemma,
                    upos=upos,
                    xpos=xpos,
                    feats=feats,
                    head=head,
                    deprel=deprel,
                    misc=misc,
                    multiword=mwt,
                    embeddings=embeddings,
                )
            )

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

    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,
        input_tokens: list[Token] | None = None,
    ) -> 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

            mwt = None
            misc = "_"
            if input_tokens is not None and i < len(input_tokens):
                mwt = input_tokens[i].multiword
                misc = input_tokens[i].misc

            embeddings = {}
            if TOKEN_EMBEDDINGS:
                if "upos_xpos_feats_embedding" in outputs:
                    embeddings["upos_xpos_feats"] = outputs["upos_xpos_feats_embedding"][0][i].cpu().tolist()
                if "deprel_embedding" in outputs:
                    embeddings["deprel"] = outputs["deprel_embedding"][0][i].cpu().tolist()

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

        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)
