"""
3D field animation utilities.

This module provides tools for visualizing 3D scalar field dynamics,
including slice views and isosurface animations.
"""

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

console = Console()


def slices(
    times: Array,
    snapshots: List[Lattice],
    potential: Optional[Callable] = None,
    interval: int = 100,
    save_path: Optional[str] = None,
    figsize: Tuple[int, int] = (14, 6),
    title: str = "3D Scalar Field Central Slices",
    cmap: str = "RdBu_r",
    vmin: Optional[float] = None,
    vmax: Optional[float] = None,
) -> FuncAnimation:
    """
    Animate central XY, XZ, YZ slices of a 3D scalar field.

    Args:
        times: Snapshot times
        snapshots: Lattice states (ndim must be 3)
        potential: Potential function (optional; if provided energy displayed)
        interval: Frame delay ms
        save_path: Optional path to save animation (GIF or MP4)
        figsize: Figure size
        title: Overall title
        cmap: Colormap
        vmin/vmax: Explicit color scale limits

    Returns:
        FuncAnimation
    """
    lattice0 = snapshots[0]
    if lattice0.ndim != 3:
        raise ValueError("animate_3d_slices requires 3D lattice snapshots")

    # Determine color limits
    if vmin is None or vmax is None:
        all_fields = [np.asarray(s.field) for s in snapshots]
        fmin = min(f.min() for f in all_fields)
        fmax = max(f.max() for f in all_fields)
        if vmin is None:
            vmin = fmin
        if vmax is None:
            vmax = fmax

    nx, ny, nz = lattice0.size
    cx, cy, cz = nx // 2, ny // 2, nz // 2

    # Prepare figure
    fig = plt.figure(figsize=figsize)
    gs = GridSpec(2, 3, figure=fig, height_ratios=[1, 0.05])
    ax_xy = fig.add_subplot(gs[0, 0])
    ax_xz = fig.add_subplot(gs[0, 1])
    ax_yz = fig.add_subplot(gs[0, 2])
    ax_cbar = fig.add_subplot(gs[1, :])

    extent_xy = (0.0, lattice0.length[0], 0.0, lattice0.length[1])
    extent_xz = (0.0, lattice0.length[0], 0.0, lattice0.length[2])
    extent_yz = (0.0, lattice0.length[1], 0.0, lattice0.length[2])

    field0 = np.asarray(lattice0.field)
    im_xy = ax_xy.imshow(
        field0[:, :, cz].T,
        origin="lower",
        cmap=cmap,
        extent=extent_xy,
        vmin=vmin,
        vmax=vmax,
        animated=True,
    )
    im_xz = ax_xz.imshow(
        field0[:, cy, :].T,
        origin="lower",
        cmap=cmap,
        extent=extent_xz,
        vmin=vmin,
        vmax=vmax,
        animated=True,
    )
    im_yz = ax_yz.imshow(
        field0[cx, :, :].T,
        origin="lower",
        cmap=cmap,
        extent=extent_yz,
        vmin=vmin,
        vmax=vmax,
        animated=True,
    )

    # Set titles & axis labels with correct coordinates per slice
    ax_xy.set_title("XY (z=mid)")
    ax_xy.set_xlabel("x")
    ax_xy.set_ylabel("y")

    ax_xz.set_title("XZ (y=mid)")
    ax_xz.set_xlabel("x")
    ax_xz.set_ylabel("z")

    ax_yz.set_title("YZ (x=mid)")
    ax_yz.set_xlabel("y")
    ax_yz.set_ylabel("z")

    cbar = plt.colorbar(im_xy, cax=ax_cbar, orientation="horizontal")
    cbar.set_label("φ")

    time_text = fig.text(
        0.02,
        0.95,
        "",
        fontsize=12,
        bbox=dict(boxstyle="round", facecolor="white", alpha=0.6),
    )

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

    def update(frame):
        snap = snapshots[frame]
        f = np.asarray(snap.field)
        im_xy.set_array(f[:, :, cz].T)
        im_xz.set_array(f[:, cy, :].T)
        im_yz.set_array(f[cx, :, :].T)
        time_text.set_text(f"t = {times[frame]:.3f}")
        return im_xy, im_xz, im_yz, time_text

    anim = FuncAnimation(
        fig,
        update,
        frames=len(snapshots),
        interval=interval,
        blit=True,
        repeat=True,
    )

    if save_path is not None:
        console.print(
            f"[cyan]💾 Saving 3D field 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 field(
    times: Array,
    snapshots: List[Lattice],
    potential: Optional[Callable] = None,
    interval: int = 100,
    save_path: Optional[str] = None,
    figsize: Tuple[int, int] = (10, 8),
    title: str = "3D Scalar Field Evolution",
    isosurface_levels: Optional[List[float]] = None,
    fps: int = 20,
) -> FuncAnimation:
    """
    Create a 3D isosurface animation of scalar field evolution.

    This function creates volumetric visualizations showing isosurfaces of constant
    field value, which is ideal for visualizing topological structures like bubbles,
    domain walls, and defects in 3D field configurations.

    Args:
        times: Array of snapshot times
        snapshots: List of Lattice objects (must be 3D)
        potential: Potential function (optional)
        interval: Delay between frames in milliseconds
        save_path: If provided, save animation to this path
        figsize: Figure size (width, height)
        title: Animation title
        isosurface_levels: List of field values to show as isosurfaces
        fps: Frames per second for saved animation

    Returns:
        FuncAnimation object

    Note:
        Requires matplotlib with 3D support. Isosurface rendering uses
        contour slices with moderate detail for performance.
    """
    from mpl_toolkits.mplot3d import Axes3D
    from matplotlib import cm

    lattice = snapshots[0]
    if lattice.ndim != 3:
        raise ValueError("animate_3d_field requires 3D lattice")

    # Default isosurface levels
    if isosurface_levels is None:
        # Use ±1 sigma from mean as default
        all_fields = [np.asarray(snap.field) for snap in snapshots]
        field_mean = np.mean([f.mean() for f in all_fields])
        field_std = np.mean([f.std() for f in all_fields])
        isosurface_levels = [field_mean + field_std, field_mean - field_std]

    nx, ny, nz = lattice.size
    lx, ly, lz = lattice.length

    # Create coordinate grids
    x = np.linspace(0, lx, nx)
    y = np.linspace(0, ly, ny)
    z = np.linspace(0, lz, nz)
    X, Y, Z = np.meshgrid(x, y, z, indexing="ij")

    # Setup figure
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(111, projection="3d")

    # Set viewing angle
    ax.view_init(elev=20, azim=45)

    # Set axis limits
    ax.set_xlim(0, lx)
    ax.set_ylim(0, ly)
    ax.set_zlim(0, lz)
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")

    time_text = fig.text(
        0.02,
        0.95,
        "",
        fontsize=12,
        bbox=dict(boxstyle="round", facecolor="white", alpha=0.7),
    )

    fig.suptitle(title, fontsize=14, fontweight="bold")

    def update(frame):
        """Update animation frame."""
        ax.clear()

        snap = snapshots[frame]
        t = times[frame]
        field = np.asarray(snap.field)

        # Set axis properties again after clear
        ax.set_xlim(0, lx)
        ax.set_ylim(0, ly)
        ax.set_zlim(0, lz)
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_zlabel("z")
        ax.view_init(elev=20, azim=45 + frame * 0.5)  # Slow rotation

        # Plot isosurfaces using contour slices
        # We'll show XY, XZ, and YZ slices with contours
        colors = ["blue", "red", "green", "orange"]

        for idx, level in enumerate(isosurface_levels[:4]):  # Max 4 levels
            color = colors[idx % len(colors)]

            # XY slices (at various Z positions)
            for zi in range(0, nz, max(1, nz // 5)):
                field_slice = field[:, :, zi]
                if np.any(field_slice >= level) and np.any(field_slice <= level):
                    ax.contour(
                        X[:, :, zi],
                        Y[:, :, zi],
                        field_slice,
                        levels=[level],
                        colors=color,
                        alpha=0.3,
                        offset=z[zi],
                        zdir="z",
                    )

            # XZ slices (at various Y positions)
            for yi in range(0, ny, max(1, ny // 5)):
                field_slice = field[:, yi, :]
                if np.any(field_slice >= level) and np.any(field_slice <= level):
                    ax.contour(
                        X[:, yi, :],
                        field_slice,
                        Z[:, yi, :],
                        levels=[level],
                        colors=color,
                        alpha=0.3,
                        offset=y[yi],
                        zdir="y",
                    )

            # YZ slices (at various X positions)
            for xi in range(0, nx, max(1, nx // 5)):
                field_slice = field[xi, :, :]
                if np.any(field_slice >= level) and np.any(field_slice <= level):
                    ax.contour(
                        Y[xi, :, :],
                        field_slice,
                        Z[xi, :, :],
                        levels=[level],
                        colors=color,
                        alpha=0.3,
                        offset=x[xi],
                        zdir="x",
                    )

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

        return ax, time_text

    # Create animation
    anim = FuncAnimation(
        fig,
        update,
        frames=len(snapshots),
        interval=interval,
        blit=False,  # 3D animations can't use blitting
        repeat=True,
    )

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

    return anim
