src/data/tokenizer.py

"""Deterministic value-to-token mapping for structured transaction features.

Each feature has its own tokenizer that maps raw values to token IDs using
the reserved token convention (D6): 0=MASK, 1=OOV, 2=NULL, 3+=real values.
Categorical features map directly; bucketed features use quantile or uniform
boundaries computed during fit().

OOV handling (D3): values outside the known range map to OOV token (1) and
increment a per-feature counter logged during eval.
"""

import hashlib
import json
from pathlib import Path
from typing import Any

import numpy as np

from src.data.schema import (
    FeatureSchema,
    SchemaConfig,
    MASK_TOKEN,
    NULL_TOKEN,
    OOV_TOKEN,
    VALUES_START,
    load_schema,
)


class FeatureTokenizer:
    """Tokenizer for a single feature. Handles encode, decode, and OOV tracking."""

    def __init__(self, schema: FeatureSchema) -> None:
        self._schema = schema
        self._boundaries: np.ndarray | None = None
        self._oov_count: int = 0
        self._fitted: bool = schema.type != "bucketed"

    @property
    def name(self) -> str:
        return self._schema.name

    @property
    def vocab_size(self) -> int:
        return self._schema.vocab_size

    @property
    def is_fitted(self) -> bool:
        return self._fitted

    @property
    def oov_count(self) -> int:
        return self._oov_count

    def reset_oov_count(self) -> None:
        self._oov_count = 0

    def fit(self, values: np.ndarray) -> None:
        """Compute bucket boundaries from data. Only valid for bucketed features.

        Args:
            values: 1-D array of raw continuous values to compute boundaries from.
        """
        assert self._schema.type == "bucketed", (
            f"fit() only applies to bucketed features, not {self._schema.type}"
        )
        assert self._schema.bucket_range is not None
        assert self._schema.bucket_method is not None

        num_buckets = self._schema.num_values

        if self._schema.bucket_method == "quantile":
            quantiles = np.linspace(0.0, 1.0, num_buckets + 1)
            self._boundaries = np.quantile(values, quantiles).astype(np.float64)
        elif self._schema.bucket_method == "uniform":
            lo, hi = self._schema.bucket_range
            self._boundaries = np.linspace(lo, hi, num_buckets + 1, dtype=np.float64)
        else:
            raise ValueError(f"Unknown bucket_method: {self._schema.bucket_method}")

        self._fitted = True

    def fit_uniform_from_range(self) -> None:
        """Compute uniform boundaries directly from the schema's bucket_range."""
        assert self._schema.type == "bucketed"
        assert self._schema.bucket_range is not None
        lo, hi = self._schema.bucket_range
        self._boundaries = np.linspace(
            lo, hi, self._schema.num_values + 1, dtype=np.float64
        )
        self._fitted = True

    def encode(self, values: np.ndarray) -> np.ndarray:
        """Map raw values to token IDs.

        Args:
            values: array of raw feature values (any shape).

        Returns:
            int16 array of token IDs, same shape as input.
        """
        assert self._fitted, f"Feature '{self.name}' not fitted. Call fit() first."
        original_shape = values.shape
        flat = values.ravel()

        if self._schema.type == "bucketed":
            token_ids, oov_count = self._encode_bucketed(flat)
        else:
            token_ids, oov_count = self._encode_categorical(flat)

        self._oov_count += oov_count
        return token_ids.reshape(original_shape)

    def decode(self, token_ids: np.ndarray) -> np.ndarray:
        """Map token IDs back to values. Bucketed features return bucket centers.

        Special tokens (MASK, OOV, NULL) decode to NaN.

        Args:
            token_ids: array of token IDs (any shape).

        Returns:
            float64 array of decoded values, same shape as input.
        """
        original_shape = token_ids.shape
        flat = token_ids.ravel().astype(np.int64)
        result = np.full(len(flat), np.nan, dtype=np.float64)

        value_mask = flat >= VALUES_START
        value_indices = flat[value_mask] - VALUES_START

        if self._schema.type == "bucketed" and self._boundaries is not None:
            centers = (self._boundaries[:-1] + self._boundaries[1:]) / 2.0
            valid = value_indices < len(centers)
            result_positions = np.where(value_mask)[0]
            result[result_positions[valid]] = centers[value_indices[valid]]
        else:
            result[value_mask] = value_indices.astype(np.float64)

        return result.reshape(original_shape)

    def _encode_categorical(self, values: np.ndarray) -> tuple[np.ndarray, int]:
        """Encode categorical/binary values. Returns (token_ids, oov_count)."""
        int_values = values.astype(np.int64)
        oov_mask = (int_values < 0) | (int_values >= self._schema.num_values)
        token_ids = (int_values + VALUES_START).astype(np.int16)
        token_ids[oov_mask] = OOV_TOKEN
        return token_ids, int(oov_mask.sum())

    def _encode_bucketed(self, values: np.ndarray) -> tuple[np.ndarray, int]:
        """Encode bucketed values using pre-computed boundaries.

        Returns (token_ids, oov_count). Uses np.digitize on internal boundaries
        so bucket index i covers [boundaries[i], boundaries[i+1]).
        The last bucket includes its upper bound: [boundaries[-2], boundaries[-1]].
        """
        assert self._boundaries is not None
        internal_bins = self._boundaries[1:-1]
        bucket_idx = np.digitize(values, internal_bins)
        # bucket_idx in [0, num_values - 1] for values within boundary range.
        # Can be num_values for values above the last internal bin, but np.digitize
        # with N-1 internal bins returns at most N-1 for values < boundaries[-1].
        # Values exactly at boundaries[-1] get num_values - 1 (last bucket).
        bucket_idx = np.clip(bucket_idx, 0, self._schema.num_values - 1)

        oov_mask = (values < self._boundaries[0]) | (values > self._boundaries[-1])
        token_ids = (bucket_idx + VALUES_START).astype(np.int16)
        token_ids[oov_mask] = OOV_TOKEN
        return token_ids, int(oov_mask.sum())

    def get_state(self) -> dict[str, Any]:
        """Serializable state for fingerprinting and persistence."""
        state: dict[str, Any] = {
            "name": self._schema.name,
            "type": self._schema.type,
            "num_values": self._schema.num_values,
            "vocab_size": self._schema.vocab_size,
        }
        if self._boundaries is not None:
            state["boundaries"] = self._boundaries.tolist()
        return state


class TransactionTokenizer:
    """Orchestrates tokenization across all features in the schema.

    Usage:
        schema = load_schema("data/schema.yaml")
        tokenizer = TransactionTokenizer(schema)

        # Fit bucketed features from raw data
        tokenizer.fit_feature("amount", raw_amounts)
        tokenizer.fit_feature("days_since_last", raw_days)

        # Encode
        token_ids = tokenizer.encode_feature("amount", raw_values)

        # Save state and compute fingerprint
        tokenizer.save_state("data/synthetic/tokenizer_state.json")
        fp = tokenizer.compute_fingerprint("data/schema.yaml")
    """

    def __init__(self, schema: SchemaConfig) -> None:
        self._schema = schema
        self._tokenizers: dict[str, FeatureTokenizer] = {}
        for feature in schema.features:
            self._tokenizers[feature.name] = FeatureTokenizer(feature)

    @property
    def feature_names(self) -> list[str]:
        return self._schema.feature_names()

    @property
    def num_features(self) -> int:
        return self._schema.num_features

    def get_feature_tokenizer(self, name: str) -> FeatureTokenizer:
        return self._tokenizers[name]

    def fit_feature(self, name: str, values: np.ndarray) -> None:
        """Fit bucket boundaries for a single bucketed feature."""
        self._tokenizers[name].fit(values)

    def is_all_fitted(self) -> bool:
        return all(t.is_fitted for t in self._tokenizers.values())

    def encode_feature(self, name: str, values: np.ndarray) -> np.ndarray:
        """Encode raw values for one feature. Returns int16 token IDs."""
        return self._tokenizers[name].encode(values)

    def decode_feature(self, name: str, token_ids: np.ndarray) -> np.ndarray:
        """Decode token IDs for one feature. Returns float64 values."""
        return self._tokenizers[name].decode(token_ids)

    def inject_nulls(
        self, token_ids: np.ndarray, null_mask: np.ndarray
    ) -> np.ndarray:
        """Replace positions where null_mask is True with NULL token.

        Args:
            token_ids: int16 array of token IDs.
            null_mask: boolean array, same shape as token_ids.

        Returns:
            Copy of token_ids with NULLs injected.
        """
        result = token_ids.copy()
        result[null_mask] = NULL_TOKEN
        return result

    @property
    def oov_counts(self) -> dict[str, int]:
        return {name: t.oov_count for name, t in self._tokenizers.items()}

    def reset_oov_counts(self) -> None:
        for t in self._tokenizers.values():
            t.reset_oov_count()

    def get_state(self) -> dict[str, Any]:
        """Full tokenizer state for persistence and fingerprinting."""
        return {
            "num_features": self._schema.num_features,
            "num_transactions": self._schema.num_transactions,
            "features": [
                self._tokenizers[name].get_state()
                for name in self._schema.feature_names()
            ],
        }

    def save_state(self, path: str | Path) -> None:
        """Save tokenizer state to JSON for fingerprint computation (D4)."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        state = self.get_state()
        with open(path, "w") as fh:
            json.dump(state, fh, indent=2, sort_keys=True)

    @classmethod
    def from_state(cls, state_path: str | Path, schema: SchemaConfig) -> "TransactionTokenizer":
        """Reconstruct tokenizer from saved state."""
        with open(state_path) as fh:
            state = json.load(fh)

        tokenizer = cls(schema)
        for feat_state in state["features"]:
            name = feat_state["name"]
            ft = tokenizer._tokenizers[name]
            if "boundaries" in feat_state:
                ft._boundaries = np.array(feat_state["boundaries"], dtype=np.float64)
                ft._fitted = True
        return tokenizer

    def compute_fingerprint(self, *config_paths: str | Path) -> str:
        """SHA256 fingerprint of tokenizer state plus config files (D4).

        The fingerprint includes the tokenizer's bucket boundaries and vocab
        sizes, plus the raw bytes of any additional config files (schema,
        generator config, split indices). Eval refuses to run if fingerprints
        don't match.

        Args:
            config_paths: paths to config files to include in the hash.

        Returns:
            Hex-encoded SHA256 digest.
        """
        hasher = hashlib.sha256()

        state_bytes = json.dumps(self.get_state(), sort_keys=True).encode("utf-8")
        hasher.update(state_bytes)

        for path in sorted(str(p) for p in config_paths):
            with open(path, "rb") as fh:
                hasher.update(fh.read())

        return hasher.hexdigest()