"""
Visualization functions for Woseco analysis.
"""
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import networkx as nx
from typing import Dict, List, Optional, Tuple


def plot_network_stats(
    stats_dict: Dict[str, List[float]],
    title: str = "Network Statistics",
    ylabel: str = "Value",
    xlabel: str = "Rating",
    figsize: Tuple[int, int] = (12, 8),
    save_path: Optional[str] = None
) -> plt.Figure:
    """
    Create grouped bar charts for network statistics.

    Args:
        stats_dict: Dictionary mapping metric names to lists of values
        title: Plot title
        ylabel: Y-axis label
        xlabel: X-axis label
        figsize: Figure size (width, height)
        save_path: Optional path to save figure

    Returns:
        Matplotlib Figure object
    """
    # Prepare data
    metric_names = list(stats_dict.keys())
    num_metrics = len(metric_names)
    num_groups = len(stats_dict[metric_names[0]])

    # Create positions for bars
    x = np.arange(num_groups)
    width = 0.8 / num_metrics

    # Create figure
    fig, ax = plt.subplots(figsize=figsize)

    # Plot bars for each metric
    for i, metric_name in enumerate(metric_names):
        values = stats_dict[metric_name]
        offset = (i - num_metrics / 2) * width + width / 2
        ax.bar(x + offset, values, width, label=metric_name)

    # Customize plot
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_xticks(x)
    ax.set_xticklabels([f"{i+1}" for i in range(num_groups)])
    ax.legend()
    ax.grid(axis='y', alpha=0.3)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')

    return fig


def plot_rating_distribution(
    df,
    rating_columns: List[str],
    figsize: Tuple[int, int] = (12, 6),
    save_path: Optional[str] = None
) -> plt.Figure:
    """
    Plot distribution of ratings across multiple raters.

    Args:
        df: DataFrame with rating columns
        rating_columns: List of rating column names
        figsize: Figure size
        save_path: Optional path to save figure

    Returns:
        Matplotlib Figure object
    """
    fig, axes = plt.subplots(1, len(rating_columns), figsize=figsize)

    if len(rating_columns) == 1:
        axes = [axes]

    for ax, col in zip(axes, rating_columns):
        df[col].hist(ax=ax, bins=5, edgecolor='black')
        ax.set_title(f"Distribution: {col}")
        ax.set_xlabel("Rating")
        ax.set_ylabel("Frequency")
        ax.set_xticks([1, 2, 3, 4, 5])

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')

    return fig


def plot_graph(
    graph: nx.Graph,
    layout: str = "spring",
    figsize: Tuple[int, int] = (12, 10),
    node_size: int = 500,
    font_size: int = 10,
    save_path: Optional[str] = None
) -> plt.Figure:
    """
    Visualize a concept network graph.

    Args:
        graph: NetworkX graph
        layout: Layout algorithm ("spring", "circular", "kamada_kawai", "shell")
        figsize: Figure size
        node_size: Size of nodes
        font_size: Font size for labels
        save_path: Optional path to save figure

    Returns:
        Matplotlib Figure object
    """
    fig, ax = plt.subplots(figsize=figsize)

    # Choose layout
    if layout == "spring":
        pos = nx.spring_layout(graph, k=0.5, iterations=50)
    elif layout == "circular":
        pos = nx.circular_layout(graph)
    elif layout == "kamada_kawai":
        pos = nx.kamada_kawai_layout(graph)
    elif layout == "shell":
        pos = nx.shell_layout(graph)
    else:
        pos = nx.spring_layout(graph)

    # Draw graph
    nx.draw_networkx_nodes(graph, pos, node_size=node_size, node_color='lightblue', ax=ax)
    nx.draw_networkx_edges(graph, pos, alpha=0.5, ax=ax)
    nx.draw_networkx_labels(graph, pos, font_size=font_size, ax=ax)

    ax.set_title("Concept Network")
    ax.axis('off')

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')

    return fig


def plot_metrics_comparison(
    rater_metrics: Dict[str, Dict[int, Dict[str, List[float]]]],
    metric_name: str,
    figsize: Tuple[int, int] = (12, 6),
    save_path: Optional[str] = None
) -> plt.Figure:
    """
    Compare a specific metric across raters and ratings.

    Args:
        rater_metrics: Nested dict (rater -> rating -> metrics)
        metric_name: Name of metric to plot
        figsize: Figure size
        save_path: Optional path to save figure

    Returns:
        Matplotlib Figure object
    """
    fig, ax = plt.subplots(figsize=figsize)

    for rater, rating_metrics in rater_metrics.items():
        ratings = sorted(rating_metrics.keys())
        values = []

        for rating in ratings:
            metric_values = rating_metrics[rating][metric_name]
            values.append(np.mean(metric_values))

        ax.plot(ratings, values, marker='o', label=rater, linewidth=2)

    ax.set_xlabel("Rating")
    ax.set_ylabel(metric_name.upper())
    ax.set_title(f"{metric_name.upper()} Comparison Across Raters")
    ax.legend()
    ax.grid(alpha=0.3)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')

    return fig


def plot_pca(
    pca_result: Dict[str, np.ndarray],
    labels: Optional[List[str]] = None,
    figsize: Tuple[int, int] = (10, 8),
    save_path: Optional[str] = None
) -> plt.Figure:
    """
    Plot PCA results.

    Args:
        pca_result: Dictionary from pca_analysis function
        labels: Optional labels for points
        figsize: Figure size
        save_path: Optional path to save figure

    Returns:
        Matplotlib Figure object
    """
    transformed = pca_result["transformed"]
    explained_variance = pca_result["explained_variance"]

    fig, ax = plt.subplots(figsize=figsize)

    if transformed.shape[1] >= 2:
        scatter = ax.scatter(
            transformed[:, 0],
            transformed[:, 1],
            c=range(len(transformed)),
            cmap='viridis',
            s=100,
            alpha=0.6
        )

        if labels:
            for i, label in enumerate(labels):
                ax.annotate(label, (transformed[i, 0], transformed[i, 1]))

        ax.set_xlabel(f"PC1 ({explained_variance[0]:.2%} variance)")
        ax.set_ylabel(f"PC2 ({explained_variance[1]:.2%} variance)")
        ax.set_title("PCA of Network Metrics")
        plt.colorbar(scatter, ax=ax, label="Sample Index")

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')

    return fig


def plot_unexpectedness(
    unexpectedness_scores: np.ndarray,
    ratings: Optional[np.ndarray] = None,
    figsize: Tuple[int, int] = (10, 6),
    save_path: Optional[str] = None
) -> plt.Figure:
    """
    Plot unexpectedness scores.

    Args:
        unexpectedness_scores: Array of unexpectedness scores
        ratings: Optional ratings for coloring
        figsize: Figure size
        save_path: Optional path to save figure

    Returns:
        Matplotlib Figure object
    """
    fig, ax = plt.subplots(figsize=figsize)

    if ratings is not None:
        scatter = ax.scatter(
            range(len(unexpectedness_scores)),
            unexpectedness_scores,
            c=ratings,
            cmap='RdYlGn',
            s=50,
            alpha=0.6
        )
        plt.colorbar(scatter, ax=ax, label="Rating")
    else:
        ax.scatter(
            range(len(unexpectedness_scores)),
            unexpectedness_scores,
            s=50,
            alpha=0.6
        )

    ax.set_xlabel("Story Index")
    ax.set_ylabel("Unexpectedness Score")
    ax.set_title("Story Unexpectedness (1 - Similarity to Prompt)")
    ax.grid(alpha=0.3)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')

    return fig


def plot_heatmap(
    data_matrix: np.ndarray,
    x_labels: Optional[List[str]] = None,
    y_labels: Optional[List[str]] = None,
    title: str = "Heatmap",
    figsize: Tuple[int, int] = (10, 8),
    save_path: Optional[str] = None
) -> plt.Figure:
    """
    Create a heatmap visualization.

    Args:
        data_matrix: 2D array to visualize
        x_labels: Labels for x-axis
        y_labels: Labels for y-axis
        title: Plot title
        figsize: Figure size
        save_path: Optional path to save figure

    Returns:
        Matplotlib Figure object
    """
    fig, ax = plt.subplots(figsize=figsize)

    sns.heatmap(
        data_matrix,
        annot=True,
        fmt=".2f",
        cmap="YlOrRd",
        xticklabels=x_labels,
        yticklabels=y_labels,
        ax=ax
    )

    ax.set_title(title)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')

    return fig
