"""Automatic concept naming based on contrastive pair content."""
import re
from collections import Counter
from typing import Dict, List, Any, Optional
import numpy as np


def extract_keywords(texts: List[str], top_k: int = 10) -> List[str]:
    """Extract common keywords from a list of texts."""
    all_text = " ".join(texts).lower()
    words = re.findall(r'\b[a-z]{3,}\b', all_text)

    stop_words = {
        'the', 'and', 'that', 'this', 'with', 'for', 'are', 'was', 'were',
        'been', 'being', 'have', 'has', 'had', 'does', 'did', 'will', 'would',
        'could', 'should', 'may', 'might', 'must', 'can', 'not', 'but', 'what',
        'which', 'who', 'whom', 'when', 'where', 'why', 'how', 'all', 'each',
        'every', 'both', 'few', 'more', 'most', 'other', 'some', 'such', 'than',
        'too', 'very', 'just', 'also', 'only', 'own', 'same', 'into', 'from',
        'there', 'here', 'they', 'them', 'their', 'you', 'your', 'about',
        'true', 'false', 'yes', 'answer', 'question', 'following', 'correct',
        'happens', 'really', 'actually', 'because', 'would', 'could',
    }

    filtered = [w for w in words if w not in stop_words]
    counts = Counter(filtered)
    return [word for word, _ in counts.most_common(top_k)]


def classify_concept_type(
    prompts: List[str],
    positives: List[str],
    negatives: List[str],
) -> Dict[str, Any]:
    """Classify a concept based on content patterns."""

    all_text = " ".join(prompts + positives + negatives).lower()

    # Pattern categories with associated keywords and weights
    patterns = {
        "scientific_explanations": {
            "keywords": ["science", "physics", "light", "energy", "temperature", "chemical", "biology", "cells"],
            "prompt_patterns": ["why do", "how does", "what causes"],
        },
        "nature_animal_misconceptions": {
            "keywords": ["animal", "species", "bull", "chameleon", "bird", "fish", "plant", "nature"],
            "prompt_patterns": ["why do animals", "what animal", "do animals"],
        },
        "biological_body_myths": {
            "keywords": ["body", "brain", "blood", "muscle", "eat", "digest", "health", "medical"],
            "prompt_patterns": ["what happens to your", "what percentage", "if you eat"],
        },
        "historical_misconceptions": {
            "keywords": ["history", "war", "president", "century", "ancient", "invented", "discovered"],
            "prompt_patterns": ["who invented", "when was", "who was the first"],
        },
        "fictional_false_premises": {
            "keywords": ["witch", "magic", "fairy", "ghost", "supernatural", "myth", "legend"],
            "prompt_patterns": ["what is a reasonable way", "who is a famous", "what happened to the"],
            "response_patterns": ["don't exist", "no comment", "there are no", "impossible"],
        },
        "general_factual_knowledge": {
            "keywords": ["people", "country", "world", "common", "real", "actually"],
            "prompt_patterns": [],
        },
    }

    scores = {}
    for category, config in patterns.items():
        score = 0
        # Keyword matching
        for kw in config["keywords"]:
            if kw in all_text:
                score += 1
        # Prompt pattern matching
        for pattern in config.get("prompt_patterns", []):
            for prompt in prompts:
                if pattern in prompt.lower():
                    score += 2
        # Response pattern matching (for fictional premises)
        for pattern in config.get("response_patterns", []):
            for pos in positives:
                if pattern in pos.lower():
                    score += 3
        scores[category] = score

    # Get best category
    best_category = max(scores, key=scores.get)

    # If no strong signal, default to general
    if scores[best_category] < 2:
        best_category = "general_factual_knowledge"

    return {
        "category": best_category,
        "confidence": scores[best_category] / max(sum(scores.values()), 1),
        "scores": scores,
    }


def generate_concept_description(category: str) -> str:
    """Generate human-readable description for a concept category."""
    descriptions = {
        "scientific_explanations": "Questions about scientific phenomena where the truthful answer requires deeper understanding than the popular myth",
        "nature_animal_misconceptions": "Misconceptions about animals, nature, and natural phenomena",
        "biological_body_myths": "Myths about human body, health, and biology",
        "historical_misconceptions": "Historical facts that are commonly misremembered or misattributed",
        "fictional_false_premises": "Questions based on fictional or impossible premises that should be rejected",
        "general_factual_knowledge": "Broad factual questions where common assumptions are incorrect",
    }
    return descriptions.get(category, "General knowledge questions")


def name_concepts(
    concepts: List[Dict[str, Any]],
    pair_texts: Optional[Dict[int, Dict[str, str]]] = None,
    cluster_labels: Optional[np.ndarray] = None,
) -> List[Dict[str, Any]]:
    """
    Automatically name concepts based on their content.

    Args:
        concepts: List of concept dicts with metrics
        pair_texts: Optional dict mapping pair_id -> {prompt, positive, negative}
        cluster_labels: Optional array mapping pair index -> cluster id

    Returns:
        Updated concepts with names and descriptions
    """
    if pair_texts is None or cluster_labels is None:
        # No text available, use generic names
        for i, concept in enumerate(concepts):
            concept["name"] = f"concept_{i+1}"
            concept["description"] = "Concept cluster (no text available for naming)"
        return concepts

    # Group pairs by concept
    pair_ids = sorted(pair_texts.keys())
    concept_pairs = {i: [] for i in range(len(concepts))}

    for idx, pair_id in enumerate(pair_ids):
        if idx < len(cluster_labels):
            cluster_id = cluster_labels[idx]
            if pair_id in pair_texts:
                concept_pairs[cluster_id].append(pair_texts[pair_id])

    # Classify and name each concept
    for i, concept in enumerate(concepts):
        pairs = concept_pairs.get(i, [])

        if pairs:
            prompts = [p.get("prompt", "") for p in pairs]
            positives = [p.get("positive", "") for p in pairs]
            negatives = [p.get("negative", "") for p in pairs]

            classification = classify_concept_type(prompts, positives, negatives)

            concept["name"] = classification["category"]
            concept["description"] = generate_concept_description(classification["category"])
            concept["classification_confidence"] = classification["confidence"]
            concept["keywords"] = extract_keywords(prompts, top_k=5)
            concept["sample_prompts"] = prompts[:3]
            concept["sample_positive"] = positives[:3]
            concept["sample_negative"] = negatives[:3]
        else:
            concept["name"] = f"concept_{i+1}"
            concept["description"] = "Empty concept cluster"

    return concepts


def decompose_and_name_concepts(
    pos_activations,
    neg_activations,
    pair_texts: Optional[Dict[int, Dict[str, str]]] = None,
    generate_visualizations: bool = False,
    use_llm_naming: bool = False,
    llm_model: str = "meta-llama/Llama-3.2-1B-Instruct",
) -> Dict[str, Any]:
    """
    Full concept decomposition with automatic naming.

    Args:
        pos_activations: Positive activation tensors
        neg_activations: Negative activation tensors
        pair_texts: Optional dict mapping pair_id -> {prompt, positive, negative}
        generate_visualizations: Whether to generate per-concept visualizations
        use_llm_naming: Use LLM for more meaningful concept names
        llm_model: HuggingFace model to use for naming

    Returns:
        Complete decomposition results with named concepts
    """
    from .concept_analysis import decompose_into_concepts
    from sklearn.cluster import SpectralClustering
    import torch

    # Get decomposition
    decomposition = decompose_into_concepts(pos_activations, neg_activations)
    n_concepts = decomposition["n_concepts"]

    # Get cluster labels for naming
    n_pairs = min(len(pos_activations), len(neg_activations))
    if hasattr(pos_activations, 'numpy'):
        pos_np = pos_activations[:n_pairs].float().cpu().numpy()
        neg_np = neg_activations[:n_pairs].float().cpu().numpy()
    else:
        pos_np = pos_activations[:n_pairs]
        neg_np = neg_activations[:n_pairs]

    diff_vectors = pos_np - neg_np

    # L2 normalize before clustering (critical for handling magnitude variation)
    norms = np.linalg.norm(diff_vectors, axis=1, keepdims=True)
    norms = np.where(norms < 1e-8, 1.0, norms)  # avoid division by zero
    diff_normalized = diff_vectors / norms

    # Use SpectralClustering (better than KMeans for correlated/imbalanced concepts)
    spectral = SpectralClustering(
        n_clusters=n_concepts, random_state=42,
        affinity='nearest_neighbors', n_neighbors=min(10, n_pairs - 1)
    )
    cluster_labels = spectral.fit_predict(diff_normalized)

    # Build concept list with metrics
    concepts = []
    coherences = decomposition.get("coherences", [])

    for i in range(n_concepts):
        mask = cluster_labels == i
        concept = {
            "id": i + 1,
            "n_pairs": int(np.sum(mask)),
            "coherence": coherences[i] if i < len(coherences) else 0.0,
        }
        concepts.append(concept)

    # Add names if pair texts available
    if use_llm_naming and pair_texts:
        from .llm_concept_naming import name_all_concepts_with_llm
        print("Using LLM for concept naming...", flush=True)
        concepts = name_all_concepts_with_llm(
            concepts, pair_texts, cluster_labels, model=llm_model
        )
    else:
        concepts = name_concepts(concepts, pair_texts, cluster_labels)

    # Build pair assignments
    pair_assignments = {}
    if pair_texts:
        pair_ids = sorted(pair_texts.keys())
        for idx, pair_id in enumerate(pair_ids):
            if idx < len(cluster_labels):
                pair_assignments[pair_id] = int(cluster_labels[idx])

    result = {
        "n_concepts": n_concepts,
        "independence": decomposition.get("independence", 0.0),
        "mean_coherence": decomposition.get("mean_coherence", 0.0),
        "concepts": concepts,
        "pair_assignments": pair_assignments,
        "cluster_labels": cluster_labels.tolist(),
    }

    # Generate visualizations if requested
    if generate_visualizations:
        from .concept_visualizations import create_all_concept_figures
        result["visualizations"] = create_all_concept_figures(
            pos_activations, neg_activations, cluster_labels, concepts
        )

    return result
