src/next_action/train.py

"""
NextActionPredictor training script.

Data source: ~6 000 synthetic (domain, entity_flags, prior_contact → action)
             examples generated in-memory from DOMAIN_ACTION_PRIORS.
             No download required.  Trains in < 30 seconds on CPU.

Label assignment rules (derived from DOMAIN_ACTION_PRIORS):
    prior_contact = 0  →  company_support  (always try company first)
    prior_contact = 1  →  priors[domain][1]  (domain regulatory authority)
    prior_contact = 1
      + has_AMOUNT + has_REF_ID
      + domain has a tertiary priors[2]
      + 20% chance     →  priors[domain][2]  (escalate to legal / ombudsman)

This creates a distribution the MLP learns naturally: company_support is the
dominant class (~50%) while the 5 escalation authorities split the rest.

CLI usage:
    python -m src.next_action.train --output_path models/next_action/model.pt
"""

from __future__ import annotations

import argparse
import logging
import os
import random

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

from src.next_action.model import (
    ACTION2ID,
    DOMAIN_LABELS_ORDERED,
    ENTITY_NAMES,
    FEATURE_DIM,
    GUIDE_MLP,
    NUM_ACTIONS,
    ACTION_LABELS,
    build_feature_vector,
)
from src.next_action.priors import DOMAIN_ACTION_PRIORS

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Label assignment
# ---------------------------------------------------------------------------

def _label_for(
    domain: str,
    entity_flags: list[float],
    prior_contact: float,
    rng: random.Random,
) -> int:
    """Assign a training label from the priors and feature context."""
    priors = DOMAIN_ACTION_PRIORS[domain]   # ordered list of action strings

    if prior_contact < 0.5:
        action = priors[0]                  # always "company_support"
    else:
        has_amount  = entity_flags[1] > 0.5  # ENTITY_NAMES[1] = "AMOUNT"
        has_ref_id  = entity_flags[3] > 0.5  # ENTITY_NAMES[3] = "REF_ID"
        strong_case = has_amount and has_ref_id

        if strong_case and len(priors) >= 3 and rng.random() < 0.20:
            action = priors[2]              # tertiary (legal / secondary ombudsman)
        else:
            action = priors[1]              # primary regulatory authority

    return ACTION2ID[action]


# ---------------------------------------------------------------------------
# Dataset builder
# ---------------------------------------------------------------------------

def build_synthetic_dataset(
    n_samples: int = 6000,
    seed: int = 42,
) -> tuple[list[list[float]], list[int]]:
    """
    Generate *n_samples* (feature_vector, label_id) pairs in memory.

    Distribution:
      ~50%  prior_contact = 0  →  company_support label
      ~50%  prior_contact = 1  →  domain authority or legal label
    Entity flags are sampled independently with p=0.55 of being set (biased
    toward having evidence, which is realistic for complaint scenarios).
    """
    rng = random.Random(seed)
    X: list[list[float]] = []
    y: list[int] = []

    for _ in range(n_samples):
        domain = rng.choice(DOMAIN_LABELS_ORDERED)
        entity_flags = [float(rng.random() < 0.55) for _ in ENTITY_NAMES]
        prior_contact = float(rng.random() < 0.50)

        fv = build_feature_vector(domain, entity_flags, prior_contact)
        label = _label_for(domain, entity_flags, prior_contact, rng)

        X.append(fv)
        y.append(label)

    # Log class distribution for transparency
    counts: dict[int, int] = {}
    for lbl in y:
        counts[lbl] = counts.get(lbl, 0) + 1
    dist = ", ".join(
        f"{ACTION_LABELS[k]}={v}" for k, v in sorted(counts.items())
    )
    logger.info("Synthetic dataset: %d examples — %s", n_samples, dist)

    return X, y


# ---------------------------------------------------------------------------
# Training entry point
# ---------------------------------------------------------------------------

def train(args: argparse.Namespace) -> None:
    """Train GUIDE_MLP and save checkpoint to args.output_path."""
    logging.basicConfig(level=logging.INFO)

    # 1. Data
    X, y = build_synthetic_dataset(n_samples=args.n_samples)
    X_t = torch.tensor(X, dtype=torch.float32)
    y_t = torch.tensor(y, dtype=torch.long)

    dataset = TensorDataset(X_t, y_t)
    loader  = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)

    # 2. Model, loss, optimiser
    model     = GUIDE_MLP()
    criterion = nn.CrossEntropyLoss()
    optimiser = torch.optim.Adam(
        model.parameters(), lr=args.lr, weight_decay=1e-4
    )

    # 3. Training loop  (pure PyTorch — no Trainer — keeps startup time minimal)
    model.train()
    for epoch in range(1, args.epochs + 1):
        epoch_loss = 0.0
        correct    = 0

        for xb, yb in loader:
            optimiser.zero_grad()
            logits = model(xb)
            loss   = criterion(logits, yb)
            loss.backward()
            optimiser.step()
            epoch_loss += loss.item() * len(yb)
            correct    += (logits.argmax(1) == yb).sum().item()

        if epoch % 10 == 0 or epoch == args.epochs:
            acc = correct / len(dataset)
            logger.info(
                "Epoch %3d/%d — loss: %.4f  acc: %.3f",
                epoch, args.epochs,
                epoch_loss / len(dataset),
                acc,
            )

    # 4. Evaluate final accuracy
    model.eval()
    with torch.no_grad():
        all_logits = model(X_t)
    final_acc = (all_logits.argmax(1) == y_t).float().mean().item()
    logger.info("Training complete — final train acc: %.3f", final_acc)

    # 5. Save checkpoint
    os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True)
    torch.save(
        {
            "state_dict":    model.state_dict(),
            "feature_dim":   FEATURE_DIM,
            "num_actions":   NUM_ACTIONS,
            "action_labels": ACTION_LABELS,
            "domain_labels": DOMAIN_LABELS_ORDERED,
            "entity_names":  ENTITY_NAMES,
        },
        args.output_path,
    )
    logger.info("NextActionPredictor checkpoint saved to %s", args.output_path)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Train NextActionPredictor MLP")
    p.add_argument(
        "--output_path", default="models/next_action/model.pt",
        help="Path to save the .pt checkpoint",
    )
    p.add_argument(
        "--n_samples", type=int, default=6000,
        help="Number of synthetic training examples",
    )
    p.add_argument(
        "--epochs", type=int, default=60,
        help="Training epochs (default 60 → well under 30 s on CPU)",
    )
    p.add_argument(
        "--batch_size", type=int, default=128,
        help="Mini-batch size",
    )
    p.add_argument(
        "--lr", type=float, default=3e-3,
        help="Adam learning rate",
    )
    return p.parse_args()


if __name__ == "__main__":
    train(parse_args())