evaluation/eval_metrics.py

"""Shared evaluation helpers: calibration, curves, one-hot encoding, temperature scaling."""
from __future__ import annotations

import json
import os

import numpy as np
import pandas as pd
from sklearn.metrics import (
    accuracy_score,
    average_precision_score,
    brier_score_loss,
    classification_report,
    confusion_matrix,
    log_loss,
    matthews_corrcoef,
    precision_recall_curve,
    roc_auc_score,
    roc_curve,
)


def apply_temperature_scaling(y_prob: np.ndarray, temperature: float) -> np.ndarray:
    """p_cal = sigmoid(logit(p) / T)."""
    eps = 1e-15
    p = np.clip(y_prob, eps, 1 - eps)
    logit_p = np.log(p / (1 - p))
    scaled_logit = logit_p / temperature
    p_cal = 1.0 / (1.0 + np.exp(-scaled_logit))
    return np.clip(p_cal, 0.0, 1.0)


def compute_ece(y_true: np.ndarray, y_prob: np.ndarray, n_bins: int = 10) -> float:
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    for i in range(n_bins):
        in_bin = (y_prob > bin_boundaries[i]) & (y_prob <= bin_boundaries[i + 1])
        prop_in_bin = np.mean(in_bin)
        if prop_in_bin > 0:
            avg_confidence_in_bin = np.mean(y_prob[in_bin])
            avg_accuracy_in_bin = np.mean(y_true[in_bin])
            ece += prop_in_bin * np.abs(avg_accuracy_in_bin - avg_confidence_in_bin)
    return float(ece)


def save_confusion_matrix(y_true, y_pred, model_name: str, output_dir: str) -> None:
    cm = confusion_matrix(y_true, y_pred)
    if cm.size == 4:
        tn, fp, fn, tp = cm.ravel()
    elif cm.size == 1:
        unique_pred = np.unique(y_pred)
        if len(unique_pred) == 1:
            if unique_pred[0] == 0:
                tn = cm[0, 0] if 0 in np.unique(y_true) else 0
                fp, fn, tp = 0, int(np.sum(y_true == 1)), 0
            else:
                tn, fp, fn = 0, int(np.sum(y_true == 0)), 0
                tp = cm[0, 0] if 1 in np.unique(y_true) else 0
        else:
            tn = fp = fn = tp = 0
    else:
        tn = cm[0, 0] if cm.shape[0] > 0 and cm.shape[1] > 0 else 0
        fp = cm[0, 1] if cm.shape[0] > 0 and cm.shape[1] > 1 else 0
        fn = cm[1, 0] if cm.shape[0] > 1 and cm.shape[1] > 0 else 0
        tp = cm[1, 1] if cm.shape[0] > 1 and cm.shape[1] > 1 else 0
    mcc = matthews_corrcoef(y_true, y_pred)
    n = tp + tn + fp + fn
    metrics = {
        'true_negatives': tn,
        'false_positives': fp,
        'false_negatives': fn,
        'true_positives': tp,
        'accuracy': (tp + tn) / n if n > 0 else 0,
        'sensitivity': tp / (tp + fn) if (tp + fn) > 0 else 0,
        'specificity': tn / (tn + fp) if (tn + fp) > 0 else 0,
        'precision': tp / (tp + fp) if (tp + fp) > 0 else 0,
        'f1_score': 2 * tp / (2 * tp + fp + fn) if (2 * tp + fp + fn) > 0 else 0,
        'mcc': mcc,
    }
    pd.DataFrame(
        cm,
        columns=['Predicted Negative', 'Predicted Positive'],
        index=['Actual Negative', 'Actual Positive'],
    ).to_csv(os.path.join(output_dir, f'{model_name}_confusion_matrix.csv'))
    pd.DataFrame([metrics]).to_csv(
        os.path.join(output_dir, f'{model_name}_classification_metrics.csv'), index=False
    )


def save_roc_curve(y_true, y_pred_proba, model_name: str, output_dir: str) -> None:
    fpr, tpr, thresholds = roc_curve(y_true, y_pred_proba)
    roc_auc = roc_auc_score(y_true, y_pred_proba)
    pd.DataFrame({'fpr': fpr, 'tpr': tpr, 'thresholds': thresholds}).to_csv(
        os.path.join(output_dir, f'{model_name}_roc_curve.csv'), index=False
    )
    pd.DataFrame({'auc_score': [roc_auc]}).to_csv(
        os.path.join(output_dir, f'{model_name}_auc_score.csv'), index=False
    )


def save_precision_recall_curve(y_true, y_pred_proba, model_name: str, output_dir: str) -> None:
    precision, recall, thresholds = precision_recall_curve(y_true, y_pred_proba)
    avg_precision = average_precision_score(y_true, y_pred_proba)
    pd.DataFrame(
        {'precision': precision, 'recall': recall, 'thresholds': np.append(thresholds, 1)}
    ).to_csv(os.path.join(output_dir, f'{model_name}_precision_recall_curve.csv'), index=False)
    pd.DataFrame({'avg_precision': [avg_precision]}).to_csv(
        os.path.join(output_dir, f'{model_name}_avg_precision.csv'), index=False
    )


def compute_calibration_metrics(y_true, y_prob, model_name: str, output_dir: str) -> dict:
    eps = 1e-15
    y_prob_clip = np.clip(y_prob, eps, 1 - eps)
    ll = log_loss(y_true, np.column_stack([1 - y_prob_clip, y_prob_clip]), labels=[0, 1])
    brier = brier_score_loss(y_true, y_prob)
    ece = compute_ece(y_true, y_prob)
    pd.DataFrame({'metric': ['log_loss', 'brier_score', 'ece'], 'value': [ll, brier, ece]}).to_csv(
        os.path.join(output_dir, f'{model_name}_calibration_metrics.csv'), index=False
    )
    return {'log_loss': ll, 'brier_score': brier, 'ece': ece}


def run_validation_diagnostics(
    y_true, y_prob, model_name: str, output_dir: str, n_bins: int = 10, n_thresholds: int = 101
) -> float:
    thresholds = np.linspace(0, 1, n_thresholds)
    precisions, recalls, f1s = [], [], []
    for t in thresholds:
        y_pred = (y_prob >= t).astype(int)
        tp = np.sum((y_pred == 1) & (y_true == 1))
        fp = np.sum((y_pred == 1) & (y_true == 0))
        fn = np.sum((y_pred == 0) & (y_true == 1))
        p = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        r = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        f1 = 2 * p * r / (p + r) if (p + r) > 0 else 0.0
        precisions.append(p)
        recalls.append(r)
        f1s.append(f1)
    best_idx = int(np.argmax(f1s))
    best_threshold = float(thresholds[best_idx])
    pd.DataFrame(
        {'threshold': thresholds, 'precision': precisions, 'recall': recalls, 'f1': f1s}
    ).to_csv(os.path.join(output_dir, f'{model_name}_metric_vs_threshold.csv'), index=False)

    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    bin_centers, mean_predicted, mean_actual, counts = [], [], [], []
    for i in range(n_bins):
        low, high = bin_boundaries[i], bin_boundaries[i + 1]
        in_bin = (
            (y_prob >= low) & (y_prob < high)
            if i < n_bins - 1
            else (y_prob >= low) & (y_prob <= high)
        )
        if np.sum(in_bin) > 0:
            bin_centers.append((low + high) / 2)
            mean_predicted.append(np.mean(y_prob[in_bin]))
            mean_actual.append(np.mean(y_true[in_bin]))
            counts.append(int(np.sum(in_bin)))
    pd.DataFrame(
        {
            'bin_center': bin_centers,
            'mean_predicted': mean_predicted,
            'mean_actual': mean_actual,
            'count': counts,
        }
    ).to_csv(os.path.join(output_dir, f'{model_name}_reliability_diagram.csv'), index=False)

    y_pred = (y_prob >= best_threshold).astype(int)
    cm = confusion_matrix(y_true, y_pred)
    cm_norm = cm.astype(float) / cm.sum(axis=1, keepdims=True) if cm.size == 4 else cm.astype(float)
    pd.DataFrame(
        cm_norm,
        columns=['Predicted Negative', 'Predicted Positive'],
        index=['Actual Negative', 'Actual Positive'],
    ).to_csv(os.path.join(output_dir, f'{model_name}_normalized_confusion_matrix.csv'))
    pd.DataFrame(
        cm,
        columns=['Predicted Negative', 'Predicted Positive'],
        index=['Actual Negative', 'Actual Positive'],
    ).to_csv(os.path.join(output_dir, f'{model_name}_confusion_matrix_at_threshold.csv'))
    return best_threshold


def resolve_id_column(df: pd.DataFrame) -> str:
    for c in ('variant', 'sequence_id', 'id'):
        if c in df.columns:
            return c
    return str(df.columns[0])


def csv_sequence_ids(df: pd.DataFrame) -> set[str]:
    """Sequence IDs from a CSV id column (variant / sequence_id / id / first column)."""
    id_col = resolve_id_column(df)
    return set(df[id_col].astype(str).str.strip())


def resolve_id_and_label_columns(df: pd.DataFrame) -> tuple[str, str]:
    id_col = resolve_id_column(df)
    if 'label' in df.columns:
        label_col = 'label'
    elif 'target' in df.columns:
        label_col = 'target'
    else:
        raise ValueError("CSV must include a 'label' or 'target' column.")
    return id_col, label_col


def build_label_map_int(df: pd.DataFrame) -> tuple[str, str, dict[str, int]]:
    """Map sequence id string -> 0/1 for lowFRET/highFRET."""
    id_col, label_col = resolve_id_and_label_columns(df)
    id_series = df[id_col].astype(str).str.strip()
    label_series = df[label_col].astype(str).str.strip()
    label_map: dict[str, int] = {}
    for k, v in zip(id_series, label_series):
        label_map[k] = 1 if v.lower() == 'highfret' else 0
    return id_col, label_col, label_map


def print_brief_metrics(model_name: str, y_true, y_prob, y_pred, threshold: float) -> None:
    acc = accuracy_score(y_true, y_pred)
    try:
        roc_auc = roc_auc_score(y_true, y_prob)
    except Exception:
        roc_auc = float('nan')
    mcc = matthews_corrcoef(y_true, y_pred)
    print(
        f"{model_name} (n={len(y_true)}): thr={threshold:.4f} acc={acc:.4f} roc_auc={roc_auc:.4f} mcc={mcc:.4f}"
    )


def run_supervised_evaluation(
    y_true: np.ndarray,
    y_prob: np.ndarray,
    y_pred: np.ndarray,
    model_name: str,
    output_dir: str,
    threshold: float,
    detailed_metrics: bool,
) -> None:
    """If detailed_metrics, write diagnostic CSVs and verbose prints; else one short line."""
    if not detailed_metrics:
        print_brief_metrics(model_name, y_true, y_prob, y_pred, threshold)
        return
    cal = compute_calibration_metrics(y_true, y_prob, model_name, output_dir)
    save_confusion_matrix(y_true, y_pred, model_name, output_dir)
    save_roc_curve(y_true, y_prob, model_name, output_dir)
    save_precision_recall_curve(y_true, y_prob, model_name, output_dir)
    run_validation_diagnostics(y_true, y_prob, model_name, output_dir)
    acc = accuracy_score(y_true, y_pred)
    try:
        roc_auc = roc_auc_score(y_true, y_prob)
    except Exception:
        roc_auc = float('nan')
    cm = confusion_matrix(y_true, y_pred)
    if cm.size == 4:
        tn, fp, fn, tp = cm.ravel()
    else:
        tn = fp = fn = tp = 0
    sensitivity = tp / (tp + fn) if (tp + fn) > 0 else 0
    specificity = tn / (tn + fp) if (tn + fp) > 0 else 0
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    f1 = 2 * tp / (2 * tp + fp + fn) if (2 * tp + fp + fn) > 0 else 0
    mcc = matthews_corrcoef(y_true, y_pred)
    print('\n' + '=' * 60)
    print(f'METRICS ({model_name}, n={len(y_true)})')
    print('=' * 60)
    print(f'Threshold used: {threshold:.4f}')
    print(f'Accuracy: {acc:.4f}  ROC AUC: {roc_auc:.4f}  MCC: {mcc:.4f}')
    print(
        f'Sensitivity: {sensitivity:.4f}  Specificity: {specificity:.4f}  '
        f'Precision: {precision:.4f}  F1: {f1:.4f}'
    )
    print(
        f"Log Loss: {cal['log_loss']:.4f}  Brier: {cal['brier_score']:.4f}  ECE: {cal['ece']:.4f}"
    )
    print('Classification Report:\n', classification_report(y_true, y_pred, zero_division=0))
    print('Confusion Matrix:\n', cm)
    print('=' * 60)


# --- One-hot encoding (used by evaluate_*_one_hot.py) ---

ONE_HOT_AA_DIM = 20
CANONICAL_AAS = 'ACDEFGHIKLMNPQRSTVWY'
AA_TO_INDEX = {aa: i for i, aa in enumerate(CANONICAL_AAS)}

# 'truncate' or 'skip' (str alias for callers that import LongSequencePolicy)
LongSequencePolicy = str


def infer_one_hot_target_length_from_sklearn(model) -> int | None:
    """Positions (L) from a trained one-hot RF: n_features == L * ONE_HOT_AA_DIM."""
    n_feat = getattr(model, 'n_features_in_', None)
    if n_feat is None:
        n_feat = getattr(model, 'n_features_', None)
    if n_feat is None:
        return None
    n_feat = int(n_feat)
    if n_feat % ONE_HOT_AA_DIM != 0:
        raise ValueError(
            f'RF n_features={n_feat} is not divisible by {ONE_HOT_AA_DIM} (one-hot width)'
        )
    return n_feat // ONE_HOT_AA_DIM


def read_one_hot_target_length_json(model_dir: str, key: str = 'one_hot_rf_model') -> int | None:
    path = os.path.join(model_dir, 'model_parameters.json')
    if not os.path.isfile(path):
        return None
    with open(path, 'r') as f:
        params = json.load(f)
    block = params.get(key) or {}
    if 'input_shape_3d' in block:
        return int(block['input_shape_3d'][0])
    n_flat = block.get('n_features_flat')
    if n_flat is not None and int(n_flat) % ONE_HOT_AA_DIM == 0:
        return int(n_flat) // ONE_HOT_AA_DIM
    return None


def encode_aa_sequence(sequence: str, target_length: int) -> np.ndarray:
    """Shape (target_length, 20): pad short sequences; encode only first target_length residues."""
    encoding = np.zeros((target_length, 20))
    for i, aa in enumerate(sequence):
        if i >= target_length:
            break
        idx = AA_TO_INDEX.get(aa.upper())
        if idx is not None:
            encoding[i, idx] = 1
    return encoding


def encode_aa_sequences(sequences: list[str], target_length: int) -> np.ndarray:
    """Stack encodings: (n_seq, target_length, 20)."""
    out = np.zeros((len(sequences), target_length, 20))
    for i, seq in enumerate(sequences):
        out[i] = encode_aa_sequence(seq, target_length)
    return out


def prepare_sequences_from_fasta(
    records: list[tuple[str, str]],
    target_length: int,
    long_sequence_policy: LongSequencePolicy = 'truncate',
) -> tuple[list[str], list[str], list[int], list[bool], list[tuple[str, int, str]]]:
    """
    Filter FASTA records for inference.

    Returns:
        sequence_ids, sequences (possibly truncated), original_lengths, truncated_flags,
        skipped (id, length, reason) for empty sequences or policy=skip on long seqs.
    """
    sequence_ids: list[str] = []
    sequences: list[str] = []
    original_lengths: list[int] = []
    truncated_flags: list[bool] = []
    skipped: list[tuple[str, int, str]] = []

    for header, seq in records:
        orig_len = len(seq)
        if orig_len == 0:
            skipped.append((header, 0, 'empty'))
            continue
        if orig_len > target_length:
            if long_sequence_policy == 'skip':
                skipped.append((header, orig_len, 'too_long'))
                continue
            sequence_ids.append(header)
            sequences.append(seq[:target_length])
            original_lengths.append(orig_len)
            truncated_flags.append(True)
            continue
        sequence_ids.append(header)
        sequences.append(seq)
        original_lengths.append(orig_len)
        truncated_flags.append(False)

    return sequence_ids, sequences, original_lengths, truncated_flags, skipped