"""
Core lattice structures for scalar and gauge fields.

This module defines the fundamental lattice data structures:
- Lattice: Simple scalar field (1D, 2D, 3D)
- CoupledLattice: Scalar + U(1) gauge (Abelian Higgs model)
- GaugeLattice: Pure U(1) gauge field
"""

from typing import Tuple, Optional
import jax.numpy as jnp
from jax import Array
from dataclasses import dataclass


@dataclass
class Lattice:
    """
    Represents a scalar field on a periodic lattice (supports 1D, 2D, 3D).

    Attributes:
        field: Scalar field values (shape matches size tuple)
        field_dot: Time derivative of the field (same shape)
        dx: Lattice spacing (assumed uniform, taken from first dimension)
        size: Tuple of lattice grid counts per dimension (always a tuple internally)
        length: Tuple of physical lengths per dimension (always a tuple internally)
        ndim: Number of spatial dimensions
    """

    field: Array
    field_dot: Array
    dx: float
    size: Tuple[int, ...]
    length: Tuple[float, ...]
    ndim: int

    def __init__(
        self,
        size: int | Tuple[int, ...],
        length: float | Tuple[float, ...],
        field: Optional[Array] = None,
        field_dot: Optional[Array] = None,
    ):
        """
        Initialize a lattice with given size and physical length.

        Args:
            size: Number of grid points (int for 1D, tuple for 2D)
            length: Physical length (float for 1D, tuple for 2D)
            field: Initial field values (if None, initialized to zeros)
            field_dot: Initial time derivative (if None, initialized to zeros)
        """
        # Handle size and length for 1D and 2D
        if isinstance(size, int):
            self.size = (size,)
            self.ndim = 1
        else:
            self.size = size
            self.ndim = len(size)

        if isinstance(length, (int, float)):
            # Replicate scalar length across all spatial dimensions
            if isinstance(size, int):
                self.length = (float(length),)
            else:
                self.length = tuple(float(length) for _ in range(len(size)))
        else:
            # If a tuple was provided but length dimensionality doesn't match size, replicate or raise
            if isinstance(size, tuple) and len(length) == 1 and len(size) > 1:
                self.length = tuple(float(length[0]) for _ in range(len(size)))
            elif isinstance(size, tuple) and len(length) != len(size):
                raise ValueError(
                    f"Length tuple dimensionality {len(length)} does not match size {len(size)}"
                )
            else:
                self.length = tuple(float(val) for val in length)

        # Compute lattice spacing
        # Assume uniform spacing; use first dimension for dx
        self.dx = self.length[0] / self.size[0]

        # Initialize fields
        if field is None:
            self.field = jnp.zeros(self.size)
        else:
            self.field = field

        if field_dot is None:
            self.field_dot = jnp.zeros(self.size)
        else:
            self.field_dot = field_dot

    def copy(self) -> "Lattice":
        """Create a copy of the lattice."""
        return Lattice(
            size=self.size,
            length=self.length,
            field=self.field.copy(),
            field_dot=self.field_dot.copy(),
        )

    def update(self, field: Array, field_dot: Array) -> "Lattice":
        """
        Create a new Lattice with updated field values.

        Args:
            field: New field values
            field_dot: New field time derivative

        Returns:
            New Lattice instance with updated values
        """
        return Lattice(
            size=self.size,
            length=self.length,
            field=field,
            field_dot=field_dot,
        )

    @property
    def volume(self) -> float:
        """Total volume of the lattice."""
        vol = 1.0
        for length_dim in self.length:
            vol *= length_dim
        return vol

    @property
    def dV(self) -> float:
        """Volume element (grid cell volume)."""
        return float(self.volume / jnp.prod(jnp.array(self.size)))

    def __repr__(self) -> str:
        return (
            f"Lattice(ndim={self.ndim}, size={self.size}, "
            f"length={self.length}, dx={self.dx:.4f})"
        )


@dataclass
class CoupledLattice:
    """
    Combined scalar and gauge field lattice for Abelian Higgs model.

    Fields:
    -------
    phi : complex array (N, N, N)
        Complex scalar field (Higgs field)
    pi : complex array (N, N, N)
        Conjugate momentum to phi
    links : complex array (3, N, N, N)
        U(1) gauge links U_i(n) = e^{i g A_i(n) dx}
    E : float array (3, N, N, N)
        Electric field (conjugate to links)

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

    phi: jnp.ndarray
    pi: jnp.ndarray
    links: jnp.ndarray
    E: jnp.ndarray
    m: float
    lambda_: float
    g: float
    dx: float
    size: Tuple[int, int, int]
    length: Tuple[float, float, float]

    def copy(self):
        """Create a copy of the lattice."""
        return CoupledLattice(
            phi=self.phi.copy(),
            pi=self.pi.copy(),
            links=self.links.copy(),
            E=self.E.copy(),
            m=self.m,
            lambda_=self.lambda_,
            g=self.g,
            dx=self.dx,
            size=self.size,
            length=self.length,
        )

    def replace(self, **kwargs) -> "CoupledLattice":
        """Create a new CoupledLattice with specified fields replaced."""
        return CoupledLattice(
            phi=kwargs.get("phi", self.phi),
            pi=kwargs.get("pi", self.pi),
            links=kwargs.get("links", self.links),
            E=kwargs.get("E", self.E),
            m=kwargs.get("m", self.m),
            lambda_=kwargs.get("lambda_", self.lambda_),
            g=kwargs.get("g", self.g),
            dx=kwargs.get("dx", self.dx),
            size=kwargs.get("size", self.size),
            length=kwargs.get("length", self.length),
        )

    def __repr__(self):
        return (
            f"CoupledLattice(\n"
            f"  size={self.size}, dx={self.dx:.4f}\n"
            f"  m={self.m}, lambda={self.lambda_}, g={self.g}\n"
            f"  phi: {self.phi.shape} {self.phi.dtype}\n"
            f"  pi: {self.pi.shape} {self.pi.dtype}\n"
            f"  links: {self.links.shape} {self.links.dtype}\n"
            f"  E: {self.E.shape} {self.E.dtype}\n"
            f")"
        )


@dataclass
class GaugeLattice:
    """
    U(1) gauge fields on a 3D periodic lattice.

    Uses compact formulation with link variables U_i ∈ U(1).

    Attributes:
        links: Complex link variables (3, N, N, N)
        E: Electric field components (3, N, N, N)
        dx: Lattice spacing
        size: Lattice dimensions
        length: Physical box size
        g: Gauge coupling
        a: Scale factor (cosmology)
        H: Hubble parameter (cosmology)
    """

    links: Array
    E: Array
    dx: float
    size: Tuple[int, int, int]
    length: Tuple[float, float, float]
    g: float = 1.0
    a: float = 1.0
    H: float = 0.0

    def __init__(
        self,
        size: Tuple[int, int, int],
        length: float,
        g: float = 1.0,
        links: Optional[Array] = None,
        E: Optional[Array] = None,
        a: float = 1.0,
        H: float = 0.0,
    ):
        """Initialize a gauge lattice."""
        self.size = size
        self.length = (length, length, length)
        self.dx = length / size[0]
        self.g = g
        self.a = a
        self.H = H

        if links is None:
            self.links = jnp.ones((3, *size), dtype=jnp.complex64)
        else:
            self.links = links.astype(jnp.complex64)

        if E is None:
            self.E = jnp.zeros((3, *size), dtype=jnp.float32)
        else:
            self.E = E.astype(jnp.float32)

    def copy(self) -> "GaugeLattice":
        """Create a copy of the gauge lattice."""
        return GaugeLattice(
            size=self.size,
            length=self.length[0],
            g=self.g,
            links=self.links.copy(),
            E=self.E.copy(),
            a=self.a,
            H=self.H,
        )

    def update(
        self,
        links: Array,
        E: Array,
        a: Optional[float] = None,
        H: Optional[float] = None,
    ) -> "GaugeLattice":
        """Create a new GaugeLattice with updated field values."""
        return GaugeLattice(
            size=self.size,
            length=self.length[0],
            g=self.g,
            links=links,
            E=E,
            a=a if a is not None else self.a,
            H=H if H is not None else self.H,
        )

    def replace(self, **kwargs) -> "GaugeLattice":
        """Create a new GaugeLattice with specified fields replaced."""
        return GaugeLattice(
            size=kwargs.get("size", self.size),
            length=kwargs.get("length", self.length[0]),
            g=kwargs.get("g", self.g),
            links=kwargs.get("links", self.links),
            E=kwargs.get("E", self.E),
            a=kwargs.get("a", self.a),
            H=kwargs.get("H", self.H),
        )

    @property
    def volume(self) -> float:
        """Total physical volume."""
        return self.length[0] ** 3

    @property
    def dV(self) -> float:
        """Volume element."""
        return self.dx**3

    def __repr__(self) -> str:
        return (
            f"GaugeLattice(size={self.size}, length={self.length[0]:.2f}, "
            f"dx={self.dx:.4f}, g={self.g:.2f}, a={self.a:.3f}, H={self.H:.3e})"
        )
