"""Difficulty scoring for error-aware sample weighting.

Computes per-sample weights based on:
1. Span factor — wide spans lose tone resolution (physical difficulty)
2. SNR factor — low SNR signals are harder to classify
3. Misclassification factor — up-weight samples the current model gets wrong
   (only meaningful when a base model is available for evaluation)
4. Label trust — all saved captures are trusted (user reviewed them)

Usage:
    from scripts.difficulty_scorer import compute_difficulty_scores

    weights, report = compute_difficulty_scores(
        dataset, model, signal_labels,
        misclassification_weight=3.0,
        enable_misclassification=True,
        enable_span=True,
        enable_snr=True,
    )
"""

from typing import List, Tuple, Dict
import numpy as np


def _span_factor(span_norm: float) -> float:
    """Physical difficulty multiplier based on span (normalized log10(span)/6).

    Wide spans lose tone resolution — at 100kHz span, FT8's 6.25Hz tones
    are sub-pixel width and collapse to a single spike.
    """
    if span_norm < 0.62:
        return 1.0
    elif span_norm < 0.74:
        return 1.1
    elif span_norm < 0.83:
        return 1.3
    else:
        return 1.5


def _snr_factor(snr_norm: float) -> float:
    """Physical difficulty multiplier based on SNR (normalized /40).

    Low SNR signals have noisier features and are harder to classify.
    """
    if snr_norm < 0.25:
        return 1.2
    return 1.0


def _misclassification_factor(
    model,
    patch: np.ndarray,
    aux: np.ndarray,
    true_label_idx: int,
    signal_labels: list,
    misclassification_weight: float,
) -> Tuple[float, str, str, float]:
    """Check if model prediction differs from training label.

    Returns (factor, predicted_label, true_label, predicted_confidence).
    Factor is misclassification_weight if wrong, 1.0 if correct.
    """
    import torch

    patch_t = torch.from_numpy(
        patch.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32)
    )
    aux_t = torch.from_numpy(aux[np.newaxis, ...].astype(np.float32))
    with torch.no_grad():
        output = model(patch_t, aux=aux_t)
    probs = torch.softmax(output[0], dim=0).numpy()
    pred_idx = int(np.argmax(probs))
    pred_conf = float(probs[pred_idx])

    pred_label = signal_labels[pred_idx] if pred_idx < len(signal_labels) else "Unknown"
    true_label = signal_labels[true_label_idx] if true_label_idx < len(signal_labels) else "Unknown"

    if pred_idx != true_label_idx:
        return misclassification_weight, pred_label, true_label, pred_conf
    return 1.0, pred_label, true_label, pred_conf


def compute_difficulty_scores(
    dataset: list,
    model,
    signal_labels: list,
    misclassification_weight: float = 3.0,
    enable_misclassification: bool = True,
    enable_span: bool = True,
    enable_snr: bool = True,
    max_weight: float = 5.0,
) -> Tuple[List[float], Dict]:
    """Compute per-sample difficulty weights.

    Parameters
    ----------
    dataset : list
        List of (patch, aux, label_idx, is_ai_labeled, ai_confidence, filename)
        tuples from load_captures().
    model : torch.nn.Module or None
        PyTorch model for pre-training evaluation. If None,
        misclassification detection is skipped (span+SNR only).
    signal_labels : list
        Label list for mapping indices to names.
    misclassification_weight : float
        Multiplier for samples the model misclassifies (default 3.0).
    enable_misclassification : bool
        Whether to evaluate model predictions (requires model).
    enable_span : bool
        Whether to apply span difficulty factor.
    enable_snr : bool
        Whether to apply SNR difficulty factor.
    max_weight : float
        Cap on total weight (default 5.0).

    Returns
    -------
    (weights, report) where weights is a list of floats and report
    is a dict with statistics for logging.
    """
    from collections import Counter

    weights = []
    misclassified = []
    correct_count = 0
    total = len(dataset)
    weight_components = {"span": [], "snr": [], "misclass": []}

    for patch, aux, label_idx, is_ai, ai_conf, fname in dataset:
        w = 1.0

        sf = _span_factor(float(aux[0])) if enable_span else 1.0
        w *= sf
        weight_components["span"].append(sf)

        snr_f = _snr_factor(float(aux[5])) if enable_snr else 1.0
        w *= snr_f
        weight_components["snr"].append(snr_f)

        if enable_misclassification and model is not None:
            mf, pred_label, true_label, pred_conf = _misclassification_factor(
                model, patch, aux, label_idx, signal_labels, misclassification_weight
            )
            w *= mf
            weight_components["misclass"].append(mf)
            if mf > 1.0:
                misclassified.append({
                    "true": true_label,
                    "pred": pred_label,
                    "conf": pred_conf,
                    "filename": fname,
                })
            else:
                correct_count += 1
        else:
            weight_components["misclass"].append(1.0)

        w = min(w, max_weight)
        weights.append(w)

    wrong_count = len(misclassified)
    report = {
        "total_samples": total,
        "correct": correct_count,
        "misclassified": wrong_count,
        "accuracy": correct_count / total if total > 0 else 0,
        "misclassification_weight": misclassification_weight,
        "enable_misclassification": enable_misclassification and model is not None,
        "enable_span": enable_span,
        "enable_snr": enable_snr,
        "weight_stats": {
            "min": min(weights) if weights else 0,
            "max": max(weights) if weights else 0,
            "mean": float(np.mean(weights)) if weights else 0,
            "median": float(np.median(weights)) if weights else 0,
        },
        "misclassified_samples": misclassified,
        "misclassified_by_true_label": dict(Counter(
            m["true"] for m in misclassified
        )),
        "confused_pairs": dict(Counter(
            f"{m['true']} -> {m['pred']}" for m in misclassified
        )),
    }

    return weights, report


def print_difficulty_report(report: Dict) -> None:
    """Print a human-readable difficulty analysis report."""
    print(f"\n{'=' * 60}")
    print("=== Difficulty Analysis ===")
    print(f"{'=' * 60}")
    print(f"Total samples: {report['total_samples']}")
    print(f"Correct: {report['correct']} ({100 * report['accuracy']:.1f}%)")
    print(f"Misclassified: {report['misclassified']}")
    print(f"Misclassification weight: {report['misclassification_weight']}x")

    if report["misclassified"] > 0:
        print(f"\nPer-class misclassification:")
        for label, count in sorted(
            report["misclassified_by_true_label"].items(), key=lambda x: -x[1]
        ):
            print(f"  {label:6s}: {count}")

        print(f"\nTop confused pairs (true -> predicted):")
        for pair, count in sorted(
            report["confused_pairs"].items(), key=lambda x: -x[1]
        )[:10]:
            print(f"  {pair:20s}: {count}")

    ws = report["weight_stats"]
    print(f"\nWeight distribution:")
    print(f"  min={ws['min']:.2f}, max={ws['max']:.2f}, "
          f"mean={ws['mean']:.2f}, median={ws['median']:.2f}")
    print(f"{'=' * 60}\n")