"""
Field creation utilities for scalar, gauge, and coupled lattices.

This module consolidates all field initialization functions for creating
various lattice configurations.
"""

from typing import Tuple
import jax.numpy as jnp
from jax import random

from .lattice import CoupledLattice, GaugeLattice


# =============================================================================
# Coupled Scalar-Gauge Lattice Creation
# =============================================================================


def create_vacuum_coupled_lattice(
    size: Tuple[int, int, int],
    length: float | Tuple[float, float, float],
    m: float = 1.0,
    lambda_: float = 1.0,
    g: float = 1.0,
) -> CoupledLattice:
    """
    Create a coupled lattice in vacuum state.

    Vacuum: phi = 0, pi = 0, U = 1, E = 0

    Parameters:
    -----------
    size : tuple
        Lattice dimensions (Nx, Ny, Nz)
    length : float or tuple
        Physical box size
    m : float
        Scalar mass
    lambda_ : float
        Scalar self-coupling
    g : float
        Gauge coupling

    Returns:
    --------
    CoupledLattice in vacuum state
    """
    if isinstance(length, (int, float)):
        length = (length, length, length)

    dx = length[0] / size[0]

    phi = jnp.zeros(size, dtype=jnp.complex64)
    pi = jnp.zeros(size, dtype=jnp.complex64)
    links = jnp.ones((3,) + size, dtype=jnp.complex64)
    E = jnp.zeros((3,) + size, dtype=jnp.float32)

    return CoupledLattice(
        phi=phi,
        pi=pi,
        links=links,
        E=E,
        m=m,
        lambda_=lambda_,
        g=g,
        dx=dx,
        size=size,
        length=length,
    )


def create_random_coupled_lattice(
    key: random.PRNGKey,
    size: Tuple[int, int, int],
    length: float | Tuple[float, float, float],
    m: float = 1.0,
    lambda_: float = 1.0,
    g: float = 1.0,
    amplitude: float = 0.1,
) -> CoupledLattice:
    """
    Create a coupled lattice with random initial conditions.

    Parameters:
    -----------
    key : PRNGKey
        Random number generator key
    size : tuple
        Lattice dimensions (Nx, Ny, Nz)
    length : float or tuple
        Physical box size
    m : float
        Scalar mass
    lambda_ : float
        Scalar self-coupling
    g : float
        Gauge coupling
    amplitude : float
        Amplitude of random perturbations

    Returns:
    --------
    CoupledLattice with random fields
    """
    if isinstance(length, (int, float)):
        length = (length, length, length)

    dx = length[0] / size[0]

    # Split random key
    k1, k2, k3, k4, k5, k6 = random.split(key, 6)

    # Random scalar field (complex)
    phi_re = amplitude * random.normal(k1, size, dtype=jnp.float32)
    phi_im = amplitude * random.normal(k2, size, dtype=jnp.float32)
    phi = phi_re + 1j * phi_im

    # Random conjugate momentum (complex)
    pi_re = amplitude * random.normal(k3, size, dtype=jnp.float32)
    pi_im = amplitude * random.normal(k4, size, dtype=jnp.float32)
    pi = pi_re + 1j * pi_im

    # Random gauge links (phases on unit circle)
    theta = amplitude * random.normal(k5, (3,) + size, dtype=jnp.float32)
    links = jnp.exp(1j * theta)

    # Random electric field
    E = amplitude * random.normal(k6, (3,) + size, dtype=jnp.float32)

    return CoupledLattice(
        phi=phi,
        pi=pi,
        links=links,
        E=E,
        m=m,
        lambda_=lambda_,
        g=g,
        dx=dx,
        size=size,
        length=length,
    )


def create_higgs_vev_lattice(
    size: Tuple[int, int, int],
    length: float | Tuple[float, float, float],
    m: float = 1.0,
    lambda_: float = 1.0,
    g: float = 1.0,
    vev_amplitude: float = None,
) -> CoupledLattice:
    """
    Create a coupled lattice with Higgs field at vacuum expectation value.

    For symmetry breaking potential V = -m²|φ|²/2 + λ|φ|⁴/4,
    the VEV is |φ| = sqrt(m²/λ).

    Parameters:
    -----------
    size : tuple
        Lattice dimensions (Nx, Ny, Nz)
    length : float or tuple
        Physical box size
    m : float
        Scalar mass parameter
    lambda_ : float
        Scalar self-coupling
    g : float
        Gauge coupling
    vev_amplitude : float, optional
        VEV amplitude. If None, uses sqrt(m²/λ)

    Returns:
    --------
    CoupledLattice with Higgs at VEV
    """
    if isinstance(length, (int, float)):
        length = (length, length, length)

    dx = length[0] / size[0]

    # Vacuum expectation value
    if vev_amplitude is None:
        if lambda_ > 0 and m > 0:
            vev_amplitude = jnp.sqrt(m**2 / lambda_)
        else:
            vev_amplitude = 0.0

    # Constant field at VEV (real for simplicity)
    phi = jnp.full(size, vev_amplitude, dtype=jnp.complex64)
    pi = jnp.zeros(size, dtype=jnp.complex64)
    links = jnp.ones((3,) + size, dtype=jnp.complex64)
    E = jnp.zeros((3,) + size, dtype=jnp.float32)

    return CoupledLattice(
        phi=phi,
        pi=pi,
        links=links,
        E=E,
        m=m,
        lambda_=lambda_,
        g=g,
        dx=dx,
        size=size,
        length=length,
    )


# =============================================================================
# Pure Gauge Lattice Creation
# =============================================================================


def create_vacuum_gauge_lattice(
    size: Tuple[int, int, int],
    length: float,
    g: float = 1.0,
) -> GaugeLattice:
    """
    Create a gauge lattice in vacuum state (links=1, E=0).

    Args:
        size: Grid dimensions
        length: Physical box size
        g: Gauge coupling

    Returns:
        GaugeLattice in vacuum configuration
    """
    return GaugeLattice(size=size, length=length, g=g)


def create_random_gauge_lattice(
    key: random.PRNGKey,
    size: Tuple[int, int, int],
    length: float,
    g: float = 1.0,
    amplitude: float = 0.1,
) -> GaugeLattice:
    """
    Create a gauge lattice with small random perturbations.

    Initializes links with small random phases and electric field with
    small random values, suitable for testing dynamics.

    Args:
        key: JAX random key
        size: Grid dimensions
        length: Physical box size
        g: Gauge coupling
        amplitude: Amplitude of random fluctuations

    Returns:
        GaugeLattice with random initial conditions
    """
    key1, key2 = random.split(key)

    # Random phases for links (small perturbation around identity)
    phases = amplitude * random.normal(key1, shape=(3, *size))
    links = jnp.exp(1j * phases).astype(jnp.complex64)

    # Random electric field
    E = amplitude * random.normal(key2, shape=(3, *size)).astype(jnp.float32)

    return GaugeLattice(size=size, length=length, g=g, links=links, E=E)
