"""PyTorch Dataset for Universal Dependencies."""

from functools import partial
from typing import Any, Optional

import torch
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset
from transformers import PreTrainedTokenizer

from .conllu_reader import Document, Sentence
from .label_encoder import MorphoEncoder, LabelEncoder, CharVocab
from .tokenization import AlignedTokens, tokenize_and_align


class UDDataset(Dataset):
    """PyTorch Dataset for UD treebanks."""

    def __init__(
        self,
        documents: list[Document],
        tokenizer: PreTrainedTokenizer,
        morpho_encoder: MorphoEncoder,
        deprel_encoder: LabelEncoder,
        char_vocab: CharVocab,
        max_length: int = 512,
    ):
        """Initialize dataset.

        Args:
            documents: List of UD documents.
            tokenizer: HuggingFace tokenizer.
            morpho_encoder: Decomposed morphological label encoder.
            deprel_encoder: Encoder for dependency relations.
            char_vocab: Character vocabulary for lemmatization.
            max_length: Maximum sequence length.
        """
        all_sentences: list[Sentence] = []
        for doc in documents:
            all_sentences.extend(doc.sentences)

        self.tokenizer = tokenizer
        self.morpho_encoder = morpho_encoder
        self.deprel_encoder = deprel_encoder
        self.char_vocab = char_vocab
        self.max_length = max_length

        # Filter out sentences that would be truncated (words lost due to
        # max_length).  Truncated sentences have invalid head indices so they
        # cannot be used for training or evaluation.
        self.sentences: list[Sentence] = []
        self.skipped_count = 0
        self.skipped_words = 0
        for sentence in all_sentences:
            words = sentence.words
            aligned = tokenize_and_align(
                self.tokenizer, [t.form for t in words], self.max_length
            )
            n_aligned = len(set(
                wid for wid in aligned.word_ids if wid is not None
            ))
            if n_aligned < len(words):
                self.skipped_count += 1
                self.skipped_words += len(words) - n_aligned
            else:
                self.sentences.append(sentence)

    def scan_truncation(self) -> dict[str, int]:
        """Return stats about sentences skipped during init due to max_length.

        Returns:
            Dict with skipped_count, skipped_words, total_sentences
            (total includes skipped).
        """
        return {
            "truncated_count": self.skipped_count,
            "truncated_words": self.skipped_words,
            "total_sentences": len(self.sentences) + self.skipped_count,
        }

    @property
    def truncated_count(self) -> int:
        return self.skipped_count

    @property
    def truncated_words(self) -> int:
        return self.skipped_words

    def __len__(self) -> int:
        return len(self.sentences)

    def __getitem__(self, idx: int) -> dict[str, Any]:
        sentence = self.sentences[idx]
        words = sentence.words

        # Tokenize and align
        aligned = tokenize_and_align(
            self.tokenizer, [t.form for t in words], self.max_length
        )

        # Encode morphological labels (separate UPOS, XPOS, FEATS)
        upos_labels = [self.morpho_encoder.encode_upos(t.upos) for t in words]
        xpos_labels = [self.morpho_encoder.encode_xpos(t.xpos) for t in words]
        feats_labels = [self.morpho_encoder.encode_feats(t.feats) for t in words]

        # Encode dependency labels
        head_labels = [t.head if t.head is not None else 0 for t in words]
        deprel_labels = [self.deprel_encoder.encode(t.deprel) for t in words]

        # Encode characters for lemmatization
        # For each word: encode form (input) and lemma (target), with EOS appended
        form_chars = [self.char_vocab.encode(t.form) + [self.char_vocab.eos_id] for t in words]
        lemma_chars = [self.char_vocab.encode(t.lemma) + [self.char_vocab.eos_id] for t in words]

        return {
            "input_ids": aligned.input_ids,
            "attention_mask": aligned.attention_mask,
            "word_start_mask": aligned.word_start_mask,
            "upos_labels": upos_labels,
            "xpos_labels": xpos_labels,
            "feats_labels": feats_labels,
            "head_labels": head_labels,
            "deprel_labels": deprel_labels,
            "form_chars": form_chars,
            "lemma_chars": lemma_chars,
            "sentence": sentence,  # Keep for evaluation
        }


def collate_fn(
    batch: list[dict[str, Any]], pad_token_id: int = 1, char_pad_id: int = 0
) -> dict[str, Any]:
    """Collate function for DataLoader.

    Args:
        batch: List of samples from dataset.
        pad_token_id: Padding token ID for input_ids.
        char_pad_id: Padding ID for character sequences.

    Returns:
        Batched and padded tensors.
    """
    # Separate sentences (not tensors)
    sentences = [item["sentence"] for item in batch]

    # Pad input_ids and attention_mask
    input_ids = [torch.tensor(item["input_ids"], dtype=torch.long) for item in batch]
    attention_mask = [torch.tensor(item["attention_mask"], dtype=torch.long) for item in batch]
    word_start_mask = [torch.tensor(item["word_start_mask"], dtype=torch.bool) for item in batch]

    input_ids_padded = pad_sequence(input_ids, batch_first=True, padding_value=pad_token_id)
    attention_mask_padded = pad_sequence(attention_mask, batch_first=True, padding_value=0)
    word_start_mask_padded = pad_sequence(word_start_mask, batch_first=True, padding_value=False)

    # Get max word count for this batch
    word_counts = [sum(item["word_start_mask"]) for item in batch]
    max_words = max(word_counts)

    # Pad word-level labels
    upos_labels = []
    xpos_labels = []
    feats_labels = []
    head_labels = []
    deprel_labels = []

    # Determine number of feat categories from first sample
    num_feat_cats = len(batch[0]["feats_labels"][0]) if batch[0]["feats_labels"] else 0

    for item in batch:
        n_words = len(item["upos_labels"])
        pad_len = max_words - n_words

        upos_labels.append(item["upos_labels"] + [0] * pad_len)
        xpos_labels.append(item["xpos_labels"] + [0] * pad_len)
        feats_labels.append(
            item["feats_labels"] + [[0] * num_feat_cats] * pad_len
        )
        head_labels.append(item["head_labels"] + [0] * pad_len)
        deprel_labels.append(item["deprel_labels"] + [0] * pad_len)

    upos_labels_tensor = torch.tensor(upos_labels, dtype=torch.long)
    xpos_labels_tensor = torch.tensor(xpos_labels, dtype=torch.long)
    feats_labels_tensor = torch.tensor(feats_labels, dtype=torch.long)
    head_labels_tensor = torch.tensor(head_labels, dtype=torch.long)
    deprel_labels_tensor = torch.tensor(deprel_labels, dtype=torch.long)

    # Create word mask (which positions in max_words are valid)
    word_mask = torch.zeros(len(batch), max_words, dtype=torch.bool)
    for i, count in enumerate(word_counts):
        word_mask[i, :count] = True

    # Pad character sequences
    # CNN lemmatizer needs form_chars and lemma_chars padded to same length
    max_form_len = max(
        max(len(chars) for chars in item["form_chars"]) if item["form_chars"] else 1
        for item in batch
    )
    max_lemma_len = max(
        max(len(chars) for chars in item["lemma_chars"]) if item["lemma_chars"] else 1
        for item in batch
    )
    max_char_len = max(max_form_len, max_lemma_len)

    form_chars_padded = torch.full(
        (len(batch), max_words, max_char_len), char_pad_id, dtype=torch.long
    )
    lemma_chars_padded = torch.full(
        (len(batch), max_words, max_char_len), char_pad_id, dtype=torch.long
    )

    for b, item in enumerate(batch):
        for w, chars in enumerate(item["form_chars"]):
            form_chars_padded[b, w, : len(chars)] = torch.tensor(chars, dtype=torch.long)
        for w, chars in enumerate(item["lemma_chars"]):
            lemma_chars_padded[b, w, : len(chars)] = torch.tensor(chars, dtype=torch.long)

    return {
        "input_ids": input_ids_padded,
        "attention_mask": attention_mask_padded,
        "word_start_mask": word_start_mask_padded,
        "word_mask": word_mask,
        "upos_labels": upos_labels_tensor,
        "xpos_labels": xpos_labels_tensor,
        "feats_labels": feats_labels_tensor,
        "head_labels": head_labels_tensor,
        "deprel_labels": deprel_labels_tensor,
        "form_chars": form_chars_padded,
        "lemma_chars": lemma_chars_padded,
        "sentences": sentences,
    }


def create_dataloaders(
    train_dataset: UDDataset,
    dev_dataset: UDDataset,
    train_batch_size: int = 16,
    eval_batch_size: int = 32,
    num_workers: int = 0,
    pad_token_id: int = 1,
    char_pad_id: int = 0,
) -> tuple[DataLoader, DataLoader]:
    """Create train and dev dataloaders.

    Args:
        train_dataset: Training dataset.
        dev_dataset: Development dataset.
        train_batch_size: Batch size for training.
        eval_batch_size: Batch size for evaluation.
        num_workers: Number of data loading workers.
        pad_token_id: Padding token ID.
        char_pad_id: Character padding ID.

    Returns:
        Tuple of (train_loader, dev_loader).
    """

    collate = partial(collate_fn, pad_token_id=pad_token_id, char_pad_id=char_pad_id)
    use_pin_memory = torch.cuda.is_available()

    train_loader = DataLoader(
        train_dataset,
        batch_size=train_batch_size,
        shuffle=True,
        num_workers=num_workers,
        collate_fn=collate,
        pin_memory=use_pin_memory,
    )

    dev_loader = DataLoader(
        dev_dataset,
        batch_size=eval_batch_size,
        shuffle=False,
        num_workers=num_workers,
        collate_fn=collate,
        pin_memory=use_pin_memory,
    )

    return train_loader, dev_loader
