"""
Lattice and Brillouin zone handling for periodic systems.

This module provides functionality for:

- Real-space and reciprocal-space lattice vectors
- Brillouin zone construction and k-point sampling
- High-symmetry path generation
"""

from __future__ import annotations

import jax
import jax.numpy as jnp
from beartype.typing import List, NamedTuple, Tuple
from jax import Array

# Type aliases
RealVec = Array  # Shape (3,)
ReciprocalVec = Array  # Shape (3,)
LatticeMatrix = Array  # Shape (3, 3)


class Lattice(NamedTuple):
    """
    Crystal lattice specification.

    Attributes
    ----------
    vectors : LatticeMatrix
        Real-space lattice vectors as rows, shape (3, 3).
        vectors[i] is the i-th lattice vector.
    reciprocal : LatticeMatrix
        Reciprocal lattice vectors as rows, shape (3, 3).
        Computed automatically from real-space vectors.
    """

    vectors: LatticeMatrix
    reciprocal: LatticeMatrix

    @classmethod
    def from_vectors(cls, a1: RealVec, a2: RealVec, a3: RealVec) -> Lattice:
        """
        Create lattice from three lattice vectors.

        Parameters
        ----------
        a1 : RealVec
            First lattice vector in Angstroms.
        a2 : RealVec
            Second lattice vector in Angstroms.
        a3 : RealVec
            Third lattice vector in Angstroms.

        Returns
        -------
        lattice : Lattice
            Lattice object with computed reciprocal vectors.
        """
        vectors = jnp.stack([a1, a2, a3], axis=0)
        reciprocal = compute_reciprocal_lattice(vectors)
        return cls(vectors=vectors, reciprocal=reciprocal)

    @classmethod
    def from_parameters(
        cls,
        a: float,
        b: float,
        c: float,
        alpha: float,
        beta: float,
        gamma: float,
    ) -> Lattice:
        """
        Create lattice from lattice parameters.

        Parameters
        ----------
        a : float
            Lattice constant a in Angstroms.
        b : float
            Lattice constant b in Angstroms.
        c : float
            Lattice constant c in Angstroms.
        alpha : float
            Angle alpha in degrees.
        beta : float
            Angle beta in degrees.
        gamma : float
            Angle gamma in degrees.

        Returns
        -------
        lattice : Lattice
            Lattice object.
        """
        vectors = lattice_vectors_from_parameters(a, b, c, alpha, beta, gamma)
        reciprocal = compute_reciprocal_lattice(vectors)
        return cls(vectors=vectors, reciprocal=reciprocal)

    @property
    def volume(self) -> float:
        """Unit cell volume in Angstrom^3."""
        return jnp.abs(jnp.linalg.det(self.vectors))


class BrillouinZone(NamedTuple):
    """
    Brillouin zone with k-point sampling.

    Attributes
    ----------
    lattice : Lattice
        Associated crystal lattice.
    kpoints : Array
        Array of k-points in fractional coordinates, shape (nk, 3).
    weights : Array
        Integration weights for each k-point, shape (nk,).
    """

    lattice: Lattice
    kpoints: Array
    weights: Array

    @classmethod
    def monkhorst_pack(
        cls,
        lattice: Lattice,
        grid: Tuple[int, int, int],
        shift: Tuple[float, float, float] = (0.0, 0.0, 0.0),
    ) -> BrillouinZone:
        """
        Generate Monkhorst-Pack k-point grid.

        Parameters
        ----------
        lattice : Lattice
            Crystal lattice.
        grid : Tuple[int, int, int]
            Number of k-points along each reciprocal lattice vector.
        shift : Tuple[float, float, float], optional
            Optional shift of the grid (in fractional coordinates).
            Default is (0.0, 0.0, 0.0).

        Returns
        -------
        bz : BrillouinZone
            BrillouinZone with uniform k-point grid.
        """
        kpoints, weights = monkhorst_pack_grid(grid, shift)
        return cls(lattice=lattice, kpoints=kpoints, weights=weights)

    @classmethod
    def from_path(
        cls,
        lattice: Lattice,
        path_points: List[Tuple[str, Array]],
        num_points: int = 50,
    ) -> BrillouinZone:
        """
        Generate k-points along a high-symmetry path.

        Parameters
        ----------
        lattice : Lattice
            Crystal lattice.
        path_points : List[Tuple[str, Array]]
            List of (label, k-point) tuples defining the path.
        num_points : int, optional
            Number of points between each pair of high-symmetry points.
            Default is 50.

        Returns
        -------
        bz : BrillouinZone
            BrillouinZone with k-points along the path.
        """
        kpoints = generate_kpath(path_points, num_points)
        weights = jnp.ones(len(kpoints)) / len(kpoints)
        return cls(lattice=lattice, kpoints=kpoints, weights=weights)

    @property
    def nkpoints(self) -> int:
        """Number of k-points."""
        return len(self.kpoints)

    def to_cartesian(self) -> Array:
        """
        Convert k-points to Cartesian coordinates (1/Angstrom).

        Returns
        -------
        kpoints_cart : Array
            k-points in Cartesian coordinates, shape (nk, 3).
        """
        return kpoints_to_cartesian(self.kpoints, self.lattice.reciprocal)


# =============================================================================
# Pure functions for lattice operations
# =============================================================================


@jax.jit
def compute_reciprocal_lattice(vectors: LatticeMatrix) -> LatticeMatrix:
    """
    Compute reciprocal lattice vectors.

    The reciprocal lattice vectors b_i satisfy: a_i · b_j = 2π δ_ij.

    Parameters
    ----------
    vectors : LatticeMatrix
        Real-space lattice vectors as rows, shape (3, 3).

    Returns
    -------
    reciprocal : LatticeMatrix
        Reciprocal lattice vectors as rows, shape (3, 3).
    """
    # b = 2π (a^T)^{-1}
    return 2 * jnp.pi * jnp.linalg.inv(vectors).T


def lattice_vectors_from_parameters(
    a: float, b: float, c: float, alpha: float, beta: float, gamma: float
) -> LatticeMatrix:
    """
    Convert lattice parameters to lattice vectors.

    Uses convention where a is along x, b is in xy-plane.

    Parameters
    ----------
    a : float
        Lattice constant a in Angstroms.
    b : float
        Lattice constant b in Angstroms.
    c : float
        Lattice constant c in Angstroms.
    alpha : float
        Angle alpha in degrees.
    beta : float
        Angle beta in degrees.
    gamma : float
        Angle gamma in degrees.

    Returns
    -------
    vectors : LatticeMatrix
        Lattice vectors as rows, shape (3, 3).
    """
    # Convert to radians
    alpha_r = jnp.deg2rad(alpha)
    beta_r = jnp.deg2rad(beta)
    gamma_r = jnp.deg2rad(gamma)

    # First vector along x
    a1 = jnp.array([a, 0.0, 0.0])

    # Second vector in xy-plane
    a2 = jnp.array([b * jnp.cos(gamma_r), b * jnp.sin(gamma_r), 0.0])

    # Third vector
    cx = c * jnp.cos(beta_r)
    cy = (
        c
        * (jnp.cos(alpha_r) - jnp.cos(beta_r) * jnp.cos(gamma_r))
        / jnp.sin(gamma_r)
    )
    cz = jnp.sqrt(c**2 - cx**2 - cy**2)
    a3 = jnp.array([cx, cy, cz])

    return jnp.stack([a1, a2, a3], axis=0)


def monkhorst_pack_grid(
    grid: Tuple[int, int, int],
    shift: Tuple[float, float, float] = (0.0, 0.0, 0.0),
) -> Tuple[Array, Array]:
    """
    Generate Monkhorst-Pack k-point grid.

    Parameters
    ----------
    grid : Tuple[int, int, int]
        (n1, n2, n3) grid dimensions.
    shift : Tuple[float, float, float], optional
        Shift in fractional coordinates. Default is (0.0, 0.0, 0.0).

    Returns
    -------
    kpoints : Array
        Array of shape (n1*n2*n3, 3).
    weights : Array
        Array of shape (n1*n2*n3,).
    """
    n1, n2, n3 = grid
    s1, s2, s3 = shift

    # Generate grid points
    k1 = (2 * jnp.arange(n1) - n1 + 1 + 2 * s1) / (2 * n1)
    k2 = (2 * jnp.arange(n2) - n2 + 1 + 2 * s2) / (2 * n2)
    k3 = (2 * jnp.arange(n3) - n3 + 1 + 2 * s3) / (2 * n3)

    # Create meshgrid and reshape
    K1, K2, K3 = jnp.meshgrid(k1, k2, k3, indexing="ij")
    kpoints = jnp.stack([K1.ravel(), K2.ravel(), K3.ravel()], axis=-1)

    # Uniform weights
    nk = n1 * n2 * n3
    weights = jnp.ones(nk) / nk

    return kpoints, weights


def generate_kpath(
    path_points: List[Tuple[str, Array]], num_points: int = 50
) -> Array:
    """
    Generate k-points along a high-symmetry path.

    Parameters
    ----------
    path_points : List[Tuple[str, Array]]
        List of (label, k-point) tuples.
    num_points : int, optional
        Points between each pair of high-symmetry points. Default is 50.

    Returns
    -------
    kpoints : Array
        Array of k-points, shape (n_total, 3).
    """
    kpoints_list = []

    for i in range(len(path_points) - 1):
        k_start = jnp.asarray(path_points[i][1])
        k_end = jnp.asarray(path_points[i + 1][1])

        # Linear interpolation
        t = jnp.linspace(
            0, 1, num_points, endpoint=(i == len(path_points) - 2)
        )
        segment = k_start[None, :] + t[:, None] * (k_end - k_start)[None, :]
        kpoints_list.append(segment)

    return jnp.concatenate(kpoints_list, axis=0)


@jax.jit
def kpoints_to_cartesian(kpoints: Array, reciprocal: LatticeMatrix) -> Array:
    """
    Convert k-points from fractional to Cartesian coordinates.

    Parameters
    ----------
    kpoints : Array
        k-points in fractional coordinates, shape (..., 3).
    reciprocal : LatticeMatrix
        Reciprocal lattice vectors as rows, shape (3, 3).

    Returns
    -------
    kpoints_cart : Array
        k-points in Cartesian coordinates (1/Angstrom), shape (..., 3).
    """
    return kpoints @ reciprocal


# =============================================================================
# Common high-symmetry points for specific lattices
# =============================================================================


def hexagonal_path() -> List[Tuple[str, Array]]:
    """High-symmetry path for hexagonal lattice (e.g., Bi2Te3)."""
    return [
        ("Γ", jnp.array([0.0, 0.0, 0.0])),
        ("M", jnp.array([0.5, 0.0, 0.0])),
        ("K", jnp.array([1 / 3, 1 / 3, 0.0])),
        ("Γ", jnp.array([0.0, 0.0, 0.0])),
        ("A", jnp.array([0.0, 0.0, 0.5])),
        ("L", jnp.array([0.5, 0.0, 0.5])),
        ("H", jnp.array([1 / 3, 1 / 3, 0.5])),
        ("A", jnp.array([0.0, 0.0, 0.5])),
    ]


def rhombohedral_path() -> List[Tuple[str, Array]]:
    """High-symmetry path for rhombohedral lattice."""
    return [
        ("Γ", jnp.array([0.0, 0.0, 0.0])),
        ("Z", jnp.array([0.5, 0.5, 0.5])),
        ("F", jnp.array([0.5, 0.5, 0.0])),
        ("Γ", jnp.array([0.0, 0.0, 0.0])),
        ("L", jnp.array([0.5, 0.0, 0.0])),
    ]
