encoder/src/model/transaction_fm_multisurface.py

"""Orchestrator: encoder + projector + multi-surface LFM wrapper + 3 heads.

Sibling of `transaction_fm.build_transaction_fm`. Reuses the encoder
and projector from the existing demo so feature-token handling is
identical. Replaces:

  - `LfmPseudoTokenBackbone` → `LfmMultiSurfaceBackbone` (loads
    Lfm2ForCausalLM, exposes LM head + text path).
  - The four classification heads → one probability head + one
    attribution head; the LM head from the backbone provides
    reasoning generation.

The forward signatures are NOT the parent's MultiHeadModel contract
(`(B, T, F) int -> dict[head_name, logits]`). This model takes a
`MixedModalityBatch` and dispatches based on `batch.head_target`.
This makes per-batch homogeneous-head sampling explicit at the type
level: a batch is tagged for one head, and the model only computes
that head's loss.

Per-surface LoRA is implemented by loading a different LoRA artifact
per surface (each `LfmMultiSurfaceBackbone` instance carries one
adapter). A multi-surface inference server keeps one backbone in
memory with a LoRA registry that hot-swaps via PEFT's `set_adapter`.
That's a serving concern; this module trains one surface at a time.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F

from src.data.schema import SchemaConfig

from encoder.src.data.mixed_modality import (
    MixedModalityBatch,
    build_combined_attention_mask,
    build_lm_target_layout,
)
from encoder.src.model.heads.attribution_head import (
    AttributionHead,
    AttributionHeadConfig,
)
from encoder.src.model.heads.fraud_pattern_head import (
    FraudPatternHead,
    FraudPatternHeadConfig,
)
from encoder.src.model.heads.multi_treatment_probability_head import (
    MultiTreatmentProbabilityHead,
    MultiTreatmentProbabilityHeadConfig,
)
from encoder.src.model.heads.probability_head import (
    ProbabilityHead,
    ProbabilityHeadConfig,
)
from encoder.src.model.lfm_multisurface_wrapper import (
    LfmMultiSurfaceBackbone,
    build_multisurface_lora_config,
)
from encoder.src.model.projection_adapter import ProjectionAdapter
from encoder.src.model.transaction_encoder import TransactionEncoder


# Index of the "headline" class for the probability head. By convention
# class 2 ({unlikely=0, ambiguous=1, likely=2}) is the high-severity
# class whose softmax probability gets surfaced in the UI as the score.
HEADLINE_CLASS_INDEX = 2

# LM-loss ignore index. Matches PyTorch CrossEntropyLoss default and
# the value used in `build_lm_target_layout`.
LM_IGNORE_INDEX = -100


class TransactionMultiSurfaceModel(nn.Module):
    """Encoder + projector + multi-surface LFM backbone + 3 heads.

    Forward dispatches based on `batch.head_target`:
        - "probability" → run encoder + projector + backbone (no LM
          head) + probability head; loss = CE(logits, labels).
        - "attribution" → run encoder + projector + backbone (no LM
          head) + attribution head; loss = weighted BCE.
        - "lm"          → run encoder + projector + backbone (WITH LM
          head) + LM next-token loss on the text positions.

    All three paths share the same encoder + projector + backbone
    weights. The only difference is which head's loss is computed,
    enforcing the per-batch homogeneous-head sampling doctrine.

    Args:
        encoder: TransactionEncoder (compress mode).
        projector: ProjectionAdapter mapping (B, 64, d_encoder) → (B, 64, d_lfm).
        backbone: LfmMultiSurfaceBackbone.
        prob_head: ProbabilityHead.
        attr_head: AttributionHead.
    """

    def __init__(
        self,
        encoder: TransactionEncoder,
        projector: ProjectionAdapter,
        backbone: LfmMultiSurfaceBackbone,
        prob_head: ProbabilityHead | MultiTreatmentProbabilityHead | FraudPatternHead,
        attr_head: AttributionHead,
    ) -> None:
        super().__init__()
        self.encoder = encoder
        self.projector = projector
        self.backbone = backbone
        self.prob_head = prob_head
        self.attr_head = attr_head

    # ---- shared forward components ----

    def _build_combined_embeds(
        self,
        feature_ids: torch.Tensor,
        text_input_ids: torch.Tensor,
        disputed_idx: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Encoder + projector + concat with SEP + text embeddings.

        Args:
            feature_ids: (B, 64, 15) int64 tx feature ids.
            text_input_ids: (B, T_txt) int64 text tokens.
            disputed_idx: optional (B,) int64. When present, the encoder
                adds its learnable disputed marker at that position so
                the backbone knows which transaction is being asked about.

        Returns:
            (B, 64 + 1 + T_txt, d_lfm) combined embedding sequence,
            cast to the backbone's dtype.
        """
        # (B, 64, d_encoder)
        tx_encoded = self.encoder(feature_ids, disputed_idx=disputed_idx)
        # (B, 64, d_lfm)
        tx_pseudo = self.projector(tx_encoded)

        # Cast to backbone dtype (bf16 on GPU, fp32 on CPU).
        target_dtype = next(self.backbone.base.parameters()).dtype
        if tx_pseudo.dtype != target_dtype:
            tx_pseudo = tx_pseudo.to(target_dtype)

        batch_size = feature_ids.shape[0]
        device = feature_ids.device

        # (B, T_txt, d_lfm)
        text_embeds = self.backbone.embed_text(text_input_ids).to(target_dtype)
        # (B, 1, d_lfm)
        sep_embeds = self.backbone.embed_sep(batch_size, device).to(target_dtype)

        # (B, 64 + 1 + T_txt, d_lfm)
        return torch.cat([tx_pseudo, sep_embeds, text_embeds], dim=1)

    # ---- per-head forward + loss ----

    def forward_probability(
        self,
        batch: MixedModalityBatch,
    ) -> dict[str, torch.Tensor]:
        """Probability head path. Computes CE loss against labels_probability."""
        if batch.labels_probability is None:
            raise ValueError(
                "Batch tagged for probability head has no labels_probability.",
            )

        combined = self._build_combined_embeds(
            batch.feature_ids, batch.text_input_ids,
            disputed_idx=batch.disputed_idx,
        )
        attn_mask = build_combined_attention_mask(
            batch_size=batch.batch_size,
            num_tx_positions=batch.num_tx_positions,
            text_attention_mask=batch.text_attention_mask,
            device=combined.device,
        )
        outputs = self.backbone.forward_mixed(
            combined_embeds=combined,
            attention_mask=attn_mask,
            compute_lm_logits=False,
        )
        # (B, 3)
        logits = self.prob_head(outputs["hidden_states"])
        loss = self.prob_head.compute_loss(logits, batch.labels_probability)
        return {"loss": loss, "logits": logits}

    def forward_attribution(
        self,
        batch: MixedModalityBatch,
    ) -> dict[str, torch.Tensor]:
        """Attribution head path. Computes BCE loss per tx position."""
        if batch.labels_attribution is None:
            raise ValueError(
                "Batch tagged for attribution head has no labels_attribution.",
            )

        combined = self._build_combined_embeds(
            batch.feature_ids, batch.text_input_ids,
            disputed_idx=batch.disputed_idx,
        )
        attn_mask = build_combined_attention_mask(
            batch_size=batch.batch_size,
            num_tx_positions=batch.num_tx_positions,
            text_attention_mask=batch.text_attention_mask,
            device=combined.device,
        )
        outputs = self.backbone.forward_mixed(
            combined_embeds=combined,
            attention_mask=attn_mask,
            compute_lm_logits=False,
        )
        # (B, 64)
        logits = self.attr_head(outputs["hidden_states"])
        loss = self.attr_head.compute_loss(logits, batch.labels_attribution)
        return {"loss": loss, "logits": logits}

    def forward_lm(
        self,
        batch: MixedModalityBatch,
    ) -> dict[str, torch.Tensor]:
        """LM head path. Teacher-forced next-token CE on text positions.

        The text portion of `combined` is what we condition AND
        predict. The LM target layout puts the text token ids at
        positions [num_tx_positions + 1, T_total) and -100 elsewhere.
        Standard next-token shift: logits[..., :-1, :] predict
        targets[..., 1:].
        """
        combined = self._build_combined_embeds(
            batch.feature_ids, batch.text_input_ids,
            disputed_idx=batch.disputed_idx,
        )
        attn_mask = build_combined_attention_mask(
            batch_size=batch.batch_size,
            num_tx_positions=batch.num_tx_positions,
            text_attention_mask=batch.text_attention_mask,
            device=combined.device,
        )
        outputs = self.backbone.forward_mixed(
            combined_embeds=combined,
            attention_mask=attn_mask,
            compute_lm_logits=True,
        )
        # (B, T_total, vocab_size)
        lm_logits = outputs["lm_logits"]

        # Targets aligned to combined sequence; ignore_index at non-text
        # positions and at text padding.
        targets = build_lm_target_layout(
            batch_size=batch.batch_size,
            num_tx_positions=batch.num_tx_positions,
            text_input_ids=batch.text_input_ids,
            text_attention_mask=batch.text_attention_mask,
            ignore_index=LM_IGNORE_INDEX,
        )
        # Shift for next-token prediction: predict pos t+1 from pos t.
        # (B, T_total - 1, vocab_size)
        shift_logits = lm_logits[..., :-1, :].contiguous()
        # (B, T_total - 1)
        shift_targets = targets[..., 1:].contiguous()

        # Flatten and compute CE. ignore_index handles non-text and padding.
        loss = F.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_targets.view(-1),
            ignore_index=LM_IGNORE_INDEX,
        )
        return {"loss": loss, "lm_logits": lm_logits}

    def forward(
        self,
        batch: MixedModalityBatch,
    ) -> dict[str, torch.Tensor]:
        """Dispatch based on `batch.head_target`. The single entry point.

        The surface_trainer's batch sampler tags each batch with one
        head target; this forward enforces the doctrine that each
        minibatch updates exactly one head's loss.

        Contract gate: every dispute-legitimacy batch must carry
        `disputed_idx`. The label, attribution, and reasoning are all
        conditional on which transaction is being disputed; a batch
        without disputed_idx makes the task unidentifiable and the
        model silently degenerates to memorization. Fail loudly here
        so the bug is caught at the first forward pass rather than
        after a full training run.
        """
        if batch.disputed_idx is None:
            raise ValueError(
                "MixedModalityBatch is missing disputed_idx. The dispute "
                "legitimacy task's labels are conditional on disputed_idx; "
                "the model cannot learn without it. Check the dataset "
                "collate function.",
            )
        head = batch.head_target
        if head == "probability":
            return self.forward_probability(batch)
        if head == "attribution":
            return self.forward_attribution(batch)
        if head == "lm":
            return self.forward_lm(batch)
        raise ValueError(
            f"Unknown head_target: {head!r}. "
            f"Expected one of {{'probability', 'attribution', 'lm'}}.",
        )

    # ---- inference helpers ----

    @torch.no_grad()
    def predict(
        self,
        batch: MixedModalityBatch,
    ) -> dict[str, torch.Tensor]:
        """Compute all three heads' outputs in one forward (eval-time).

        Used by the demo UI: one inference returns the probability
        score, the top-k attributed transactions, and (if requested)
        the generated reasoning. No loss computed.
        """
        combined = self._build_combined_embeds(
            batch.feature_ids, batch.text_input_ids,
            disputed_idx=batch.disputed_idx,
        )
        attn_mask = build_combined_attention_mask(
            batch_size=batch.batch_size,
            num_tx_positions=batch.num_tx_positions,
            text_attention_mask=batch.text_attention_mask,
            device=combined.device,
        )
        outputs = self.backbone.forward_mixed(
            combined_embeds=combined,
            attention_mask=attn_mask,
            compute_lm_logits=False,
        )
        hidden = outputs["hidden_states"]
        prob_logits = self.prob_head(hidden)
        attr_logits = self.attr_head(hidden)

        return {
            "probability_logits": prob_logits,
            "probability_score": self.prob_head.score(prob_logits),
            "probability_band": self.prob_head.predict_band(prob_logits),
            "attribution_logits": attr_logits,
            "attribution_probabilities": self.attr_head.probabilities(attr_logits),
            "attribution_top_k": self.attr_head.top_k_positions(attr_logits, k=8),
        }

    # ---- introspection ----

    def num_parameters(self) -> dict[str, int]:
        """Parameter breakdown across components."""
        return {
            "encoder": sum(p.numel() for p in self.encoder.parameters()),
            "projector": sum(p.numel() for p in self.projector.parameters()),
            "backbone_total": self.backbone.total_parameters(),
            "backbone_trainable": self.backbone.trainable_parameters(),
            "prob_head": self.prob_head.num_parameters(),
            "attr_head": self.attr_head.num_parameters(),
            "trainable_total": sum(
                p.numel() for p in self.parameters() if p.requires_grad
            ),
            "total": sum(p.numel() for p in self.parameters()),
        }


def build_transaction_multisurface(
    schema: SchemaConfig,
    model_path: str | Path,
    encoder_cfg: dict[str, Any] | None = None,
    projector_cfg: dict[str, Any] | None = None,
    head_cfg: dict[str, Any] | None = None,
    lora_cfg: dict[str, Any] | None = None,
    dtype: torch.dtype = torch.bfloat16,
    device_map: str | None = "auto",
) -> TransactionMultiSurfaceModel:
    """Construct the multi-surface model from configs.

    Args:
        schema: SchemaConfig from the parent's data module.
        model_path: HF-format directory or repo ID for LFM2.5.
        encoder_cfg: optional overrides for the TransactionEncoder
            (d_feat, d_encoder, mlp_hidden).
        projector_cfg: optional overrides for the ProjectionAdapter
            (hidden, use_layernorm).
        head_cfg: dict with `probability` and `attribution` sub-configs.
            See the YAML configs for shape.
        lora_cfg: per-surface LoRA configuration. {r, alpha, dropout,
            target_modules}.
        dtype: bfloat16 (GPU) or float32 (CPU).
        device_map: "auto" for GPU, None for CPU.

    Returns:
        TransactionMultiSurfaceModel ready for training or inference.
    """
    encoder_cfg = encoder_cfg or {}
    projector_cfg = projector_cfg or {}
    head_cfg = head_cfg or {}
    lora_cfg = lora_cfg or {}

    # 1) Backbone — loads weights, applies LoRA. We build this first
    # because the projector + heads need d_lfm.
    lora = None
    if lora_cfg.get("enabled", True):
        lora = build_multisurface_lora_config(
            r=lora_cfg.get("r", 16),
            alpha=lora_cfg.get("alpha", 32),
            dropout=lora_cfg.get("dropout", 0.05),
            target_modules=lora_cfg.get("target_modules"),
        )
    backbone = LfmMultiSurfaceBackbone(
        model_path,
        lora=lora,
        dtype=dtype,
        device_map=device_map,
        freeze_base=True,  # encoder pattern; LoRA + heads carry the gradient
    )

    # 2) Encoder + projector — compress mode (one pseudo-token per tx).
    encoder = TransactionEncoder(
        schema,
        d_feat=encoder_cfg.get("d_feat", 32),
        d_encoder=encoder_cfg.get("d_encoder", 256),
        mlp_hidden=encoder_cfg.get("mlp_hidden", 384),
        enable_collections_markers=encoder_cfg.get(
            "enable_collections_markers", False,
        ),
        enable_fraud_markers=encoder_cfg.get(
            "enable_fraud_markers", False,
        ),
    )
    projector = ProjectionAdapter(
        d_encoder=encoder_cfg.get("d_encoder", 256),
        d_lfm=backbone.d_lfm,
        hidden=projector_cfg.get("hidden", 2 * backbone.d_lfm),
        use_layernorm=projector_cfg.get("use_layernorm", True),
    )

    # 3) Heads.
    prob_cfg_dict = head_cfg.get("probability", {})
    head_type = prob_cfg_dict.get("type", "probability")
    prob_head: ProbabilityHead | MultiTreatmentProbabilityHead | FraudPatternHead
    if head_type == "multi_treatment_probability":
        prob_head = MultiTreatmentProbabilityHead(
            MultiTreatmentProbabilityHeadConfig(
                name=prob_cfg_dict.get("name", "collections_treatment"),
                num_treatments=prob_cfg_dict.get("num_treatments", 4),
                num_bands=prob_cfg_dict.get("num_bands", 3),
                hidden_dim=backbone.d_lfm,
                mlp_hidden=prob_cfg_dict.get("mlp_hidden", 512),
                dropout=prob_cfg_dict.get("dropout", 0.1),
                num_tx_positions=prob_cfg_dict.get("num_tx_positions", 64),
                band_class_weights=prob_cfg_dict.get("band_class_weights"),
            )
        )
    elif head_type == "fraud_pattern":
        prob_head = FraudPatternHead(
            FraudPatternHeadConfig(
                name=prob_cfg_dict.get("name", "fraud_pattern"),
                num_stages=prob_cfg_dict.get("num_stages", 5),
                num_types=prob_cfg_dict.get("num_types", 4),
                hidden_dim=backbone.d_lfm,
                mlp_hidden=prob_cfg_dict.get("mlp_hidden", 256),
                dropout=prob_cfg_dict.get("dropout", 0.1),
                num_tx_positions=prob_cfg_dict.get("num_tx_positions", 64),
                stage_class_weights=prob_cfg_dict.get("stage_class_weights"),
                type_class_weights=prob_cfg_dict.get("type_class_weights"),
            )
        )
    elif head_type == "probability":
        prob_head = ProbabilityHead(
            ProbabilityHeadConfig(
                name=prob_cfg_dict.get("name", "dispute_legitimacy"),
                num_classes=prob_cfg_dict.get("num_classes", 3),
                hidden_dim=backbone.d_lfm,
                mlp_hidden=prob_cfg_dict.get("mlp_hidden", 256),
                dropout=prob_cfg_dict.get("dropout", 0.1),
                num_tx_positions=prob_cfg_dict.get("num_tx_positions", 64),
            )
        )
    else:
        raise ValueError(
            f"Unknown probability head type: {head_type!r}. Expected one of "
            f"{{'probability', 'multi_treatment_probability', 'fraud_pattern'}}.",
        )

    attr_cfg_dict = head_cfg.get("attribution", {})
    attr_head = AttributionHead(
        AttributionHeadConfig(
            name=attr_cfg_dict.get("name", "behavioral_attribution"),
            hidden_dim=backbone.d_lfm,
            mlp_hidden=attr_cfg_dict.get("mlp_hidden", 64),
            dropout=attr_cfg_dict.get("dropout", 0.1),
            num_tx_positions=attr_cfg_dict.get("num_tx_positions", 64),
            pos_weight=attr_cfg_dict.get("pos_weight", 5.0),
        )
    )

    return TransactionMultiSurfaceModel(
        encoder=encoder,
        projector=projector,
        backbone=backbone,
        prob_head=prob_head,
        attr_head=attr_head,
    )