"""Minimal self-contained inference model for the Booru prompt generator release."""
import math
import re
from typing import Dict, List, Optional, Set

import numpy as np
import torch
import torch.nn as nn


PAD_ID = 0
BOS_ID = 1
EOS_ID = 2
RATING_G = 3
RATING_S = 4
RATING_Q = 5
RATING_E = 6
OFFSET = 7

_RATING_TOKENS = {"g": RATING_G, "s": RATING_S, "q": RATING_Q, "e": RATING_E}


def _apply_top_k_top_p(probs: torch.Tensor, top_k: int, top_p: float) -> torch.Tensor:
    """Filter a probability distribution with top-k and/or nucleus (top-p) sampling."""
    if top_k > 0:
        k = min(top_k, probs.size(0))
        threshold = torch.topk(probs, k).values[-1]
        probs = probs.where(probs >= threshold, torch.zeros_like(probs))
    if top_p < 1.0:
        sorted_probs, sorted_idx = torch.sort(probs, descending=True)
        cumsum = torch.cumsum(sorted_probs, dim=0)
        nucleus_mask = cumsum <= top_p
        if nucleus_mask.any():
            nucleus_mask[0] = True
        kept = torch.zeros_like(probs, dtype=torch.bool)
        kept.scatter_(0, sorted_idx, nucleus_mask)
        probs = probs.where(kept, torch.zeros_like(probs))
    return probs


class Vocab:
    def __init__(self, tags: List[str], counts: List[int]):
        self.tags = tags
        self.tag_to_idx = {tag: i for i, tag in enumerate(tags)}
        self.counts = np.array(counts, dtype=np.int64)
        self.total = int(self.counts.sum())
        self.freqs = self.counts.astype(np.float64) / max(self.total, 1)

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


class SimpleGraph:
    def __init__(self, vocab: Vocab, mutex: List[List[int]]):
        self.vocab = vocab
        self.mutex = mutex


class TagTransformer(nn.Module):
    """Permutation-equivariant Transformer for tag-set generation."""

    def __init__(
        self,
        vocab_size: int,
        d_model: int = 256,
        nhead: int = 4,
        num_layers: int = 4,
        dim_feedforward: int = 1024,
        dropout: float = 0.1,
        max_len: int = 256,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.max_len = max_len
        self.embedding = nn.Embedding(vocab_size + OFFSET, d_model, padding_idx=PAD_ID)
        layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True,
            norm_first=True,
        )
        self.transformer = nn.TransformerEncoder(layer, num_layers=num_layers, enable_nested_tensor=False)
        self.output = nn.Linear(d_model, vocab_size + OFFSET)
        self._init_weights()

    def _init_weights(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        seq_len = input_ids.size(1)
        mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=input_ids.device).bool()
        x = self.embedding(input_ids)
        padding_mask = input_ids == PAD_ID
        x = self.transformer(
            x,
            mask=mask,
            src_key_padding_mask=padding_mask,
            is_causal=True,
        )
        return self.output(x)


class NeuralPromptGenerator:
    GROUP_SUFFIXES = ["_hair", "_eyes"]

    def __init__(
        self,
        model: TagTransformer,
        graph: SimpleGraph,
        device: Optional[torch.device] = None,
        seed: Optional[int] = None,
        distribution_weight: float = 0.75,
        fallback_top_k: int = 20,
    ):
        self.model = model
        self.graph = graph
        self.vocab = graph.vocab
        self.device = device or torch.device("cpu")
        self.model.to(self.device)
        self.model.eval()
        self.rng = np.random.default_rng(seed)
        self.distribution_weight = distribution_weight
        self.fallback_top_k = fallback_top_k

        self._n_real = len(self.vocab)
        self._real_token_ids = torch.arange(self._n_real, device=self.device) + OFFSET
        self._log_uniform = -np.log(max(self._n_real, 1))
        self._special_ids = {PAD_ID, BOS_ID, EOS_ID}

        self._group_ids = np.full(self._n_real, -1, dtype=np.int32)
        for i, tag in enumerate(self.vocab.tags):
            for gid, suffix in enumerate(self.GROUP_SUFFIXES):
                if tag.endswith(suffix):
                    self._group_ids[i] = gid
                    break
        self._group_ids_tensor = torch.from_numpy(self._group_ids).to(self.device)

    def _target_log_prob(self, idx: int, alpha: float) -> float:
        log_emp = np.log(max(self.vocab.freqs[idx], 1e-12))
        return (1.0 - 2.0 * alpha) * log_emp

    def _num_people(self, tag_idxs: Set[int]) -> int:
        n = 0
        for idx in tag_idxs:
            tag = self.vocab.tags[idx]
            if tag == "solo":
                n = max(n, 1)
            elif tag in ("1girl", "1boy", "1other"):
                n = max(n, 1)
            elif tag in ("2girls", "2boys", "2others"):
                n = max(n, 2)
            elif tag in ("3girls", "3boys", "3others"):
                n = max(n, 3)
            elif tag in ("4girls", "4boys", "4others"):
                n = max(n, 4)
            elif tag in ("5girls", "5boys", "5others"):
                n = max(n, 5)
            elif tag in (
                "6+girls", "6+boys", "6+others",
                "multiple_girls", "multiple_boys", "multiple_others",
            ):
                n = max(n, 2)
        return max(n, 1)

    def generate(
        self,
        alpha: float,
        count: int,
        length: int = 30,
        anchor: Optional[List[str]] = None,
        blacklist: Optional[List[str]] = None,
        min_prob: float = 0.0,
        temperature: float = 1.0,
        top_k: int = 0,
        top_p: float = 1.0,
        rating: Optional[str] = None,
    ) -> List[List[str]]:
        anchor = anchor or []
        anchor_indices = [self.vocab.tag_to_idx[t] for t in anchor if t in self.vocab.tag_to_idx]
        blacklist = blacklist or []
        blacklist_indices = {self.vocab.tag_to_idx[t] for t in blacklist if t in self.vocab.tag_to_idx}
        rating_token = _RATING_TOKENS.get((rating or "g").lower(), RATING_G)

        target_bias = torch.tensor(
            [self._target_log_prob(i, alpha) for i in range(self._n_real)],
            dtype=torch.float32,
            device=self.device,
        )

        results: List[List[str]] = []
        with torch.no_grad():
            for _ in range(count):
                prompt_tokens: List[int] = [BOS_ID, rating_token] + [idx + OFFSET for idx in anchor_indices]
                present_tag_idxs: Set[int] = set(anchor_indices)
                excluded_tag_idxs: Set[int] = set(blacklist_indices)
                group_counts: Dict[int, int] = {}
                for idx in anchor_indices:
                    excluded_tag_idxs.update(self.graph.mutex[idx])
                    gid = self._group_ids[idx]
                    if gid >= 0:
                        group_counts[gid] = group_counts.get(gid, 0) + 1

                target_len = max(length, len(anchor_indices))
                max_len = getattr(self.model, "max_len", 256)
                if target_len > max_len - 2:
                    target_len = max_len - 2

                while len(prompt_tokens) - 1 < target_len:
                    input_ids = torch.tensor([prompt_tokens], dtype=torch.long, device=self.device)
                    logits = self.model(input_ids)[:, -1, :]

                    model_probs = torch.softmax(logits / max(temperature, 1e-6), dim=-1).squeeze(0)
                    allowed = model_probs >= min_prob

                    real_allowed = allowed[self._real_token_ids]
                    if not real_allowed.any():
                        k = min(self.fallback_top_k, self._n_real)
                        topk = torch.topk(model_probs[self._real_token_ids], k=k).indices
                        real_allowed = torch.zeros(self._n_real, dtype=torch.bool, device=self.device)
                        real_allowed[topk] = True
                        allowed = allowed.clone()
                        allowed[self._real_token_ids] = real_allowed

                    max_people = self._num_people(present_tag_idxs)
                    full_group_ids = [gid for gid, c in group_counts.items() if c >= max_people]

                    biased_logits = logits.squeeze(0).clone()
                    biased_logits[self._real_token_ids] += self.distribution_weight * target_bias

                    biased_logits[~allowed] = -float("inf")
                    for tid in self._special_ids:
                        biased_logits[tid] = -float("inf")
                    for idx in present_tag_idxs:
                        biased_logits[idx + OFFSET] = -float("inf")
                    for idx in excluded_tag_idxs:
                        biased_logits[idx + OFFSET] = -float("inf")
                    if full_group_ids:
                        full_groups_tensor = torch.tensor(full_group_ids, dtype=torch.int32, device=self.device)
                        group_full_mask = torch.isin(self._group_ids_tensor, full_groups_tensor)
                        biased_logits[self._real_token_ids[group_full_mask]] = -float("inf")

                    probs = torch.softmax(biased_logits / max(temperature, 1e-6), dim=0)
                    probs = _apply_top_k_top_p(probs, top_k, top_p)
                    if not torch.isfinite(probs).all() or probs.sum() <= 0:
                        break
                    probs = probs / probs.sum()
                    probs = probs.cpu().numpy()
                    probs = probs / probs.sum()

                    token_id = int(self.rng.choice(len(probs), p=probs))
                    if token_id in self._special_ids:
                        break
                    tag_idx = token_id - OFFSET
                    if tag_idx < 0 or tag_idx >= self._n_real:
                        break
                    if tag_idx in present_tag_idxs or tag_idx in excluded_tag_idxs:
                        break

                    prompt_tokens.append(token_id)
                    present_tag_idxs.add(tag_idx)
                    excluded_tag_idxs.update(self.graph.mutex[tag_idx])
                    gid = self._group_ids[tag_idx]
                    if gid >= 0:
                        group_counts[gid] = group_counts.get(gid, 0) + 1

                results.append([
                    self.vocab.tags[i - OFFSET]
                    for i in prompt_tokens
                    if i >= OFFSET
                ])

        return results