"""
1D field animation utilities.

This module provides tools for visualizing 1D scalar field dynamics,
including field configurations and phase space trajectories.
"""

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] = (14, 5),
    show_energy: bool = True,
    title: str = "1D Scalar Field Evolution",
    ylim: Optional[Tuple[float, float]] = None,
) -> FuncAnimation:
    """
    Create an animation of 1D 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 (e.g., 'anim.gif')
        figsize: Figure size (width, height)
        show_energy: Whether to show energy evolution subplot
        title: Animation title
        ylim: Y-axis limits for field plot (auto if None)

    Returns:
        FuncAnimation object

    Example:
        >>> from jaxlatt.potentials import quadratic_potential
        >>> anim = animate_1d_field(times, snapshots, quadratic_potential(m=1.0))
        >>> # In Jupyter: from IPython.display import HTML; HTML(anim.to_jshtml())
    """
    # Validate inputs
    if show_energy and potential is None:
        raise ValueError("Must provide potential function if show_energy=True")

    lattice = snapshots[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)

    # Predefine energy arrays/handles for type checking
    E_total = E_kinetic = E_gradient = E_potential = []  # type: ignore
    line_E_total = line_E_kin = line_E_grad = line_E_pot = energy_marker = None  # type: ignore
    ax_energy = None  # type: ignore

    if show_energy:
        assert potential is not None
        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 y-limits for field
    if ylim is None:
        all_fields = [np.asarray(snap.field) for snap in snapshots]
        fmin = min(f.min() for f in all_fields)
        fmax = max(f.max() for f in all_fields)
        margin = 0.1 * (fmax - fmin)
        ylim = (fmin - margin, fmax + margin)

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

    (line_field,) = ax_field.plot([], [], "b-", linewidth=2, label="φ(x)")
    (line_velocity,) = ax_field.plot(
        [], [], "r--", linewidth=1.5, alpha=0.7, label="φ̇(x)"
    )
    ax_field.set_xlim(0, float(length))
    ax_field.set_ylim(ylim)
    ax_field.set_xlabel("x", fontsize=11)
    ax_field.set_ylabel("Field value", fontsize=11)
    ax_field.grid(True, alpha=0.3)
    ax_field.legend(loc="upper right")
    time_text = ax_field.text(
        0.02,
        0.95,
        "",
        transform=ax_field.transAxes,
        fontsize=11,
        verticalalignment="top",
        bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
    )

    if show_energy:
        (line_E_total,) = ax_energy.plot([], [], "k-", linewidth=2, label="Total")  # type: ignore
        (line_E_kin,) = ax_energy.plot(
            [], [], "r-", linewidth=1.5, alpha=0.7, label="Kinetic"
        )  # type: ignore
        (line_E_grad,) = ax_energy.plot(
            [], [], "g-", linewidth=1.5, alpha=0.7, label="Gradient"
        )  # type: ignore
        (line_E_pot,) = ax_energy.plot(
            [], [], "b-", linewidth=1.5, alpha=0.7, label="Potential"
        )  # type: ignore
        ax_energy.set_xlim(float(times[0]), float(times[-1]))  # type: ignore
        Emin = min(min(E_kinetic), min(E_gradient), min(E_potential)) * 0.9
        Emax = max(E_total) * 1.1
        ax_energy.set_ylim(Emin, Emax)  # type: ignore
        ax_energy.set_xlabel("Time", fontsize=11)  # type: ignore
        ax_energy.set_ylabel("Energy", fontsize=11)  # type: ignore
        ax_energy.legend(loc="upper right", fontsize=9)  # type: ignore
        ax_energy.grid(True, alpha=0.3)  # type: ignore
        (energy_marker,) = ax_energy.plot([], [], "ro", markersize=8)  # type: ignore

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

    def init():
        """Initialize animation."""
        line_field.set_data([], [])
        line_velocity.set_data([], [])
        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 (
                line_field,
                line_velocity,
                time_text,
                line_E_total,
                line_E_kin,
                line_E_grad,
                line_E_pot,
                energy_marker,
            )

        return line_field, line_velocity, time_text

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

        # Update field plot
        field = np.asarray(snap.field)
        field_dot = np.asarray(snap.field_dot)
        line_field.set_data(x, field)
        line_velocity.set_data(x, field_dot)
        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 (
                line_field,
                line_velocity,
                time_text,
                line_E_total,
                line_E_kin,
                line_E_grad,
                line_E_pot,
                energy_marker,
            )

        return line_field, line_velocity, 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 phase space 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 phase_space(
    times: Array,
    snapshots: List[Lattice],
    x_index: Optional[int] = None,
    interval: int = 50,
    save_path: Optional[str] = None,
    figsize: Tuple[int, int] = (10, 8),
    title: str = "Phase Space Evolution",
) -> FuncAnimation:
    """
    Create a phase space animation showing φ vs φ̇ at a specific spatial point.

    Useful for visualizing oscillatory dynamics and energy exchange.

    Args:
        times: Array of snapshot times
        snapshots: List of Lattice objects
        x_index: Spatial index to track (center if None)
        interval: Delay between frames in milliseconds
        save_path: If provided, save animation to this path
        figsize: Figure size
        title: Animation title

    Returns:
        FuncAnimation object
    """
    lattice = snapshots[0]

    # Default to center of lattice
    if x_index is None:
        x_index = lattice.size[0] // 2

    # Extract trajectory in phase space
    field_vals = [
        float(np.asarray(snap.field).flatten()[x_index]) for snap in snapshots
    ]
    field_dot_vals = [
        float(np.asarray(snap.field_dot).flatten()[x_index]) for snap in snapshots
    ]

    # Setup figure
    fig, (ax_phase, ax_time) = plt.subplots(2, 1, figsize=figsize)

    # Phase space plot
    (line_phase,) = ax_phase.plot([], [], "b-", linewidth=2, alpha=0.7)
    (point_phase,) = ax_phase.plot([], [], "ro", markersize=10)

    phi_min, phi_max = min(field_vals), max(field_vals)
    phi_dot_min, phi_dot_max = min(field_dot_vals), max(field_dot_vals)
    phi_margin = 0.1 * (phi_max - phi_min)
    phi_dot_margin = 0.1 * (phi_dot_max - phi_dot_min)

    ax_phase.set_xlim(phi_min - phi_margin, phi_max + phi_margin)
    ax_phase.set_ylim(phi_dot_min - phi_dot_margin, phi_dot_max + phi_dot_margin)
    ax_phase.set_xlabel("φ", fontsize=12)
    ax_phase.set_ylabel("φ̇", fontsize=12)
    ax_phase.set_title("Phase Space Trajectory", fontsize=12)
    ax_phase.grid(True, alpha=0.3)
    ax_phase.axhline(y=0, color="k", linewidth=0.5, alpha=0.5)
    ax_phase.axvline(x=0, color="k", linewidth=0.5, alpha=0.5)

    # Time evolution plot
    (line_phi,) = ax_time.plot([], [], "b-", linewidth=2, label="φ(t)")
    (line_phi_dot,) = ax_time.plot([], [], "r-", linewidth=2, label="φ̇(t)")
    (marker_current,) = ax_time.plot([], [], "ko", markersize=8)

    ax_time.set_xlim(times[0], times[-1])
    y_min = min(phi_min, phi_dot_min) - max(phi_margin, phi_dot_margin)
    y_max = max(phi_max, phi_dot_max) + max(phi_margin, phi_dot_margin)
    ax_time.set_ylim(y_min, y_max)
    ax_time.set_xlabel("Time", fontsize=12)
    ax_time.set_ylabel("Value", fontsize=12)
    ax_time.legend(loc="upper right")
    ax_time.grid(True, alpha=0.3)
    ax_time.axhline(y=0, color="k", linewidth=0.5, alpha=0.5)

    time_text = ax_time.text(
        0.02,
        0.98,
        "",
        transform=ax_time.transAxes,
        fontsize=11,
        verticalalignment="top",
        bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
    )

    fig.suptitle(f"{title} (x_index = {x_index})", fontsize=14, fontweight="bold")
    plt.subplots_adjust(left=0.12, right=0.95, top=0.92, bottom=0.12)

    def init():
        """Initialize animation."""
        line_phase.set_data([], [])
        point_phase.set_data([], [])
        line_phi.set_data([], [])
        line_phi_dot.set_data([], [])
        marker_current.set_data([], [])
        time_text.set_text("")
        return (
            line_phase,
            point_phase,
            line_phi,
            line_phi_dot,
            marker_current,
            time_text,
        )

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

        # Update phase space trajectory
        line_phase.set_data(field_vals[: frame + 1], field_dot_vals[: frame + 1])
        point_phase.set_data([field_vals[frame]], [field_dot_vals[frame]])

        # Update time series
        line_phi.set_data(times[: frame + 1], field_vals[: frame + 1])
        line_phi_dot.set_data(times[: frame + 1], field_dot_vals[: frame + 1])
        marker_current.set_data([t], [field_vals[frame]])

        time_text.set_text(f"t = {t:.3f}")

        return (
            line_phase,
            point_phase,
            line_phi,
            line_phi_dot,
            marker_current,
            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 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
