src/clause_classifier.py

"""
clause_classifier.py — Multi-label legal clause classifier.

Supports three backbone models:
  - Legal-BERT:       nlpaueb/legal-bert-base-uncased
  - DeBERTa-v3:       microsoft/deberta-v3-base
  - Legal-RoBERTa-L:  lexlms/legal-roberta-large

All backbones share identical training logic (AsymmetricLoss, AdamW,
linear warmup, early stopping). Model selection and per-class threshold
optimisation are handled by ModelSelector, which evaluates all single
backbones and every 2- and 3-model ensemble combination. Production
inference uses the saved production_config.json for deterministic,
tuning-free deployment.

Usage:
    # Train individual backbones
    python src/clause_classifier.py --mode train --backbone legalbert
    python src/clause_classifier.py --mode train --backbone deberta
    python src/clause_classifier.py --mode train --backbone legalroberta

    # Select best model / ensemble and optimise thresholds
    python src/clause_classifier.py --mode select

    # Final evaluation on test set (no tuning)
    python src/clause_classifier.py --mode evaluate

    # Single-clause prediction
    python src/clause_classifier.py --mode predict --text "The Licensor shall..."
"""

import argparse
import json
import math
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from loguru import logger
from sklearn.metrics import (
    classification_report,
    f1_score,
    precision_score,
    recall_score,
)
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
)

sys.path.insert(0, str(Path(__file__).parent.parent))
import config

logger.remove()
logger.add(config.LOGS_DIR / "classifier.log", rotation="10 MB", level="DEBUG")
logger.add(sys.stderr, level="INFO")

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

BACKBONE_MODEL_MAP: Dict[str, str] = {
    "legalbert":    config.LEGAL_BERT_MODEL,
    "deberta":      config.DEBERTA_MODEL,
    "legalroberta": config.LEGAL_ROBERTA_MODEL,
}

BACKBONE_DIR_MAP: Dict[str, Path] = {
    "legalbert":    config.CLASSIFIER_LEGALBERT_DIR,
    "deberta":      config.CLASSIFIER_DEBERTA_DIR,
    "legalroberta": config.CLASSIFIER_LEGALROBERTA_DIR,
}

# Per-backbone training hyperparameters
# DeBERTa-v3 NaN root causes: disentangled attention sensitivity to high LR,
# head gradient explosion, and AdamW default eps=1e-8 (must be 1e-6 per paper).
# Legal-RoBERTa-large uses the same settings as Legal-BERT — RoBERTa's pooler
# is stable at 2e-5 and the large model handles fp16 without NaN issues.
BACKBONE_LR_MAP: Dict[str, float] = {
    "legalbert":    2e-5,
    "deberta":      3e-6,   # 2e-6 was stable but too slow; 5e-6 caused NaN — 3e-6 is the midpoint
    "legalroberta": 1e-5,   # slightly lower than BERT-base due to larger model size
}
BACKBONE_EPOCHS_MAP: Dict[str, int] = {
    "legalbert":    90,
    "deberta":      90,
    "legalroberta": 90,
}
BACKBONE_HEAD_MULT_MAP: Dict[str, int] = {
    "legalbert":    10,
    "deberta":      1,    # any head multiplier causes NaN for DeBERTa
    "legalroberta": 10,
}
BACKBONE_ADAM_EPS_MAP: Dict[str, float] = {
    "legalbert":    1e-8,
    "deberta":      1e-6,  # required by DeBERTa paper — prevents NaN
    "legalroberta": 1e-8,
}
BACKBONE_GRAD_CLIP_MAP: Dict[str, float] = {
    "legalbert":    1.0,
    "deberta":      0.5,
    "legalroberta": 1.0,
}

# DeBERTa-v3: use mean-pooling over last hidden states instead of the default
# pooler. The AutoModelForSequenceClassification pooler adds instability at
# LR=3e-6 and can produce all-zero predictions for many epochs.
# Legal-RoBERTa-large: standard CLS pooler via AutoModelForSequenceClassification
# is stable — no mean-pooling needed.
BACKBONE_USE_MEAN_POOL: Dict[str, bool] = {
    "legalbert":    False,
    "deberta":      True,
    "legalroberta": False,
}
# DeBERTa at LR=2e-6 produces near-zero F1 for the first ~15 epochs while the
# encoder representations stabilise. Early stopping must not fire during warmup.
BACKBONE_MIN_EPOCHS_MAP: Dict[str, int] = {
    "legalbert":    10,
    "deberta":      20,
    "legalroberta": 10,
}
# Per-backbone early stopping patience (epochs without val F1 improvement).
BACKBONE_PATIENCE_MAP: Dict[str, int] = {
    "legalbert":    6,
    "deberta":      8,
    "legalroberta": 6,
}
# Legal-RoBERTa-large (355M params) needs smaller batch to fit in 8GB VRAM.
# Gradient accumulation steps compensate so effective batch = CLASSIFIER_BATCH_SIZE.
BACKBONE_BATCH_SIZE_MAP: Dict[str, int] = {
    "legalbert":    config.CLASSIFIER_BATCH_SIZE,      # 16
    "deberta":      config.CLASSIFIER_BATCH_SIZE,      # 16
    "legalroberta": config.CLASSIFIER_BATCH_SIZE // 2, # 8 — large model, 8GB VRAM
}
BACKBONE_GRAD_ACCUM_MAP: Dict[str, int] = {
    "legalbert":    2,
    "deberta":      2,
    "legalroberta": 4,  # 8 × 4 = effective batch 32 (slightly larger than 16×2 — helps large model)
}


# ---------------------------------------------------------------------------
# 1. DATASET
# ---------------------------------------------------------------------------

class ClauseDataset(Dataset):
    """PyTorch Dataset with pre-tokenization and sliding-window chunking.

    Short clauses (≤ max_length tokens) are batch-tokenized upfront — zero
    CPU work per batch during training.

    Long clauses (> max_length tokens) are split into overlapping chunks of
    max_length tokens (stride = max_length // 2). Each chunk inherits the
    parent clause's labels. This ensures the model sees the full text of long
    clauses including exceptions, limitations, and conditions that often appear
    beyond the first max_length tokens.

    Construction is a one-time cost. After that __getitem__ is pure tensor
    indexing.

    Args:
        texts: List of clause text strings.
        labels: List of multi-hot label vectors (length = NUM_CLAUSE_TYPES).
        tokenizer: HuggingFace tokenizer instance.
        max_length: Maximum token sequence length per chunk.
        stride: Overlap between consecutive chunks (tokens).
    """

    def __init__(
        self,
        texts: List[str],
        labels: List[List[int]],
        tokenizer: AutoTokenizer,
        max_length: int = config.CLASSIFIER_MAX_LENGTH,
        stride: int = 128,
    ):
        logger.info(f"Pre-tokenizing {len(texts):,} texts (runs once)...")

        cls_id    = tokenizer.cls_token_id
        sep_id    = tokenizer.sep_token_id
        pad_id    = tokenizer.pad_token_id
        max_chunk = max_length - 2  # reserve 2 positions for [CLS] and [SEP]

        # ── Step 1: tokenize all texts without padding to get lengths ─────
        raw = tokenizer(
            texts,
            add_special_tokens=True,
            truncation=False,
            padding=False,
            return_attention_mask=False,
        )
        raw_ids = raw["input_ids"]  # list of variable-length lists

        # ── Step 2: partition into short and long ─────────────────────────
        short_texts:  List[str]       = []
        short_labels: List[List[int]] = []
        long_ids_list:   List[List[int]]       = []
        long_labels_list: List[List[int]]      = []

        for i, ids in enumerate(raw_ids):
            if len(ids) <= max_length:
                short_texts.append(texts[i])
                short_labels.append(labels[i])
            else:
                long_ids_list.append(ids)
                long_labels_list.append(labels[i])

        # ── Step 3: batch-tokenize short texts with padding (fast) ────────
        parts_ids:   List[torch.Tensor] = []
        parts_masks: List[torch.Tensor] = []
        parts_lbls:  List[torch.Tensor] = []

        if short_texts:
            enc = tokenizer(
                short_texts,
                max_length=max_length,
                padding="max_length",
                truncation=True,
                return_tensors="pt",
            )
            parts_ids.append(enc["input_ids"])
            parts_masks.append(enc["attention_mask"])
            parts_lbls.append(torch.tensor(short_labels, dtype=torch.float32))

        # ── Step 4: sliding-window chunk long texts ────────────────────────
        chunk_ids:   List[List[int]] = []
        chunk_masks: List[List[int]] = []
        chunk_lbls:  List[List[int]] = []
        n_long_chunks = 0

        for ids, label in zip(long_ids_list, long_labels_list):
            body  = ids[1:-1]  # strip [CLS] and [SEP] added by tokenizer
            start = 0
            while start < len(body):
                chunk  = body[start : start + max_chunk]
                seq    = [cls_id] + chunk + [sep_id]
                pad_n  = max_length - len(seq)
                chunk_ids.append(seq + [pad_id] * pad_n)
                chunk_masks.append([1] * len(seq) + [0] * pad_n)
                chunk_lbls.append(label)
                n_long_chunks += 1
                next_start = start + max_chunk - stride
                if next_start >= len(body):
                    break
                start = next_start

        if chunk_ids:
            parts_ids.append(torch.tensor(chunk_ids,   dtype=torch.long))
            parts_masks.append(torch.tensor(chunk_masks, dtype=torch.long))
            parts_lbls.append(torch.tensor(chunk_lbls,  dtype=torch.float32))

        self.input_ids      = torch.cat(parts_ids,   dim=0)
        self.attention_mask = torch.cat(parts_masks, dim=0)
        self.labels         = torch.cat(parts_lbls,  dim=0)

        logger.info(
            f"Pre-tokenization complete. "
            f"{len(texts):,} clauses → {len(self.labels):,} training examples "
            f"({len(long_ids_list):,} long clauses expanded into "
            f"{n_long_chunks:,} chunks via sliding window, stride={stride})."
        )

    def __len__(self) -> int:
        return len(self.labels)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        return {
            "input_ids":      self.input_ids[idx],
            "attention_mask": self.attention_mask[idx],
            "labels":         self.labels[idx],
        }


# ---------------------------------------------------------------------------
# 2. MODEL
# ---------------------------------------------------------------------------

class ClauseClassifierModel(nn.Module):
    """Standard backbone classifier using AutoModelForSequenceClassification.

    Used for Legal-BERT and any backbone where the default pooler is stable.

    Args:
        model_name: HuggingFace model identifier or local checkpoint path.
        num_labels: Number of output labels (default: NUM_CLAUSE_TYPES).
    """

    def __init__(
        self,
        model_name: str,
        num_labels: int = config.NUM_CLAUSE_TYPES,
    ):
        super().__init__()
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_name,
            num_labels=num_labels,
            problem_type="multi_label_classification",
            ignore_mismatched_sizes=True,
        )

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Return raw logits (batch, num_labels)."""
        return self.model(input_ids=input_ids, attention_mask=attention_mask).logits


class MeanPoolClassifierModel(nn.Module):
    """Mean-pooling classifier — used for DeBERTa-v3.

    DeBERTa-v3's AutoModelForSequenceClassification pooler is unstable at
    the low learning rates required to avoid NaN (LR ≤ 3e-6). Mean-pooling
    the last hidden states produces a more stable representation for the
    classification head.

    Architecture:
        DeBERTa encoder → mean-pool last hidden states over non-pad tokens
        → LayerNorm → Dropout(0.1) → Linear(hidden, num_labels)

    Args:
        model_name: HuggingFace model identifier or local checkpoint path.
        num_labels: Number of output labels.
    """

    def __init__(
        self,
        model_name: str,
        num_labels: int = config.NUM_CLAUSE_TYPES,
    ):
        super().__init__()
        from transformers import AutoConfig, AutoModel
        cfg = AutoConfig.from_pretrained(model_name, ignore_mismatched_sizes=True)
        self.encoder   = AutoModel.from_pretrained(model_name, config=cfg)
        hidden         = cfg.hidden_size
        self.norm      = nn.LayerNorm(hidden)
        self.dropout   = nn.Dropout(0.1)
        self.classifier = nn.Linear(hidden, num_labels)
        # Initialise head to small values — prevents large logits on epoch 1
        nn.init.normal_(self.classifier.weight, std=0.02)
        nn.init.zeros_(self.classifier.bias)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Return raw logits (batch, num_labels) via mean-pooled encoder output."""
        outputs     = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        hidden      = outputs.last_hidden_state          # (B, T, H)
        # Mask padding tokens before mean-pooling
        mask_exp    = attention_mask.unsqueeze(-1).float()
        summed      = (hidden * mask_exp).sum(dim=1)
        lengths     = mask_exp.sum(dim=1).clamp(min=1e-9)
        pooled      = summed / lengths                   # (B, H)
        pooled      = self.dropout(self.norm(pooled))
        return self.classifier(pooled)                   # (B, num_labels)

    def save_pretrained(self, save_dir: str) -> None:
        """Save encoder + head weights in a HuggingFace-compatible layout."""
        import os, json as _json
        os.makedirs(save_dir, exist_ok=True)
        self.encoder.save_pretrained(save_dir)
        torch.save(
            {"norm": self.norm.state_dict(),
             "dropout": self.dropout.state_dict(),
             "classifier": self.classifier.state_dict()},
            os.path.join(save_dir, "head_weights.pt"),
        )

    @classmethod
    def from_pretrained(cls, save_dir: str, num_labels: int = config.NUM_CLAUSE_TYPES):
        """Load a saved MeanPoolClassifierModel checkpoint."""
        import os
        model = cls(save_dir, num_labels=num_labels)
        head_path = os.path.join(save_dir, "head_weights.pt")
        if os.path.exists(head_path):
            state = torch.load(head_path, map_location="cpu")
            model.norm.load_state_dict(state["norm"])
            model.classifier.load_state_dict(state["classifier"])
        return model


def build_model(backbone: str, model_name: str) -> nn.Module:
    """Instantiate the correct model class for a given backbone.

    Args:
        backbone: 'legalbert' or 'deberta'.
        model_name: HuggingFace identifier or local checkpoint path.

    Returns:
        Initialised model (ClauseClassifierModel or MeanPoolClassifierModel).
    """
    if BACKBONE_USE_MEAN_POOL.get(backbone, False):
        return MeanPoolClassifierModel(model_name)
    return ClauseClassifierModel(model_name)


# ---------------------------------------------------------------------------
# Asymmetric Loss for multi-label classification
# ---------------------------------------------------------------------------

class AsymmetricLoss(nn.Module):
    """Asymmetric Loss (ASL) for multi-label classification.

    Reference: Ben-Baruch et al. (2021) "Asymmetric Loss For Multi-Label
    Classification" (ICCV 2021).

    ASL addresses the extreme positive/negative imbalance in multi-label
    problems by applying:
      - A higher focusing parameter (gamma_neg) to easy negatives to
        down-weight their contribution
      - A lower focusing parameter (gamma_pos) to positives to preserve
        gradient from rare labels
      - An optional probability margin (clip) that shifts negative
        probabilities by 'm', effectively ignoring very easy negatives

    Default parameters follow the paper's recommended values for
    fine-grained multi-label classification.

    Args:
        gamma_neg: Focusing exponent for negative samples (default 4).
        gamma_pos: Focusing exponent for positive samples (default 1).
        clip: Probability shift for easy-negative clipping (default 0.05).
        reduction: 'mean' or 'sum'.
    """

    def __init__(
        self,
        gamma_neg: float = 4.0,
        gamma_pos: float = 1.0,
        clip: float = 0.05,
        reduction: str = "mean",
    ):
        super().__init__()
        self.gamma_neg = gamma_neg
        self.gamma_pos = gamma_pos
        self.clip      = clip
        self.reduction = reduction

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        """Compute ASL.

        Args:
            logits: Raw model logits (batch, num_labels).
            targets: Soft or hard targets in [0, 1] (batch, num_labels).

        Returns:
            Scalar loss.
        """
        probs    = torch.sigmoid(logits)
        probs_m  = probs                                 # negative branch

        if self.clip is not None and self.clip > 0:
            probs_m = (probs_m + self.clip).clamp(max=1)

        # Binary cross-entropy for both branches
        loss_pos = targets       * torch.log(probs.clamp(min=1e-8))
        loss_neg = (1 - targets) * torch.log((1 - probs_m).clamp(min=1e-8))

        # Asymmetric focusing
        with torch.no_grad():
            pt0 = probs_m.detach()
            pt1 = probs.detach()
            pt  = pt1 * targets + pt0 * (1 - targets)
            w_p = torch.pow(1 - pt1, self.gamma_pos)
            w_n = torch.pow(pt0,     self.gamma_neg)
            w   = w_p * targets + w_n * (1 - targets)

        loss = -w * (loss_pos + loss_neg)

        return loss.mean() if self.reduction == "mean" else loss.sum()


# Backward-compatible alias so any existing imports keep working.
LegalBERTClassifier = ClauseClassifierModel


# ---------------------------------------------------------------------------
# 3. DATA LOADING UTILITIES
# ---------------------------------------------------------------------------

def _load_split(split_name: str) -> Tuple[List[str], List[List[int]]]:
    """Load a train/val/test CSV and convert clause_type to multi-hot vectors.

    Args:
        split_name: One of 'train', 'val', 'test'.

    Returns:
        Tuple of (texts, labels) where labels are multi-hot integer lists.

    Raises:
        FileNotFoundError: If the split CSV does not exist.
    """
    csv_path = config.PROCESSED_DIR / f"{split_name}.csv"
    if not csv_path.exists():
        raise FileNotFoundError(
            f"{split_name}.csv not found. Run data_pipeline.py --mode split first."
        )
    df     = pd.read_csv(csv_path)
    texts  = df["clause_text"].fillna("").tolist()
    labels = [_encode_labels(row) for row in df["clause_type"].fillna("").tolist()]
    return texts, labels


def _encode_labels(clause_type_str: str) -> List[int]:
    """Convert a pipe-separated clause type string to a multi-hot vector.

    Args:
        clause_type_str: e.g. 'Indemnification|Governing Law'.

    Returns:
        Multi-hot integer list of length NUM_CLAUSE_TYPES.
    """
    vec = [0] * config.NUM_CLAUSE_TYPES
    if not clause_type_str.strip():
        return vec
    type_to_idx = {ct: i for i, ct in enumerate(config.CUAD_CLAUSE_TYPES)}
    for ct in clause_type_str.split("|"):
        ct = ct.strip()
        if ct in type_to_idx:
            vec[type_to_idx[ct]] = 1
    return vec




def _build_optimizer_param_groups(
    hf_model: nn.Module,
    learning_rate: float,
    no_decay: List[str],
    head_mult: int = 10,
) -> List[Dict]:
    """Build backbone-agnostic optimizer parameter groups.

    Works for any HuggingFace model regardless of internal attribute names
    (bert, deberta, roberta, etc.) by identifying the classification head
    via standard attribute names and treating everything else as the encoder.

    Applies:
      - head_mult × LR to the classification head (10 for Legal-BERT, 1 for DeBERTa)
      - Weight decay only to non-bias / non-LayerNorm encoder params

    Args:
        hf_model: The AutoModelForSequenceClassification instance.
        learning_rate: Base learning rate for encoder layers.
        no_decay: Parameter name substrings that should have zero weight decay.
        head_mult: LR multiplier for the classification head. Use 1 for DeBERTa
            (any higher multiplier causes NaN during DeBERTa training).

    Returns:
        List of param group dicts for AdamW.
    """
    all_named = list(hf_model.named_parameters())

    head_ids: set = set()
    for attr in ("classifier", "pooler"):
        sub = getattr(hf_model, attr, None)
        if sub is not None:
            for p in sub.parameters():
                head_ids.add(id(p))

    enc_decay    = [p for n, p in all_named if id(p) not in head_ids and not any(nd in n for nd in no_decay)]
    enc_no_decay = [p for n, p in all_named if id(p) not in head_ids and     any(nd in n for nd in no_decay)]
    head_params  = [p for n, p in all_named if id(p) in head_ids]

    return [
        {"params": enc_decay,    "lr": learning_rate,             "weight_decay": 0.01},
        {"params": enc_no_decay, "lr": learning_rate,             "weight_decay": 0.0},
        {"params": head_params,  "lr": learning_rate * head_mult, "weight_decay": 0.01},
    ]


# ---------------------------------------------------------------------------
# 4. TRAINER
# ---------------------------------------------------------------------------

class ClauseClassifierTrainer:
    """Manages the full training loop for a single classifier backbone.

    Both Legal-BERT and DeBERTa use identical training logic:
      - AsymmetricLoss (gamma_neg=4, gamma_pos=1, clip=0.05)
      - AdamW with backbone-agnostic layer-wise LR
      - Linear warmup + linear decay scheduler
      - Gradient accumulation (effective batch = BATCH_SIZE × 2)
      - Best checkpoint saved by validation F1 macro
      - Per-backbone early stopping patience (legalbert=6, deberta=8)
      - Min epochs before early stopping (legalbert=10, deberta=20)
      - Auto-resume: saves training_state.json + training_state.pt after every
        epoch; if these files exist at startup, training resumes automatically
        from the next epoch with optimizer and scheduler state fully restored.

    Args:
        backbone: Either 'legalbert' or 'deberta'.
    """

    def __init__(self, backbone: str = "legalbert"):
        if backbone not in BACKBONE_MODEL_MAP:
            raise ValueError(
                f"backbone must be one of {list(BACKBONE_MODEL_MAP)}, got '{backbone}'"
            )
        self.backbone   = backbone
        self.model_name = BACKBONE_MODEL_MAP[backbone]
        self.ckpt_dir   = BACKBONE_DIR_MAP[backbone]
        self.device     = "cuda" if torch.cuda.is_available() else "cpu"

        # Per-backbone hyperparameters
        self.lr            = BACKBONE_LR_MAP[backbone]
        self.epochs        = BACKBONE_EPOCHS_MAP[backbone]
        self.head_mult     = BACKBONE_HEAD_MULT_MAP[backbone]
        self.adam_eps      = BACKBONE_ADAM_EPS_MAP[backbone]
        self.clip_norm     = BACKBONE_GRAD_CLIP_MAP[backbone]
        self.patience      = BACKBONE_PATIENCE_MAP[backbone]
        self.min_epochs    = BACKBONE_MIN_EPOCHS_MAP[backbone]
        self.batch_size    = BACKBONE_BATCH_SIZE_MAP[backbone]
        self.grad_accum    = BACKBONE_GRAD_ACCUM_MAP[backbone]
        # label_smooth removed — ASL handles noisy labels via gamma focusing;
        # soft targets (label smoothing) cause fp16 log overflow → Train Loss: inf

        logger.info(
            f"Backbone: {backbone} ({self.model_name}) | "
            f"lr={self.lr}  epochs={self.epochs}  eps={self.adam_eps}  clip={self.clip_norm}"
        )
        logger.info(f"Training device: {self.device}")

        # Legal-BERT requires use_fast=False in certain environments
        use_fast = (backbone != "legalbert")
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_fast=use_fast)
        self.model: Optional[ClauseClassifierModel] = None

    def train(self) -> Dict:
        """Run the full training pipeline and return final evaluation metrics.

        Returns:
            Dict with 'backbone', 'best_val_f1', 'best_epoch', 'epochs_trained',
            and per-epoch 'history'.
        """
        state_json = self.ckpt_dir / "training_state.json"
        state_pt   = self.ckpt_dir / "training_state.pt"
        best_dir   = self.ckpt_dir / "best"
        self.ckpt_dir.mkdir(parents=True, exist_ok=True)

        # ── Auto-resume detection ──────────────────────────────────────────
        resuming = state_json.exists() and best_dir.exists()
        if resuming:
            saved       = json.loads(state_json.read_text())
            start_epoch = saved["epoch"] + 1
            best_val_f1 = saved["best_val_f1"]
            best_epoch  = saved["best_epoch"]
            no_improve  = saved["no_improve"]
            history     = saved["history"]
            logger.info(
                f"[{self.backbone}] RESUMING from epoch {start_epoch} "
                f"(best F1={best_val_f1:.4f} at epoch {best_epoch})"
            )
            use_fast = (self.backbone != "legalbert")
            self.tokenizer = AutoTokenizer.from_pretrained(str(best_dir), use_fast=use_fast)
            # MeanPoolClassifierModel has a custom from_pretrained that reloads
            # the head weights stored in head_weights.pt alongside the encoder.
            if BACKBONE_USE_MEAN_POOL.get(self.backbone, False):
                self.model = MeanPoolClassifierModel.from_pretrained(str(best_dir)).to(self.device)
            else:
                self.model = ClauseClassifierModel(str(best_dir)).to(self.device)
        else:
            start_epoch = 1
            best_val_f1 = 0.0
            best_epoch  = 0
            no_improve  = 0
            history: List[Dict] = []
            logger.info(f"[{self.backbone}] Loading training and validation data...")
            self.model = build_model(self.backbone, self.model_name).to(self.device)
            # Gradient checkpointing: recompute activations during backward instead
            # of storing them — cuts activation memory ~60% for large models.
            # Only enabled for legalroberta-large which is VRAM-bound on 8GB.
            if self.backbone == "legalroberta":
                if hasattr(self.model, "model"):
                    self.model.model.gradient_checkpointing_enable()
                elif hasattr(self.model, "encoder"):
                    self.model.encoder.gradient_checkpointing_enable()

        train_texts, train_labels = _load_split("train")
        val_texts,   val_labels   = _load_split("val")

        # Training: ClauseDataset chunks long clauses into overlapping windows.
        # Each chunk is an independent training example — correct for training.
        train_dataset = ClauseDataset(train_texts, train_labels, self.tokenizer)
        train_loader  = DataLoader(
            train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=0,
            pin_memory=(self.device == "cuda"),
        )
        # Validation: raw texts + labels kept separately so _evaluate_aggregated
        # can run sliding window per clause and max-pool predictions back into
        # one prediction per original clause — matching the inference path exactly.

        # AsymmetricLoss (ASL) handles positive/negative imbalance via asymmetric
        # focusing (gamma_neg > gamma_pos) — replaces BCEWithLogitsLoss+pos_weight.
        criterion = AsymmetricLoss(gamma_neg=4.0, gamma_pos=1.0, clip=0.05)

        no_decay = ["bias", "LayerNorm.weight"]
        # MeanPoolClassifierModel exposes the head directly on self.model;
        # ClauseClassifierModel wraps AutoModelForSequenceClassification at self.model.model.
        hf_model_for_params = (
            self.model
            if BACKBONE_USE_MEAN_POOL.get(self.backbone, False)
            else self.model.model
        )
        optimizer_params = _build_optimizer_param_groups(
            hf_model_for_params, self.lr, no_decay, head_mult=self.head_mult
        )
        optimizer = torch.optim.AdamW(optimizer_params, eps=self.adam_eps)

        total_steps = (
            len(train_loader) // self.grad_accum
        ) * self.epochs

        scheduler = get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=config.CLASSIFIER_WARMUP_STEPS,
            num_training_steps=total_steps,
        )

        # Restore optimiser + scheduler state if resuming
        if resuming and state_pt.exists():
            pt = torch.load(str(state_pt), map_location=self.device)
            optimizer.load_state_dict(pt["optimizer"])
            scheduler.load_state_dict(pt["scheduler"])
            logger.info(f"[{self.backbone}] Optimizer and scheduler state restored.")

        for epoch in range(start_epoch, self.epochs + 1):
            train_loss = self._train_epoch(train_loader, criterion, optimizer, scheduler, epoch)

            # NaN abort: _train_epoch returns float('nan') when >50 NaN batches hit.
            # Skip val evaluation and no_improve increment — model state is corrupt for
            # this epoch; the best checkpoint from a prior epoch is still valid.
            if math.isnan(train_loss):
                logger.warning(
                    f"[{self.backbone}] Epoch {epoch} aborted (NaN) — "
                    "skipping val eval, no_improve unchanged."
                )
                history.append({"epoch": epoch, "train_loss": train_loss,
                                 "f1_macro": float("nan"), "precision_macro": float("nan"),
                                 "recall_macro": float("nan")})
                continue

            val_metrics = self._evaluate_aggregated(val_texts, val_labels)
            val_f1      = val_metrics["f1_macro"]

            logger.info(
                f"Epoch {epoch}/{self.epochs} — "
                f"Train Loss: {train_loss:.4f} | "
                f"Val F1: {val_f1:.4f} | "
                f"Val Precision: {val_metrics['precision_macro']:.4f} | "
                f"Val Recall: {val_metrics['recall_macro']:.4f}"
            )
            history.append({"epoch": epoch, "train_loss": train_loss, **val_metrics})

            if val_f1 > best_val_f1:
                best_val_f1 = val_f1
                best_epoch  = epoch
                no_improve  = 0
                self._save_checkpoint(is_best=True)
                logger.info(f"  → New best checkpoint saved (F1={val_f1:.4f})")
            else:
                no_improve += 1

            # Save full training state after every epoch (enables resume after crash)
            torch.save(
                {"optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict()},
                str(state_pt),
            )
            state_json.write_text(json.dumps({
                "epoch":       epoch,
                "best_val_f1": best_val_f1,
                "best_epoch":  best_epoch,
                "no_improve":  no_improve,
                "history":     history,
            }))

            if epoch >= self.min_epochs and no_improve >= self.patience:
                logger.info(
                    f"Early stopping at epoch {epoch} "
                    f"(patience={self.patience}, min_epochs={self.min_epochs})"
                )
                break

        logger.info(
            f"Training complete [{self.backbone}]. "
            f"Best val F1: {best_val_f1:.4f} at epoch {best_epoch}"
        )
        # Clean up resume files — training finished normally
        for fp in [state_json, state_pt]:
            if fp.exists():
                fp.unlink()

        return {
            "backbone":       self.backbone,
            "best_val_f1":    best_val_f1,
            "best_epoch":     best_epoch,
            "epochs_trained": len(history),
            "history":        history,
        }

    def _train_epoch(
        self,
        loader: DataLoader,
        criterion: nn.Module,
        optimizer: torch.optim.Optimizer,
        scheduler,
        epoch: int,
    ) -> float:
        """Run one training epoch and return average loss.

        Args:
            loader: Training DataLoader.
            criterion: AsymmetricLoss instance.
            optimizer: AdamW optimiser.
            scheduler: Linear warmup scheduler.
            epoch: Current epoch number (for progress bar label).

        Returns:
            Average training loss over the epoch.
        """
        self.model.train()
        total_loss  = 0.0
        nan_batches = 0
        optimizer.zero_grad()

        use_amp = (self.device == "cuda")
        scaler  = torch.amp.GradScaler(device="cuda", enabled=use_amp, init_scale=2**10)

        pbar = tqdm(loader, desc=f"Epoch {epoch} [{self.backbone}]", leave=False)
        for step, batch in enumerate(pbar):
            input_ids      = batch["input_ids"].to(self.device)
            attention_mask = batch["attention_mask"].to(self.device)
            labels         = batch["labels"].to(self.device)

            with torch.amp.autocast(device_type="cuda", enabled=use_amp):
                logits = self.model(input_ids, attention_mask)

            # NaN guard — critical for DeBERTa-v3 early in training
            if torch.isnan(logits).any():
                nan_batches += 1
                optimizer.zero_grad()
                if nan_batches > 50:
                    logger.error(
                        f"[{self.backbone}] >50 NaN batches in epoch {epoch} — aborting."
                    )
                    return float("nan")  # sentinel: caller skips val + no_improve
                continue

            # ASL loss in fp32 — log ops in ASL overflow in fp16; cast logits/labels
            # to float32 before loss computation regardless of autocast context
            loss = criterion(logits.float(), labels.float()) / self.grad_accum

            # Guard against residual NaN/inf loss
            if not torch.isfinite(loss):
                nan_batches += 1
                optimizer.zero_grad()
                if nan_batches > 50:
                    logger.error(
                        f"[{self.backbone}] >50 NaN/inf loss batches in epoch {epoch} — aborting."
                    )
                    return float("nan")
                continue

            scaler.scale(loss).backward()
            total_loss += loss.item() * self.grad_accum

            if (step + 1) % self.grad_accum == 0:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip_norm)
                scaler.step(optimizer)
                scaler.update()
                scheduler.step()
                optimizer.zero_grad()

            pbar.set_postfix({"loss": f"{loss.item() * self.grad_accum:.4f}"})

        if nan_batches > 0:
            logger.warning(f"[{self.backbone}] Epoch {epoch}: skipped {nan_batches} NaN batches")

        return total_loss / max(len(loader) - nan_batches, 1)

    @torch.no_grad()
    def _evaluate_aggregated(
        self,
        texts: List[str],
        labels: List[List[int]],
        batch_size: int = 64,
    ) -> Dict:
        """Evaluate on raw clause texts with sliding-window aggregation.

        Each clause — regardless of length — produces exactly ONE prediction.
        Long clauses are split into overlapping 256-token chunks, all chunks
        are run through the model, and per-class probabilities are max-pooled
        back into a single prediction per clause.

        This matches the inference path exactly, ensuring F1 is measured per
        original clause and not inflated by repeated chunk predictions.

        Args:
            texts:      Raw clause text strings.
            labels:     Ground-truth multi-hot label lists.
            batch_size: Number of chunks to process per GPU batch.

        Returns:
            Dict with f1_macro, precision_macro, recall_macro.
        """
        self.model.eval()
        use_amp   = (self.device == "cuda")
        max_len   = config.CLASSIFIER_MAX_LENGTH
        stride    = 128
        max_chunk = max_len - 2
        cls_id    = self.tokenizer.cls_token_id
        sep_id    = self.tokenizer.sep_token_id
        pad_id    = self.tokenizer.pad_token_id

        all_preds:  List[np.ndarray] = []
        all_labels: List[np.ndarray] = []

        for text, label in tqdm(
            zip(texts, labels), total=len(texts), desc="Evaluating", leave=False
        ):
            token_ids = self.tokenizer.encode(text, add_special_tokens=False)

            if len(token_ids) <= max_len - 2:
                # Short clause — single forward pass
                enc = self.tokenizer(
                    text,
                    max_length=max_len,
                    padding="max_length",
                    truncation=True,
                    return_tensors="pt",
                ).to(self.device)
                with torch.amp.autocast(device_type="cuda", enabled=use_amp):
                    logits = self.model(enc["input_ids"], enc["attention_mask"])
                probs = torch.sigmoid(logits).float().cpu().numpy()[0]
            else:
                # Long clause — sliding window, max-pool across chunks
                chunks_ids:  List[List[int]] = []
                chunks_mask: List[List[int]] = []
                start = 0
                while start < len(token_ids):
                    chunk = token_ids[start : start + max_chunk]
                    seq   = [cls_id] + chunk + [sep_id]
                    pad_n = max_len - len(seq)
                    chunks_ids.append(seq + [pad_id] * pad_n)
                    chunks_mask.append([1] * len(seq) + [0] * pad_n)
                    next_start = start + max_chunk - stride
                    if next_start >= len(token_ids):
                        break
                    start = next_start

                chunk_probs: List[np.ndarray] = []
                for i in range(0, len(chunks_ids), batch_size):
                    ids_t  = torch.tensor(chunks_ids[i:i+batch_size],  dtype=torch.long).to(self.device)
                    mask_t = torch.tensor(chunks_mask[i:i+batch_size], dtype=torch.long).to(self.device)
                    with torch.amp.autocast(device_type="cuda", enabled=use_amp):
                        logits = self.model(ids_t, mask_t)
                    chunk_probs.append(torch.sigmoid(logits).float().cpu().numpy())

                probs = np.max(np.vstack(chunk_probs), axis=0)  # max-pool across chunks

            all_preds.append((probs >= 0.5).astype(int))
            all_labels.append(label)

        y_pred = np.vstack(all_preds)
        y_true = np.vstack(all_labels)

        return {
            "f1_macro":        f1_score(y_true, y_pred, average="macro", zero_division=0),
            "precision_macro": precision_score(y_true, y_pred, average="macro", zero_division=0),
            "recall_macro":    recall_score(y_true, y_pred, average="macro", zero_division=0),
        }

    def _save_checkpoint(self, is_best: bool = False) -> None:
        """Save model weights and tokeniser to the backbone-specific checkpoint dir.

        Args:
            is_best: If True, saves to backbone/best/; otherwise backbone/last/.
        """
        save_dir = self.ckpt_dir / ("best" if is_best else "last")
        save_dir.mkdir(parents=True, exist_ok=True)
        # MeanPoolClassifierModel has a custom save_pretrained that also writes
        # head_weights.pt alongside the encoder files.
        if BACKBONE_USE_MEAN_POOL.get(self.backbone, False):
            self.model.save_pretrained(str(save_dir))
        else:
            self.model.model.save_pretrained(str(save_dir))
        self.tokenizer.save_pretrained(str(save_dir))
        logger.debug(f"Checkpoint saved to {save_dir}")


# ---------------------------------------------------------------------------
# 5. MODEL SELECTOR
# ---------------------------------------------------------------------------

class ModelSelector:
    """Selects the best production model and optimises per-class thresholds.

    Evaluation is performed exclusively on the validation set.
    The test set is never touched during selection.

    Process:
      1. Load available checkpoints (Legal-BERT, DeBERTa, or both).
      2. Compute sigmoid probabilities on the full validation set.
      3. Search ensemble weights 0.0–1.0 (step 0.1) when both are available.
      4. For each candidate, optimise per-class F1 thresholds on val.
      5. Pick the winner by highest macro F1 on val.
      6. Save production_config.json for deterministic inference.
    """

    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

    def select(self) -> Dict:
        """Run full model selection and save production config.

        Returns:
            Saved production config dict.
        """
        val_texts, val_labels = _load_split("val")
        y_true = np.array(val_labels, dtype=np.float32)

        # --- Collect probabilities from available checkpoints ---------------
        available: Dict[str, np.ndarray] = {}
        for backbone in ("legalbert", "deberta", "legalroberta"):
            ckpt = BACKBONE_DIR_MAP[backbone] / "best"
            if not ckpt.exists():
                logger.info(f"No checkpoint for {backbone} — skipping")
                continue
            logger.info(f"Computing val probabilities: {backbone}")
            available[backbone] = self._get_probs(backbone, val_texts)

        if not available:
            raise RuntimeError(
                "No trained checkpoints found. Run --mode train --backbone <name> first."
            )

        # --- Evaluate each single backbone ----------------------------------
        results: Dict[str, Dict] = {}

        for backbone, probs in available.items():
            temperature = self._find_temperature(probs, y_true)
            cal_probs   = self._calibrate(probs, temperature)
            thresholds  = self._optimise_thresholds(cal_probs, y_true)
            f1 = f1_score(y_true, cal_probs >= thresholds, average="macro", zero_division=0)
            results[backbone] = {
                "thresholds":  thresholds,
                "probs":       probs,
                "temperature": temperature,
                "f1":          f1,
                "w_lb":        1.0 if backbone == "legalbert"    else 0.0,
                "w_db":        1.0 if backbone == "deberta"      else 0.0,
                "w_lr":        1.0 if backbone == "legalroberta" else 0.0,
            }
            logger.info(
                f"  {backbone} val F1 (T={temperature:.1f}, per-class thr): {f1:.4f}"
            )

        # --- Evaluate all ensemble combinations when ≥2 checkpoints exist --
        # Weight search step = 0.1; all weights ≥ 0 and sum to 1.0.
        # For 2-model ensembles: w_a ∈ {0.0, 0.1, …, 1.0}, w_b = 1 - w_a.
        # For 3-model ensemble:  w_lb, w_db ∈ grid with w_lb+w_db ≤ 1,
        #                        w_lr = 1 - w_lb - w_db.
        weight_steps = [round(x * 0.1, 1) for x in range(0, 11)]

        def _eval_combo(probs_combined: np.ndarray) -> Tuple[float, np.ndarray, float]:
            """Return (f1, thresholds, temperature) for a blended prob array."""
            temp = self._find_temperature(probs_combined, y_true)
            cal  = self._calibrate(probs_combined, temp)
            thr  = self._optimise_thresholds(cal, y_true)
            f1   = f1_score(y_true, cal >= thr, average="macro", zero_division=0)
            return f1, thr, temp

        # 2-model pairs
        pairs = [
            ("legalbert",    "deberta",      "w_lb", "w_db"),
            ("legalbert",    "legalroberta", "w_lb", "w_lr"),
            ("deberta",      "legalroberta", "w_db", "w_lr"),
        ]
        for bk_a, bk_b, field_a, field_b in pairs:
            if bk_a not in available or bk_b not in available:
                continue
            best_f1, best_wa = -1.0, 0.5
            for wa in weight_steps:
                wb = round(1.0 - wa, 1)
                combined = wa * available[bk_a] + wb * available[bk_b]
                f1, _, _ = _eval_combo(combined)
                if f1 > best_f1:
                    best_f1, best_wa = f1, wa
            best_wb = round(1.0 - best_wa, 1)
            combined = best_wa * available[bk_a] + best_wb * available[bk_b]
            f1, thr, temp = _eval_combo(combined)
            key = f"ensemble_{bk_a}_{bk_b}"
            results[key] = {
                "thresholds":  thr,
                "probs":       combined,
                "temperature": temp,
                "f1":          f1,
                "w_lb":        best_wa if bk_a == "legalbert"    else (best_wb if bk_b == "legalbert"    else 0.0),
                "w_db":        best_wa if bk_a == "deberta"      else (best_wb if bk_b == "deberta"      else 0.0),
                "w_lr":        best_wa if bk_a == "legalroberta" else (best_wb if bk_b == "legalroberta" else 0.0),
            }
            logger.info(
                f"  {key} ({field_a}={best_wa:.1f}, {field_b}={best_wb:.1f}) val F1: {f1:.4f}"
            )

        # 3-model ensemble
        if all(b in available for b in ("legalbert", "deberta", "legalroberta")):
            best_f1 = -1.0
            best_weights = (0.34, 0.33, 0.33)
            for w_lb in weight_steps:
                for w_db in weight_steps:
                    w_lr = round(1.0 - w_lb - w_db, 1)
                    if w_lr < 0:
                        continue
                    combined = (
                        w_lb * available["legalbert"] +
                        w_db * available["deberta"] +
                        w_lr * available["legalroberta"]
                    )
                    f1, _, _ = _eval_combo(combined)
                    if f1 > best_f1:
                        best_f1 = f1
                        best_weights = (w_lb, w_db, w_lr)
            w_lb, w_db, w_lr = best_weights
            combined = (
                w_lb * available["legalbert"] +
                w_db * available["deberta"] +
                w_lr * available["legalroberta"]
            )
            f1, thr, temp = _eval_combo(combined)
            results["ensemble"] = {
                "thresholds":  thr,
                "probs":       combined,
                "temperature": temp,
                "f1":          f1,
                "w_lb":        w_lb,
                "w_db":        w_db,
                "w_lr":        w_lr,
            }
            logger.info(
                f"  ensemble (w_lb={w_lb:.1f}, w_db={w_db:.1f}, w_lr={w_lr:.1f}) val F1: {f1:.4f}"
            )

        # --- Pick winner ----------------------------------------------------
        winner_key = max(results, key=lambda k: results[k]["f1"])
        winner     = results[winner_key]
        temperature = winner["temperature"]
        logger.info(
            f"Winner: {winner_key}  (val F1 = {winner['f1']:.4f}, T={temperature:.1f})"
        )

        prod_config = {
            "model_type":           winner_key,
            "weight_legalbert":     winner["w_lb"],
            "weight_deberta":       winner["w_db"],
            "weight_legalroberta":  winner.get("w_lr", 0.0),
            "thresholds":           {
                ct: float(winner["thresholds"][i])
                for i, ct in enumerate(config.CUAD_CLAUSE_TYPES)
            },
            "temperature":          temperature,
            "unknown_threshold":    config.UNKNOWN_PROB_THRESHOLD,
            "val_f1_macro":         round(winner["f1"], 6),
        }

        config.CLASSIFIER_DIR.mkdir(parents=True, exist_ok=True)
        with open(config.CLASSIFIER_PRODUCTION_CONFIG, "w") as f:
            json.dump(prod_config, f, indent=2)
        logger.info(f"Production config saved → {config.CLASSIFIER_PRODUCTION_CONFIG}")
        return prod_config

    def _get_probs(self, backbone: str, texts: List[str]) -> np.ndarray:
        """Load a backbone checkpoint and return sigmoid probs for all texts.

        Args:
            backbone: 'legalbert' or 'deberta'.
            texts: List of clause strings.

        Returns:
            Array of shape (N, NUM_CLAUSE_TYPES).
        """
        ckpt_dir  = BACKBONE_DIR_MAP[backbone] / "best"
        tokenizer = AutoTokenizer.from_pretrained(str(ckpt_dir))
        use_mean_pool = BACKBONE_USE_MEAN_POOL.get(backbone, False)
        if use_mean_pool:
            model = MeanPoolClassifierModel.from_pretrained(str(ckpt_dir)).to(self.device)
        else:
            model = ClauseClassifierModel(str(ckpt_dir)).to(self.device)
        model.eval()

        all_probs: List[np.ndarray] = []
        batch_size = config.CLASSIFIER_BATCH_SIZE * 2

        with torch.no_grad():
            for start in tqdm(
                range(0, len(texts), batch_size),
                desc=f"  {backbone} probs",
                leave=False,
            ):
                batch   = texts[start : start + batch_size]
                encoded = tokenizer(
                    batch,
                    max_length=config.CLASSIFIER_MAX_LENGTH,
                    padding=True,
                    truncation=True,
                    return_tensors="pt",
                ).to(self.device)
                # Both model classes take (input_ids, attention_mask) and return
                # raw logits directly — no .logits attribute access needed.
                logits = model(encoded["input_ids"], encoded["attention_mask"])
                all_probs.append(torch.sigmoid(logits).cpu().numpy())

        return np.vstack(all_probs)

    @staticmethod
    def _optimise_thresholds(probs: np.ndarray, y_true: np.ndarray) -> np.ndarray:
        """Find the per-class threshold that maximises per-class F1.

        Searches thresholds in [0.1, 0.2, ..., 0.9] independently for each
        of the 100 classes. Falls back to 0.5 if no positive examples exist.

        Args:
            probs: Probability array (N, NUM_CLAUSE_TYPES).
            y_true: Ground-truth multi-hot array (N, NUM_CLAUSE_TYPES).

        Returns:
            Threshold array of shape (NUM_CLAUSE_TYPES,).
        """
        search     = [round(t * 0.1, 1) for t in range(1, 10)]
        thresholds = np.full(config.NUM_CLAUSE_TYPES, 0.5)

        for i in range(config.NUM_CLAUSE_TYPES):
            if y_true[:, i].sum() == 0:
                continue  # no positives — leave at 0.5
            best_f1 = -1.0
            for t in search:
                preds = (probs[:, i] >= t).astype(int)
                f1 = f1_score(y_true[:, i], preds, average="binary", zero_division=0)
                if f1 > best_f1:
                    best_f1        = f1
                    thresholds[i]  = t

        return thresholds

    @staticmethod
    def _find_temperature(probs: np.ndarray, y_true: np.ndarray) -> float:
        """Find temperature T that minimises binary cross-entropy on the val set.

        Temperature scaling: calibrated_prob = sigmoid(logit(p) / T)
          T > 1  →  softens probabilities (model was overconfident)
          T < 1  →  sharpens probabilities (model was underconfident)
          T = 1  →  no change (already well-calibrated)

        Searches T ∈ {0.5, 0.6, …, 2.0} and picks the value that minimises
        mean binary cross-entropy across all classes and all val examples.

        Args:
            probs: Sigmoid probability array (N, NUM_CLAUSE_TYPES).
            y_true: Ground-truth multi-hot array (N, NUM_CLAUSE_TYPES).

        Returns:
            Optimal temperature scalar (float).
        """
        import torch.nn.functional as F

        probs_clipped = np.clip(probs, 1e-7, 1 - 1e-7)
        logits_t = torch.tensor(
            np.log(probs_clipped / (1 - probs_clipped)), dtype=torch.float32
        )
        labels_t = torch.tensor(y_true, dtype=torch.float32)

        best_T, best_nll = 1.0, float("inf")
        for T in [round(t * 0.1, 1) for t in range(5, 21)]:  # 0.5 … 2.0
            nll = F.binary_cross_entropy_with_logits(
                logits_t / T, labels_t
            ).item()
            if nll < best_nll:
                best_nll, best_T = nll, T

        return best_T

    @staticmethod
    def _calibrate(probs: np.ndarray, temperature: float) -> np.ndarray:
        """Apply temperature scaling to a probability array (numpy).

        Args:
            probs: Raw probability array, any shape.
            temperature: Temperature scalar T.

        Returns:
            Calibrated probability array, same shape.
        """
        if temperature == 1.0:
            return probs
        clipped = np.clip(probs, 1e-7, 1 - 1e-7)
        logits  = np.log(clipped / (1 - clipped))
        return 1.0 / (1.0 + np.exp(-logits / temperature))


# ---------------------------------------------------------------------------
# 6. INFERENCE
# ---------------------------------------------------------------------------

class ClauseClassifierInference:
    """Production inference using saved production_config.json.

    Applies in order for each clause:
      1. Low-confidence check: if max(probs) < unknown_threshold → "Other"
      2. Per-class thresholds from production config
      3. No-label fallback: if nothing passes thresholds → "Other"
      4. Return multi-label list (or ["Other"])

    Long clauses (>CLASSIFIER_MAX_LENGTH tokens) are handled via
    sliding window with max-pooling aggregation.

    Args:
        checkpoint_dir: Legacy override — loads this directory directly as
            a Legal-BERT model, bypassing production config. Preserves
            backward compatibility with any existing callers.
    """

    def __init__(self, checkpoint_dir: Optional[Path] = None):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # ------------------------------------------------------------------
        # Resolve production config vs legacy direct-checkpoint path
        # ------------------------------------------------------------------
        if checkpoint_dir is not None:
            # Legacy path: treat as a single Legal-BERT checkpoint
            self.model_type  = "legalbert"
            self.w_lb        = 1.0
            self.w_db        = 0.0
            self.w_lr        = 0.0
            self.thresholds  = np.full(config.NUM_CLAUSE_TYPES, config.CLASSIFIER_THRESHOLD)
            self.temperature = 1.0
            self.unknown_thr = config.UNKNOWN_PROB_THRESHOLD
            self._load_backbone("legalbert", checkpoint_dir)
            return

        if config.CLASSIFIER_PRODUCTION_CONFIG.exists():
            with open(config.CLASSIFIER_PRODUCTION_CONFIG) as f:
                prod = json.load(f)
            self.model_type  = prod["model_type"]
            self.w_lb        = prod.get("weight_legalbert",    0.0)
            self.w_db        = prod.get("weight_deberta",      0.0)
            self.w_lr        = prod.get("weight_legalroberta", 0.0)
            thr = prod["thresholds"]
            if isinstance(thr, dict):
                self.thresholds = np.array([thr[ct] for ct in config.CUAD_CLAUSE_TYPES])
            else:
                self.thresholds = np.array(thr)
            self.temperature = float(prod.get("temperature", 1.0))
            self.unknown_thr = prod.get("unknown_threshold", config.UNKNOWN_PROB_THRESHOLD)
            logger.info(
                f"Loaded production config: {self.model_type} "
                f"(val F1 = {prod.get('val_f1_macro', 'n/a')})"
            )
        else:
            # Fallback: Legal-BERT best checkpoint with default thresholds
            logger.warning(
                "No production_config.json found. "
                "Run --mode select after training to optimise thresholds. "
                "Falling back to Legal-BERT best checkpoint with threshold=0.5."
            )
            self.model_type  = "legalbert"
            self.w_lb        = 1.0
            self.w_db        = 0.0
            self.w_lr        = 0.0
            self.thresholds  = np.full(config.NUM_CLAUSE_TYPES, config.CLASSIFIER_THRESHOLD)
            self.temperature = 1.0
            self.unknown_thr = config.UNKNOWN_PROB_THRESHOLD

        # ------------------------------------------------------------------
        # Load required backbone(s)
        # ------------------------------------------------------------------
        self.model_lb: Optional[nn.Module] = None
        self.tok_lb:   Optional[AutoTokenizer] = None
        self.model_db: Optional[nn.Module] = None
        self.tok_db:   Optional[AutoTokenizer] = None
        self.model_lr: Optional[nn.Module] = None
        self.tok_lr:   Optional[AutoTokenizer] = None

        if self.w_lb > 0:
            ckpt = config.CLASSIFIER_LEGALBERT_DIR / "best"
            if not ckpt.exists():
                raise FileNotFoundError(
                    f"Legal-BERT checkpoint not found at {ckpt}. "
                    "Run: python src/clause_classifier.py --mode train --backbone legalbert"
                )
            self._load_backbone("legalbert", ckpt)

        if self.w_db > 0:
            ckpt = config.CLASSIFIER_DEBERTA_DIR / "best"
            if not ckpt.exists():
                raise FileNotFoundError(
                    f"DeBERTa checkpoint not found at {ckpt}. "
                    "Run: python src/clause_classifier.py --mode train --backbone deberta"
                )
            self._load_backbone("deberta", ckpt)

        if self.w_lr > 0:
            ckpt = config.CLASSIFIER_LEGALROBERTA_DIR / "best"
            if not ckpt.exists():
                raise FileNotFoundError(
                    f"Legal-RoBERTa checkpoint not found at {ckpt}. "
                    "Run: python src/clause_classifier.py --mode train --backbone legalroberta"
                )
            self._load_backbone("legalroberta", ckpt)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    @torch.no_grad()
    def predict(
        self,
        texts: List[str],
        threshold: Optional[float] = None,
    ) -> List[Dict]:
        """Predict clause types for a list of clause texts.

        Args:
            texts: List of clause text strings.
            threshold: Single threshold override (uses per-class if None).

        Returns:
            List of dicts, each with:
                'clause_types': list of predicted labels (may be ["Other"])
                'probabilities': dict mapping each clause type to its probability
        """
        probs_lb = self._backbone_probs("legalbert",    texts) if self.model_lb is not None else None
        probs_db = self._backbone_probs("deberta",      texts) if self.model_db is not None else None
        probs_lr = self._backbone_probs("legalroberta", texts) if self.model_lr is not None else None

        active = [
            (p, w) for p, w in (
                (probs_lb, self.w_lb),
                (probs_db, self.w_db),
                (probs_lr, self.w_lr),
            ) if p is not None and w > 0
        ]
        if not active:
            raise RuntimeError("No backbone loaded — check production_config.json weights.")
        probs = sum(w * p for p, w in active)

        probs = self._apply_temperature(probs)

        results = []
        for prob_vec in probs:
            clause_types = self._decode(prob_vec, threshold)
            prob_dict    = {
                config.CUAD_CLAUSE_TYPES[i]: float(prob_vec[i])
                for i in range(config.NUM_CLAUSE_TYPES)
            }
            results.append({"clause_types": clause_types, "probabilities": prob_dict})

        return results

    def predict_single(
        self, text: str, threshold: Optional[float] = None
    ) -> Dict:
        """Convenience wrapper for single-clause prediction.

        Args:
            text: Single clause text string.
            threshold: Optional single-value threshold override.

        Returns:
            Prediction dict with 'clause_types' and 'probabilities'.
        """
        return self.predict([text], threshold=threshold)[0]

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _apply_temperature(self, probs: np.ndarray) -> np.ndarray:
        """Apply temperature scaling using the production temperature.

        Delegates to ModelSelector._calibrate for consistent behaviour.
        """
        return ModelSelector._calibrate(probs, self.temperature)

    def _decode(
        self, prob_vec: np.ndarray, threshold_override: Optional[float]
    ) -> List[str]:
        """Apply open-set detection and thresholds to one probability vector.

        Logic (evaluated in order):
          1. max(prob) < unknown_threshold  → ["Other"]  (too uncertain overall)
          2. Apply per-class thresholds
          3. No label passed               → ["Other"]  (no confident prediction)
          4. Return the list of passing labels

        Args:
            prob_vec: Probability array of length NUM_CLAUSE_TYPES.
            threshold_override: If set, replaces per-class thresholds uniformly.

        Returns:
            List of predicted clause type strings.
        """
        # Step 1: global confidence gate
        if float(prob_vec.max()) < self.unknown_thr:
            return ["Other"]

        # Step 2: per-class (or override) thresholds
        if threshold_override is not None:
            mask = prob_vec >= threshold_override
        else:
            mask = prob_vec >= self.thresholds

        # Step 3: no-label fallback
        if not mask.any():
            return ["Other"]

        return [config.CUAD_CLAUSE_TYPES[i] for i, hit in enumerate(mask) if hit]

    def _load_backbone(self, backbone: str, ckpt_dir: Path) -> None:
        """Load tokeniser and model for one backbone into instance attributes.

        Args:
            backbone: 'legalbert' or 'deberta'.
            ckpt_dir: Path to the HuggingFace checkpoint directory.
        """
        logger.info(f"Loading {backbone} checkpoint from {ckpt_dir}")
        tok   = AutoTokenizer.from_pretrained(str(ckpt_dir))
        if BACKBONE_USE_MEAN_POOL.get(backbone, False):
            model = MeanPoolClassifierModel.from_pretrained(str(ckpt_dir)).to(self.device)
        else:
            model = ClauseClassifierModel(str(ckpt_dir)).to(self.device)
        model.eval()

        if backbone == "legalbert":
            self.tok_lb   = tok
            self.model_lb = model
        elif backbone == "deberta":
            self.tok_db   = tok
            self.model_db = model
        else:  # legalroberta
            self.tok_lr   = tok
            self.model_lr = model

    def _backbone_probs(self, backbone: str, texts: List[str]) -> np.ndarray:
        """Run a backbone over all texts, using sliding window when needed.

        Args:
            backbone: 'legalbert' or 'deberta'.
            texts: List of clause text strings.

        Returns:
            Probability array of shape (N, NUM_CLAUSE_TYPES).
        """
        if backbone == "legalbert":
            model, tok = self.model_lb, self.tok_lb
        elif backbone == "deberta":
            model, tok = self.model_db, self.tok_db
        else:
            model, tok = self.model_lr, self.tok_lr

        all_probs: List[np.ndarray] = []
        for text in texts:
            n_tokens = len(tok.encode(text, add_special_tokens=True))
            if n_tokens > config.CLASSIFIER_MAX_LENGTH:
                p = self._sliding_window(text, model, tok)
            else:
                p = self._forward_single(text, model, tok)
            all_probs.append(p)

        return np.array(all_probs)

    def _forward_single(
        self,
        text: str,
        model: nn.Module,
        tok: AutoTokenizer,
    ) -> np.ndarray:
        """Single forward pass for one clause text.

        Args:
            text: Clause string.
            model: Loaded HuggingFace model.
            tok: Corresponding tokeniser.

        Returns:
            Probability vector of length NUM_CLAUSE_TYPES.
        """
        encoded = tok(
            text,
            max_length=config.CLASSIFIER_MAX_LENGTH,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        ).to(self.device)
        logits = model(encoded["input_ids"], encoded["attention_mask"])
        return torch.sigmoid(logits).cpu().numpy()[0]

    def _sliding_window(
        self,
        text: str,
        model: nn.Module,
        tok: AutoTokenizer,
        stride: int = 128,
    ) -> np.ndarray:
        """Sliding window inference for clauses that exceed max token length.

        Splits the token sequence into overlapping chunks (stride=128),
        runs each chunk through the model, and aggregates with max pooling
        (highest probability per class across all chunks).

        Args:
            text: Full clause text.
            model: Loaded HuggingFace model.
            tok: Corresponding tokeniser.
            stride: Overlap between consecutive windows (tokens).

        Returns:
            Aggregated probability vector of length NUM_CLAUSE_TYPES.
        """
        cls_id    = tok.cls_token_id
        sep_id    = tok.sep_token_id
        token_ids = tok.encode(text, add_special_tokens=False)
        max_chunk = config.CLASSIFIER_MAX_LENGTH - 2  # reserve [CLS] + [SEP]

        chunk_probs: List[np.ndarray] = []
        start = 0
        while start < len(token_ids):
            chunk     = token_ids[start : start + max_chunk]
            input_ids = torch.tensor([[cls_id] + chunk + [sep_id]], device=self.device)
            attn_mask = torch.ones_like(input_ids)
            logits    = model(input_ids=input_ids, attention_mask=attn_mask)
            chunk_probs.append(torch.sigmoid(logits).cpu().numpy()[0])

            next_start = start + max_chunk - stride
            if next_start >= len(token_ids):
                break
            start = next_start

        if not chunk_probs:
            return np.zeros(config.NUM_CLAUSE_TYPES)

        return np.max(chunk_probs, axis=0)


# ---------------------------------------------------------------------------
# 7. EVALUATOR
# ---------------------------------------------------------------------------

class ClauseClassifierEvaluator:
    """Evaluates production inference on the held-out test set.

    Uses the model and thresholds from production_config.json.
    The test set must not be used for any tuning — evaluation only.
    """

    def __init__(self):
        self.inference = ClauseClassifierInference()

    def evaluate(self) -> Dict:
        """Run evaluation on the test set and save metrics to evaluation_results.json.

        Returns:
            Dict with macro metrics, per-class metrics, and model_type.
        """
        test_texts, test_labels_hot = _load_split("test")
        logger.info(f"Evaluating on {len(test_texts)} test clauses...")

        predictions = self.inference.predict(test_texts)
        type_to_idx = {ct: i for i, ct in enumerate(config.CUAD_CLAUSE_TYPES)}

        pred_hot = []
        for pred in predictions:
            vec = [0] * config.NUM_CLAUSE_TYPES
            for ct in pred["clause_types"]:
                if ct in type_to_idx:
                    vec[type_to_idx[ct]] = 1
            pred_hot.append(vec)

        y_true = np.array(test_labels_hot)
        y_pred = np.array(pred_hot)

        report = classification_report(
            y_true, y_pred,
            target_names=config.CUAD_CLAUSE_TYPES,
            output_dict=True,
            zero_division=0,
        )

        results = {
            "classifier": {
                "model_type":      self.inference.model_type,
                "f1_macro":        f1_score(y_true, y_pred, average="macro", zero_division=0),
                "precision_macro": precision_score(y_true, y_pred, average="macro", zero_division=0),
                "recall_macro":    recall_score(y_true, y_pred, average="macro", zero_division=0),
                "per_class":       report,
                "test_size":       len(test_texts),
            }
        }

        existing: Dict = {}
        if config.EVAL_RESULTS_PATH.exists():
            with open(config.EVAL_RESULTS_PATH) as f:
                existing = json.load(f)
        existing.update(results)
        with open(config.EVAL_RESULTS_PATH, "w") as f:
            json.dump(existing, f, indent=2)

        logger.info(f"Evaluation results saved → {config.EVAL_RESULTS_PATH}")
        logger.info(
            f"F1 Macro: {results['classifier']['f1_macro']:.4f}  |  "
            f"Precision: {results['classifier']['precision_macro']:.4f}  |  "
            f"Recall: {results['classifier']['recall_macro']:.4f}"
        )
        return results


# ---------------------------------------------------------------------------
# 8. CLI
# ---------------------------------------------------------------------------

def main() -> None:
    """Command-line interface for training, model selection, evaluation, and prediction."""
    parser = argparse.ArgumentParser(
        description="Legal Clause Classifier",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python src/clause_classifier.py --mode train --backbone legalbert
  python src/clause_classifier.py --mode train --backbone deberta
  python src/clause_classifier.py --mode train --backbone legalroberta
  python src/clause_classifier.py --mode select
  python src/clause_classifier.py --mode evaluate
  python src/clause_classifier.py --mode predict --text "The Licensor shall..."
""",
    )
    parser.add_argument(
        "--mode",
        choices=["train", "select", "evaluate", "predict"],
        required=True,
        help="Pipeline stage to run",
    )
    parser.add_argument(
        "--backbone",
        choices=["legalbert", "deberta", "legalroberta"],
        default="legalbert",
        help="Backbone for --mode train (default: legalbert)",
    )
    parser.add_argument(
        "--text",
        type=str,
        help="Clause text for --mode predict",
    )
    args = parser.parse_args()

    if args.mode == "train":
        trainer = ClauseClassifierTrainer(backbone=args.backbone)
        metrics = trainer.train()
        print(f"\nTraining complete [{args.backbone}]. Best val F1: {metrics['best_val_f1']:.4f}")

    elif args.mode == "select":
        selector    = ModelSelector()
        prod_config = selector.select()
        print(f"\nModel selection complete.")
        print(f"  Winner:         {prod_config['model_type']}")
        print(f"  Val F1 (macro): {prod_config['val_f1_macro']:.4f}")
        print(f"  Config saved:   {config.CLASSIFIER_PRODUCTION_CONFIG}")

    elif args.mode == "evaluate":
        evaluator = ClauseClassifierEvaluator()
        results   = evaluator.evaluate()
        c = results["classifier"]
        print(f"\nF1 Macro:        {c['f1_macro']:.4f}")
        print(f"Precision Macro: {c['precision_macro']:.4f}")
        print(f"Recall Macro:    {c['recall_macro']:.4f}")
        print(f"Model type:      {c['model_type']}")

    elif args.mode == "predict":
        if not args.text:
            parser.error("--text is required with --mode predict")
        inference = ClauseClassifierInference()
        result    = inference.predict_single(args.text)
        print(f"\nPredicted clause types: {result['clause_types']}")
        top5 = sorted(result["probabilities"].items(), key=lambda x: x[1], reverse=True)[:5]
        print("Top 5 probabilities:")
        for ct, prob in top5:
            print(f"  {ct}: {prob:.4f}")


if __name__ == "__main__":
    main()