encoder/src/demo/copilot_inference_collections.py

"""Inference plumbing for the Collections Co-Pilot demo.

Mirror of `copilot_inference.py` (dispute) adapted to the Collections
surface:
  - The model uses MultiTreatmentProbabilityHead (K=4 treatments × 3
    bands) instead of the 3-class dispute head.
  - The "verdict" is now (per-treatment LIKELY-band probability,
    dominant treatment, per-treatment band argmax).
  - The reasoning template grounds in the cross-position signals
    (velocity / subscription burden / merchant diversity / large
    amount count) rather than dispute-specific anomalies.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator

import numpy as np
import torch
import yaml

from src.data.schema import SchemaConfig, load_schema

from encoder.src.data.mixed_modality import MixedModalityBatch, tokenize_texts
from encoder.src.data.synthetic_collections import (
    BAND_LIKELY,
    BAND_NAMES,
    CONTEXT_IDX_DEFAULT,
    NUM_BANDS,
    NUM_TREATMENTS,
    TREATMENT_NAMES,
    TREATMENT_NO_OFFER,
    TREATMENT_PAYMENT_PLAN,
    TREATMENT_SETTLEMENT,
    TREATMENT_SOFT_TOUCH,
    attribution_for_treatment,
    signal_large_amount_count,
    signal_merchant_diversity,
    signal_recent_velocity,
    signal_spending_volatility,
    signal_subscription_burden,
)
from encoder.src.model.transaction_fm_multisurface import (
    TransactionMultiSurfaceModel,
    build_transaction_multisurface,
)


DEMO_SEED = 42

# Color band per treatment band. Same green/yellow/red palette as
# dispute so the audience picks up cross-surface visual consistency.
BAND_COLORS: dict[int, str] = {
    0: "#22c55e",  # green — unlikely_respond
    1: "#eab308",  # yellow — ambiguous
    2: "#ef4444",  # red — likely_respond (note: in Collections, red is
                   # not negative — it's the model's HIGH-confidence
                   # recommendation. We override colors per treatment
                   # below to keep cross-surface semantics intact.)
}

# Per-treatment LIKELY-band display color — these are the "positive"
# brand colors of each option, distinct from the band semantics above.
TREATMENT_COLORS: dict[int, str] = {
    TREATMENT_SETTLEMENT: "#2563eb",    # blue — settle (financial close)
    TREATMENT_PAYMENT_PLAN: "#16a34a",  # green — plan (preserved relationship)
    TREATMENT_SOFT_TOUCH: "#ca8a04",    # amber — soft (light touch)
    TREATMENT_NO_OFFER: "#737373",      # gray — no offer (write-off)
}


@dataclass
class CollectionsCastMember:
    """One entry from `encoder/data/collections_cast.json`.

    Attributes map 1:1 with the JSON fields. The renderer reads
    display_name, pattern, context_text, treatment_label_names,
    dominant_treatment_name. The inference module reads customer_idx,
    context_idx, context_text.
    """

    pattern: str
    display_name: str
    customer_idx: int
    context_idx: int
    treatment_labels: list[int]
    treatment_label_names: list[str]
    dominant_treatment: int
    dominant_treatment_name: str
    description: str
    context_text: str


@dataclass
class CollectionsResult:
    """Result of one Collections inference call.

    Attributes:
        likely_scores: (K,) float — P(likely_respond) per treatment.
        predicted_bands: (K,) int — argmax band per treatment.
        dominant_treatment: int — treatment with the highest
            LIKELY-band probability (the model's recommendation).
        attribution_probs: (64,) float — per-position contribution
            probabilities (per blueprint, the same single attribution
            head is shared across treatments; it attends to the
            position that drove the dominant verdict).
        top_k_positions: (k,) int — indices of the top-contributing tx.
        context_idx: position the model's bias markers landed on.
    """

    likely_scores: np.ndarray            # (K,) float
    predicted_bands: list[int]           # (K,) int
    dominant_treatment: int
    attribution_probs: np.ndarray        # (64,) float
    top_k_positions: np.ndarray          # (k,) int
    context_idx: int


class CollectionsCopilotModel:
    """Encapsulates the Collections multi-surface model + cast + tokenizer."""

    def __init__(
        self,
        model: TransactionMultiSurfaceModel,
        schema: SchemaConfig,
        histories: np.ndarray,
        cast: list[CollectionsCastMember],
        device: torch.device,
    ) -> None:
        self.model = model
        self.schema = schema
        self.histories = histories
        self.cast = cast
        self.device = device
        self.tokenizer = model.backbone.tokenizer
        self.model.eval()

    @classmethod
    def from_paths(
        cls,
        checkpoint_path: Path,
        model_config_path: Path,
        schema_path: Path,
        histories_path: Path,
        cast_path: Path,
        device: torch.device = torch.device("cpu"),
    ) -> "CollectionsCopilotModel":
        """Build inference stack from on-disk artifacts.

        Args:
            checkpoint_path: slim checkpoint (.pt) from slim_checkpoint.py.
            model_config_path: same YAML the trainer used.
            schema_path: parent's data/schema.yaml.
            histories_path: parent's data/synthetic/token_ids.npy.
            cast_path: encoder/data/collections_cast.json.
            device: torch.device. CPU is the default for demo replay.

        Returns:
            CollectionsCopilotModel ready for predict / stream_reasoning.
        """
        schema = load_schema(schema_path)
        histories = np.load(histories_path, mmap_mode="r")
        cast = _load_cast(cast_path)

        with model_config_path.open() as f:
            mcfg = yaml.safe_load(f)

        dtype = torch.float32 if device.type == "cpu" else torch.bfloat16

        model = build_transaction_multisurface(
            schema=schema,
            model_path=mcfg["backbone"]["hf_path"],
            encoder_cfg=mcfg.get("encoder"),
            projector_cfg=mcfg.get("projector"),
            head_cfg=mcfg.get("heads"),
            lora_cfg=mcfg["backbone"].get("lora"),
            dtype=dtype,
            device_map=None if device.type == "cpu" else "auto",
        )

        ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
        if not ckpt.get("model_state_dict_slim"):
            raise ValueError(
                f"Expected a slim checkpoint at {checkpoint_path}. "
                f"Run encoder/scripts/slim_checkpoint.py first.",
            )
        state = {
            k: v.to(dtype) if v.is_floating_point() else v
            for k, v in ckpt["model_state_dict"].items()
        }
        missing, unexpected = model.load_state_dict(state, strict=False)
        if unexpected:
            raise RuntimeError(f"Unexpected state_dict keys: {unexpected[:5]} ...")

        model.to(device)
        return cls(
            model=model, schema=schema, histories=histories,
            cast=cast, device=device,
        )

    # ----- inference -----

    @torch.inference_mode()
    def predict(
        self,
        member: CollectionsCastMember,
        top_k: int = 5,
    ) -> CollectionsResult:
        """Run one inference: probability head (K×3) + attribution head."""
        torch.manual_seed(DEMO_SEED)
        batch = self._build_batch(member)
        out = self.model.predict(batch)

        # prob_logits: (1, K, num_bands). Softmax along band axis only.
        prob_logits = out["probability_logits"][0].float().cpu()  # (K, bands)
        probs = torch.softmax(prob_logits, dim=-1)
        likely_scores = probs[..., BAND_LIKELY].numpy()           # (K,)
        predicted_bands = prob_logits.argmax(dim=-1).tolist()
        dominant = int(np.argmax(likely_scores))

        attr_logits = out["attribution_logits"][0].float().cpu()   # (64,)
        attr_probs = torch.sigmoid(attr_logits).numpy()
        top_k_positions = (
            torch.topk(attr_logits, k=top_k, dim=-1).indices.numpy()
        )

        return CollectionsResult(
            likely_scores=likely_scores,
            predicted_bands=predicted_bands,
            dominant_treatment=dominant,
            attribution_probs=attr_probs,
            top_k_positions=top_k_positions,
            context_idx=member.context_idx,
        )

    def build_reasoning_text(
        self,
        member: CollectionsCastMember,
        result: CollectionsResult,
    ) -> str:
        """Render analyst-facing reasoning deterministically.

        Same Ottoguard pattern as dispute: the model contributes
        (predicted_bands, likely_scores, dominant_treatment, attribution).
        The history contributes ground-truth cross-position signals
        (velocity, subscription burden, merchant diversity, large
        amount count, volatility). This function fuses both into a
        coherent paragraph.
        """
        history = np.asarray(self.histories[member.customer_idx])
        return _render_reasoning(
            history=history,
            result=result,
        )

    def stream_reasoning(
        self,
        member: CollectionsCastMember,
        result: CollectionsResult,
        chunk_chars: int = 6,
    ) -> Iterator[str]:
        """Yield the reasoning text in cumulative chunks for the UI."""
        text = self.build_reasoning_text(member, result)
        for i in range(chunk_chars, len(text) + chunk_chars, chunk_chars):
            yield text[: min(i, len(text))]

    # ----- internals -----

    def _build_batch(
        self,
        member: CollectionsCastMember,
    ) -> MixedModalityBatch:
        history = np.asarray(self.histories[member.customer_idx]).copy()
        feature_ids = torch.from_numpy(history).long().unsqueeze(0)
        feature_ids = feature_ids.to(self.device)

        input_ids, attention_mask, lengths = tokenize_texts(
            self.tokenizer, [member.context_text], max_length=256,
        )
        input_ids = input_ids.to(self.device)
        attention_mask = attention_mask.to(self.device)
        lengths = lengths.to(self.device)

        # Collections always anchors the encoder bias at the most-recent
        # position; member.context_idx is 63 in v1.
        context_idx = torch.tensor(
            [member.context_idx], dtype=torch.long, device=self.device,
        )

        return MixedModalityBatch(
            feature_ids=feature_ids,
            text_input_ids=input_ids,
            text_attention_mask=attention_mask,
            text_lengths=lengths,
            head_target="probability",
            disputed_idx=context_idx,
        )


def _load_cast(cast_path: Path) -> list[CollectionsCastMember]:
    payload = json.loads(cast_path.read_text())
    return [
        CollectionsCastMember(
            pattern=m["pattern"],
            display_name=m["display_name"],
            customer_idx=int(m["customer_idx"]),
            context_idx=int(m["context_idx"]),
            treatment_labels=list(m["treatment_labels"]),
            treatment_label_names=list(m["treatment_label_names"]),
            dominant_treatment=int(m["dominant_treatment"]),
            dominant_treatment_name=m["dominant_treatment_name"],
            description=m["description"],
            context_text=m["context_text"],
        )
        for m in payload["cast"]
    ]


def _render_reasoning(
    history: np.ndarray,
    result: CollectionsResult,
) -> str:
    """Build the analyst-facing reasoning paragraph from cross-position signals."""
    velocity = signal_recent_velocity(history)
    sub_burden = signal_subscription_burden(history)
    unique_merch = signal_merchant_diversity(history)
    large_amt = signal_large_amount_count(history)
    volatility = signal_spending_volatility(history)

    dom = result.dominant_treatment
    dom_name = TREATMENT_NAMES[dom]
    dom_score = float(result.likely_scores[dom])

    # Per-treatment one-line scoreboard.
    score_lines = []
    for t in range(NUM_TREATMENTS):
        band = result.predicted_bands[t]
        score = float(result.likely_scores[t])
        marker = "★" if t == dom else " "
        score_lines.append(
            f"{marker} {TREATMENT_NAMES[t]}: P(respond)={score:.2f} "
            f"band={BAND_NAMES[band]}"
        )
    scoreboard = "\n".join(score_lines)

    # Pattern-specific reasoning grounded in the cross-position signals.
    if dom == TREATMENT_SETTLEMENT:
        rationale = (
            f"Settlement is the model's recommendation (P={dom_score:.2f}). "
            f"The customer has {large_amt} large-amount transactions "
            f"(discretionary capacity), recent activity is steady "
            f"(velocity {velocity:.1f}, lower = more active), and "
            f"spending volatility is moderate ({volatility:.1f}). "
            f"This profile suggests the customer can muster a one-time "
            f"lump sum without compromising day-to-day cash flow."
        )
    elif dom == TREATMENT_PAYMENT_PLAN:
        rationale = (
            f"Payment plan is the model's recommendation (P={dom_score:.2f}). "
            f"The customer carries {sub_burden} recurring obligations "
            f"in their history — they already tolerate auto-debits, "
            f"so a structured monthly payment aligns with their behavioral "
            f"pattern. Recent activity (velocity {velocity:.1f}) suggests "
            f"the customer is engaged enough to manage a multi-month plan."
        )
    elif dom == TREATMENT_SOFT_TOUCH:
        rationale = (
            f"Soft-touch is the model's recommendation (P={dom_score:.2f}). "
            f"The customer spans {unique_merch} unique merchants with "
            f"healthy spending breadth, and the recent velocity "
            f"({velocity:.1f}) suggests active engagement, not distress. "
            f"This profile typically self-resolves with light contact "
            f"and a small concession rather than aggressive collections."
        )
    elif dom == TREATMENT_NO_OFFER:
        rationale = (
            f"No-offer is the model's recommendation (P={dom_score:.2f}). "
            f"Recent activity is sparse (velocity {velocity:.1f}, well "
            f"above the active threshold), merchant diversity is low "
            f"({unique_merch}), and there are no large-amount transactions "
            f"({large_amt}) suggesting discretionary capacity. The "
            f"behavioral signature is dormant — analyst hours are better "
            f"spent on accounts with response signal."
        )
    else:
        rationale = (
            f"The model recommends {dom_name} (P={dom_score:.2f}). "
            f"Cross-position signature — velocity {velocity:.1f}, subscription "
            f"burden {sub_burden}, merchant diversity {unique_merch}, "
            f"large-amount count {large_amt}, volatility {volatility:.1f}."
        )

    top_positions_str = ", ".join(str(int(p)) for p in result.top_k_positions[:5])
    return (
        f"{rationale} The model's top contributing transactions are "
        f"positions {top_positions_str}.\n\n"
        f"Per-treatment scoreboard:\n{scoreboard}"
    )