src/next_action/model.py

"""
NextActionPredictor model definition.

Architecture: 2-hidden-layer MLP (12-dim → 64 → 64 → 6)
Input layout (12 dims, indices):
    [0-5]  domain one-hot  — ecommerce | telecom | banking | cibil | insurance | general
    [6-10] entity flags    — has_ORG | has_AMOUNT | has_DATE | has_REF_ID | has_ACCOUNT
    [11]   prior_contact   — 1.0 if user has already contacted the company, else 0.0

Output: all 6 EscalationActions sorted by confidence descending.

Fallback: if no checkpoint exists, DOMAIN_ACTION_PRIORS rule-based mapping is
          used so the pipeline never crashes.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn

from src.next_action.priors import ACTION_METADATA, DOMAIN_ACTION_PRIORS

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Label / feature constants — single source of truth for all modules
# ---------------------------------------------------------------------------

ACTION_LABELS: list[str] = [
    "company_support", "nch", "trai", "rbi_ombudsman", "irdai", "legal"
]
ACTION2ID: dict[str, int] = {a: i for i, a in enumerate(ACTION_LABELS)}

# Domain ordering must match the one-hot encoding used during training
DOMAIN_LABELS_ORDERED: list[str] = [
    "ecommerce", "telecom", "banking", "cibil", "insurance", "general"
]

# The 5 NER entity types used as routing signals (PERSON excluded — role reference)
ENTITY_NAMES: list[str] = ["ORG", "AMOUNT", "DATE", "REF_ID", "ACCOUNT"]

FEATURE_DIM: int = 12   # 6 + 5 + 1
NUM_ACTIONS: int = 6


# ---------------------------------------------------------------------------
# Feature vector builder — used by train.py and predict.py
# ---------------------------------------------------------------------------

def build_feature_vector(
    domain: str,
    entity_flags: list[float],
    prior_contact: float,
) -> list[float]:
    """
    Construct the 12-dim feature vector.

    Args:
        domain:        one of DOMAIN_LABELS_ORDERED
        entity_flags:  5 floats (0.0/1.0) in ENTITY_NAMES order
        prior_contact: 1.0 if user already contacted the company, else 0.0

    Returns: list[float] of length FEATURE_DIM (12)
    """
    domain_oh = [1.0 if d == domain else 0.0 for d in DOMAIN_LABELS_ORDERED]
    return domain_oh + list(entity_flags) + [float(prior_contact)]


def entities_dict_to_flags(entities: dict) -> list[float]:
    """
    Convert an EvidenceNER entities dict {entity_type: value} → 5-dim flag vector.

    Example: {"ORG": "Flipkart", "AMOUNT": "₹4,299"} → [1.0, 1.0, 0.0, 0.0, 0.0]
    """
    return [1.0 if name in entities and entities[name] else 0.0
            for name in ENTITY_NAMES]


# ---------------------------------------------------------------------------
# Public output type
# ---------------------------------------------------------------------------

@dataclass
class EscalationAction:
    """A single recommended escalation step."""
    action: str
    authority: str
    url: str
    confidence: float


# ---------------------------------------------------------------------------
# Raw PyTorch module
# ---------------------------------------------------------------------------

class GUIDE_MLP(nn.Module):
    """
    2-hidden-layer MLP: 12 → 64 → ReLU → 64 → ReLU → 6.

    Returns raw logits; softmax is applied at inference time.
    """

    def __init__(
        self,
        input_dim: int = FEATURE_DIM,
        hidden_dim: int = 64,
        num_classes: int = NUM_ACTIONS,
    ) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:  # (B, 12) → (B, 6)
        return self.net(x)


# ---------------------------------------------------------------------------
# Rule-based fallback
# ---------------------------------------------------------------------------

def _rule_based_predict(feature_vector: list[float]) -> list[EscalationAction]:
    """
    Return a ranked EscalationAction list from DOMAIN_ACTION_PRIORS.

    When prior_contact=1, company_support is deprioritised — the user already
    tried that path.
    """
    domain_oh = feature_vector[:6]
    prior_contact = feature_vector[11]

    domain_idx = int(max(range(6), key=lambda i: domain_oh[i]))
    domain = DOMAIN_LABELS_ORDERED[domain_idx]

    ordered: list[str] = list(DOMAIN_ACTION_PRIORS[domain])  # e.g. ["company_support","nch","legal"]
    if prior_contact > 0.5 and ordered[0] == "company_support":
        # Already tried company — rotate company_support to last position
        ordered = ordered[1:] + ordered[:1]

    # Rank-based exponential score: position 0 → highest
    decay = 0.55
    score_map: dict[str, float] = {}
    s = 1.0
    for action in ordered:
        score_map[action] = s
        s *= decay
    for action in ACTION_LABELS:
        if action not in score_map:
            score_map[action] = decay ** len(ordered)  # small residual

    total = sum(score_map.values())

    return sorted(
        [
            EscalationAction(
                action=action,
                authority=ACTION_METADATA[action]["authority"],
                url=ACTION_METADATA[action]["url"],
                confidence=round(score_map[action] / total, 4),
            )
            for action in ACTION_LABELS
        ],
        key=lambda e: -e.confidence,
    )


# ---------------------------------------------------------------------------
# NextActionPredictor — public wrapper
# ---------------------------------------------------------------------------

class NextActionPredictor:
    """
    MLP escalation router with rule-based fallback.

    Pass model_path=None (or a path that does not exist) to use the rule-based
    fallback only.  The pipeline never crashes: if the checkpoint is missing
    a warning is logged and DOMAIN_ACTION_PRIORS is used instead.
    """

    def __init__(self, model_path: Optional[str] = None) -> None:
        self._mlp: Optional[GUIDE_MLP] = None
        self._device = torch.device(
            "cuda" if torch.cuda.is_available()
            else "mps" if torch.backends.mps.is_available()
            else "cpu"
        )

        if model_path and os.path.isfile(model_path):
            try:
                ckpt = torch.load(model_path, map_location=self._device, weights_only=True)
                self._mlp = GUIDE_MLP()
                self._mlp.load_state_dict(ckpt["state_dict"])
                self._mlp.to(self._device)
                self._mlp.eval()
                logger.info("NextActionPredictor loaded from %s on %s", model_path, self._device)
            except Exception:
                logger.warning(
                    "Failed to load NextActionPredictor from %s — using rule-based fallback.",
                    model_path, exc_info=True,
                )
        else:
            logger.info(
                "No NextActionPredictor checkpoint at '%s' — using rule-based fallback.",
                model_path,
            )

    @property
    def uses_fallback(self) -> bool:
        return self._mlp is None

    def predict(self, feature_vector: list[float]) -> list[EscalationAction]:
        """Return all 6 EscalationActions ranked by confidence (highest first)."""
        if self._mlp is None:
            return _rule_based_predict(feature_vector)

        x = torch.tensor(feature_vector, dtype=torch.float32).unsqueeze(0).to(self._device)
        with torch.no_grad():
            logits = self._mlp(x)[0].cpu()    # (6,)
        probs: list[float] = torch.softmax(logits, dim=-1).tolist()

        return sorted(
            [
                EscalationAction(
                    action=ACTION_LABELS[i],
                    authority=ACTION_METADATA[ACTION_LABELS[i]]["authority"],
                    url=ACTION_METADATA[ACTION_LABELS[i]]["url"],
                    confidence=round(probs[i], 4),
                )
                for i in range(NUM_ACTIONS)
            ],
            key=lambda e: -e.confidence,
        )