encoder/src/data/synthetic_collections.py

"""Programmatic synthesizer for the collections-treatment surface.

Produces `(customer_history, context_idx, context_text, treatment_labels,
attribution_labels, reasoning_text)` examples. Like the dispute
synthesizer, this reuses the parent's existing customer histories
(`data/synthetic/token_ids.npy`) — no fresh transaction generation.

Task (per ADR 0003 + blueprint §5):

    Given a customer's 64-transaction history at the point of
    delinquency, output the probability the customer responds
    favorably to each of K=4 treatments:

        treatment 0 = settlement      (one-time reduced lump sum)
        treatment 1 = payment_plan    (3-6 month installments)
        treatment 2 = soft_touch      (light contact, no offer yet)
        treatment 3 = no_offer        (write-off; don't burn analyst hours)

    Each treatment is rated on a 3-band categorical:

        band 0 = unlikely_respond
        band 1 = ambiguous
        band 2 = likely_respond

    The model output is a (K=4) × (bands=3) softmax distribution.

Cross-position signals (mirror-able in TransactionEncoder):

  - `recent_velocity_score`  mean days_since_last across the last 16
    transactions; high = dormant.
  - `subscription_burden_score` count of `is_recurring=1` transactions;
    high = recurring obligations, plan-tolerant.
  - `merchant_diversity_score` unique merchant_id count; high =
    discretionary breadth.
  - `large_amount_score` count of large-amount transactions; high =
    discretionary capacity (settlement-able).
  - `spending_volatility_score` std of amount tokens; high = bouncy.

These signals are computed at corpus generation time (numpy) and at
inference time (torch in the encoder — see §encoder markers in the
blueprint).

Doctrine compliance (data-distribution-doctrine §6):

  - Per-treatment class balance: aim for each of the 12 (treatment, band)
    cells to receive ≥ 10% of the corpus. We sample treatment-balanced
    rather than label-balanced because no single treatment carries the
    headline signal.
  - Hard negatives: ~10% adversarial examples where one signal
    contradicts the dominant pattern (high large_amt but very dormant,
    high subscription burden but high volatility, etc.).
  - Length distribution: context_text p10=15 / p50=35 / p90=80 words.
"""

from __future__ import annotations

import json
import random
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import numpy as np


# --- Feature column indices (must match data/schema.yaml) ---
FEATURE_DAYS_SINCE_LAST = 2
FEATURE_IS_RECURRING = 3
FEATURE_MERCHANT_ID = 5
FEATURE_AMOUNT = 8
FEATURE_CUSTOMER_TENURE = 14

# Reserved token offset (MASK=0, OOV=1, NULL=2, values start at 3).
RESERVED_OFFSET = 3

# is_recurring schema: 0 = not recurring, 1 = recurring. Tokens: 3 = not, 4 = yes.
IS_RECURRING_TRUE = RESERVED_OFFSET + 1  # = 4

# Amount bucket threshold for "large". The amount feature is
# quantile-bucketed into 256 bins over [0.01, 25000]. From the
# empirical distribution (200K customers, 64 tx each):
#   p50 ~ token 100 (~$30), p90 ~ token 180 (~$300)
# Tokens >= 150 represent the discretionary band (~$150+).
AMOUNT_LARGE_THRESH = RESERVED_OFFSET + 150  # = 153

# Recent window for velocity computation.
RECENT_WINDOW = 16

# Cast / context placement: the verdict is rendered "as of the most
# recent transaction." We always set context_idx = 63 so the encoder's
# bias markers land on the most-recent position.
CONTEXT_IDX_DEFAULT = 63

# Treatment ids — match the head layout exactly.
TREATMENT_SETTLEMENT = 0
TREATMENT_PAYMENT_PLAN = 1
TREATMENT_SOFT_TOUCH = 2
TREATMENT_NO_OFFER = 3
NUM_TREATMENTS = 4
TREATMENT_NAMES = ["settlement", "payment_plan", "soft_touch", "no_offer"]

# Band labels — match ProbabilityHead convention (highest-severity last).
BAND_UNLIKELY = 0
BAND_AMBIGUOUS = 1
BAND_LIKELY = 2
NUM_BANDS = 3
BAND_NAMES = ["unlikely_respond", "ambiguous", "likely_respond"]


# --- Cross-position signal helpers (numpy) ---
#
# These MUST mirror the torch computations in TransactionEncoder so
# corpus labels and inference-time encoder biases agree. If you change
# a threshold here, change it in the encoder too.


def signal_recent_velocity(history: np.ndarray) -> float:
    """Mean days_since_last across the last RECENT_WINDOW transactions.

    High value = the customer's recent activity is sparse (dormant).
    Low value = active recent activity.

    The `days_since_last` feature is bucketed [0, 365] → 30 bins, then
    +3 for reserved tokens. We use the raw token mean as the velocity
    proxy; higher token = more days between transactions = more dormant.
    """
    return float(history[-RECENT_WINDOW:, FEATURE_DAYS_SINCE_LAST].mean())


def signal_subscription_burden(history: np.ndarray) -> int:
    """Count of `is_recurring == 1` transactions in the history.

    Customers with many recurring obligations already tolerate
    auto-debits → payment plans are realistic.
    """
    return int(np.sum(history[:, FEATURE_IS_RECURRING] == IS_RECURRING_TRUE))


def signal_merchant_diversity(history: np.ndarray) -> int:
    """Number of unique merchant_id values in the history.

    High diversity = discretionary breadth (entertainment, travel,
    dining). Low diversity = subsistence (groceries, gas, utilities).
    """
    return int(np.unique(history[:, FEATURE_MERCHANT_ID]).size)


def signal_large_amount_count(history: np.ndarray) -> int:
    """Count of large-amount transactions (token >= AMOUNT_LARGE_THRESH)."""
    return int(np.sum(history[:, FEATURE_AMOUNT] >= AMOUNT_LARGE_THRESH))


def signal_spending_volatility(history: np.ndarray) -> float:
    """Standard deviation of amount tokens across the full history."""
    return float(history[:, FEATURE_AMOUNT].astype(np.float32).std())


def signal_customer_tenure(history: np.ndarray) -> int:
    """Customer tenure bucket at the most-recent transaction (token id)."""
    return int(history[-1, FEATURE_CUSTOMER_TENURE])


# --- Label rule (per treatment) ---
#
# Thresholds are empirical from data/synthetic/token_ids.npy distribution:
#   recent_velocity:     p50=6.4, p90=11.0 (higher = more dormant)
#   sub_burden:          p50=10,  p90=18
#   unique_merchants:    p50=15,  p90=19
#   large_amount_count:  p50=6,   p90=23
#   spending_volatility: p50=44,  p90=70
#
# The rule is intentionally conservative on the LIKELY band — we want
# the headline scores to peak only for textbook patterns. The model is
# evaluated on 5/6 cast accuracy with calibration, not 100% confidence.


def classify_settlement(history: np.ndarray) -> int:
    """Likelihood the customer responds to a settlement offer.

    Settlement requires: discretionary capacity (can muster lump sum) +
    stable behavior (won't bail mid-negotiation) + active engagement.
    """
    large_amt = signal_large_amount_count(history)
    velocity = signal_recent_velocity(history)
    volatility = signal_spending_volatility(history)

    if large_amt >= 12 and velocity <= 8.0 and volatility <= 60.0:
        return BAND_LIKELY
    if large_amt >= 5 and velocity <= 11.0:
        return BAND_AMBIGUOUS
    return BAND_UNLIKELY


def classify_payment_plan(history: np.ndarray) -> int:
    """Likelihood the customer accepts a payment plan.

    Plans work for customers with recurring obligations (they already
    tolerate auto-debits) and reasonably active recent behavior.
    """
    sub_burden = signal_subscription_burden(history)
    velocity = signal_recent_velocity(history)

    if sub_burden >= 14 and velocity <= 10.0:
        return BAND_LIKELY
    if sub_burden >= 6 and velocity <= 13.0:
        return BAND_AMBIGUOUS
    return BAND_UNLIKELY


def classify_soft_touch(history: np.ndarray) -> int:
    """Likelihood the customer self-resolves with a light contact.

    Soft-touch works for diversified, moderately active spenders whose
    delinquency is more likely a temporary cash crunch than a structural
    problem.
    """
    unique_merch = signal_merchant_diversity(history)
    velocity = signal_recent_velocity(history)
    volatility = signal_spending_volatility(history)

    if (
        unique_merch >= 17
        and velocity <= 10.0
        and 30.0 <= volatility <= 70.0
    ):
        return BAND_LIKELY
    if unique_merch >= 12 and velocity <= 13.0:
        return BAND_AMBIGUOUS
    return BAND_UNLIKELY


def classify_no_offer(history: np.ndarray) -> int:
    """Likelihood that no-offer (write-off) is the right answer.

    No-offer is "right" when the behavioral signature suggests the
    customer is dormant + subsistence-only: sparse recent activity,
    low merchant diversity, no large-amount transactions. The bank
    saves analyst hours by not pursuing.
    """
    velocity = signal_recent_velocity(history)
    unique_merch = signal_merchant_diversity(history)
    large_amt = signal_large_amount_count(history)

    if velocity >= 11.0 and unique_merch <= 12 and large_amt <= 4:
        return BAND_LIKELY
    if velocity >= 9.5 and large_amt <= 6:
        return BAND_AMBIGUOUS
    return BAND_UNLIKELY


def classify_all_treatments(history: np.ndarray) -> list[int]:
    """Per-treatment labels in the canonical order.

    Returns:
        list[int] of length NUM_TREATMENTS = 4, ordered
        [settlement, payment_plan, soft_touch, no_offer], each in
        {BAND_UNLIKELY, BAND_AMBIGUOUS, BAND_LIKELY}.
    """
    return [
        classify_settlement(history),
        classify_payment_plan(history),
        classify_soft_touch(history),
        classify_no_offer(history),
    ]


def dominant_treatment(treatment_labels: list[int]) -> int:
    """Return the treatment index with the highest band.

    Tie-break order: settlement > payment_plan > soft_touch > no_offer.
    The order reflects "lighter intervention preferred when tied" —
    i.e., the analyst would prefer a settlement to a plan when both
    score the same.
    """
    best = 0
    best_band = treatment_labels[0]
    for i in range(1, NUM_TREATMENTS):
        if treatment_labels[i] > best_band:
            best = i
            best_band = treatment_labels[i]
    return best


# --- Attribution labels ---
#
# Per-position contributions to the dominant treatment's verdict.
# Different treatments highlight different positions:
#   settlement   → recent large-amount transactions
#   payment_plan → recurring (is_recurring=1) positions
#   soft_touch   → diverse merchants (positions whose merchant_id
#                  is "unique to this position" within history)
#   no_offer     → high-days_since_last sparse cluster at the end


def attribution_for_treatment(
    history: np.ndarray,
    treatment_idx: int,
) -> np.ndarray:
    """Per-position attribution for the specified treatment."""
    attr = np.zeros(64, dtype=np.float32)
    if treatment_idx == TREATMENT_SETTLEMENT:
        # Large-amount positions
        mask = history[:, FEATURE_AMOUNT] >= AMOUNT_LARGE_THRESH
        attr[mask] = 1.0
    elif treatment_idx == TREATMENT_PAYMENT_PLAN:
        # Recurring positions
        mask = history[:, FEATURE_IS_RECURRING] == IS_RECURRING_TRUE
        attr[mask] = 1.0
    elif treatment_idx == TREATMENT_SOFT_TOUCH:
        # Diverse-merchant positions: pick positions whose merchant_id
        # appears only once in the history (the "exploration" positions).
        merchants = history[:, FEATURE_MERCHANT_ID]
        unique_vals, counts = np.unique(merchants, return_counts=True)
        singletons = unique_vals[counts == 1]
        mask = np.isin(merchants, singletons)
        attr[mask] = 1.0
    elif treatment_idx == TREATMENT_NO_OFFER:
        # Recent dormancy: positions in the last 16 with high
        # days_since_last (sparse).
        recent = np.arange(48, 64)
        dsl = history[recent, FEATURE_DAYS_SINCE_LAST]
        # Mark positions whose dsl is at or above the customer's median.
        median_dsl = float(np.median(history[:, FEATURE_DAYS_SINCE_LAST]))
        for pos in recent:
            if float(history[pos, FEATURE_DAYS_SINCE_LAST]) >= median_dsl:
                attr[pos] = 1.0
    # Always mark the context position (the "now" anchor).
    attr[CONTEXT_IDX_DEFAULT] = max(attr[CONTEXT_IDX_DEFAULT], 1.0)
    return attr


# --- Context text bank (analyst-facing delinquency context) ---
#
# Five tones × 4 templates = 20 surface variants. Following the
# dispute synthesizer's pattern.

CONTEXT_TEMPLATES: dict[str, list[str]] = {
    "formal": [
        "Customer is {dpd} days past due with an outstanding balance of ${balance:,}. Please assess treatment options.",
        "Account flagged for collections review. {dpd} dpd, ${balance:,} outstanding. Recommend an appropriate treatment from the available catalog.",
        "Delinquency notice: {dpd} days past due, balance ${balance:,}. Requesting model-recommended collections action.",
    ],
    "casual": [
        "Hey — this customer is {dpd} days behind, ${balance:,} owed. What's the right treatment?",
        "Account at {dpd} dpd, ${balance:,} on the books. Treatment recommendation?",
        "Customer is past due {dpd} days for ${balance:,}. Best path forward?",
    ],
    "terse": [
        "{dpd}dpd ${balance:,}. Treatment?",
        "Past due {dpd}d, ${balance:,}.",
        "Collections review. {dpd}/${balance:,}.",
    ],
    "detailed": [
        "This customer is {dpd} days past due on a ${balance:,} balance. Last successful payment was {last_pay} days ago. Account opened {tenure} months ago. Recommend treatment.",
        "Delinquency summary: {dpd} dpd, ${balance:,} outstanding, last payment {last_pay} days back. Customer tenure {tenure} months. What's the model's call?",
    ],
    "urgent": [
        "FYI — this account is {dpd} days past due, ${balance:,} outstanding. Need a treatment call before EOD.",
        "Time-sensitive: {dpd}dpd, ${balance:,} balance. Pre-charge-off window closes soon. Treatment?",
    ],
}


def _build_context_text(
    history: np.ndarray,
    rng: random.Random,
) -> tuple[str, str, dict[str, int]]:
    """Render a delinquency context string for the analyst.

    Returns (text, tone, vars) where vars captures the dpd/balance/etc.
    used in the template for downstream auditing.
    """
    # Synthesize plausible delinquency parameters. These are NOT in the
    # tokenized history (the schema is transaction-level, not account-
    # level), so we generate them from a distribution that matches what
    # an analyst would see in production. The model conditions on these
    # as text context but the label is computed from history alone.
    # Wide ranges produce ~unique strings per example; tight finite
    # buckets here produce too many duplicate context strings and trip
    # the contamination gate.
    dpd = rng.randint(30, 180)
    balance = rng.randint(500, 25000)
    last_pay = rng.randint(20, 180)
    tenure = max(1, signal_customer_tenure(history) - RESERVED_OFFSET) * 12  # months

    tone = rng.choice(list(CONTEXT_TEMPLATES.keys()))
    template = rng.choice(CONTEXT_TEMPLATES[tone])
    text = template.format(
        dpd=dpd,
        balance=balance,
        last_pay=last_pay,
        tenure=tenure,
    )
    return text, tone, {
        "dpd": dpd,
        "balance": balance,
        "last_pay": last_pay,
        "tenure_months": tenure,
    }


# --- Reasoning text (templated, lesson 4) ---


def _decode_treatment_name(treatment_idx: int) -> str:
    return TREATMENT_NAMES[treatment_idx]


def _decode_band_name(band: int) -> str:
    return BAND_NAMES[band]


def build_reasoning_text(
    history: np.ndarray,
    treatment_labels: list[int],
) -> str:
    """Programmatic reasoning grounded in the cross-position signals.

    The text is fully deterministic from `treatment_labels` + history.
    Following lesson 4, this is NOT generated by the LM head; the model
    renders it from its (per-treatment band + ground-truth signals)
    output at inference time.
    """
    dom_idx = dominant_treatment(treatment_labels)
    dom_band = treatment_labels[dom_idx]
    dom_name = _decode_treatment_name(dom_idx)
    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)

    parts: list[str] = []
    parts.append(
        f"Recommended treatment: {dom_name} (band: {_decode_band_name(dom_band)})."
    )
    parts.append(
        f"Behavioral signature — recent velocity {velocity:.1f} "
        f"(lower=more active), subscription burden {sub_burden}, "
        f"merchant diversity {unique_merch}, large-amount transactions {large_amt}, "
        f"spending volatility {volatility:.1f}."
    )
    if dom_idx == TREATMENT_SETTLEMENT:
        parts.append(
            f"Settlement is the dominant option because the customer has "
            f"{large_amt} discretionary-band transactions and stable spending "
            f"behavior — consistent with the capacity to muster a lump sum."
        )
    elif dom_idx == TREATMENT_PAYMENT_PLAN:
        parts.append(
            f"Payment plan is the dominant option because the customer "
            f"already maintains {sub_burden} recurring obligations — "
            f"a structured monthly debit aligns with their existing pattern."
        )
    elif dom_idx == TREATMENT_SOFT_TOUCH:
        parts.append(
            f"Soft-touch is the dominant option because the customer "
            f"shows {unique_merch} unique merchants and active recent "
            f"behavior — pattern suggests temporary cash crunch, not "
            f"structural distress."
        )
    elif dom_idx == TREATMENT_NO_OFFER:
        parts.append(
            f"No-offer is the dominant option because the recent activity "
            f"pattern is sparse (velocity {velocity:.1f}) and the merchant "
            f"diversity is low — behavioral signature suggests the customer "
            f"is dormant. Continued outreach is unlikely to convert."
        )
    return " ".join(parts)


# --- Example dataclass ---


@dataclass
class CollectionsExample:
    """One synthesized collections-treatment training example."""

    customer_idx: int
    context_idx: int                # always 63 in v1 (most-recent position)
    context_text: str
    treatment_labels: list[int]     # length NUM_TREATMENTS = 4
    attribution_labels: list[float] # length 64
    reasoning_text: str
    tone: str
    is_adversarial: bool
    context_vars: dict[str, int]    # dpd, balance, last_pay, tenure for audit

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


def synthesize_one(
    history: np.ndarray,
    customer_idx: int,
    rng: random.Random,
    adversarial: bool = False,
) -> CollectionsExample:
    """Produce one collections example for a given customer history.

    The label is fully determined by `history`; the context_text is
    rendered with rng-driven tone variation. Adversarial flag adds
    surface-form perturbation to the context text (lesson 3 boundary
    coverage) but does NOT change the label.
    """
    treatment_labels = classify_all_treatments(history)
    dom_idx = dominant_treatment(treatment_labels)
    attribution = attribution_for_treatment(history, dom_idx)
    context_text, tone, context_vars = _build_context_text(history, rng)
    if adversarial:
        context_text = _apply_adversarial_perturbation(context_text, rng)
    reasoning = build_reasoning_text(history, treatment_labels)
    return CollectionsExample(
        customer_idx=customer_idx,
        context_idx=CONTEXT_IDX_DEFAULT,
        context_text=context_text,
        treatment_labels=treatment_labels,
        attribution_labels=attribution.tolist(),
        reasoning_text=reasoning,
        tone=tone,
        is_adversarial=adversarial,
        context_vars=context_vars,
    )


def _apply_adversarial_perturbation(text: str, rng: random.Random) -> str:
    """Same shape as the dispute synthesizer's perturbation: light case-flip,
    occasional double-space. Preserves readability while training the
    model against surface-form noise (doctrine §3 production drift)."""
    chars = list(text)
    for i in range(len(chars)):
        if chars[i].isalpha() and rng.random() < 0.10:
            chars[i] = chars[i].swapcase()
    if rng.random() < 0.3 and " " in text:
        idx = text.index(" ")
        chars.insert(idx, " ")
    return "".join(chars)


# --- Corpus generation ---


def generate_corpus(
    histories: np.ndarray,
    train_indices: np.ndarray,
    target_size: int = 4000,
    seed: int = 42,
    adversarial_fraction: float = 0.10,
) -> list[CollectionsExample]:
    """Generate a class-balanced collections corpus.

    Sampling strategy: we draw customers uniformly from train_indices,
    classify, and accept up to a per-(dominant_treatment) cap so the
    final corpus has reasonable balance across the four treatments.

    This is intentionally less hand-tuned than the dispute synthesizer
    because the rule has 4 treatments × 3 bands = 12 cells; trying to
    target each cell is overkill at 4000 examples. The cast curator
    handles per-pattern coverage separately.

    Args:
        histories: (N, 64, 15) full corpus.
        train_indices: indices into histories for the training pool.
        target_size: total corpus size.
        seed: RNG seed.
        adversarial_fraction: fraction of examples with text perturbation.

    Returns:
        List of CollectionsExample, length ~target_size.
    """
    rng = random.Random(seed)
    np_rng = np.random.RandomState(seed)
    examples: list[CollectionsExample] = []

    # Treatment cap: each dominant treatment gets at most this many
    # examples. 4 treatments × 1500 cap = 6000 ceiling, well above 4000.
    per_treatment_cap = target_size // 2
    treatment_counts = [0] * NUM_TREATMENTS

    n_adv_target = int(target_size * adversarial_fraction)
    n_adv = 0
    attempts = 0
    max_attempts = target_size * 8

    while len(examples) < target_size and attempts < max_attempts:
        attempts += 1
        customer_idx = int(np_rng.choice(train_indices))
        history = histories[customer_idx]
        treatment_labels = classify_all_treatments(history)
        dom_idx = dominant_treatment(treatment_labels)
        if treatment_counts[dom_idx] >= per_treatment_cap:
            continue
        adversarial = (n_adv < n_adv_target) and (rng.random() < 0.15)
        example = synthesize_one(
            history, customer_idx, rng, adversarial=adversarial,
        )
        examples.append(example)
        treatment_counts[dom_idx] += 1
        if adversarial:
            n_adv += 1

    rng.shuffle(examples)
    return examples


def write_jsonl(
    examples: list[CollectionsExample],
    output_path: Path | str,
) -> None:
    """Write the corpus to JSONL, one example per line."""
    output_path = Path(output_path)
    with output_path.open("w") as f:
        for example in examples:
            f.write(json.dumps(example.to_dict()) + "\n")