"""
2D field animation utilities.

This module provides tools for visualizing 2D scalar field dynamics,
including field configurations and phase portraits.
"""

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 field(
    times: Array,
    snapshots: List[Lattice],
    potential: Optional[Callable] = None,
    interval: int = 50,
    save_path: Optional[str] = None,
    figsize: Tuple[int, int] = (5, 5),
    show_energy: bool = False,
    title: Optional[str] = None,
    cmap: str = "RdBu_r",
    vmin: Optional[float] = None,
    vmax: Optional[float] = None,
    dark: bool = False,
) -> FuncAnimation:
    """
    Create an animation of 2D scalar field evolution.

    Args:
        times: Array of snapshot times
        snapshots: List of Lattice objects at each snapshot time
        potential: Potential function (needed if show_energy=True)
        interval: Delay between frames in milliseconds
        save_path: If provided, save animation to this path
        figsize: Figure size (width, height)
        show_energy: Whether to show energy evolution subplot
        title: Animation title
        cmap: Colormap for field visualization
        vmin: Minimum value for colormap (auto if None)
        vmax: Maximum value for colormap (auto if None)
        dark: Whether to use dark mode styling (white text/labels)

    Returns:
        FuncAnimation object
    """
    import jaxlatt.plotting as plot

    # Validate inputs
    if show_energy and potential is None:
        raise ValueError("Must provide potential function if show_energy=True")

    lattice = snapshots[0]

    # Compute energy if needed
    if show_energy:
        energy_data = [
            energy_components_integrated(snap, potential) for snap in snapshots
        ]
        E_total = [e["total"] for e in energy_data]
        E_kinetic = [e["kinetic"] for e in energy_data]
        E_gradient = [e["gradient"] for e in energy_data]
        E_potential = [e["potential"] for e in energy_data]

    # Determine colormap limits
    if vmin is None or vmax is None:
        all_fields = [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)
        if vmin is None:
            vmin = field_min
        if vmax is None:
            vmax = field_max

    # Create figure
    if show_energy:
        fig = plt.figure(figsize=figsize)
        gs = GridSpec(1, 2, width_ratios=[1, 1], figure=fig)
        ax_field = fig.add_subplot(gs[0])
        ax_energy = fig.add_subplot(gs[1])
    else:
        fig, ax_field = plt.subplots(figsize=figsize)

    # Transparent background for dark mode
    if dark:
        fig.patch.set_alpha(0.0)
        ax_field.set_facecolor("none")
        if show_energy:
            ax_energy.set_facecolor("none")

    # # Initialize field plot
    im = plot.snapshot2d(
        snapshots[0], ax=ax_field, vmin=vmin, vmax=vmax, cmap=cmap, dark=dark
    )
    time_text = ax_field.text(
        0.02,
        0.98,
        "",
        transform=ax_field.transAxes,
        fontsize=11,
        verticalalignment="top",
        color="white",
        fontweight="bold",
        bbox=dict(boxstyle="round", facecolor="black", alpha=0.7),
    )

    # Initialize energy plot if needed
    if show_energy:
        (line_E_total,) = ax_energy.plot([], [], "k-", linewidth=2, label="Total")
        (line_E_kin,) = ax_energy.plot(
            [], [], "r-", linewidth=1.5, alpha=0.7, label="Kinetic"
        )
        (line_E_grad,) = ax_energy.plot(
            [], [], "g-", linewidth=1.5, alpha=0.7, label="Gradient"
        )
        (line_E_pot,) = ax_energy.plot(
            [], [], "b-", linewidth=1.5, alpha=0.7, label="Potential"
        )

        ax_energy.set_xlim(times[0], times[-1])
        E_min = min(min(E_kinetic), min(E_gradient), min(E_potential)) * 0.9
        E_max = max(E_total) * 1.1
        ax_energy.set_ylim(E_min, E_max)
        ax_energy.set_xlabel("Time", fontsize=11)
        ax_energy.set_ylabel("Energy", fontsize=11)
        ax_energy.legend(loc="upper right", fontsize=9)
        ax_energy.grid(True, alpha=0.3)

        (energy_marker,) = ax_energy.plot([], [], "ro", markersize=8)

        if dark:
            plot.apply_dark_style(ax_energy)

    title_color = "white" if dark else "black"
    fig.suptitle(title, fontsize=14, fontweight="bold", color=title_color)
    # plt.subplots_adjust(left=0.1, right=0.95, top=0.92, bottom=0.1)

    def init():
        """Initialize animation."""
        im.set_array(np.asarray(snapshots[0].field).T)
        time_text.set_text("")

        if show_energy:
            line_E_total.set_data([], [])
            line_E_kin.set_data([], [])
            line_E_grad.set_data([], [])
            line_E_pot.set_data([], [])
            energy_marker.set_data([], [])
            return (
                im,
                time_text,
                line_E_total,
                line_E_kin,
                line_E_grad,
                line_E_pot,
                energy_marker,
            )

        return im, time_text

    def update(frame):
        """Update animation frame."""
        snap = snapshots[frame]
        t = times[frame]

        # Update field
        field = np.asarray(snap.field)
        im.set_array(field.T)
        time_text.set_text(f"t = {t:.3f}")

        # Update energy plot if needed
        if show_energy:
            line_E_total.set_data(times[: frame + 1], E_total[: frame + 1])
            line_E_kin.set_data(times[: frame + 1], E_kinetic[: frame + 1])
            line_E_grad.set_data(times[: frame + 1], E_gradient[: frame + 1])
            line_E_pot.set_data(times[: frame + 1], E_potential[: frame + 1])
            energy_marker.set_data([t], [E_total[frame]])

            return (
                im,
                time_text,
                line_E_total,
                line_E_kin,
                line_E_grad,
                line_E_pot,
                energy_marker,
            )

        return im, time_text

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

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

    return anim
