"""
Analysis visualization utilities.
"""

from __future__ import annotations

from typing import Any

import numpy as np

from sltk.analysis.clustering import ClusteringResult
from sltk.analysis.datasets import LexicalStatistics


def plot_vocabulary_distribution(
    stats: LexicalStatistics | dict[str, int],
    ax=None,
    figsize=(12, 6),
    top_n: int = 30,
    log_scale: bool = True,
    title: str | None = None,
):
    """
    Plot vocabulary frequency distribution.

    Args:
        stats: LexicalStatistics or frequency dictionary
        ax: Matplotlib axes
        figsize: Figure size
        top_n: Number of top items to show
        log_scale: Whether to use log scale for y-axis
        title: Plot title

    Returns:
        Matplotlib axes
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        raise ImportError("matplotlib required: pip install sltk[vis]")

    if isinstance(stats, LexicalStatistics):
        freq = stats.frequency_distribution
    else:
        freq = stats

    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=figsize)

    # Sort by frequency
    sorted_items = sorted(freq.items(), key=lambda x: x[1], reverse=True)
    top_items = sorted_items[:top_n]

    labels, counts = zip(*top_items) if top_items else ([], [])

    ax.bar(range(len(labels)), counts, color="steelblue", alpha=0.7)
    ax.set_xticks(range(len(labels)))
    ax.set_xticklabels(labels, rotation=45, ha="right", fontsize=8)

    if log_scale:
        ax.set_yscale("log")
        ax.set_ylabel("Frequency (log scale)")
    else:
        ax.set_ylabel("Frequency")

    ax.set_xlabel("Gloss")
    ax.set_title(title or f"Top {top_n} Vocabulary Items")

    plt.tight_layout()
    return ax


def plot_zipf_distribution(
    stats: LexicalStatistics | dict[str, int],
    ax=None,
    figsize=(8, 6),
    title: str | None = None,
):
    """
    Plot Zipf's law distribution (log-log rank vs frequency).

    Args:
        stats: LexicalStatistics or frequency dictionary
        ax: Matplotlib axes
        figsize: Figure size
        title: Plot title

    Returns:
        Matplotlib axes
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        raise ImportError("matplotlib required: pip install sltk[vis]")

    if isinstance(stats, LexicalStatistics):
        freq = stats.frequency_distribution
        zipf_coef = stats.zipf_coefficient
    else:
        freq = stats
        zipf_coef = None

    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=figsize)

    # Sort by frequency
    sorted_freq = sorted(freq.values(), reverse=True)
    ranks = np.arange(1, len(sorted_freq) + 1)

    ax.loglog(ranks, sorted_freq, "b.", alpha=0.5, markersize=3)

    # Add Zipf reference line if coefficient available
    if zipf_coef is not None:
        max_freq = sorted_freq[0]
        reference = max_freq * (ranks**-zipf_coef)
        ax.loglog(ranks, reference, "r-", alpha=0.7, label=f"Zipf (α={zipf_coef:.2f})")
        ax.legend()

    ax.set_xlabel("Rank")
    ax.set_ylabel("Frequency")
    ax.set_title(title or "Zipf's Law Distribution")

    return ax


def plot_embedding_clusters(
    result: ClusteringResult,
    ax=None,
    figsize=(10, 8),
    point_size: int = 20,
    show_legend: bool = True,
    title: str | None = None,
):
    """
    Plot embedding clusters from ClusteringResult.

    Args:
        result: ClusteringResult with projections and cluster IDs
        ax: Matplotlib axes
        figsize: Figure size
        point_size: Size of scatter points
        show_legend: Whether to show cluster legend
        title: Plot title

    Returns:
        Matplotlib axes
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        raise ImportError("matplotlib required: pip install sltk[vis]")

    projections = result.projections
    cluster_ids = result.cluster_ids

    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=figsize)

    # Get unique clusters
    unique_clusters = np.unique(cluster_ids)
    colors = plt.cm.tab20(np.linspace(0, 1, len(unique_clusters)))  # type: ignore[attr-defined]

    for i, cluster_id in enumerate(unique_clusters):
        mask = cluster_ids == cluster_id
        label = f"Cluster {cluster_id}" if cluster_id >= 0 else "Noise"
        color = "gray" if cluster_id < 0 else colors[i]
        alpha = 0.3 if cluster_id < 0 else 0.7

        ax.scatter(
            projections[mask, 0],
            projections[mask, 1],
            c=[color],
            s=point_size,
            alpha=alpha,
            label=label,
        )

    ax.set_xlabel("Dimension 1")
    ax.set_ylabel("Dimension 2")

    if show_legend and len(unique_clusters) <= 15:
        ax.legend(loc="best", fontsize=8)

    ax.set_title(title or f"Embedding Clusters ({result.num_clusters} clusters)")

    return ax


def plot_duration_histogram(
    durations: list[float],
    ax=None,
    figsize=(10, 6),
    bins: int = 50,
    title: str | None = None,
):
    """
    Plot histogram of segment/sample durations.

    Args:
        durations: List of durations in seconds
        ax: Matplotlib axes
        figsize: Figure size
        bins: Number of histogram bins
        title: Plot title

    Returns:
        Matplotlib axes
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        raise ImportError("matplotlib required: pip install sltk[vis]")

    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=figsize)

    ax.hist(durations, bins=bins, color="steelblue", alpha=0.7, edgecolor="black")

    # Add statistics
    mean_dur = np.mean(durations)
    median_dur = np.median(durations)
    std_dur = np.std(durations)

    ax.axvline(mean_dur, color="red", linestyle="--", label=f"Mean: {mean_dur:.2f}s")
    ax.axvline(median_dur, color="green", linestyle="--", label=f"Median: {median_dur:.2f}s")

    ax.set_xlabel("Duration (seconds)")
    ax.set_ylabel("Count")
    ax.legend()

    stats_text = f"N={len(durations)}, Std={std_dur:.2f}s"
    ax.text(0.95, 0.95, stats_text, transform=ax.transAxes, ha="right", va="top", fontsize=10)

    ax.set_title(title or "Duration Distribution")

    return ax


def plot_vocabulary_growth(
    growth_data: dict[str, list],
    ax=None,
    figsize=(10, 6),
    title: str | None = None,
):
    """
    Plot vocabulary growth over samples.

    Args:
        growth_data: Dictionary with 'sample_indices', 'vocabulary_sizes', 'hapax_counts'
        ax: Matplotlib axes
        figsize: Figure size
        title: Plot title

    Returns:
        Matplotlib axes
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        raise ImportError("matplotlib required: pip install sltk[vis]")

    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=figsize)

    # Accept either 'sample_indices' or 'sample_counts' as key
    indices = growth_data.get("sample_indices", growth_data.get("sample_counts", []))
    vocab_sizes = growth_data["vocabulary_sizes"]
    hapax_counts = growth_data.get("hapax_counts", [])

    ax.plot(indices, vocab_sizes, "b-", label="Vocabulary Size", linewidth=2)

    if hapax_counts:
        ax2 = ax.twinx()
        ax2.plot(indices, hapax_counts, "r--", label="Hapax Count", linewidth=2)
        ax2.set_ylabel("Hapax Count", color="red")
        ax2.tick_params(axis="y", labelcolor="red")

    ax.set_xlabel("Number of Samples")
    ax.set_ylabel("Vocabulary Size", color="blue")
    ax.tick_params(axis="y", labelcolor="blue")
    ax.legend(loc="upper left")

    if hapax_counts:
        ax2.legend(loc="upper right")

    ax.set_title(title or "Vocabulary Growth")

    return ax


def plot_pos_distribution(
    pos_stats: dict[str, Any],
    ax=None,
    figsize=(10, 6),
    title: str | None = None,
):
    """
    Plot parts of speech distribution.

    Args:
        pos_stats: Output from pos_analysis()
        ax: Matplotlib axes
        figsize: Figure size
        title: Plot title

    Returns:
        Matplotlib axes
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        raise ImportError("matplotlib required: pip install sltk[vis]")

    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=figsize)

    distribution = pos_stats["distribution"]

    # Sort by count
    sorted_pos = sorted(distribution.items(), key=lambda x: x[1]["count"], reverse=True)
    labels = [p[0] for p in sorted_pos]
    counts = [p[1]["count"] for p in sorted_pos]

    bars = ax.bar(labels, counts, color="steelblue", alpha=0.7)

    # Add percentage labels
    total = sum(counts)
    for bar, count in zip(bars, counts):
        percentage = count / total * 100
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height(),
            f"{percentage:.1f}%",
            ha="center",
            va="bottom",
            fontsize=8,
        )

    ax.set_xlabel("Part of Speech")
    ax.set_ylabel("Count")
    ax.set_title(title or "POS Distribution")

    plt.xticks(rotation=45, ha="right")
    plt.tight_layout()

    return ax


# Alias for plot_vocabulary_distribution
plot_frequency_distribution = plot_vocabulary_distribution


def plot_confusion_matrix(
    matrix: np.ndarray,
    labels: list[str],
    ax=None,
    figsize=(10, 8),
    cmap: str = "Blues",
    title: str | None = None,
    normalize: bool = False,
):
    """
    Plot a confusion matrix.

    Args:
        matrix: Confusion matrix array (N x N)
        labels: Class labels
        ax: Matplotlib axes
        figsize: Figure size
        cmap: Colormap name
        title: Plot title
        normalize: Whether to normalize rows to sum to 1

    Returns:
        Matplotlib axes
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        raise ImportError("matplotlib required: pip install sltk[vis]")

    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=figsize)

    if normalize:
        matrix = matrix.astype(float) / matrix.sum(axis=1, keepdims=True)
        matrix = np.nan_to_num(matrix)

    im = ax.imshow(matrix, cmap=cmap, aspect="auto")

    # Add colorbar
    plt.colorbar(im, ax=ax)

    # Add labels
    ax.set_xticks(range(len(labels)))
    ax.set_yticks(range(len(labels)))
    ax.set_xticklabels(labels, rotation=45, ha="right", fontsize=8)
    ax.set_yticklabels(labels, fontsize=8)

    # Add text annotations
    fmt = ".2f" if normalize else "d"
    thresh = matrix.max() / 2.0
    for i in range(len(labels)):
        for j in range(len(labels)):
            value = matrix[i, j]
            color = "white" if value > thresh else "black"
            text = format(value, fmt)
            ax.text(j, i, text, ha="center", va="center", color=color, fontsize=8)

    ax.set_xlabel("Predicted")
    ax.set_ylabel("True")
    ax.set_title(title or "Confusion Matrix")

    plt.tight_layout()
    return ax
