"""
Comparison animation utilities.

This module provides tools for creating side-by-side comparisons of
different simulations, dimensions, and time snapshots.
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
from matplotlib.gridspec import GridSpec
from typing import List, Optional, Callable, Tuple
from jax import Array
from rich.console import Console

from jaxlatt.core import Lattice
from jaxlatt.observables.energy import energy_components_integrated

console = Console()


def create_comparison_animation(
    times_list: List[Array],
    snapshots_list: List[List[Lattice]],
    labels: List[str],
    potential: Optional[Callable] = None,
    interval: int = 50,
    save_path: Optional[str] = None,
    figsize: Tuple[int, int] = (16, 5),
    title: str = "Field Comparison",
) -> FuncAnimation:
    """
    Create side-by-side comparison animation for multiple simulations.

    Useful for comparing different initial conditions, potentials, or parameters.

    Args:
        times_list: List of time arrays for each simulation
        snapshots_list: List of snapshot lists for each simulation
        labels: Labels for each simulation
        potential: Potential function (for energy computation)
        interval: Delay between frames in milliseconds
        save_path: If provided, save animation to this path
        figsize: Figure size
        title: Animation title

    Returns:
        FuncAnimation object
    """
    n_sims = len(snapshots_list)

    if len(times_list) != n_sims or len(labels) != n_sims:
        raise ValueError("times_list, snapshots_list, and labels must have same length")

    # Verify all are 1D
    for snapshots in snapshots_list:
        if snapshots[0].ndim != 1:
            raise ValueError("Comparison animation currently only supports 1D fields")

    # Setup figure
    fig, axes = plt.subplots(1, n_sims, figsize=figsize, sharey=True)
    if n_sims == 1:
        axes = [axes]

    # Get spatial grid (assume all same)
    lattice = snapshots_list[0][0]
    length = lattice.length[0] if isinstance(lattice.length, tuple) else lattice.length
    size = lattice.size[0] if isinstance(lattice.size, tuple) else lattice.size
    x = np.linspace(0, float(length), int(size), endpoint=False)

    # Determine y-limits
    all_fields = []
    for snapshots in snapshots_list:
        all_fields.extend([np.asarray(snap.field) for snap in snapshots])
    field_min = min(f.min() for f in all_fields)
    field_max = max(f.max() for f in all_fields)
    margin = 0.1 * (field_max - field_min)
    ylim = (field_min - margin, field_max + margin)

    # Initialize plots
    lines = []
    time_texts = []

    for i, (ax, label) in enumerate(zip(axes, labels)):
        (line,) = ax.plot([], [], "b-", linewidth=2)
        lines.append(line)

        ax.set_xlim(0, float(length))
        ax.set_ylim(ylim)
        ax.set_xlabel("x", fontsize=11)
        if i == 0:
            ax.set_ylabel("φ(x)", fontsize=11)
        ax.set_title(label, fontsize=12)
        ax.grid(True, alpha=0.3)

        time_text = ax.text(
            0.02,
            0.95,
            "",
            transform=ax.transAxes,
            fontsize=10,
            verticalalignment="top",
            bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
        )
        time_texts.append(time_text)

    fig.suptitle(title, fontsize=14, fontweight="bold")
    plt.subplots_adjust(left=0.08, right=0.98, top=0.92, bottom=0.08)

    def init():
        """Initialize animation."""
        for line, time_text in zip(lines, time_texts):
            line.set_data([], [])
            time_text.set_text("")
        return lines + time_texts

    def update(frame):
        """Update animation frame."""
        for i, (line, time_text, snapshots, times) in enumerate(
            zip(lines, time_texts, snapshots_list, times_list)
        ):
            if frame < len(snapshots):
                snap = snapshots[frame]
                t = times[frame]
                field = np.asarray(snap.field)
                line.set_data(x, field)
                time_text.set_text(f"t = {t:.3f}")

        return lines + time_texts

    # Use minimum number of frames across all simulations
    max_frames = max(len(snapshots) for snapshots in snapshots_list)

    # Create animation
    anim = FuncAnimation(
        fig,
        update,
        init_func=init,
        frames=max_frames,
        interval=interval,
        blit=True,
        repeat=True,
    )

    # Save if requested
    if save_path is not None:
        console.print(
            f"[cyan]💾 Saving comparison animation to[/cyan] [bold]{save_path}[/bold]..."
        )
        if save_path.endswith(".gif"):
            writer = PillowWriter(fps=1000 // interval)
            anim.save(save_path, writer=writer)
        else:
            anim.save(save_path, fps=1000 // interval)
        console.print("[green]✓ Animation saved![/green]")

    return anim


def animate_multi_snapshot_grid(
    times: Array,
    snapshots: List[Lattice],
    potential: Optional[Callable] = None,
    n_snapshots: int = 6,
    interval: int = 100,
    save_path: Optional[str] = None,
    figsize: Tuple[int, int] = (16, 10),
    title: str = "Field Evolution Snapshots",
    cmap: str = "RdBu_r",
    show_energy: bool = True,
) -> FuncAnimation:
    """
    Create an animation showing a grid of snapshots at different times.

    This provides a "time-lapse" view where you can see multiple moments
    of the evolution simultaneously. Particularly useful for understanding
    the overall dynamics and identifying key transition moments.

    Args:
        times: Array of snapshot times
        snapshots: List of Lattice objects
        potential: Potential function for energy computation
        n_snapshots: Number of time snapshots to show in grid (2, 4, 6, or 9)
        interval: Delay between frames in milliseconds
        save_path: If provided, save animation to this path
        figsize: Figure size
        title: Animation title
        cmap: Colormap for 2D/3D visualizations
        show_energy: Whether to include energy evolution plot

    Returns:
        FuncAnimation object

    Example:
        >>> # Show 6 snapshots evolving together
        >>> anim = animate_multi_snapshot_grid(
        ...     times, snapshots, potential=pot,
        ...     n_snapshots=6, save_path="timelapse.gif"
        ... )
    """
    lattice = snapshots[0]
    ndim = lattice.ndim

    # Determine grid layout
    if n_snapshots == 2:
        nrows, ncols = 1, 2
    elif n_snapshots == 4:
        nrows, ncols = 2, 2
    elif n_snapshots == 6:
        nrows, ncols = 2, 3
    elif n_snapshots == 9:
        nrows, ncols = 3, 3
    else:
        raise ValueError("n_snapshots must be 2, 4, 6, or 9")

    # Create figure
    if show_energy and potential is not None:
        fig = plt.figure(figsize=figsize)
        gs = GridSpec(nrows + 1, ncols, figure=fig, height_ratios=[1] * nrows + [0.4])
        axes_grid = [
            [fig.add_subplot(gs[i, j]) for j in range(ncols)] for i in range(nrows)
        ]
        ax_energy = fig.add_subplot(gs[nrows, :])

        # Compute energy
        energy_data = [energy_components_integrated(s, potential) for s in snapshots]
        E_total = [e["total"] for e in energy_data]
    else:
        fig, axes_grid = plt.subplots(nrows, ncols, figsize=figsize)
        if nrows == 1:
            axes_grid = [axes_grid]
        ax_energy = None

    axes = [ax for row in axes_grid for ax in row]

    # Initialize plots based on dimensionality
    plot_objects = []
    time_texts = []

    if ndim == 1:
        # 1D: Line plots
        length = (
            lattice.length[0] if isinstance(lattice.length, tuple) else lattice.length
        )
        size = lattice.size[0] if isinstance(lattice.size, tuple) else lattice.size
        x = np.linspace(0, float(length), int(size), endpoint=False)

        for ax in axes:
            (line,) = ax.plot([], [], "b-", linewidth=2)
            plot_objects.append(line)
            ax.set_xlim(0, float(length))
            ax.grid(True, alpha=0.3)

            time_text = ax.text(
                0.02,
                0.95,
                "",
                transform=ax.transAxes,
                fontsize=10,
                verticalalignment="top",
                bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
            )
            time_texts.append(time_text)

    elif ndim == 2:
        # 2D: Heatmaps
        extent = [0, lattice.length[0], 0, lattice.length[1]]

        for ax in axes:
            im = ax.imshow(
                np.asarray(snapshots[0].field).T,
                origin="lower",
                cmap=cmap,
                extent=extent,
                animated=True,
            )
            plot_objects.append(im)
            ax.set_aspect("equal")

            time_text = ax.text(
                0.02,
                0.98,
                "",
                transform=ax.transAxes,
                fontsize=10,
                verticalalignment="top",
                color="white",
                fontweight="bold",
                bbox=dict(boxstyle="round", facecolor="black", alpha=0.7),
            )
            time_texts.append(time_text)

    elif ndim == 3:
        # 3D: Central slices
        mid_z = lattice.size[2] // 2
        extent = [0, lattice.length[0], 0, lattice.length[1]]

        for ax in axes:
            im = ax.imshow(
                np.asarray(snapshots[0].field)[:, :, mid_z].T,
                origin="lower",
                cmap=cmap,
                extent=extent,
                animated=True,
            )
            plot_objects.append(im)
            ax.set_aspect("equal")
            ax.set_title("XY slice (z=mid)", fontsize=9)

            time_text = ax.text(
                0.02,
                0.98,
                "",
                transform=ax.transAxes,
                fontsize=10,
                verticalalignment="top",
                color="white",
                fontweight="bold",
                bbox=dict(boxstyle="round", facecolor="black", alpha=0.7),
            )
            time_texts.append(time_text)

    # Energy plot initialization
    energy_marker = None
    if ax_energy is not None:
        (line_E,) = ax_energy.plot([], [], "k-", linewidth=2, label="Total Energy")
        (energy_marker,) = ax_energy.plot([], [], "ro", markersize=8)
        ax_energy.set_xlim(times[0], times[-1])
        ax_energy.set_ylim(min(E_total) * 0.9, max(E_total) * 1.1)
        ax_energy.set_xlabel("Time", fontsize=11)
        ax_energy.set_ylabel("Energy", fontsize=11)
        ax_energy.legend()
        ax_energy.grid(True, alpha=0.3)

    fig.suptitle(title, fontsize=14, fontweight="bold")
    plt.subplots_adjust(
        left=0.05, right=0.98, top=0.92, bottom=0.08, hspace=0.3, wspace=0.3
    )

    def update(frame):
        """Update all snapshot panels."""
        # Calculate which snapshots to show
        total_frames = len(snapshots)
        indices = [
            int(i * (total_frames - 1) / (n_snapshots - 1)) for i in range(n_snapshots)
        ]

        # Update each panel
        for i, (idx, obj, txt) in enumerate(zip(indices, plot_objects, time_texts)):
            snap = snapshots[idx]
            t = times[idx]

            if ndim == 1:
                field = np.asarray(snap.field)
                obj.set_data(x, field)
            elif ndim == 2:
                field = np.asarray(snap.field)
                obj.set_array(field.T)
            elif ndim == 3:
                field = np.asarray(snap.field)
                obj.set_array(field[:, :, mid_z].T)

            txt.set_text(f"t = {t:.2f}")

        # Update energy plot
        if ax_energy is not None:
            current_idx = indices[-1]  # Track the last snapshot's energy
            line_E.set_data(times[: current_idx + 1], E_total[: current_idx + 1])
            energy_marker.set_data([times[current_idx]], [E_total[current_idx]])

        return (
            plot_objects + time_texts + ([line_E, energy_marker] if ax_energy else [])
        )

    # Create animation
    anim = FuncAnimation(
        fig, update, frames=len(snapshots), interval=interval, blit=False, repeat=True
    )

    # Save if requested
    if save_path is not None:
        console.print(
            f"[cyan]💾 Saving multi-snapshot animation to[/cyan] [bold]{save_path}[/bold]..."
        )
        if save_path.endswith(".gif"):
            writer = PillowWriter(fps=1000 // interval)
            anim.save(save_path, writer=writer)
        else:
            anim.save(save_path, fps=1000 // interval)
        console.print("[green]✓ Animation saved![/green]")

    return anim
