"""
Potential functions for scalar field theory.

This module defines various potential functions V(phi) and uses JAX's
automatic differentiation to compute derivatives dV/dphi.

The ScalarPotential class provides a unified interface that:
1. Encapsulates potential parameters for hashability (enables JIT caching)
2. Provides autodiff-derived force F = -dV/dphi automatically
3. Supports both real and complex fields via Wirtinger derivatives
"""

from dataclasses import dataclass
from functools import lru_cache
from typing import Callable, Tuple

import jax.numpy as jnp
from jax import Array, grad, jit


# Type alias for potential functions
PotentialFunction = Callable[[Array], Array]


# =============================================================================
# ScalarPotential Class - Unified Potential Interface
# =============================================================================


@dataclass(frozen=True)
class ScalarPotential:
    """
    Unified potential representation with autodiff-derived forces.

    This class encapsulates a potential function V(phi) along with its parameters,
    providing automatic computation of forces via JAX autodiff. The frozen
    dataclass ensures hashability, enabling efficient JIT compilation caching.

    Attributes:
        name: Identifier for the potential type (e.g., "quadratic", "quartic")
        params: Tuple of parameters (must be hashable for caching)

    Example:
        >>> V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
        >>> energy_density = V(field)           # V(phi) at each point
        >>> total_energy = V.total_energy(field) # sum of V(phi)
        >>> force = V.force(field)              # F = -dV/dphi (autodiff)
    """

    name: str
    params: Tuple

    def __call__(self, field: Array) -> Array:
        """
        Evaluate potential energy density V(phi) at each lattice point.

        Args:
            field: Scalar field configuration (real or complex)

        Returns:
            Potential energy density array (same shape as field)
        """
        return _get_potential_fn(self)(field)

    def total_energy(self, field: Array) -> Array:
        """
        Compute total potential energy sum(V(phi)).

        Args:
            field: Scalar field configuration

        Returns:
            Total potential energy (scalar)
        """
        return jnp.sum(self(field)).real

    def force(self, field: Array) -> Array:
        """
        Compute force F = -dV/dphi using autodiff.

        For complex fields, computes the Wirtinger derivative -dV/dphi*.

        Args:
            field: Scalar field configuration

        Returns:
            Force array (same shape as field)
        """
        return _get_cached_force(self)(field)

    # -------------------------------------------------------------------------
    # Factory Methods
    # -------------------------------------------------------------------------

    @staticmethod
    def quadratic(m: float = 1.0) -> "ScalarPotential":
        """
        Create quadratic potential V(phi) = (1/2) m^2 |phi|^2.

        This is the simplest non-trivial potential, commonly used in
        chaotic inflation models and free field theory.

        Args:
            m: Mass parameter

        Returns:
            ScalarPotential instance
        """
        return ScalarPotential(name="quadratic", params=(m,))

    @staticmethod
    def quartic(m: float, lambda_: float) -> "ScalarPotential":
        """
        Create quartic potential V(phi) = (1/2) m^2 |phi|^2 + (1/4) lambda |phi|^4.

        Standard scalar field potential with mass term and self-interaction.
        Used in Higgs-like models and preheating studies.

        Args:
            m: Mass parameter
            lambda_: Quartic self-coupling

        Returns:
            ScalarPotential instance
        """
        return ScalarPotential(name="quartic", params=(m, lambda_))

    @staticmethod
    def double_well(mu2: float = 1.0, lam: float = 1.0) -> "ScalarPotential":
        """
        Create double-well potential V(phi) = -(1/2) mu^2 phi^2 + (1/4) lambda phi^4.

        Has minima at phi = +/- sqrt(mu^2/lambda), used for symmetry breaking studies.

        Args:
            mu2: Negative mass squared coefficient
            lam: Quartic coupling

        Returns:
            ScalarPotential instance
        """
        return ScalarPotential(name="double_well", params=(mu2, lam))

    @staticmethod
    def mexican_hat(lam: float = 1.0, v: float = 1.0) -> "ScalarPotential":
        """
        Create Mexican hat potential V(phi) = lambda (|phi|^2 - v^2)^2.

        Classic symmetry-breaking potential with circular minimum at |phi| = v.

        Args:
            lam: Coupling constant
            v: Vacuum expectation value

        Returns:
            ScalarPotential instance
        """
        return ScalarPotential(name="mexican_hat", params=(lam, v))

    @staticmethod
    def from_function(
        potential_fn: PotentialFunction, name: str = "custom"
    ) -> "ScalarPotential":
        """
        Create ScalarPotential from an arbitrary function.

        Note: Custom functions may not cache as efficiently since they
        are identified by object id rather than parameters.

        Args:
            potential_fn: Function V(phi) -> energy density
            name: Identifier for this potential

        Returns:
            ScalarPotential instance
        """
        # Use id of function as unique identifier
        return ScalarPotential(name=name, params=(id(potential_fn), potential_fn))


# =============================================================================
# Internal Implementation - Cached JIT Functions
# =============================================================================


@lru_cache(maxsize=64)
def _get_potential_fn(potential: ScalarPotential) -> Callable[[Array], Array]:
    """Get or create the JIT-compiled potential function."""
    name = potential.name
    params = potential.params

    if name == "quadratic":
        (m,) = params

        @jit
        def V(field: Array) -> Array:
            return 0.5 * m**2 * jnp.abs(field) ** 2

        return V

    elif name == "quartic":
        m, lambda_ = params

        @jit
        def V(field: Array) -> Array:
            phi_sq = jnp.abs(field) ** 2
            return 0.5 * m**2 * phi_sq + 0.25 * lambda_ * phi_sq**2

        return V

    elif name == "double_well":
        mu2, lam = params

        @jit
        def V(field: Array) -> Array:
            return -0.5 * mu2 * field**2 + 0.25 * lam * field**4

        return V

    elif name == "mexican_hat":
        lam, v = params

        @jit
        def V(field: Array) -> Array:
            return lam * (jnp.abs(field) ** 2 - v**2) ** 2

        return V

    elif name == "custom":
        # params is (id, function)
        _, potential_fn = params
        return jit(potential_fn)

    else:
        raise ValueError(f"Unknown potential type: {name}")


@lru_cache(maxsize=64)
def _get_cached_force(potential: ScalarPotential) -> Callable[[Array], Array]:
    """
    Get or create a cached, JIT-compiled force function.

    The force is computed as F = -dV/dphi using autodiff.
    For complex fields, this computes the Wirtinger derivative.
    """

    def total_energy(field: Array) -> Array:
        return jnp.sum(_get_potential_fn(potential)(field)).real

    @jit
    def force(field: Array) -> Array:
        # holomorphic=False gives Wirtinger derivative for complex fields
        return -grad(total_energy, holomorphic=False)(field)

    return force
