"""Keyword-based text encoder for conditioning the diffusion model."""

import csv
from pathlib import Path

import torch
import torch.nn as nn


def build_vocab(metadata_csv: Path, min_count: int = 5) -> list[str]:
    """Build vocabulary from metadata CSV keywords, filtering rare tokens."""
    counts: dict[str, int] = {}
    with open(metadata_csv) as f:
        reader = csv.DictReader(f)
        for row in reader:
            for kw in row["keywords"].split(","):
                kw = kw.strip().lower()
                if kw:
                    counts[kw] = counts.get(kw, 0) + 1

    vocab = sorted(kw for kw, c in counts.items() if c >= min_count)
    return vocab


class KeywordEncoder(nn.Module):
    """Encodes a set of keyword indices into a fixed-size conditioning vector."""

    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 64,
        cond_dim: int = 256,
    ) -> None:
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.proj = nn.Linear(embed_dim, cond_dim)
        self.null_embedding = nn.Parameter(torch.randn(cond_dim))

    def forward(self, token_ids: list[list[int]], device: torch.device) -> torch.Tensor:
        """Encode batch of keyword lists into conditioning vectors.

        Args:
            token_ids: List of lists of vocab indices per sample.
            device: Target device.

        Returns:
            (batch, cond_dim) conditioning tensor.
        """
        batch_embs = []
        for ids in token_ids:
            if len(ids) == 0:
                batch_embs.append(self.null_embedding)
            else:
                idx = torch.tensor(ids, device=device)
                emb = self.embedding(idx).mean(dim=0)
                batch_embs.append(self.proj(emb))
        return torch.stack(batch_embs)