"""LSTM model definitions for sequence classification."""

from __future__ import annotations

from typing import Any

import numpy as np


class Vocabulary:
    """Build vocabulary from tokenized corpus; map words to indices with UNK and PAD."""

    def __init__(self, min_freq: int = 1, unk_token: str = "<UNK>", pad_token: str = "<PAD>"):
        self.min_freq = min_freq
        self.unk_token = unk_token
        self.pad_token = pad_token
        self.word2idx: dict[str, int] = {}
        self.idx2word: list[str] = []
        self._built = False

    def build(self, tokenized_texts: list[list[str]]) -> "Vocabulary":
        """Build vocabulary from list of tokenized documents."""
        from collections import Counter
        counter: Counter = Counter()
        for tokens in tokenized_texts:
            counter.update(tokens)
        self.word2idx = {self.pad_token: 0, self.unk_token: 1}
        for w, c in counter.items():
            if c >= self.min_freq:
                self.word2idx[w] = len(self.word2idx)
        self.idx2word = [""] * len(self.word2idx)
        for w, i in self.word2idx.items():
            self.idx2word[i] = w
        self._built = True
        return self

    def encode(self, tokens: list[str]) -> list[int]:
        """Encode tokens to indices."""
        return [self.word2idx.get(t, self.word2idx.get(self.unk_token, 1)) for t in tokens]

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


def SentimentLSTM(
    vocab_size: int,
    embedding_dim: int = 128,
    hidden_dim: int = 256,
    num_layers: int = 2,
    num_classes: int = 3,
    dropout: float = 0.3,
    padding_idx: int = 0,
) -> Any:
    """Simple LSTM for sequence classification. Returns nn.Module. Requires torch."""
    try:
        import torch
        from torch import nn
    except ImportError as e:
        raise ImportError("SentimentLSTM requires PyTorch. Install with: pip install torch") from e

    class _LSTM(nn.Module):
        def __init__(self):
            super().__init__()
            self.embed = nn.Embedding(vocab_size, embedding_dim, padding_idx=padding_idx)
            self.lstm = nn.LSTM(
                embedding_dim,
                hidden_dim,
                num_layers=num_layers,
                batch_first=True,
                dropout=dropout if num_layers > 1 else 0,
            )
            self.fc = nn.Linear(hidden_dim, num_classes)
            self.dropout = nn.Dropout(dropout)

        def forward(self, x: Any) -> Any:
            emb = self.embed(x)
            out, (h_n, _) = self.lstm(emb)
            last = h_n[-1]
            return self.fc(self.dropout(last))

    return _LSTM()
