"""Non-compressing transaction encoder.

Mirrors the parent repo's `StructuredEmbedding` pattern exactly:
per-feature value tables + feature-type table, summed, expanded to a
flat sequence of length T_tx * F. The only difference is the output
dimension — we project directly to `d_lfm` (1024 for LFM2.5-350M)
instead of parent's `hidden_dim=256`.

Why this exists:
    The compressing encoder (TransactionEncoder) outputs 1 pseudo-token
    per transaction (sequence length 64 from 64 transactions). That is
    "encoder-style" — a small MLP collapses 15 features per tx into a
    single 256-dim vector, then a projector lifts to 1024-dim. Fraud
    quality is poor (ROC-AUC 0.53 at 100% labels) — likely because
    intra-transaction feature combinations get averaged into the
    256-dim bottleneck.

    This non-compressing variant keeps the full 64*15 = 960-token
    stream. No MLP-level compression. The LFM2.5 base sees the same
    sequence shape parent's structured-feature backbone sees — but on
    its own 1024-dim hidden space, with frozen text-pretrained weights
    plus LoRA. Tests whether the compression was the binding constraint
    on fraud quality.

    What we GIVE UP relative to the compressing variant:
    - The 15× sequence-length latency advantage
    - The "modality-token" framing (now we're really doing per-feature-
      token, like parent — the recipe is closer to "rebuild parent's
      embedding layer on a bigger backbone" than to LFM2-VL).

    What we keep:
    - Frozen base (cross-customer shared base story)
    - LoRA on attention
    - The pretrained text backbone

Shape contract:
    (B, T_tx, F) int64  →  (B, T_tx * F, d_lfm) float
    For default (T_tx=64, F=15): (B, 64, 15) → (B, 960, 1024)
"""

from __future__ import annotations

import torch
import torch.nn as nn

from src.data.schema import SchemaConfig


class StructuredEncoder(nn.Module):
    """Per-feature value embeddings + feature-type embeddings, summed.

    This is the parent repo's `StructuredEmbedding` at `hidden_dim=d_lfm`.
    Drop-in input-side replacement for the LFM2's `embed_tokens` table —
    same shape contract.

    Args:
        schema: parent's SchemaConfig.
        d_lfm: LFM2 backbone hidden size. 1024 for LFM2.5-350M, 2048 for
            LFM2.5-1.2B. Must match the LFM wrapper's d_lfm.

    Forward:
        token_ids: (B, T_tx, F) int64
        returns:   (B, T_tx * F, d_lfm) float, dtype matching the
                   embedding tables (fp32 by default).
    """

    def __init__(self, schema: SchemaConfig, d_lfm: int) -> None:
        super().__init__()
        self.num_features = schema.num_features
        self.num_transactions = schema.num_transactions
        self.d_lfm = d_lfm

        # Per-feature value tables, each sized to that feature's full vocab.
        # No padding_idx — reserved tokens (MASK/OOV/NULL) get learned
        # embeddings like everything else. Total params dominated by
        # merchant_id (10003 * d_lfm = 10.2M for d_lfm=1024).
        self.value_tables = nn.ModuleList(
            [nn.Embedding(f.vocab_size, d_lfm) for f in schema.features],
        )

        # Feature-type table. Distinguishes "the 5th token is an mcc"
        # from "the 5th token is a merchant_id" — RoPE alone can't carry
        # this since intra-transaction position is arbitrary. Mirrors
        # parent's pattern exactly.
        self.type_table = nn.Embedding(schema.num_features, d_lfm)

    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        # token_ids: (B, T_tx, F) int64
        B, T_tx, F = token_ids.shape
        assert F == self.num_features, (
            f"Expected {self.num_features} features, got {F}"
        )

        # Type embeddings: (F, d_lfm). Same per-feature offsets added to
        # every transaction's value embedding for that feature.
        type_indices = torch.arange(F, device=token_ids.device)
        type_emb = self.type_table(type_indices)
        # type_emb: (F, d_lfm)

        # Embed each feature column with its own table, add the type
        # offset, stack into (B, T_tx, F, d_lfm).
        feature_embeddings: list[torch.Tensor] = []
        for f_idx in range(F):
            feat_tokens = token_ids[:, :, f_idx]
            val_emb = self.value_tables[f_idx](feat_tokens)
            # val_emb: (B, T_tx, d_lfm)
            feature_embeddings.append(val_emb + type_emb[f_idx])

        stacked = torch.stack(feature_embeddings, dim=2)
        # stacked: (B, T_tx, F, d_lfm)
        return stacked.reshape(B, T_tx * F, self.d_lfm)
        # → (B, T_tx * F, d_lfm) — flat sequence the LFM can consume

    def num_embedding_params(self) -> int:
        """Total params in value + type tables (sanity check)."""
        val = sum(e.weight.numel() for e in self.value_tables)
        typ = self.type_table.weight.numel()
        return val + typ