"""
Command-line interface for Woseco.
"""
import click
import pandas as pd
import json
import pickle
from pathlib import Path
from typing import Optional

from .core import WosecoExtractor
from .graph import build_graph, export_graphs, build_graphs_from_list
from .metrics import stats, calculate_all_metrics
from .analysis import (
    unexpectedness_arrays,
    group_by_rating,
    analyze_prompt_groups,
    compare_raters,
)
from .visualization import (
    plot_network_stats,
    plot_rating_distribution,
    plot_graph,
    plot_unexpectedness,
)


@click.group()
@click.version_option(version="0.1.2")
def cli():
    """
    Woseco: WOrd SEntence COoccurrence

    Extract and analyze concept associations in creative stories.
    Supports both English and German text.
    """
    pass


@cli.command()
@click.argument('input_file', type=click.Path(exists=True))
@click.option('--output', '-o', help='Output CSV file path')
@click.option('--text-column', default='Story', help='Name of column containing text')
@click.option('--language', '-l', type=click.Choice(['en', 'de']), default='en',
              help='Language of the text (English or German)')
@click.option('--model-size', type=click.Choice(['sm', 'md', 'lg']), default='lg',
              help='Spacy model size to use')
@click.option('--similarity-threshold', type=float, default=None,
              help='Minimum similarity threshold for concept matching (0-1)')
def extract(input_file, output, text_column, language, model_size, similarity_threshold):
    """
    Extract Woseco concept associations from a CSV file.

    INPUT_FILE: Path to CSV file containing stories/text

    Example:
        woseco extract stories.csv -o results.csv --language en
    """
    click.echo(f"Loading data from {input_file}...")

    # Initialize extractor
    extractor = WosecoExtractor(language=language, model_size=model_size)

    # Process CSV
    output_path = output if output else input_file.replace('.csv', '_woseco.csv')
    df = extractor.process_csv(
        input_file,
        text_column=text_column,
        output_path=output_path,
        similarity_threshold=similarity_threshold
    )

    click.echo(f"\nProcessed {len(df)} stories")
    click.echo(f"Results saved to {output_path}")


@cli.command()
@click.argument('input_file', type=click.Path(exists=True))
@click.option('--output-dir', '-o', required=True, help='Output directory for graphs')
@click.option('--concepts-column', default='woseco_concepts',
              help='Name of column containing concept lists')
@click.option('--graph-type', type=click.Choice(['chain', 'dense']), default='chain',
              help='Type of graph to build')
@click.option('--format', '-f', type=click.Choice(['graphml', 'gexf', 'gml']), default='graphml',
              help='Graph export format')
@click.option('--prefix', default='graph', help='Prefix for output filenames')
def export(input_file, output_dir, concepts_column, graph_type, format, prefix):
    """
    Build and export concept graphs from extracted Woseco data.

    INPUT_FILE: Path to CSV file with woseco_concepts column

    Example:
        woseco export results.csv -o graphs/ --format graphml
    """
    click.echo(f"Loading data from {input_file}...")
    df = pd.read_csv(input_file)

    # Convert string representations of lists back to actual lists
    import ast
    df[concepts_column] = df[concepts_column].apply(
        lambda x: ast.literal_eval(x) if isinstance(x, str) else x
    )

    concept_lists = df[concepts_column].tolist()

    click.echo(f"Building {graph_type} graphs...")
    graphs = build_graphs_from_list(concept_lists, graph_type=graph_type)

    click.echo(f"Exporting {len(graphs)} graphs to {output_dir}...")
    export_graphs(graphs, output_dir, prefix=prefix, format=format)

    click.echo("Done!")


@cli.command()
@click.argument('input_file', type=click.Path(exists=True))
@click.option('--concepts-column', default='woseco_concepts',
              help='Name of column containing concept lists')
@click.option('--output', '-o', help='Output JSON file for metrics')
@click.option('--graph-type', type=click.Choice(['chain', 'dense']), default='chain',
              help='Type of graph to build')
def metrics(input_file, concepts_column, output, graph_type):
    """
    Calculate network metrics for Woseco concept graphs.

    INPUT_FILE: Path to CSV file with woseco_concepts column

    Metrics calculated:
    - ASPL: Average Shortest Path Length
    - MLCC: Mean Local Clustering Coefficient
    - Modularity: Community structure strength
    - Num Components: Number of connected components
    - Avg Component Size: Average size of components
    - GCC Size: Giant Connected Component size

    Example:
        woseco metrics results.csv -o metrics.json
    """
    click.echo(f"Loading data from {input_file}...")
    df = pd.read_csv(input_file)

    # Convert string representations of lists back to actual lists
    import ast
    df[concepts_column] = df[concepts_column].apply(
        lambda x: ast.literal_eval(x) if isinstance(x, str) else x
    )

    concept_lists = df[concepts_column].tolist()

    click.echo("Calculating metrics...")
    metrics_dict = stats(concept_lists, graph_type=graph_type)

    # Display summary
    click.echo("\n=== METRICS SUMMARY ===")
    for metric_name, values in metrics_dict.items():
        import numpy as np
        click.echo(f"{metric_name.upper()}: mean={np.mean(values):.4f}, std={np.std(values):.4f}")

    # Save to file if requested
    if output:
        # Convert numpy arrays to lists for JSON serialization
        output_data = {k: [float(v) for v in vals] for k, vals in metrics_dict.items()}
        with open(output, 'w') as f:
            json.dump(output_data, f, indent=2)
        click.echo(f"\nMetrics saved to {output}")


@cli.command()
@click.argument('input_file', type=click.Path(exists=True))
@click.option('--concepts-column', default='woseco_concepts',
              help='Name of column containing concept lists')
@click.option('--rating-column', required=True, help='Name of column containing ratings')
@click.option('--output', '-o', help='Output JSON file for grouped metrics')
@click.option('--graph-type', type=click.Choice(['chain', 'dense']), default='chain',
              help='Type of graph to build')
def analyze_ratings(input_file, concepts_column, rating_column, output, graph_type):
    """
    Analyze metrics grouped by creativity ratings.

    INPUT_FILE: Path to CSV file with woseco_concepts and ratings

    Example:
        woseco analyze-ratings results.csv --rating-column rating_h -o grouped.json
    """
    click.echo(f"Loading data from {input_file}...")
    df = pd.read_csv(input_file)

    # Convert string representations of lists back to actual lists
    import ast
    df[concepts_column] = df[concepts_column].apply(
        lambda x: ast.literal_eval(x) if isinstance(x, str) else x
    )

    # Group by rating
    click.echo(f"Grouping by {rating_column}...")
    grouped = group_by_rating(df, rating_column, concepts_column)

    # Calculate metrics for each rating group
    results = {}
    for rating, concept_lists in grouped.items():
        click.echo(f"  Rating {rating}: {len(concept_lists)} stories")
        metrics_dict = stats(concept_lists, graph_type=graph_type)
        results[str(rating)] = {k: [float(v) for v in vals] for k, vals in metrics_dict.items()}

    # Save to file if requested
    if output:
        with open(output, 'w') as f:
            json.dump(results, f, indent=2)
        click.echo(f"\nResults saved to {output}")


@cli.command()
@click.argument('input_file', type=click.Path(exists=True))
@click.option('--concepts-column', default='woseco_concepts',
              help='Name of column containing concept lists')
@click.option('--prompt-column', default='prompt', help='Name of column containing prompts')
@click.option('--language', '-l', type=click.Choice(['en', 'de']), default='en',
              help='Language of the text')
@click.option('--output', '-o', help='Output CSV file with unexpectedness scores')
def unexpectedness(input_file, concepts_column, prompt_column, language, output):
    """
    Calculate unexpectedness scores (semantic distance from prompts).

    INPUT_FILE: Path to CSV file with concepts and prompts

    Example:
        woseco unexpectedness results.csv --prompt-column prompt -o unexpected.csv
    """
    click.echo(f"Loading data from {input_file}...")
    df = pd.read_csv(input_file)

    # Convert string representations of lists back to actual lists
    import ast
    df[concepts_column] = df[concepts_column].apply(
        lambda x: ast.literal_eval(x) if isinstance(x, str) else x
    )

    click.echo("Calculating unexpectedness scores...")
    scores = unexpectedness_arrays(df, prompt_column, concepts_column, language)

    df['unexpectedness'] = scores

    click.echo(f"\nUnexpectedness statistics:")
    click.echo(f"  Mean: {scores.mean():.4f}")
    click.echo(f"  Std:  {scores.std():.4f}")
    click.echo(f"  Min:  {scores.min():.4f}")
    click.echo(f"  Max:  {scores.max():.4f}")

    # Save to file if requested
    output_path = output if output else input_file.replace('.csv', '_unexpected.csv')
    df.to_csv(output_path, index=False)
    click.echo(f"\nResults saved to {output_path}")


@cli.command()
@click.argument('input_file', type=click.Path(exists=True))
@click.option('--concepts-column', default='woseco_concepts',
              help='Name of column containing concept lists')
@click.option('--graph-type', type=click.Choice(['chain', 'dense']), default='chain',
              help='Type of graph to build')
@click.option('--output', '-o', required=True, help='Output image file')
@click.option('--index', type=int, default=0, help='Index of story to visualize')
def visualize(input_file, concepts_column, graph_type, output, index):
    """
    Visualize a concept network graph.

    INPUT_FILE: Path to CSV file with woseco_concepts column

    Example:
        woseco visualize results.csv -o graph.png --index 0
    """
    click.echo(f"Loading data from {input_file}...")
    df = pd.read_csv(input_file)

    # Convert string representations of lists back to actual lists
    import ast
    df[concepts_column] = df[concepts_column].apply(
        lambda x: ast.literal_eval(x) if isinstance(x, str) else x
    )

    if index >= len(df):
        click.echo(f"Error: Index {index} out of range (0-{len(df)-1})", err=True)
        return

    concepts = df.iloc[index][concepts_column]

    click.echo(f"Building graph for story {index}...")
    graph = build_graph(concepts, graph_type=graph_type)

    click.echo(f"Graph has {graph.number_of_nodes()} nodes and {graph.number_of_edges()} edges")

    click.echo(f"Saving visualization to {output}...")
    plot_graph(graph, save_path=output)

    click.echo("Done!")


@cli.command()
@click.argument('input_file', type=click.Path(exists=True))
@click.option('--rating-columns', '-r', multiple=True, required=True,
              help='Rating columns to compare (can specify multiple)')
@click.option('--concepts-column', default='woseco_concepts',
              help='Name of column containing concept lists')
@click.option('--output-dir', '-o', required=True, help='Output directory for plots')
def compare(input_file, rating_columns, concepts_column, output_dir):
    """
    Compare metrics across multiple raters.

    INPUT_FILE: Path to CSV file with ratings and concepts

    Example:
        woseco compare results.csv -r rating_h -r rating_j -r rating_k -o plots/
    """
    click.echo(f"Loading data from {input_file}...")
    df = pd.read_csv(input_file)

    # Convert string representations of lists back to actual lists
    import ast
    df[concepts_column] = df[concepts_column].apply(
        lambda x: ast.literal_eval(x) if isinstance(x, str) else x
    )

    click.echo(f"Comparing {len(rating_columns)} raters...")
    rater_metrics = compare_raters(df, list(rating_columns), concepts_column)

    # Create output directory
    import os
    os.makedirs(output_dir, exist_ok=True)

    # Plot comparisons for each metric
    from .visualization import plot_metrics_comparison

    metric_names = ['aspl', 'mlcc', 'modularity', 'num_components', 'avg_component_size', 'gcc_size']

    for metric_name in metric_names:
        output_path = os.path.join(output_dir, f"{metric_name}_comparison.png")
        click.echo(f"  Plotting {metric_name}...")
        plot_metrics_comparison(rater_metrics, metric_name, save_path=output_path)

    click.echo(f"\nPlots saved to {output_dir}")


@cli.command()
@click.option('--language', '-l', type=click.Choice(['en', 'de']), required=True,
              help='Language model to download')
@click.option('--size', '-s', type=click.Choice(['sm', 'md', 'lg']), default='lg',
              help='Model size')
def download_model(language, size):
    """
    Download spacy language models.

    Example:
        woseco download-model --language en --size lg
        woseco download-model --language de --size lg
    """
    import spacy

    if language == "en":
        model_name = f"en_core_web_{size}"
    else:
        model_name = f"de_core_news_{size}"

    click.echo(f"Downloading {model_name}...")
    spacy.cli.download(model_name)
    click.echo("Done!")


if __name__ == '__main__':
    cli()
