"""Dimensional comparison plotting utilities."""

import numpy as np
import matplotlib.pyplot as plt
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 dimensional_comparison(
    times_1d: Array,
    snapshots_1d: List[Lattice],
    times_2d: Array,
    snapshots_2d: List[Lattice],
    times_3d: Array,
    snapshots_3d: List[Lattice],
    potential: Optional[Callable] = None,
    save_path: Optional[str] = None,
    figsize: Tuple[int, int] = (18, 12),
    title: str = "Dimensional Comparison",
    vmin: float = -2.0,
    vmax: float = 2.0,
    reference_value: Optional[float] = None,
) -> plt.Figure:
    """
    Create a comprehensive comparison plot showing 1D, 2D, and 3D field evolution.

    This function generates a publication-quality comparison figure with:
    - Initial, mid, and final states for each dimension
    - Energy evolution plots for each dimension
    - Consistent color scaling and formatting
    - Optional reference lines (e.g., VEV for symmetry breaking)

    Args:
        times_1d: Time array for 1D simulation
        snapshots_1d: List of 1D Lattice snapshots
        times_2d: Time array for 2D simulation
        snapshots_2d: List of 2D Lattice snapshots
        times_3d: Time array for 3D simulation
        snapshots_3d: List of 3D Lattice snapshots
        potential: Potential function for energy computation
        save_path: If provided, save figure to this path
        figsize: Figure size (width, height)
        title: Overall figure title
        vmin: Minimum value for color scale
        vmax: Maximum value for color scale
        reference_value: Optional reference line (e.g., VEV)

    Returns:
        matplotlib Figure object

    Example:
        >>> from jaxlatt.core.potentials import double_well_potential
        >>> potential = double_well_potential(lam=0.5, mu2=1.0)
        >>> fig = dimensional_comparison(
        ...     times_1d, snaps_1d, times_2d, snaps_2d, times_3d, snaps_3d,
        ...     potential=potential, reference_value=1.414,
        ...     save_path="comparison.png"
        ... )
    """
    fig = plt.figure(figsize=figsize)
    gs = GridSpec(3, 4, figure=fig, hspace=0.3, wspace=0.3)

    # Column headers
    fig.text(0.15, 0.97, "Initial State", ha="center", fontsize=14, fontweight="bold")
    fig.text(0.39, 0.97, "Mid Evolution", ha="center", fontsize=14, fontweight="bold")
    fig.text(0.63, 0.97, "Final State", ha="center", fontsize=14, fontweight="bold")
    fig.text(
        0.87, 0.97, "Energy Evolution", ha="center", fontsize=14, fontweight="bold"
    )

    # Compute indices for snapshots
    mid_idx_1d = len(snapshots_1d) // 2
    mid_idx_2d = len(snapshots_2d) // 2
    mid_idx_3d = len(snapshots_3d) // 2

    # Compute energy if potential provided
    energy_1d = energy_2d = energy_3d = None
    if potential is not None:
        energy_1d = [energy_components_integrated(s, potential) for s in snapshots_1d]
        energy_2d = [energy_components_integrated(s, potential) for s in snapshots_2d]
        energy_3d = [energy_components_integrated(s, potential) for s in snapshots_3d]

    # =========================================================================
    # Row 1: 1D
    # =========================================================================
    ax_1d_init = fig.add_subplot(gs[0, 0])
    ax_1d_mid = fig.add_subplot(gs[0, 1])
    ax_1d_final = fig.add_subplot(gs[0, 2])
    ax_1d_energy = fig.add_subplot(gs[0, 3])

    # Get 1D coordinates
    lattice_1d = snapshots_1d[0]
    length_1d = (
        lattice_1d.length[0]
        if isinstance(lattice_1d.length, tuple)
        else lattice_1d.length
    )
    size_1d = (
        lattice_1d.size[0] if isinstance(lattice_1d.size, tuple) else lattice_1d.size
    )
    x_1d = np.linspace(0, float(length_1d), int(size_1d), endpoint=False)

    # Plot 1D fields
    ax_1d_init.plot(x_1d, np.asarray(snapshots_1d[0].field), "b-", linewidth=1.5)
    if reference_value is not None:
        ax_1d_init.axhline(reference_value, color="r", linestyle="--", alpha=0.5)
        ax_1d_init.axhline(-reference_value, color="r", linestyle="--", alpha=0.5)
    ax_1d_init.set_ylabel("1D", fontsize=12, fontweight="bold")
    ax_1d_init.set_ylim(vmin, vmax)
    ax_1d_init.grid(True, alpha=0.3)

    ax_1d_mid.plot(
        x_1d, np.asarray(snapshots_1d[mid_idx_1d].field), "b-", linewidth=1.5
    )
    if reference_value is not None:
        ax_1d_mid.axhline(reference_value, color="r", linestyle="--", alpha=0.5)
        ax_1d_mid.axhline(-reference_value, color="r", linestyle="--", alpha=0.5)
    ax_1d_mid.set_ylim(vmin, vmax)
    ax_1d_mid.grid(True, alpha=0.3)

    ax_1d_final.plot(x_1d, np.asarray(snapshots_1d[-1].field), "b-", linewidth=1.5)
    if reference_value is not None:
        ax_1d_final.axhline(reference_value, color="r", linestyle="--", alpha=0.5)
        ax_1d_final.axhline(-reference_value, color="r", linestyle="--", alpha=0.5)
    ax_1d_final.set_ylim(vmin, vmax)
    ax_1d_final.grid(True, alpha=0.3)

    # Energy plot
    if energy_1d is not None:
        E_total_1d = [e["total"] for e in energy_1d]
        ax_1d_energy.plot(times_1d, E_total_1d, "k-", linewidth=2)
        ax_1d_energy.set_ylabel("Total Energy", fontsize=10)
        ax_1d_energy.grid(True, alpha=0.3)

    # =========================================================================
    # Row 2: 2D
    # =========================================================================
    ax_2d_init = fig.add_subplot(gs[1, 0])
    ax_2d_mid = fig.add_subplot(gs[1, 1])
    ax_2d_final = fig.add_subplot(gs[1, 2])
    ax_2d_energy = fig.add_subplot(gs[1, 3])

    # Plot 2D heatmaps
    lattice_2d = snapshots_2d[0]
    extent_2d = [0, lattice_2d.length[0], 0, lattice_2d.length[1]]

    ax_2d_init.imshow(
        np.asarray(snapshots_2d[0].field).T,
        origin="lower",
        cmap="RdBu_r",
        extent=extent_2d,
        vmin=vmin,
        vmax=vmax,
    )
    ax_2d_init.set_ylabel("2D", fontsize=12, fontweight="bold")
    ax_2d_init.set_aspect("equal")

    ax_2d_mid.imshow(
        np.asarray(snapshots_2d[mid_idx_2d].field).T,
        origin="lower",
        cmap="RdBu_r",
        extent=extent_2d,
        vmin=vmin,
        vmax=vmax,
    )
    ax_2d_mid.set_aspect("equal")

    ax_2d_final.imshow(
        np.asarray(snapshots_2d[-1].field).T,
        origin="lower",
        cmap="RdBu_r",
        extent=extent_2d,
        vmin=vmin,
        vmax=vmax,
    )
    ax_2d_final.set_aspect("equal")

    # Energy plot
    if energy_2d is not None:
        E_total_2d = [e["total"] for e in energy_2d]
        ax_2d_energy.plot(times_2d, E_total_2d, "k-", linewidth=2)
        ax_2d_energy.set_ylabel("Total Energy", fontsize=10)
        ax_2d_energy.grid(True, alpha=0.3)

    # =========================================================================
    # Row 3: 3D (show central XY slice)
    # =========================================================================
    ax_3d_init = fig.add_subplot(gs[2, 0])
    ax_3d_mid = fig.add_subplot(gs[2, 1])
    ax_3d_final = fig.add_subplot(gs[2, 2])
    ax_3d_energy = fig.add_subplot(gs[2, 3])

    # Plot 3D central slices
    lattice_3d = snapshots_3d[0]
    extent_3d = [0, lattice_3d.length[0], 0, lattice_3d.length[1]]
    mid_z = lattice_3d.size[2] // 2

    ax_3d_init.imshow(
        np.asarray(snapshots_3d[0].field)[:, :, mid_z].T,
        origin="lower",
        cmap="RdBu_r",
        extent=extent_3d,
        vmin=vmin,
        vmax=vmax,
    )
    ax_3d_init.set_ylabel("3D (XY slice)", fontsize=12, fontweight="bold")
    ax_3d_init.set_xlabel("x", fontsize=10)
    ax_3d_init.set_aspect("equal")

    ax_3d_mid.imshow(
        np.asarray(snapshots_3d[mid_idx_3d].field)[:, :, mid_z].T,
        origin="lower",
        cmap="RdBu_r",
        extent=extent_3d,
        vmin=vmin,
        vmax=vmax,
    )
    ax_3d_mid.set_xlabel("x", fontsize=10)
    ax_3d_mid.set_aspect("equal")

    im_final = ax_3d_final.imshow(
        np.asarray(snapshots_3d[-1].field)[:, :, mid_z].T,
        origin="lower",
        cmap="RdBu_r",
        extent=extent_3d,
        vmin=vmin,
        vmax=vmax,
    )
    ax_3d_final.set_xlabel("x", fontsize=10)
    ax_3d_final.set_aspect("equal")

    # Energy plot
    if energy_3d is not None:
        E_total_3d = [e["total"] for e in energy_3d]
        ax_3d_energy.plot(times_3d, E_total_3d, "k-", linewidth=2)
        ax_3d_energy.set_xlabel("Time", fontsize=10)
        ax_3d_energy.set_ylabel("Total Energy", fontsize=10)
        ax_3d_energy.grid(True, alpha=0.3)

    # Add colorbar
    cbar = plt.colorbar(
        im_final,
        ax=[ax_2d_init, ax_2d_mid, ax_2d_final, ax_3d_init, ax_3d_mid, ax_3d_final],
        orientation="horizontal",
        pad=0.1,
        aspect=40,
    )
    cbar.set_label("φ(x)", fontsize=11)

    # Overall title
    fig.suptitle(title, fontsize=16, fontweight="bold", y=0.995)

    if save_path is not None:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        console.print(f"[green]✓ Saved:[/green] [bold]{save_path}[/bold]")

    return fig
