"""
Phase space plotting utilities.

This module provides tools for creating phase portraits and phase space
analysis plots for scalar field dynamics.
"""

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

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

console = Console()


def portrait(
    snapshots: List[Lattice],
    potential: Callable,
    save_path: Optional[str] = None,
    figsize: Tuple[int, int] = (14, 12),
    title: str = "Phase Portrait",
    n_sample_points: int = 100,
) -> plt.Figure:
    """
    Create a comprehensive phase portrait showing potential landscape,
    trajectories, and field statistics.

    This creates a publication-quality figure showing:
    - Potential energy landscape V(φ)
    - Phase space trajectories at multiple spatial points
    - Field distribution histograms
    - Energy evolution

    Works for any dimensionality (1D, 2D, 3D) by sampling spatial points.

    Args:
        snapshots: List of Lattice snapshots
        potential: Potential function V(φ)
        save_path: If provided, save figure to this path
        figsize: Figure size
        title: Figure title
        n_sample_points: Number of points to sample for trajectories

    Returns:
        matplotlib Figure object

    Example:
        >>> from jaxlatt.core.potentials import double_well_potential
        >>> pot = double_well_potential(lam=0.5, mu2=1.0)
        >>> fig = phase_portrait(snapshots, pot, save_path="portrait.png")
    """
    import jax.numpy as jnp

    fig = plt.figure(figsize=figsize)
    gs = GridSpec(2, 2, figure=fig, hspace=0.3, wspace=0.3)

    ax_potential = fig.add_subplot(gs[0, 0])
    ax_phase = fig.add_subplot(gs[0, 1])
    ax_histogram = fig.add_subplot(gs[1, 0])
    ax_energy = fig.add_subplot(gs[1, 1])

    # =========================================================================
    # Plot 1: Potential landscape
    # =========================================================================
    phi_range = np.linspace(-3, 3, 200)
    V_vals = [float(potential(jnp.array(phi))) for phi in phi_range]

    ax_potential.plot(phi_range, V_vals, "k-", linewidth=2)
    ax_potential.set_xlabel("φ", fontsize=12)
    ax_potential.set_ylabel("V(φ)", fontsize=12)
    ax_potential.set_title("Potential Energy", fontsize=12, fontweight="bold")
    ax_potential.grid(True, alpha=0.3)
    ax_potential.axhline(0, color="gray", linestyle="--", alpha=0.5)

    # Mark minima
    V_array = np.array(V_vals)
    if len(V_array) > 0:
        local_minima = []
        for i in range(1, len(V_array) - 1):
            if V_array[i] < V_array[i - 1] and V_array[i] < V_array[i + 1]:
                local_minima.append((phi_range[i], V_array[i]))
        for phi_min, V_min in local_minima:
            ax_potential.plot(phi_min, V_min, "ro", markersize=10, label="Vacuum")

    ax_potential.legend()

    # =========================================================================
    # Plot 2: Phase space trajectories
    # =========================================================================
    lattice = snapshots[0]
    total_points = int(np.prod(lattice.size))

    # Sample random spatial points
    n_trajectories = min(20, n_sample_points)
    sample_indices = np.random.choice(total_points, size=n_trajectories, replace=False)

    for idx in sample_indices:
        phi_traj = [float(np.asarray(s.field).flatten()[idx]) for s in snapshots]
        phi_dot_traj = [
            float(np.asarray(s.field_dot).flatten()[idx]) for s in snapshots
        ]

        ax_phase.plot(phi_traj, phi_dot_traj, alpha=0.5, linewidth=1)

    ax_phase.set_xlabel("φ", fontsize=12)
    ax_phase.set_ylabel("φ̇", fontsize=12)
    ax_phase.set_title("Phase Space Trajectories", fontsize=12, fontweight="bold")
    ax_phase.grid(True, alpha=0.3)
    ax_phase.axhline(0, color="k", linestyle="-", linewidth=0.5, alpha=0.3)
    ax_phase.axvline(0, color="k", linestyle="-", linewidth=0.5, alpha=0.3)

    # =========================================================================
    # Plot 3: Field distribution histogram
    # =========================================================================
    # Show histogram evolution (initial, mid, final)
    fields = [
        np.asarray(snapshots[0].field).flatten(),
        np.asarray(snapshots[len(snapshots) // 2].field).flatten(),
        np.asarray(snapshots[-1].field).flatten(),
    ]
    labels = ["Initial", "Mid", "Final"]
    colors = ["blue", "green", "red"]

    for field, label, color in zip(fields, labels, colors):
        ax_histogram.hist(
            field, bins=50, alpha=0.5, label=label, color=color, density=True
        )

    ax_histogram.set_xlabel("φ", fontsize=12)
    ax_histogram.set_ylabel("Probability Density", fontsize=12)
    ax_histogram.set_title(
        "Field Distribution Evolution", fontsize=12, fontweight="bold"
    )
    ax_histogram.legend()
    ax_histogram.grid(True, alpha=0.3, axis="y")

    # =========================================================================
    # Plot 4: Energy evolution
    # =========================================================================
    energy_data = [energy_components_integrated(s, potential) for s 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]

    # Generate time array (assuming uniform spacing)
    times = np.arange(len(snapshots))

    ax_energy.plot(times, E_total, "k-", linewidth=2, label="Total")
    ax_energy.plot(times, E_kinetic, "r-", linewidth=1.5, alpha=0.7, label="Kinetic")
    ax_energy.plot(times, E_gradient, "g-", linewidth=1.5, alpha=0.7, label="Gradient")
    ax_energy.plot(
        times, E_potential, "b-", linewidth=1.5, alpha=0.7, label="Potential"
    )

    ax_energy.set_xlabel("Snapshot Index", fontsize=12)
    ax_energy.set_ylabel("Energy", fontsize=12)
    ax_energy.set_title("Energy Components", fontsize=12, fontweight="bold")
    ax_energy.legend(loc="best")
    ax_energy.grid(True, alpha=0.3)

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

    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
