"""Linearity validation helpers for the Linear Representation Hypothesis."""
from typing import Tuple, Optional, Dict, List
import numpy as np
import torch
from sklearn.linear_model import LogisticRegression
from sklearn.neural_network import MLPClassifier
from sklearn.model_selection import cross_val_score, StratifiedKFold
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from wisent.core.utils.config_tools.constants import (ZERO_THRESHOLD, LINEARITY_MAX_PAIRS, LINEARITY_PCA_COMPONENTS, CV_FOLDS, DEFAULT_RANDOM_SEED, LINEARITY_HIDDEN_DIM_LARGE, LINEARITY_HIDDEN_DIM_SMALL, LINEARITY_N_CLUSTERS, LINEARITY_NONLINEAR_ERROR, LINEARITY_N_INIT, LINEARITY_N_BOOTSTRAP, LINEARITY_MLP_MAX_ITER, LINEARITY_MIN_GAPS, LINEARITY_DEGREE, LINEARITY_N_FEATURES, LINEARITY_REG_C, LINEARITY_RIDGE_ALPHA, LINEARITY_BOOTSTRAP_ITER, LINEARITY_Z_SCORE_THRESHOLD, CONFIDENCE_LEVEL)

def _prepare_data(
    pos: torch.Tensor, neg: torch.Tensor
) -> Tuple[np.ndarray, np.ndarray]:
    """Convert tensors to numpy, subsample if large, apply PCA, create labels."""
    from sklearn.decomposition import PCA
    pos_np = pos.cpu().numpy() if isinstance(pos, torch.Tensor) else pos
    neg_np = neg.cpu().numpy() if isinstance(neg, torch.Tensor) else neg
    import sys; print(f"  [TRACE] _prepare_data: {len(pos_np)} pairs, {pos_np.shape[1]} dims", file=sys.stderr, flush=True)
    if len(pos_np) > LINEARITY_MAX_PAIRS:
        idx = np.random.RandomState(DEFAULT_RANDOM_SEED).choice(len(pos_np), LINEARITY_MAX_PAIRS, replace=False)
        idx.sort()
        pos_np, neg_np = pos_np[idx], neg_np[idx]
    X = np.vstack([pos_np, neg_np])
    y = np.array([1] * len(pos_np) + [0] * len(neg_np))
    n_samples, n_features = X.shape
    pca_dims = min(n_samples - 1, n_features, LINEARITY_PCA_COMPONENTS)
    if pca_dims < n_features and pca_dims >= 2:
        X = PCA(n_components=pca_dims, random_state=DEFAULT_RANDOM_SEED).fit_transform(X)
    return X, y

def compute_probe_accuracies(
    pos: torch.Tensor,
    neg: torch.Tensor,
    n_splits: int = CV_FOLDS,
    random_state: int = DEFAULT_RANDOM_SEED,
) -> Tuple[float, float, np.ndarray, np.ndarray]:
    """Compute linear and nonlinear probe accuracies with cross-validation.

    Returns:
        Tuple of (linear_acc, nonlinear_acc, linear_scores, nonlinear_scores)
        where scores are per-fold accuracies for statistical testing.
    """
    X, y = _prepare_data(pos, neg)
    n_per_class = min(len(pos), len(neg))
    n_splits = max(2, min(n_splits, n_per_class))
    cv = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
    linear_model = LogisticRegression(solver="lbfgs", random_state=random_state)
    nonlinear_model = MLPClassifier(
        hidden_layer_sizes=(LINEARITY_HIDDEN_DIM_LARGE,), random_state=random_state
    )

    linear_scores = cross_val_score(linear_model, X, y, cv=cv, scoring="accuracy")
    nonlinear_scores = cross_val_score(nonlinear_model, X, y, cv=cv, scoring="accuracy")

    return (
        float(np.mean(linear_scores)),
        float(np.mean(nonlinear_scores)),
        linear_scores,
        nonlinear_scores,
    )

def analyze_residuals(
    pos: torch.Tensor,
    neg: torch.Tensor,
    n_clusters: int = LINEARITY_N_CLUSTERS,
    random_state: int = DEFAULT_RANDOM_SEED,
) -> Dict[str, float]:
    """Analyze linear probe residuals for systematic patterns.

    If linear model errors cluster (high silhouette), this indicates
    the linear assumption fails in predictable regions of the space.

    Returns:
        Dictionary with residual_silhouette and cluster info.
    """
    X, y = _prepare_data(pos, neg)

    model = LogisticRegression(solver="lbfgs", random_state=random_state)
    model.fit(X, y)

    probs = model.predict_proba(X)[:, 1]
    errors = np.abs(y - probs)

    error_mask = errors > LINEARITY_NONLINEAR_ERROR
    if error_mask.sum() < n_clusters * 2:
        return {
            "residual_silhouette": 0.0,
            "n_high_error": int(error_mask.sum()),
            "clusters_found": False,
        }

    high_error_X = X[error_mask]

    kmeans = KMeans(n_clusters=n_clusters, random_state=random_state, n_init=LINEARITY_N_INIT)
    cluster_labels = kmeans.fit_predict(high_error_X)

    if len(np.unique(cluster_labels)) > 1:
        sil = silhouette_score(high_error_X, cluster_labels)
    else:
        sil = 0.0

    return {
        "residual_silhouette": float(sil),
        "n_high_error": int(error_mask.sum()),
        "clusters_found": True,
    }

def bootstrap_gap_ci(
    pos: torch.Tensor,
    neg: torch.Tensor,
    n_bootstrap: int = LINEARITY_N_BOOTSTRAP,
    ci_level: float = CONFIDENCE_LEVEL,
    random_state: int = DEFAULT_RANDOM_SEED,
) -> Tuple[float, float, float]:
    """Compute bootstrap confidence interval on linear-nonlinear gap.

    Returns:
        Tuple of (gap_mean, ci_lower, ci_upper)
    """
    rng = np.random.RandomState(random_state)
    X, y = _prepare_data(pos, neg)
    n = len(X)

    gaps = []
    for _ in range(n_bootstrap):
        idx = rng.choice(n, n, replace=True)
        X_boot, y_boot = X[idx], y[idx]

        if len(np.unique(y_boot)) < 2:
            continue

        linear = LogisticRegression(solver="lbfgs", random_state=DEFAULT_RANDOM_SEED)
        nonlinear = MLPClassifier(hidden_layer_sizes=(LINEARITY_HIDDEN_DIM_SMALL,), max_iter=LINEARITY_MLP_MAX_ITER, random_state=DEFAULT_RANDOM_SEED)

        try:
            linear.fit(X_boot, y_boot)
            nonlinear.fit(X_boot, y_boot)

            linear_acc = linear.score(X_boot, y_boot)
            nonlinear_acc = nonlinear.score(X_boot, y_boot)
            gaps.append(nonlinear_acc - linear_acc)
        except Exception:
            continue

    if len(gaps) < LINEARITY_MIN_GAPS:
        return 0.0, 0.0, 0.0

    gaps = np.array(gaps)
    alpha = 1 - ci_level
    ci_lower = float(np.percentile(gaps, 100 * alpha / 2))
    ci_upper = float(np.percentile(gaps, 100 * (1 - alpha / 2)))

    return float(np.mean(gaps)), ci_lower, ci_upper

def test_cross_context_linearity(
    contexts: List[Tuple[torch.Tensor, torch.Tensor]],
    random_state: int = DEFAULT_RANDOM_SEED,
) -> Dict[str, float]:
    """Test if linear directions transfer across contexts (Lampinen-style).

    Each context is a (pos, neg) tuple from different conversational contexts.
    Trains on one context, tests on others to see if directions generalize.

    Returns:
        Dictionary with transfer_accuracy and generalization metrics.
    """
    if len(contexts) < 2:
        return {
            "transfer_accuracy": 0.0,
            "n_contexts": len(contexts),
            "sufficient_contexts": False,
        }

    transfer_scores = []

    for i, (pos_train, neg_train) in enumerate(contexts):
        X_train, y_train = _prepare_data(pos_train, neg_train)

        model = LogisticRegression(solver="lbfgs", random_state=random_state)
        model.fit(X_train, y_train)

        for j, (pos_test, neg_test) in enumerate(contexts):
            if i == j:
                continue
            X_test, y_test = _prepare_data(pos_test, neg_test)
            acc = model.score(X_test, y_test)
            transfer_scores.append(acc)

    return {
        "transfer_accuracy": float(np.mean(transfer_scores)),
        "transfer_std": float(np.std(transfer_scores)),
        "n_contexts": len(contexts),
        "sufficient_contexts": True,
    }

def ramsey_polynomial_test(
    pos: torch.Tensor,
    neg: torch.Tensor,
    degree: int = LINEARITY_DEGREE,
    n_features: int = LINEARITY_N_FEATURES,
    random_state: int = DEFAULT_RANDOM_SEED,
) -> Dict[str, float]:
    """Ramsey-style test: do polynomial features significantly improve fit?

    If polynomial features don't help, the relationship is likely linear.
    Uses PCA to reduce dimensionality before polynomial expansion.

    Returns:
        Dictionary with linear_acc, poly_acc, and improvement metrics.
    """
    from sklearn.preprocessing import PolynomialFeatures
    from sklearn.decomposition import PCA

    X, y = _prepare_data(pos, neg)

    n_components = min(n_features, X.shape[1], len(X) // 2)
    pca = PCA(n_components=n_components, random_state=random_state)
    X_reduced = pca.fit_transform(X)

    n_per_class = min(len(pos), len(neg)); _ns = max(2, min(CV_FOLDS, n_per_class))
    cv = StratifiedKFold(n_splits=_ns, shuffle=True, random_state=random_state)
    linear = LogisticRegression(solver="lbfgs", random_state=random_state)
    linear_scores = cross_val_score(linear, X_reduced, y, cv=cv, scoring="accuracy")
    linear_acc = float(np.mean(linear_scores))

    poly = PolynomialFeatures(degree=degree, include_bias=False)
    X_poly = poly.fit_transform(X_reduced)

    poly_model = LogisticRegression(solver="lbfgs", random_state=random_state, C=LINEARITY_REG_C)
    poly_scores = cross_val_score(poly_model, X_poly, y, cv=cv, scoring="accuracy")
    poly_acc = float(np.mean(poly_scores))

    improvement = poly_acc - linear_acc

    return {
        "linear_accuracy": linear_acc,
        "polynomial_accuracy": poly_acc,
        "improvement": improvement,
        "polynomial_degree": degree,
        "n_pca_components": n_components,
    }

def regression_probe(
    pos: torch.Tensor,
    neg: torch.Tensor,
    target: np.ndarray,
    n_splits: int = CV_FOLDS,
    random_state: int = DEFAULT_RANDOM_SEED,
) -> Dict[str, float]:
    """Regression probe for predicting continuous outcomes.

    Stronger than classification for predicting actual behavior magnitudes.
    Uses difference vectors to predict continuous target (e.g., toxicity score).

    Args:
        pos: Positive activations [n_samples, dim]
        neg: Negative activations [n_samples, dim]
        target: Continuous target values [n_samples]

    Returns:
        Dictionary with R², correlation, and comparison to null.
    """
    from sklearn.linear_model import Ridge
    from sklearn.model_selection import cross_val_score, KFold
    from scipy import stats

    diff = pos.cpu().numpy() - neg.cpu().numpy() if isinstance(pos, torch.Tensor) else pos - neg

    cv = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
    model = Ridge(alpha=LINEARITY_RIDGE_ALPHA)

    r2_scores = cross_val_score(model, diff, target, cv=cv, scoring="r2")
    r2_mean = float(np.mean(r2_scores))

    model.fit(diff, target)
    preds = model.predict(diff)
    corr, p_value = stats.pearsonr(preds, target)

    rng = np.random.RandomState(random_state)
    null_r2s = []
    for _ in range(LINEARITY_BOOTSTRAP_ITER):
        target_perm = rng.permutation(target)
        null_scores = cross_val_score(model, diff, target_perm, cv=cv, scoring="r2")
        null_r2s.append(np.mean(null_scores))

    null_mean = float(np.mean(null_r2s))
    null_std = float(np.std(null_r2s))
    z_score = (r2_mean - null_mean) / (null_std + ZERO_THRESHOLD)

    return {
        "r2_mean": r2_mean,
        "r2_std": float(np.std(r2_scores)),
        "correlation": float(corr),
        "correlation_p_value": float(p_value),
        "null_r2_mean": null_mean,
        "z_score_vs_null": z_score,
        "significant": z_score > LINEARITY_Z_SCORE_THRESHOLD,
    }
