"""Feature Importance (Lopez de Prado, AFML Chapter 8).

Model-agnostic feature importance methods with purged cross-validation:

- **MDA** (Mean Decrease Accuracy): permutation importance with purged k-fold.
- **SFI** (Single Feature Importance): trains on each feature individually.
- **Clustered MDA**: groups correlated features, permutes clusters.

All methods accept a generic ``score_fn(X_train, y_train, X_test, y_test)``
so they work with any model (sklearn, xgboost, a simple function, etc.).

Usage:
    def my_scorer(X_train, y_train, X_test, y_test):
        # Train a model, return accuracy on test set
        from collections import Counter
        majority = Counter(y_train).most_common(1)[0][0]
        return sum(1 for y in y_test if y == majority) / len(y_test)

    results = mda_importance(my_scorer, X, y, feature_names, n_splits=5)
    for fi in results:
        print(f"{fi.feature}: {fi.importance:.4f} +/- {fi.std:.4f}")
"""

from __future__ import annotations

import math
import random
from dataclasses import dataclass


@dataclass
class FeatureImportance:
    """Importance score for a single feature.

    Attributes:
        feature: Feature name.
        importance: Mean importance score (higher = more important).
        std: Standard deviation of importance across folds.
    """

    feature: str
    importance: float
    std: float


# ---------------------------------------------------------------------------
# Purged K-Fold
# ---------------------------------------------------------------------------


def _purged_kfold_split(
    n: int,
    n_splits: int,
    purge_gap: int,
) -> list[tuple[list[int], list[int]]]:
    """Generate purged k-fold cross-validation indices.

    Standard k-fold partitioning, but train samples within *purge_gap*
    indices of any test sample are removed. This prevents information
    leakage in time-series data where adjacent samples may be correlated.

    Args:
        n: Total number of samples.
        n_splits: Number of folds (must be >= 2).
        purge_gap: Number of indices on each side of the test set to
            exclude from the training set.

    Returns:
        List of ``(train_indices, test_indices)`` tuples, one per fold.
    """
    if n_splits < 2:
        raise ValueError("n_splits must be >= 2")
    if n < n_splits:
        raise ValueError(f"Cannot split {n} samples into {n_splits} folds")

    indices = list(range(n))
    fold_size = n // n_splits
    folds: list[tuple[list[int], list[int]]] = []

    for k in range(n_splits):
        test_start = k * fold_size
        if k == n_splits - 1:
            test_end = n
        else:
            test_end = (k + 1) * fold_size

        test_indices = indices[test_start:test_end]
        test_set = set(test_indices)

        # Build purge zone: indices within purge_gap of any test index
        purge_set: set[int] = set()
        if purge_gap > 0:
            test_min = test_start
            test_max = test_end - 1
            purge_start = max(0, test_min - purge_gap)
            purge_end = min(n - 1, test_max + purge_gap)
            for i in range(purge_start, purge_end + 1):
                purge_set.add(i)

        train_indices = [
            i for i in indices
            if i not in test_set and i not in purge_set
        ]

        folds.append((train_indices, test_indices))

    return folds


# ---------------------------------------------------------------------------
# Helper: extract column / shuffle column
# ---------------------------------------------------------------------------


def _get_column(X: list[list[float]], col: int) -> list[float]:
    """Extract a single column from a 2D list."""
    return [row[col] for row in X]


def _copy_matrix(X: list[list[float]]) -> list[list[float]]:
    """Shallow copy of a 2D list (copies row lists)."""
    return [list(row) for row in X]


def _subset_rows(X: list[list[float]], indices: list[int]) -> list[list[float]]:
    """Extract rows by index."""
    return [list(X[i]) for i in indices]


def _subset_cols(X: list[list[float]], cols: list[int]) -> list[list[float]]:
    """Extract specific columns from each row."""
    return [[row[c] for c in cols] for row in X]


def _subset_labels(y: list[float], indices: list[int]) -> list[float]:
    """Extract labels by index."""
    return [y[i] for i in indices]


# ---------------------------------------------------------------------------
# MDA Importance
# ---------------------------------------------------------------------------


def mda_importance(
    score_fn: callable,
    X: list[list[float]],
    y: list[float],
    feature_names: list[str],
    n_splits: int = 5,
    purge_gap: int = 0,
    seed: int | None = None,
) -> list[FeatureImportance]:
    """Mean Decrease Accuracy (permutation importance) with purged CV.

    For each cross-validation fold, computes a baseline score on the test
    set. Then for each feature, shuffles that feature's column in the test
    set and re-scores. The importance of a feature is the mean decrease in
    score caused by shuffling it.

    Args:
        score_fn: Callable ``(X_train, y_train, X_test, y_test) -> float``.
            Higher values indicate better performance (e.g. accuracy).
        X: Feature matrix (N x D), each inner list is one sample.
        y: Label vector (length N).
        feature_names: Names for each feature (length D).
        n_splits: Number of cross-validation folds.
        purge_gap: Number of samples to purge around test folds.
        seed: Random seed for reproducibility.

    Returns:
        List of :class:`FeatureImportance` sorted by importance descending.
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    n = len(X)
    if n == 0:
        return []

    n_features = len(feature_names)
    if n_features == 0:
        return []

    rng = random.Random(seed)
    folds = _purged_kfold_split(n, n_splits, purge_gap)

    # importance_scores[feature_idx] = list of (baseline - shuffled) per fold
    importance_scores: dict[int, list[float]] = {
        f: [] for f in range(n_features)
    }

    for train_idx, test_idx in folds:
        if not train_idx or not test_idx:
            continue

        X_train = _subset_rows(X, train_idx)
        y_train = _subset_labels(y, train_idx)
        X_test = _subset_rows(X, test_idx)
        y_test = _subset_labels(y, test_idx)

        # Baseline score
        baseline = score_fn(X_train, y_train, X_test, y_test)

        for f_idx in range(n_features):
            # Shuffle feature f_idx in test set
            X_test_shuffled = _copy_matrix(X_test)
            col_values = [row[f_idx] for row in X_test_shuffled]
            rng.shuffle(col_values)
            for i, val in enumerate(col_values):
                X_test_shuffled[i][f_idx] = val

            shuffled_score = score_fn(X_train, y_train, X_test_shuffled, y_test)
            importance_scores[f_idx].append(baseline - shuffled_score)

    results: list[FeatureImportance] = []
    for f_idx in range(n_features):
        scores = importance_scores[f_idx]
        if scores:
            mean_imp = sum(scores) / len(scores)
            if len(scores) > 1:
                variance = sum((s - mean_imp) ** 2 for s in scores) / (len(scores) - 1)
                std_imp = math.sqrt(max(variance, 0.0))
            else:
                std_imp = 0.0
        else:
            mean_imp = 0.0
            std_imp = 0.0

        results.append(
            FeatureImportance(
                feature=feature_names[f_idx],
                importance=mean_imp,
                std=std_imp,
            )
        )

    results.sort(key=lambda fi: fi.importance, reverse=True)
    return results


# ---------------------------------------------------------------------------
# SFI Importance
# ---------------------------------------------------------------------------


def sfi_importance(
    score_fn: callable,
    X: list[list[float]],
    y: list[float],
    feature_names: list[str],
    n_splits: int = 5,
) -> list[FeatureImportance]:
    """Single Feature Importance (AFML Chapter 8.6).

    Trains the model using each feature *individually* and evaluates via
    cross-validation. The importance of a feature is its cross-validated
    score when used as the sole predictor.

    Args:
        score_fn: Callable ``(X_train, y_train, X_test, y_test) -> float``.
        X: Feature matrix (N x D).
        y: Label vector (length N).
        feature_names: Names for each feature (length D).
        n_splits: Number of cross-validation folds.

    Returns:
        List of :class:`FeatureImportance` sorted by importance descending.
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    n = len(X)
    if n == 0:
        return []

    n_features = len(feature_names)
    if n_features == 0:
        return []

    folds = _purged_kfold_split(n, n_splits, purge_gap=0)
    results: list[FeatureImportance] = []

    for f_idx in range(n_features):
        fold_scores: list[float] = []

        for train_idx, test_idx in folds:
            if not train_idx or not test_idx:
                continue

            # Single-feature matrices (N x 1)
            X_train_single = [[X[i][f_idx]] for i in train_idx]
            y_train = _subset_labels(y, train_idx)
            X_test_single = [[X[i][f_idx]] for i in test_idx]
            y_test = _subset_labels(y, test_idx)

            score = score_fn(X_train_single, y_train, X_test_single, y_test)
            fold_scores.append(score)

        if fold_scores:
            mean_score = sum(fold_scores) / len(fold_scores)
            if len(fold_scores) > 1:
                variance = sum(
                    (s - mean_score) ** 2 for s in fold_scores
                ) / (len(fold_scores) - 1)
                std_score = math.sqrt(max(variance, 0.0))
            else:
                std_score = 0.0
        else:
            mean_score = 0.0
            std_score = 0.0

        results.append(
            FeatureImportance(
                feature=feature_names[f_idx],
                importance=mean_score,
                std=std_score,
            )
        )

    results.sort(key=lambda fi: fi.importance, reverse=True)
    return results


# ---------------------------------------------------------------------------
# Clustering helpers (simple agglomerative)
# ---------------------------------------------------------------------------


def _correlation_matrix(X: list[list[float]]) -> list[list[float]]:
    """Compute Pearson correlation matrix for features (columns of X).

    Returns a D x D matrix where D is the number of features.
    """
    n = len(X)
    if n == 0:
        return []

    n_features = len(X[0]) if X else 0
    if n_features == 0:
        return []

    # Compute means
    means = [0.0] * n_features
    for row in X:
        for j in range(n_features):
            means[j] += row[j]
    for j in range(n_features):
        means[j] /= n

    # Compute standard deviations and covariances
    stds = [0.0] * n_features
    for row in X:
        for j in range(n_features):
            stds[j] += (row[j] - means[j]) ** 2
    for j in range(n_features):
        stds[j] = math.sqrt(stds[j] / max(n - 1, 1))

    corr = [[0.0] * n_features for _ in range(n_features)]
    for i in range(n_features):
        corr[i][i] = 1.0
        for j in range(i + 1, n_features):
            if stds[i] < 1e-15 or stds[j] < 1e-15:
                corr[i][j] = 0.0
                corr[j][i] = 0.0
                continue
            cov_ij = 0.0
            for row in X:
                cov_ij += (row[i] - means[i]) * (row[j] - means[j])
            cov_ij /= max(n - 1, 1)
            r = cov_ij / (stds[i] * stds[j])
            r = max(-1.0, min(1.0, r))
            corr[i][j] = r
            corr[j][i] = r

    return corr


def _agglomerative_cluster(
    corr: list[list[float]],
    n_clusters: int,
) -> list[list[int]]:
    """Simple agglomerative (single-linkage) clustering on a correlation matrix.

    Uses distance = 1 - |correlation| so that highly correlated features
    (positively or negatively) are grouped together.

    Args:
        corr: D x D correlation matrix.
        n_clusters: Desired number of clusters.

    Returns:
        List of clusters, each cluster is a list of feature indices.
    """
    n = len(corr)
    if n == 0:
        return []
    if n_clusters >= n:
        return [[i] for i in range(n)]

    # Initialize: each feature is its own cluster
    clusters: list[list[int]] = [[i] for i in range(n)]

    # Distance matrix: 1 - |corr|
    dist = [[0.0] * n for _ in range(n)]
    for i in range(n):
        for j in range(i + 1, n):
            d = 1.0 - abs(corr[i][j])
            dist[i][j] = d
            dist[j][i] = d

    # Track which clusters are still active
    active = list(range(n))

    while len(active) > n_clusters:
        # Find the two closest active clusters (single linkage)
        best_dist = float("inf")
        merge_a = -1
        merge_b = -1

        for idx_a in range(len(active)):
            for idx_b in range(idx_a + 1, len(active)):
                ca = active[idx_a]
                cb = active[idx_b]
                # Single linkage: minimum distance between any pair
                min_d = float("inf")
                for i in clusters[ca]:
                    for j in clusters[cb]:
                        if dist[i][j] < min_d:
                            min_d = dist[i][j]
                if min_d < best_dist:
                    best_dist = min_d
                    merge_a = idx_a
                    merge_b = idx_b

        if merge_a < 0 or merge_b < 0:
            break

        # Merge cluster at merge_b into merge_a
        ca = active[merge_a]
        cb = active[merge_b]
        clusters[ca] = clusters[ca] + clusters[cb]
        active.pop(merge_b)

    return [clusters[i] for i in active]


# ---------------------------------------------------------------------------
# Clustered MDA
# ---------------------------------------------------------------------------


def clustered_mda(
    score_fn: callable,
    X: list[list[float]],
    y: list[float],
    feature_names: list[str],
    n_clusters: int = 3,
    n_splits: int = 5,
    seed: int | None = None,
) -> list[FeatureImportance]:
    """Clustered feature importance (AFML Chapter 8.7).

    Groups features by correlation into clusters using simple agglomerative
    clustering, then permutes entire clusters at once. This accounts for
    substitution effects between correlated features that standard MDA
    misses.

    When one feature in a correlated group is shuffled, the model can
    compensate by using the remaining correlated features. Shuffling the
    entire cluster eliminates this substitution, giving a more accurate
    picture of the group's true importance.

    Args:
        score_fn: Callable ``(X_train, y_train, X_test, y_test) -> float``.
        X: Feature matrix (N x D).
        y: Label vector (length N).
        feature_names: Names for each feature (length D).
        n_clusters: Number of feature clusters.
        n_splits: Number of cross-validation folds.
        seed: Random seed for reproducibility.

    Returns:
        List of :class:`FeatureImportance`, one per *cluster*. The
        ``feature`` field contains comma-separated names of features in
        that cluster. Sorted by importance descending.
    """
    from horizon._horizon import auth_require_pro
    auth_require_pro()

    n = len(X)
    if n == 0:
        return []

    n_features = len(feature_names)
    if n_features == 0:
        return []

    rng = random.Random(seed)

    # Build correlation matrix and cluster features
    corr = _correlation_matrix(X)
    clusters = _agglomerative_cluster(corr, n_clusters)

    folds = _purged_kfold_split(n, n_splits, purge_gap=0)
    results: list[FeatureImportance] = []

    for cluster in clusters:
        cluster_name = ", ".join(feature_names[i] for i in cluster)
        fold_importances: list[float] = []

        for train_idx, test_idx in folds:
            if not train_idx or not test_idx:
                continue

            X_train = _subset_rows(X, train_idx)
            y_train = _subset_labels(y, train_idx)
            X_test = _subset_rows(X, test_idx)
            y_test = _subset_labels(y, test_idx)

            # Baseline score
            baseline = score_fn(X_train, y_train, X_test, y_test)

            # Shuffle all features in the cluster simultaneously
            X_test_shuffled = _copy_matrix(X_test)
            n_test = len(X_test_shuffled)

            for f_idx in cluster:
                col_values = [X_test_shuffled[i][f_idx] for i in range(n_test)]
                rng.shuffle(col_values)
                for i in range(n_test):
                    X_test_shuffled[i][f_idx] = col_values[i]

            shuffled_score = score_fn(X_train, y_train, X_test_shuffled, y_test)
            fold_importances.append(baseline - shuffled_score)

        if fold_importances:
            mean_imp = sum(fold_importances) / len(fold_importances)
            if len(fold_importances) > 1:
                variance = sum(
                    (s - mean_imp) ** 2 for s in fold_importances
                ) / (len(fold_importances) - 1)
                std_imp = math.sqrt(max(variance, 0.0))
            else:
                std_imp = 0.0
        else:
            mean_imp = 0.0
            std_imp = 0.0

        results.append(
            FeatureImportance(
                feature=cluster_name,
                importance=mean_imp,
                std=std_imp,
            )
        )

    results.sort(key=lambda fi: fi.importance, reverse=True)
    return results
