"""
Pose visualization utilities.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np

from sltk.data.core import PoseSequence

# Standard skeleton connections for visualization
SKELETON_CONNECTIONS = {
    "smplx": [
        # Body
        (0, 1),
        (0, 2),
        (1, 4),
        (2, 5),
        (4, 7),
        (5, 8),  # Legs
        (0, 3),
        (3, 6),
        (6, 9),
        (9, 12),
        (12, 15),  # Spine to head
        (9, 13),
        (13, 16),
        (16, 18),
        (18, 20),  # Left arm
        (9, 14),
        (14, 17),
        (17, 19),
        (19, 21),  # Right arm
    ],
    "mediapipe": [
        # Face
        (0, 1),
        (1, 2),
        (2, 3),
        (3, 7),
        (0, 4),
        (4, 5),
        (5, 6),
        (6, 8),
        # Body
        (11, 12),
        (11, 13),
        (12, 14),
        (13, 15),
        (14, 16),
        (11, 23),
        (12, 24),
        (23, 24),
        (23, 25),
        (24, 26),
        (25, 27),
        (26, 28),
        (27, 29),
        (28, 30),
        (29, 31),
        (30, 32),
        # Arms
        (15, 17),
        (15, 19),
        (15, 21),
        (17, 19),
        (16, 18),
        (16, 20),
        (16, 22),
        (18, 20),
    ],
    "wilor": [
        # Hand connections (for each hand)
        (0, 1),
        (1, 2),
        (2, 3),
        (3, 4),  # Thumb
        (0, 5),
        (5, 6),
        (6, 7),
        (7, 8),  # Index
        (0, 9),
        (9, 10),
        (10, 11),
        (11, 12),  # Middle
        (0, 13),
        (13, 14),
        (14, 15),
        (15, 16),  # Ring
        (0, 17),
        (17, 18),
        (18, 19),
        (19, 20),  # Pinky
        (5, 9),
        (9, 13),
        (13, 17),  # Palm connections
    ],
}


def plot_pose_2d(
    poses: PoseSequence | np.ndarray,
    frame: int = 0,
    ax=None,
    connections: list[tuple[int, int]] | None = None,
    point_size: int = 50,
    line_width: float = 2.0,
    color: str = "blue",
    show_indices: bool = False,
    title: str | None = None,
):
    """
    Plot a single frame of 2D poses.

    Args:
        poses: PoseSequence or numpy array
        frame: Frame index to plot
        ax: Matplotlib axes (creates new if None)
        connections: Skeleton connections to draw
        point_size: Size of keypoint markers
        line_width: Width of connection lines
        color: Color for plotting
        show_indices: Whether to show keypoint indices
        title: Optional plot title

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

    if isinstance(poses, PoseSequence):
        data = poses.data[frame]
        pose_format = poses.format
    else:
        data = poses[frame] if poses.ndim == 3 else poses
        pose_format = None

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

    # Get connections
    if connections is None and pose_format:
        connections = SKELETON_CONNECTIONS.get(pose_format, [])

    # Plot keypoints
    x, y = data[:, 0], data[:, 1]
    ax.scatter(x, y, s=point_size, c=color, zorder=3)

    # Plot connections
    for i, j in connections or []:
        if i < len(data) and j < len(data):
            ax.plot([x[i], x[j]], [y[i], y[j]], c=color, linewidth=line_width, zorder=2)

    # Show indices
    if show_indices:
        for idx, (xi, yi) in enumerate(zip(x, y)):
            ax.annotate(str(idx), (xi, yi), fontsize=8, ha="center", va="bottom")

    ax.set_aspect("equal")
    ax.invert_yaxis()  # Image coordinates

    if title:
        ax.set_title(title)

    return ax


def plot_pose_sequence(
    poses: PoseSequence,
    frames: list[int] | None = None,
    num_frames: int = 8,
    figsize: tuple[int, int] = (16, 4),
    connections: list[tuple[int, int]] | None = None,
    title: str | None = None,
):
    """
    Plot multiple frames from a pose sequence.

    Args:
        poses: PoseSequence to visualize
        frames: Specific frame indices (auto-selects if None)
        num_frames: Number of frames to plot
        figsize: Figure size
        connections: Skeleton connections
        title: Overall title

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

    if frames is None:
        # Evenly space frames
        total = poses.num_frames
        frames = [int(i * total / num_frames) for i in range(num_frames)]
        frames = [min(f, total - 1) for f in frames]

    fig, axes = plt.subplots(1, len(frames), figsize=figsize)
    if len(frames) == 1:
        axes = [axes]

    for ax, frame in zip(axes, frames):
        plot_pose_2d(poses, frame=frame, ax=ax, connections=connections, title=f"Frame {frame}")

    if title:
        fig.suptitle(title)

    plt.tight_layout()
    return fig


def plot_hand_keypoints(
    poses: PoseSequence | np.ndarray,
    frame: int = 0,
    hand: str = "right",
    ax=None,
    figsize: tuple[int, int] = (6, 6),
    color: str = "blue",
    show_indices: bool = True,
):
    """
    Plot hand keypoints with finger labels.

    Args:
        poses: Hand pose data
        frame: Frame index
        hand: 'left' or 'right'
        ax: Matplotlib axes
        figsize: Figure size
        color: Plot color
        show_indices: Whether to show keypoint indices

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

    connections = SKELETON_CONNECTIONS.get("wilor", [])

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

    ax = plot_pose_2d(
        poses,
        frame=frame,
        ax=ax,
        connections=connections,
        color=color,
        show_indices=show_indices,
        title=f"{hand.title()} Hand",
    )

    # Add finger labels
    finger_names = ["Wrist", "Thumb", "Index", "Middle", "Ring", "Pinky"]
    finger_tips = [0, 4, 8, 12, 16, 20]

    if isinstance(poses, PoseSequence):
        data = poses.data[frame]
    else:
        data = poses[frame] if poses.ndim == 3 else poses

    for name, tip_idx in zip(finger_names, finger_tips):
        if tip_idx < len(data):
            ax.annotate(
                name,
                (data[tip_idx, 0], data[tip_idx, 1]),
                fontsize=8,
                ha="center",
                va="top",
                color="red",
            )

    return ax


def create_pose_video(
    poses: PoseSequence,
    output_path: str | Path,
    fps: float | None = None,
    connections: list[tuple[int, int]] | None = None,
    figsize: tuple[int, int] = (8, 8),
    dpi: int = 100,
):
    """
    Create a video animation from pose sequence.

    Args:
        poses: PoseSequence to animate
        output_path: Output video file path
        fps: Output FPS (uses pose fps if None)
        connections: Skeleton connections
        figsize: Figure size
        dpi: Figure DPI

    Returns:
        Path to created video
    """
    try:
        import matplotlib.pyplot as plt
        from matplotlib.animation import FFMpegWriter, FuncAnimation
    except ImportError:
        raise ImportError("matplotlib required: pip install sltk[vis]")

    output_path = Path(output_path)
    fps = fps or poses.fps

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

    # Get data bounds for consistent axes
    data = poses.data
    x_min, x_max = data[:, :, 0].min(), data[:, :, 0].max()
    y_min, y_max = data[:, :, 1].min(), data[:, :, 1].max()
    padding = max(x_max - x_min, y_max - y_min) * 0.1

    def update(frame):
        ax.clear()
        plot_pose_2d(poses, frame=frame, ax=ax, connections=connections)
        ax.set_xlim(x_min - padding, x_max + padding)
        ax.set_ylim(y_max + padding, y_min - padding)  # Inverted
        ax.set_title(f"Frame {frame} / {poses.num_frames}")
        return []

    anim = FuncAnimation(fig, update, frames=poses.num_frames, blit=True)

    writer = FFMpegWriter(fps=int(fps), bitrate=1800)
    anim.save(str(output_path), writer=writer, dpi=dpi)

    plt.close(fig)
    return output_path


# Color schemes for visualization
COLOR_SCHEMES = {
    "default": {
        "body": "#3498db",
        "left_hand": "#e74c3c",
        "right_hand": "#2ecc71",
        "face": "#9b59b6",
    },
    "colorblind": {
        "body": "#0072B2",
        "left_hand": "#D55E00",
        "right_hand": "#009E73",
        "face": "#CC79A7",
    },
}


def prepare_animation_data(
    poses: PoseSequence,
    step: int = 1,
    start: int = 0,
    end: int | None = None,
) -> list[np.ndarray]:
    """
    Prepare pose data for animation.

    Args:
        poses: PoseSequence to animate
        step: Frame step size
        start: Starting frame
        end: Ending frame (inclusive)

    Returns:
        List of frame data arrays
    """
    end = end if end is not None else poses.num_frames
    frames = []

    for i in range(start, end, step):
        frames.append(poses.data[i].copy())

    return frames
