"""
Analysis functions for Woseco concept networks.
"""
import numpy as np
import pandas as pd
from typing import List, Dict, Optional
import spacy


def unexpectedness_arrays(
    df: pd.DataFrame,
    prompt_column: str = "prompt",
    concepts_column: str = "woseco_concepts",
    language: str = "en",
    model_size: str = "lg"
) -> np.ndarray:
    """
    Measure semantic distance between story concepts and prompt words.

    This quantifies how unexpected/creative the story concepts are relative
    to the original prompt words. Higher distances indicate more creative leaps.

    Args:
        df: DataFrame with prompt and concept columns
        prompt_column: Name of column containing prompts
        concepts_column: Name of column containing concept lists
        language: Language code for spacy model
        model_size: Spacy model size

    Returns:
        Numpy array of unexpectedness scores (one per story)
    """
    # Load spacy model
    model_name = f"{language}_core_{'web' if language == 'en' else 'news'}_{model_size}"
    try:
        nlp = spacy.load(model_name)
    except OSError:
        spacy.cli.download(model_name)
        nlp = spacy.load(model_name)

    unexpectedness_scores = []

    for idx, row in df.iterrows():
        # Parse prompt words (e.g., "belief-faith-sing" -> ["belief", "faith", "sing"])
        prompt = row[prompt_column]
        if isinstance(prompt, str):
            prompt_words = [w.strip() for w in prompt.split('-')]
        else:
            prompt_words = []

        # Get concepts for this story
        concepts = row[concepts_column]
        if not isinstance(concepts, list):
            unexpectedness_scores.append(0.0)
            continue

        if not prompt_words or not concepts:
            unexpectedness_scores.append(0.0)
            continue

        # Calculate average similarity between all concepts and all prompt words
        similarities = []
        for concept in concepts:
            concept_token = nlp(concept)
            for prompt_word in prompt_words:
                prompt_token = nlp(prompt_word)
                similarity = concept_token.similarity(prompt_token)
                similarities.append(similarity)

        # Unexpectedness is inverse of similarity (lower similarity = more unexpected)
        if similarities:
            avg_similarity = np.mean(similarities)
            unexpectedness = 1.0 - avg_similarity
            unexpectedness_scores.append(unexpectedness)
        else:
            unexpectedness_scores.append(0.0)

    return np.array(unexpectedness_scores)


def group_by_rating(
    df: pd.DataFrame,
    rating_column: str,
    concepts_column: str = "woseco_concepts"
) -> Dict[int, List[List[str]]]:
    """
    Group concept lists by rating.

    Args:
        df: DataFrame with ratings and concepts
        rating_column: Name of column containing ratings
        concepts_column: Name of column containing concept lists

    Returns:
        Dictionary mapping rating values to lists of concept lists
    """
    grouped = {}

    for rating in df[rating_column].unique():
        mask = df[rating_column] == rating
        concepts = df.loc[mask, concepts_column].tolist()
        grouped[rating] = concepts

    return grouped


def filter_by_prompt(
    df: pd.DataFrame,
    prompt_value: str,
    prompt_column: str = "prompt"
) -> pd.DataFrame:
    """
    Filter dataframe by prompt value.

    Args:
        df: Input DataFrame
        prompt_value: Prompt to filter by
        prompt_column: Name of column containing prompts

    Returns:
        Filtered DataFrame
    """
    return df[df[prompt_column] == prompt_value].copy()


def calculate_rating_statistics(
    df: pd.DataFrame,
    rating_columns: List[str]
) -> pd.DataFrame:
    """
    Calculate statistics across multiple raters.

    Args:
        df: DataFrame with rating columns
        rating_columns: List of column names containing ratings

    Returns:
        DataFrame with mean, std, min, max ratings
    """
    result = df.copy()

    # Calculate statistics across raters
    ratings = df[rating_columns]
    result['rating_mean'] = ratings.mean(axis=1)
    result['rating_std'] = ratings.std(axis=1)
    result['rating_min'] = ratings.min(axis=1)
    result['rating_max'] = ratings.max(axis=1)

    return result


def pca_analysis(
    metrics_matrix: np.ndarray,
    n_components: int = 2
) -> Dict[str, np.ndarray]:
    """
    Perform PCA dimensionality reduction on metrics.

    Args:
        metrics_matrix: Matrix of metrics (samples x features)
        n_components: Number of PCA components

    Returns:
        Dictionary with transformed data and explained variance
    """
    from sklearn.decomposition import PCA

    pca = PCA(n_components=n_components)
    transformed = pca.fit_transform(metrics_matrix)

    return {
        "transformed": transformed,
        "explained_variance": pca.explained_variance_ratio_,
        "components": pca.components_,
    }


def create_metrics_matrix(
    grouped_metrics: Dict[int, Dict[str, List[float]]],
    metric_names: Optional[List[str]] = None
) -> np.ndarray:
    """
    Create a matrix from grouped metrics for statistical analysis.

    Args:
        grouped_metrics: Dictionary mapping groups to metric dictionaries
        metric_names: List of metric names to include (None = all)

    Returns:
        Numpy matrix (groups x metrics)
    """
    if not grouped_metrics:
        return np.array([])

    # Get metric names
    first_group = list(grouped_metrics.values())[0]
    if metric_names is None:
        metric_names = list(first_group.keys())

    # Build matrix
    matrix_rows = []
    for group_key in sorted(grouped_metrics.keys()):
        group_metrics = grouped_metrics[group_key]
        row = [np.mean(group_metrics[name]) for name in metric_names]
        matrix_rows.append(row)

    return np.array(matrix_rows)


def compare_raters(
    df: pd.DataFrame,
    rating_columns: List[str],
    concepts_column: str = "woseco_concepts"
) -> Dict[str, Dict[str, List[float]]]:
    """
    Compare metrics across different raters.

    Args:
        df: DataFrame with ratings and concepts
        rating_columns: List of rating column names
        concepts_column: Name of column containing concept lists

    Returns:
        Dictionary mapping rater names to grouped metrics
    """
    from .metrics import stats

    rater_metrics = {}

    for rater_col in rating_columns:
        # Group by rating for this rater
        grouped = group_by_rating(df, rater_col, concepts_column)

        # Calculate metrics for each rating group
        rating_metrics = {}
        for rating, concept_lists in grouped.items():
            metrics = stats(concept_lists)
            rating_metrics[rating] = metrics

        rater_name = rater_col.replace('rating_', '')
        rater_metrics[rater_name] = rating_metrics

    return rater_metrics


def analyze_prompt_groups(
    df: pd.DataFrame,
    prompt_column: str = "prompt",
    concepts_column: str = "woseco_concepts"
) -> Dict[str, Dict[str, float]]:
    """
    Analyze metrics for different prompt groups.

    Args:
        df: DataFrame with prompts and concepts
        prompt_column: Name of column containing prompts
        concepts_column: Name of column containing concept lists

    Returns:
        Dictionary mapping prompts to aggregated metrics
    """
    from .metrics import stats, aggregate_metrics

    prompt_metrics = {}

    for prompt in df[prompt_column].unique():
        # Filter to this prompt
        prompt_df = filter_by_prompt(df, prompt, prompt_column)

        # Get concepts
        concepts = prompt_df[concepts_column].tolist()

        # Calculate metrics
        metrics = stats(concepts)

        # Aggregate
        metrics_list = []
        for i in range(len(concepts)):
            story_metrics = {k: v[i] for k, v in metrics.items()}
            metrics_list.append(story_metrics)

        aggregated = aggregate_metrics(metrics_list)
        prompt_metrics[prompt] = aggregated

    return prompt_metrics
