"""
Segment visualization utilities.
"""

from __future__ import annotations

import numpy as np

from sltk.data.core import SegmentList


def plot_segments(
    segments: SegmentList,
    ax=None,
    figsize=(14, 4),
    colors: dict[str, str] | None = None,
    show_labels: bool = True,
    label_rotation: int = 45,
    title: str | None = None,
    tier_height: float = 0.8,
):
    """
    Plot segments as a timeline.

    Args:
        segments: SegmentList to visualize
        ax: Matplotlib axes
        figsize: Figure size
        colors: Optional color mapping by tier or label
        show_labels: Whether to show segment labels
        label_rotation: Rotation angle for labels
        title: Plot title
        tier_height: Height of each tier row

    Returns:
        Matplotlib axes
    """
    try:
        import matplotlib.patches as patches
        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)

    # Group by tier
    tiers = segments.tiers
    tier_y = {tier: i for i, tier in enumerate(tiers)}

    # Default colors
    default_colors = plt.cm.Set3(np.linspace(0, 1, len(tiers)))  # type: ignore[attr-defined]
    if colors is None:
        colors = {tier: default_colors[i] for i, tier in enumerate(tiers)}

    # Plot each segment
    for seg in segments:
        y = tier_y[seg.tier]
        color = colors.get(seg.tier or "", colors.get(seg.label or "", "steelblue"))

        rect = patches.Rectangle(
            (seg.start, y - tier_height / 2),
            seg.duration,
            tier_height,
            linewidth=1,
            edgecolor="black",
            facecolor=color,
            alpha=0.7,
        )
        ax.add_patch(rect)

        # Add label
        if show_labels and seg.label:
            ax.text(
                seg.midpoint,
                y,
                seg.label,
                ha="center",
                va="center",
                fontsize=8,
                rotation=label_rotation,
            )

    # Configure axes
    ax.set_xlim(segments.span[0] - 0.5, segments.span[1] + 0.5)
    ax.set_ylim(-0.5, len(tiers) - 0.5)
    ax.set_yticks(range(len(tiers)))
    ax.set_yticklabels(tiers)
    ax.set_xlabel("Time (seconds)")

    if title:
        ax.set_title(title)

    return ax


def plot_segment_comparison(
    reference: SegmentList,
    prediction: SegmentList,
    ax=None,
    figsize=(14, 6),
    title: str | None = None,
):
    """
    Compare reference and predicted segments.

    Args:
        reference: Ground truth segments
        prediction: Predicted segments
        ax: Matplotlib axes
        figsize: Figure size
        title: Plot title

    Returns:
        Matplotlib axes
    """
    try:
        import matplotlib.patches as patches
        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)

    # Plot reference (top row)
    for seg in reference:
        rect = patches.Rectangle(
            (seg.start, 1.5),
            seg.duration,
            0.8,
            linewidth=1,
            edgecolor="black",
            facecolor="green",
            alpha=0.6,
            label="Reference" if seg == reference[0] else None,
        )
        ax.add_patch(rect)
        if seg.label:
            ax.text(seg.midpoint, 1.9, seg.label, ha="center", fontsize=7, rotation=45)

    # Plot prediction (bottom row)
    for seg in prediction:
        rect = patches.Rectangle(
            (seg.start, 0.5),
            seg.duration,
            0.8,
            linewidth=1,
            edgecolor="black",
            facecolor="blue",
            alpha=0.6,
            label="Prediction" if seg == prediction[0] else None,
        )
        ax.add_patch(rect)
        if seg.label:
            ax.text(seg.midpoint, 0.9, seg.label, ha="center", fontsize=7, rotation=45)

    # Mark overlaps
    for ref in reference:
        for pred in prediction:
            if ref.overlaps(pred):
                inter = ref.intersection(pred)
                if inter:
                    rect = patches.Rectangle(
                        (inter.start, 1.0),
                        inter.duration,
                        0.3,
                        linewidth=0,
                        facecolor="yellow",
                        alpha=0.8,
                    )
                    ax.add_patch(rect)

    # Configure axes
    all_segs = list(reference) + list(prediction)
    if all_segs:
        min_t = min(s.start for s in all_segs)
        max_t = max(s.end for s in all_segs)
        ax.set_xlim(min_t - 0.5, max_t + 0.5)

    ax.set_ylim(0, 2.5)
    ax.set_yticks([0.9, 1.9])
    ax.set_yticklabels(["Prediction", "Reference"])
    ax.set_xlabel("Time (seconds)")

    ax.legend(loc="upper right")

    if title:
        ax.set_title(title)

    return ax


def plot_segment_timeline(
    segments_dict: dict[str, SegmentList],
    ax=None,
    figsize=(14, 6),
    title: str | None = None,
):
    """
    Plot multiple segment lists as a stacked timeline.

    Args:
        segments_dict: Dictionary mapping names to SegmentLists
        ax: Matplotlib axes
        figsize: Figure size
        title: Plot title

    Returns:
        Matplotlib axes
    """
    try:
        import matplotlib.patches as patches
        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)

    names = list(segments_dict.keys())
    colors = plt.cm.tab10(np.linspace(0, 1, len(names)))  # type: ignore[attr-defined]

    for i, (name, segments) in enumerate(segments_dict.items()):
        for seg in segments:
            rect = patches.Rectangle(
                (seg.start, i + 0.1),
                seg.duration,
                0.8,
                linewidth=1,
                edgecolor="black",
                facecolor=colors[i],
                alpha=0.7,
            )
            ax.add_patch(rect)

    # Get time bounds
    all_segs = []
    for segs in segments_dict.values():
        all_segs.extend(list(segs))

    if all_segs:
        min_t = min(s.start for s in all_segs)
        max_t = max(s.end for s in all_segs)
        ax.set_xlim(min_t - 0.5, max_t + 0.5)

    ax.set_ylim(0, len(names))
    ax.set_yticks([i + 0.5 for i in range(len(names))])
    ax.set_yticklabels(names)
    ax.set_xlabel("Time (seconds)")

    if title:
        ax.set_title(title)

    return ax


# Alias for plot_segments
plot_timeline = plot_segments


def get_tier_colors(tiers: list[str], colormap: str = "Set3") -> list[str]:
    """
    Get distinct colors for a list of tiers.

    Args:
        tiers: List of tier names
        colormap: Matplotlib colormap name

    Returns:
        List of hex color strings
    """
    try:
        import matplotlib.colormaps as colormaps
        import matplotlib.colors as mcolors
    except ImportError:
        # Return basic colors without matplotlib
        basic_colors = [
            "#1f77b4",
            "#ff7f0e",
            "#2ca02c",
            "#d62728",
            "#9467bd",
            "#8c564b",
            "#e377c2",
            "#7f7f7f",
            "#bcbd22",
            "#17becf",
        ]
        return [basic_colors[i % len(basic_colors)] for i in range(len(tiers))]

    cmap = colormaps[colormap]
    colors = []
    for i, _ in enumerate(tiers):
        rgba = cmap(i / max(len(tiers) - 1, 1))
        colors.append(mcolors.to_hex(rgba))

    return colors
