"""Decomposition metrics for Zwiad Step 3: Decomposition Test."""

from typing import List, Tuple, Dict
import torch
import numpy as np
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.metrics import silhouette_score
from wisent.core.constants import (
    DEFAULT_RANDOM_SEED, ZERO_THRESHOLD,
    DECOMP_ADAPTIVE_K_MIN, DECOMP_ADAPTIVE_K_MAX,
    DECOMP_MIN_CLUSTER_SIZE_BASE, DECOMP_CLUSTER_SIZE_RATIO,
    DECOMP_MIN_SILHOUETTE, DECOMP_MAX_SILHOUETTE_SAMPLES,
    DECOMP_MIN_CONCEPT_SAMPLES, DECOMP_MIN_CONCEPT_SAMPLES_CHOW,
    DECOMP_THRESHOLD_BASE, DECOMP_THRESHOLD_MIN, DECOMP_THRESHOLD_MAX,
    DECOMP_COSINE_SIM_THRESHOLD,
    DECOMP_PCA_DIMS_MAX,
    DECOMP_KMEANS_N_INIT_MIN, DECOMP_KMEANS_N_INIT_MAX,
    DECOMP_KMEANS_SCALING_FACTOR,
)


def _adaptive_max_k(n_samples: int) -> int:
    """Adaptive max clusters: sqrt(n/2) clamped to [2, 15]."""
    max_k = int(np.sqrt(n_samples / 2))
    return max(DECOMP_ADAPTIVE_K_MIN, min(max_k, DECOMP_ADAPTIVE_K_MAX))


def _adaptive_min_cluster_size(n_samples: int) -> int:
    """Adaptive min samples per cluster: max(5, n/20)."""
    return max(DECOMP_MIN_CLUSTER_SIZE_BASE, n_samples // DECOMP_CLUSTER_SIZE_RATIO)


def find_optimal_clustering(
    diff_vectors: torch.Tensor,
    min_silhouette: float = DECOMP_MIN_SILHOUETTE,
) -> Tuple[int, List[int], float]:
    """
    Find optimal number of clusters with adaptive parameters.

    Uses silhouette score to determine optimal k.
    Only reports multiple concepts if silhouette exceeds min_silhouette threshold.

    Args:
        diff_vectors: Difference vectors (pos - neg) to cluster
        min_silhouette: Minimum silhouette score to accept multiple clusters

    Returns:
        Tuple of (n_concepts, cluster_labels, best_silhouette)
    """
    diff_np = diff_vectors.cpu().numpy() if isinstance(diff_vectors, torch.Tensor) else diff_vectors
    n_samples, n_features = diff_np.shape

    pca_dims = min(n_samples - 1, n_features, DECOMP_PCA_DIMS_MAX)
    if pca_dims < n_features and pca_dims >= 2:
        _pca = PCA(n_components=pca_dims, random_state=DEFAULT_RANDOM_SEED)
        diff_np = _pca.fit_transform(diff_np)

    # Subsample for silhouette search (O(n²)), refit on full data with best k
    if n_samples > DECOMP_MAX_SILHOUETTE_SAMPLES:
        _idx = np.random.RandomState(DEFAULT_RANDOM_SEED).choice(n_samples, DECOMP_MAX_SILHOUETTE_SAMPLES, replace=False)
        diff_sub = diff_np[_idx]
        n_sub = DECOMP_MAX_SILHOUETTE_SAMPLES
    else:
        diff_sub = diff_np
        n_sub = n_samples

    max_k = _adaptive_max_k(n_sub)
    min_cluster_size = _adaptive_min_cluster_size(n_sub)
    max_k = min(max_k, n_sub // min_cluster_size)
    if max_k < 2:
        return 1, [0] * n_samples, 0.0

    n_init = max(DECOMP_KMEANS_N_INIT_MIN, min(DECOMP_KMEANS_N_INIT_MAX, DECOMP_KMEANS_SCALING_FACTOR // n_sub + 1))
    best_k = 1
    best_silhouette = -1.0

    for k in range(2, max_k + 1):
        km = KMeans(n_clusters=k, random_state=DEFAULT_RANDOM_SEED, n_init=n_init)
        labels = km.fit_predict(diff_sub)
        cluster_sizes = np.bincount(labels)
        if len(cluster_sizes) < 2 or cluster_sizes.min() < min_cluster_size:
            continue
        sil = silhouette_score(diff_sub, labels)
        if sil > best_silhouette:
            best_silhouette = sil
            best_k = k

    if best_silhouette < min_silhouette:
        return 1, [0] * n_samples, float(best_silhouette)

    # Refit on full data with best k to get labels for all samples
    final_km = KMeans(n_clusters=best_k, random_state=DEFAULT_RANDOM_SEED, n_init=n_init)
    best_labels = final_km.fit_predict(diff_np).tolist()
    return best_k, best_labels, float(best_silhouette)


def geometry_per_concept(
    pos: torch.Tensor,
    neg: torch.Tensor,
    cluster_labels: List[int],
    n_concepts: int,
) -> Dict[int, Dict[str, float]]:
    """
    Run geometry test (linear vs nonlinear) per concept cluster.

    This addresses: "Does the direction differ for violent vs sexual vs illegal?"
    Each concept may have different geometry (linear vs nonlinear encoding).

    Returns:
        Dict mapping concept_id -> {linear_acc, nonlinear_acc, gap, diagnosis}
    """
    from .geometry_metrics import compute_linear_nonlinear_gap

    pos_np = pos.cpu().numpy() if isinstance(pos, torch.Tensor) else pos
    neg_np = neg.cpu().numpy() if isinstance(neg, torch.Tensor) else neg
    labels = np.array(cluster_labels)

    results = {}

    for concept_id in range(n_concepts):
        mask = labels == concept_id
        if mask.sum() < DECOMP_MIN_CONCEPT_SAMPLES:
            results[concept_id] = {
                "n_samples": int(mask.sum()),
                "linear_accuracy": None,
                "nonlinear_accuracy": None,
                "gap": None,
                "diagnosis": "INSUFFICIENT_SAMPLES",
            }
            continue

        pos_concept = torch.tensor(pos_np[mask])
        neg_concept = torch.tensor(neg_np[mask])

        linear_acc, nonlinear_acc = compute_linear_nonlinear_gap(pos_concept, neg_concept)
        gap = nonlinear_acc - linear_acc

        n_samples = int(mask.sum())
        adaptive_threshold = DECOMP_THRESHOLD_BASE / np.sqrt(n_samples)
        adaptive_threshold = max(DECOMP_THRESHOLD_MIN, min(adaptive_threshold, DECOMP_THRESHOLD_MAX))

        diagnosis = "NONLINEAR" if gap > adaptive_threshold else "LINEAR"

        results[concept_id] = {
            "n_samples": n_samples,
            "linear_accuracy": linear_acc,
            "nonlinear_accuracy": nonlinear_acc,
            "gap": gap,
            "gap_threshold": adaptive_threshold,
            "diagnosis": diagnosis,
        }

    return results


def compute_decomposition_metrics(
    pos: torch.Tensor,
    neg: torch.Tensor,
    min_silhouette: float = DECOMP_MIN_SILHOUETTE,
) -> Dict:
    """
    Compute decomposition metrics: clustering and per-concept geometry analysis.

    Args:
        pos: Positive activations
        neg: Negative activations
        min_silhouette: Minimum silhouette score for multiple clusters

    Returns:
        Dict with clustering results and per-concept geometry analysis
    """
    diff_vectors = pos - neg
    n_concepts, cluster_labels, silhouette = find_optimal_clustering(diff_vectors, min_silhouette)
    per_concept = geometry_per_concept(pos, neg, cluster_labels, n_concepts)
    chow = chow_test_analog(pos, neg, cluster_labels, n_concepts)
    return {
        "n_concepts": n_concepts,
        "silhouette_score": silhouette,
        "cluster_labels": cluster_labels,
        "per_concept_geometry": per_concept,
        "chow_test": chow,
    }


def chow_test_analog(
    pos: torch.Tensor,
    neg: torch.Tensor,
    cluster_labels: List[int],
    n_concepts: int,
) -> Dict[str, float]:
    """
    Chow test analog: do probe coefficients differ significantly across concepts?

    Fits linear probe per concept and tests if coefficients are significantly different.
    If yes, the "direction" differs by content type.

    Returns:
        Dict with coefficient_variance, f_statistic_analog, diagnosis.
    """
    from sklearn.linear_model import LogisticRegression

    pos_np = pos.cpu().numpy() if isinstance(pos, torch.Tensor) else pos
    neg_np = neg.cpu().numpy() if isinstance(neg, torch.Tensor) else neg
    labels = np.array(cluster_labels)

    coefficients = []

    for concept_id in range(n_concepts):
        mask = labels == concept_id
        if mask.sum() < DECOMP_MIN_CONCEPT_SAMPLES_CHOW:
            continue

        X = np.vstack([pos_np[mask], neg_np[mask]])
        y = np.array([1] * mask.sum() + [0] * mask.sum())

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

        coef = model.coef_[0]
        coef_norm = coef / (np.linalg.norm(coef) + ZERO_THRESHOLD)
        coefficients.append(coef_norm)

    if len(coefficients) < 2:
        return {
            "n_concepts_tested": len(coefficients),
            "sufficient_concepts": False,
            "coefficients_differ": None,
        }

    cosine_similarities = []
    for i in range(len(coefficients)):
        for j in range(i + 1, len(coefficients)):
            cos_sim = np.dot(coefficients[i], coefficients[j])
            cosine_similarities.append(cos_sim)

    mean_cos = float(np.mean(cosine_similarities))
    min_cos = float(np.min(cosine_similarities))

    coefficients_differ = min_cos < DECOMP_COSINE_SIM_THRESHOLD

    return {
        "n_concepts_tested": len(coefficients),
        "sufficient_concepts": True,
        "mean_cosine_similarity": mean_cos,
        "min_cosine_similarity": min_cos,
        "coefficients_differ": coefficients_differ,
        "interpretation": "Directions differ by concept" if coefficients_differ else "Similar direction across concepts",
    }
