"""
Concept analysis for detecting and decomposing multiple concepts.

These functions analyze whether activations contain a single concept
or multiple interleaved concepts, and how to separate them.
"""

import torch
import numpy as np
from typing import Dict, Any, List, Tuple, Optional


def detect_multiple_concepts(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    max_concepts: int = 5,
) -> Dict[str, Any]:
    """
    Detect if activations contain multiple concepts.
    
    Uses clustering on difference vectors to find distinct concepts.
    """
    try:
        from sklearn.cluster import KMeans
        from sklearn.metrics import silhouette_score
        
        n_pairs = min(len(pos_activations), len(neg_activations))
        if n_pairs < 10:
            return {"n_concepts": 1, "silhouette_scores": {}}
        
        diff_vectors = (pos_activations[:n_pairs] - neg_activations[:n_pairs]).float().cpu().numpy()
        
        # Normalize
        norms = np.linalg.norm(diff_vectors, axis=1, keepdims=True)
        valid_mask = norms.squeeze() > 1e-8
        diff_normalized = diff_vectors[valid_mask] / norms[valid_mask]
        
        if len(diff_normalized) < 10:
            return {"n_concepts": 1, "silhouette_scores": {}}
        
        silhouette_scores = {}
        for k in range(2, min(max_concepts + 1, len(diff_normalized) // 3)):
            km = KMeans(n_clusters=k, random_state=42, n_init=10)
            labels = km.fit_predict(diff_normalized)
            
            try:
                score = silhouette_score(diff_normalized, labels)
                silhouette_scores[k] = float(score)
            except:
                silhouette_scores[k] = 0.0
        
        if not silhouette_scores:
            return {"n_concepts": 1, "silhouette_scores": {}}
        
        # Best k is the one with highest silhouette
        best_k = max(silhouette_scores, key=silhouette_scores.get)
        best_score = silhouette_scores[best_k]

        return {
            # Raw metrics - no threshold-based decisions
            "best_k": best_k,
            "best_silhouette": best_score,
            "silhouette_scores": silhouette_scores,
            # For backwards compatibility, but users should look at best_k and best_silhouette
            "n_concepts": best_k,
        }
    except Exception:
        return {"n_concepts": 1, "silhouette_scores": {}}


def split_by_concepts(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    n_concepts: int = 2,
) -> List[Tuple[torch.Tensor, torch.Tensor]]:
    """
    Split activations into separate concepts using clustering.
    
    Returns list of (pos, neg) pairs for each concept.
    """
    try:
        from sklearn.cluster import KMeans
        
        n_pairs = min(len(pos_activations), len(neg_activations))
        diff_vectors = (pos_activations[:n_pairs] - neg_activations[:n_pairs]).float().cpu().numpy()
        
        norms = np.linalg.norm(diff_vectors, axis=1, keepdims=True)
        valid_mask = norms.squeeze() > 1e-8
        diff_normalized = diff_vectors[valid_mask] / norms[valid_mask]
        
        km = KMeans(n_clusters=n_concepts, random_state=42, n_init=10)
        labels = km.fit_predict(diff_normalized)
        
        # Map back to original indices
        valid_indices = np.where(valid_mask)[0]
        
        result = []
        for c in range(n_concepts):
            concept_mask = labels == c
            original_indices = valid_indices[concept_mask]
            
            if len(original_indices) > 0:
                pos_c = pos_activations[original_indices]
                neg_c = neg_activations[original_indices]
                result.append((pos_c, neg_c))
        
        return result
    except Exception:
        return [(pos_activations, neg_activations)]


def analyze_concept_independence(
    concepts: List[Tuple[torch.Tensor, torch.Tensor]],
) -> Dict[str, Any]:
    """
    Analyze how independent different concepts are.
    
    Measures angle between concept directions.
    """
    try:
        if len(concepts) < 2:
            return {"n_concepts": len(concepts), "independence": 1.0}
        
        directions = []
        for pos, neg in concepts:
            n = min(len(pos), len(neg))
            if n < 2:
                continue
            diff_mean = (pos[:n] - neg[:n]).float().cpu().numpy().mean(axis=0)
            norm = np.linalg.norm(diff_mean)
            if norm > 1e-8:
                directions.append(diff_mean / norm)
        
        if len(directions) < 2:
            return {"n_concepts": len(concepts), "independence": 1.0}
        
        # Compute pairwise angles
        angles = []
        for i in range(len(directions)):
            for j in range(i + 1, len(directions)):
                cos_angle = np.dot(directions[i], directions[j])
                angle = np.degrees(np.arccos(np.clip(np.abs(cos_angle), 0, 1)))
                angles.append(angle)
        
        mean_angle = float(np.mean(angles))
        independence = mean_angle / 90.0  # Normalize: 90 degrees = fully independent
        
        return {
            "n_concepts": len(concepts),
            "mean_angle": mean_angle,
            "independence": float(np.clip(independence, 0, 1)),
            "angles": angles,
        }
    except Exception:
        return {"n_concepts": len(concepts), "independence": 0.5}


def compute_concept_coherence(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
) -> float:
    """
    Compute how coherent the concept is (single direction vs. spread).
    
    High coherence = all pairs point same direction
    Low coherence = pairs point different directions
    """
    try:
        n_pairs = min(len(pos_activations), len(neg_activations))
        if n_pairs < 3:
            return 0.0
        
        diff_vectors = (pos_activations[:n_pairs] - neg_activations[:n_pairs]).float().cpu().numpy()
        
        norms = np.linalg.norm(diff_vectors, axis=1, keepdims=True)
        valid_mask = norms.squeeze() > 1e-8
        diff_normalized = diff_vectors[valid_mask] / norms[valid_mask]
        
        if len(diff_normalized) < 3:
            return 0.0
        
        # Coherence = how much variance is explained by first PC
        U, S, Vh = np.linalg.svd(diff_normalized, full_matrices=False)
        
        if len(S) == 0 or S.sum() == 0:
            return 0.0
        
        coherence = float((S[0] ** 2) / (S ** 2).sum())
        return coherence
    except Exception:
        return 0.0


def compute_concept_stability(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    n_bootstrap: int = 20,
) -> float:
    """
    Compute stability of concept direction across bootstrap samples.
    """
    try:
        from .direction_metrics import compute_direction_stability
        
        result = compute_direction_stability(
            pos_activations, neg_activations, n_bootstrap=n_bootstrap
        )
        return result.get("stability_score", 0.0)
    except Exception:
        return 0.0


def decompose_into_concepts(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    n_concepts: int = None,
) -> Dict[str, Any]:
    """
    Full concept decomposition: detect, split, and analyze.
    """
    detection = detect_multiple_concepts(pos_activations, neg_activations)
    
    if n_concepts is None:
        n_concepts = detection.get("n_concepts", 1)
    
    if n_concepts == 1:
        return {
            "n_concepts": 1,
            "concepts": [(pos_activations, neg_activations)],
            "independence": 1.0,
            "coherence": compute_concept_coherence(pos_activations, neg_activations),
        }
    
    concepts = split_by_concepts(pos_activations, neg_activations, n_concepts)
    independence = analyze_concept_independence(concepts)
    
    coherences = [compute_concept_coherence(p, n) for p, n in concepts]
    
    return {
        "n_concepts": len(concepts),
        "concepts": concepts,
        "independence": independence.get("independence", 0.5),
        "coherences": coherences,
        "mean_coherence": float(np.mean(coherences)) if coherences else 0.0,
    }


def get_pair_concept_assignments(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    n_concepts: int = 2,
) -> Dict[str, Any]:
    """
    Get concept assignment for each pair.

    Returns which concept each pair belongs to, along with confidence scores.

    Args:
        pos_activations: [N, hidden_dim] positive activations
        neg_activations: [N, hidden_dim] negative activations
        n_concepts: Number of concepts to cluster into

    Returns:
        Dict with:
        - assignments: List[int] - concept index for each pair (0 to n_concepts-1)
        - confidences: List[float] - how strongly each pair belongs to its cluster
        - distances_to_clusters: List[List[float]] - distance to each cluster center
        - cluster_centers: np.ndarray - the cluster centers (steering directions per concept)
        - pairs_by_concept: Dict[int, List[int]] - mapping from concept to list of pair indices
    """
    try:
        from sklearn.cluster import KMeans

        n_pairs = min(len(pos_activations), len(neg_activations))
        diff_vectors = (pos_activations[:n_pairs] - neg_activations[:n_pairs]).float().cpu().numpy()

        norms = np.linalg.norm(diff_vectors, axis=1, keepdims=True)
        valid_mask = norms.squeeze() > 1e-8
        diff_normalized = diff_vectors[valid_mask] / norms[valid_mask]

        if len(diff_normalized) < n_concepts * 2:
            # Not enough samples
            return {
                "assignments": [0] * n_pairs,
                "confidences": [1.0] * n_pairs,
                "error": "not enough valid samples for clustering",
            }

        km = KMeans(n_clusters=n_concepts, random_state=42, n_init=10)
        labels = km.fit_predict(diff_normalized)

        # Distances to each cluster center
        distances = km.transform(diff_normalized)  # [n_valid, n_concepts]

        # Confidence: inverse of distance to assigned cluster, normalized
        assigned_distances = distances[np.arange(len(labels)), labels]
        max_dist = distances.max() + 1e-8
        confidences_valid = 1 - (assigned_distances / max_dist)

        # Map back to original indices
        valid_indices = np.where(valid_mask)[0]

        # Build full arrays with -1 for invalid pairs
        assignments = [-1] * n_pairs
        confidences = [0.0] * n_pairs
        distances_to_clusters = [[float('inf')] * n_concepts for _ in range(n_pairs)]

        for i, orig_idx in enumerate(valid_indices):
            assignments[orig_idx] = int(labels[i])
            confidences[orig_idx] = float(confidences_valid[i])
            distances_to_clusters[orig_idx] = distances[i].tolist()

        # Group pairs by concept
        pairs_by_concept = {c: [] for c in range(n_concepts)}
        for pair_idx, concept in enumerate(assignments):
            if concept >= 0:
                pairs_by_concept[concept].append(pair_idx)

        # Cluster centers (these are the steering directions for each concept)
        cluster_centers = km.cluster_centers_

        # Compute steering vector for each concept (unnormalized)
        concept_steering_vectors = []
        for c in range(n_concepts):
            concept_indices = pairs_by_concept[c]
            if concept_indices:
                concept_diffs = diff_vectors[concept_indices]
                steering_vec = concept_diffs.mean(axis=0)
                concept_steering_vectors.append(steering_vec)
            else:
                concept_steering_vectors.append(np.zeros(diff_vectors.shape[1]))

        return {
            "assignments": assignments,
            "confidences": confidences,
            "distances_to_clusters": distances_to_clusters,
            "cluster_centers": cluster_centers,  # Normalized directions
            "concept_steering_vectors": concept_steering_vectors,  # Actual steering vectors
            "pairs_by_concept": pairs_by_concept,
            "n_pairs": n_pairs,
            "n_valid_pairs": int(valid_mask.sum()),
            "n_invalid_pairs": int((~valid_mask).sum()),
        }

    except Exception as e:
        return {
            "assignments": [0] * min(len(pos_activations), len(neg_activations)),
            "confidences": [0.0] * min(len(pos_activations), len(neg_activations)),
            "error": str(e),
        }


def find_mixed_pairs(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    threshold: float = 0.3,
) -> List[int]:
    """
    Find pairs that don't fit well into any cluster (mixed concepts).
    """
    try:
        from sklearn.cluster import KMeans
        
        n_pairs = min(len(pos_activations), len(neg_activations))
        diff_vectors = (pos_activations[:n_pairs] - neg_activations[:n_pairs]).float().cpu().numpy()
        
        norms = np.linalg.norm(diff_vectors, axis=1, keepdims=True)
        valid_mask = norms.squeeze() > 1e-8
        diff_normalized = diff_vectors[valid_mask] / norms[valid_mask]
        
        km = KMeans(n_clusters=2, random_state=42, n_init=10)
        km.fit(diff_normalized)
        
        distances = km.transform(diff_normalized).min(axis=1)
        median_dist = np.median(distances)
        
        mixed_mask = distances > median_dist * (1 + threshold)
        valid_indices = np.where(valid_mask)[0]
        
        return list(valid_indices[mixed_mask])
    except Exception:
        return []


def get_pure_concept_pairs(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    concept_idx: int = 0,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Get pairs belonging to a specific concept cluster.
    """
    concepts = split_by_concepts(pos_activations, neg_activations, n_concepts=2)
    
    if concept_idx < len(concepts):
        return concepts[concept_idx]
    
    return pos_activations, neg_activations


def analyze_concept_structure(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
) -> Dict[str, Any]:
    """
    Analyze concept structure - returns raw metrics for steering decisions.

    Does NOT make recommendations - just returns metrics describing
    whether the data contains one or multiple concepts.
    """
    decomposition = decompose_into_concepts(pos_activations, neg_activations)

    n_concepts = decomposition["n_concepts"]
    independence = decomposition["independence"]
    mean_coherence = decomposition.get("mean_coherence", 0.0)
    coherences = decomposition.get("coherences", [])

    # Also get raw detection metrics
    detection = detect_multiple_concepts(pos_activations, neg_activations)

    return {
        # How many concepts were detected
        "n_concepts": n_concepts,
        "best_silhouette": detection.get("best_silhouette"),
        "silhouette_scores": detection.get("silhouette_scores", {}),

        # Independence: 0 = same direction, 1 = orthogonal
        "independence": independence,

        # Coherence per concept (1 = single direction within concept)
        "coherences": coherences,
        "mean_coherence": mean_coherence,

        # Sizes of each concept cluster
        "concept_sizes": [len(c[0]) for c in decomposition.get("concepts", [])],
    }


# Keep old name for backwards compatibility
def recommend_per_concept_steering(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
) -> Dict[str, Any]:
    """Deprecated: Use analyze_concept_structure instead."""
    return analyze_concept_structure(pos_activations, neg_activations)
