"""
Band structure computation for DiffTB.

This module provides:

- Band structure along high-symmetry paths
- Fermi surface calculations
- Band topology analysis
"""

from __future__ import annotations

from functools import partial

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

from difftb.core.eigensolver import solve_eigenproblem
from difftb.core.lattice import BrillouinZone


class BandStructure(NamedTuple):
    """
    Band structure results.

    Attributes
    ----------
    kpoints : Array
        k-points along path, shape (nk, 3).
    kdist : Array
        Cumulative distance along path, shape (nk,).
    eigenvalues : Array
        Band energies, shape (nk, n_bands).
    eigenvectors : Array, optional
        Band wavefunctions, shape (nk, n_orb, n_bands).
    high_symmetry_points : List[Tuple[str, float]], optional
        List of (name, kdist) for labeling.
    fermi_level : float, optional
        Fermi level if computed.
    """

    kpoints: Array
    kdist: Array
    eigenvalues: Array
    eigenvectors: Optional[Array] = None
    high_symmetry_points: Optional[List[Tuple[str, float]]] = None
    fermi_level: Optional[float] = None


def compute_bands(
    hamiltonian_fn: Callable[[Array], Tuple[Array, Array]],
    bz: BrillouinZone,
    return_eigenvectors: bool = False,
) -> BandStructure:
    """
    Compute band structure along k-path.

    Parameters
    ----------
    hamiltonian_fn : Callable[[Array], Tuple[Array, Array]]
        Function k -> (H(k), S(k)).
    bz : BrillouinZone
        BrillouinZone with k-path.
    return_eigenvectors : bool, optional
        Whether to return wavefunctions. Default is False.

    Returns
    -------
    bands : BandStructure
        BandStructure object.
    """
    kpoints = bz.kpoints
    nk = len(kpoints)

    # Compute Hamiltonian at each k-point
    eigenvalues_list = []
    eigenvectors_list = []

    for i in range(nk):
        k = kpoints[i]
        H_k, S_k = hamiltonian_fn(k)

        result = solve_eigenproblem(H_k, S_k)
        eigenvalues_list.append(result.eigenvalues)
        if return_eigenvectors:
            eigenvectors_list.append(result.eigenvectors)

    eigenvalues = jnp.stack(eigenvalues_list, axis=0)
    eigenvectors = (
        jnp.stack(eigenvectors_list, axis=0) if return_eigenvectors else None
    )

    # Compute k-distance for plotting
    kdist = compute_kdist(kpoints, bz.lattice.reciprocal)

    return BandStructure(
        kpoints=kpoints,
        kdist=kdist,
        eigenvalues=eigenvalues,
        eigenvectors=eigenvectors,
    )


@partial(jax.jit, static_argnums=(1,))
def compute_bands_vmapped(
    kpoints: Array,
    hamiltonian_fn: Callable[[Array], Tuple[Array, Array]],
) -> Tuple[Array, Array]:
    """
    Vectorized band computation using vmap.

    More efficient for large k-point sets.

    Parameters
    ----------
    kpoints : Array
        Array of k-points, shape (nk, 3).
    hamiltonian_fn : Callable[[Array], Tuple[Array, Array]]
        Function k -> (H(k), S(k)).

    Returns
    -------
    eigenvalues : Array
        Shape (nk, n_bands).
    eigenvectors : Array
        Shape (nk, n_orb, n_bands).
    """

    def single_k(k):
        H_k, S_k = hamiltonian_fn(k)
        result = solve_eigenproblem(H_k, S_k)
        return result.eigenvalues, result.eigenvectors

    return jax.vmap(single_k)(kpoints)


@jax.jit
def compute_kdist(kpoints: Array, reciprocal: Array) -> Array:
    """
    Compute cumulative distance along k-path.

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

    Returns
    -------
    kdist : Array
        Cumulative distance, shape (nk,).
    """
    # Convert to Cartesian
    kpoints_cart = kpoints @ reciprocal

    # Compute distances between consecutive points
    dk = jnp.diff(kpoints_cart, axis=0)
    distances = jnp.linalg.norm(dk, axis=1)

    # Cumulative distance
    return jnp.concatenate([jnp.array([0.0]), jnp.cumsum(distances)])


# =============================================================================
# Specialized band structure computations
# =============================================================================


def compute_surface_bands(
    hamiltonian_fn: Callable,
    bz: BrillouinZone,
    n_layers: int,
    surface_orbital_indices: Array,
) -> Tuple[BandStructure, Array]:
    """
    Compute surface-projected band structure for slab geometry.

    Parameters
    ----------
    hamiltonian_fn : Callable
        Slab Hamiltonian function.
    bz : BrillouinZone
        Brillouin zone with k-path.
    n_layers : int
        Number of layers in slab.
    surface_orbital_indices : Array
        Indices of surface orbitals.

    Returns
    -------
    bands : BandStructure
        Band structure.
    surface_weight : Array
        Surface character of each band, shape (nk, n_bands).
    """
    bands = compute_bands(hamiltonian_fn, bz, return_eigenvectors=True)

    # Compute surface weight for each band
    # weight = sum over surface orbitals of |c_i|^2
    surface_weight = jnp.sum(
        jnp.abs(bands.eigenvectors[:, surface_orbital_indices, :]) ** 2, axis=1
    )

    return bands, surface_weight


def compute_spin_texture(bands: BandStructure, spin_matrices: Array) -> Array:
    """
    Compute spin texture <S> for each band at each k-point.

    Parameters
    ----------
    bands : BandStructure
        Band structure with eigenvectors.
    spin_matrices : Array
        Pauli matrices, shape (3, n_orb, n_orb).

    Returns
    -------
    spin_texture : Array
        <S_x>, <S_y>, <S_z> for each band, shape (nk, n_bands, 3).
    """
    if bands.eigenvectors is None:
        raise ValueError("Band structure must include eigenvectors")

    # <ψ_n|S_i|ψ_n>
    return jnp.einsum(
        "iab,kbn,kan->kni",
        spin_matrices,
        bands.eigenvectors,
        jnp.conj(bands.eigenvectors),
    ).real


# =============================================================================
# Band gap and topology
# =============================================================================


def find_band_gap(
    bands: BandStructure, n_occupied: int
) -> Tuple[float, Array, Array]:
    """
    Find band gap and its location in k-space.

    Parameters
    ----------
    bands : BandStructure
        Band structure.
    n_occupied : int
        Number of occupied bands.

    Returns
    -------
    gap : float
        Band gap in eV.
    k_vbm : Array
        k-point of valence band maximum.
    k_cbm : Array
        k-point of conduction band minimum.
    """
    vb = bands.eigenvalues[:, n_occupied - 1]  # Valence band
    cb = bands.eigenvalues[:, n_occupied]  # Conduction band

    vbm = jnp.max(vb)
    cbm = jnp.min(cb)

    k_vbm = bands.kpoints[jnp.argmax(vb)]
    k_cbm = bands.kpoints[jnp.argmin(cb)]

    gap = cbm - vbm

    return float(gap), k_vbm, k_cbm


def is_direct_gap(
    bands: BandStructure, n_occupied: int, tolerance: float = 0.01
) -> bool:
    """
    Check if band gap is direct (VBM and CBM at same k).

    Parameters
    ----------
    bands : BandStructure
        Band structure.
    n_occupied : int
        Number of occupied bands.
    tolerance : float, optional
        k-space tolerance for "same point". Default is 0.01.

    Returns
    -------
    is_direct : bool
        True if direct gap.
    """
    gap, k_vbm, k_cbm = find_band_gap(bands, n_occupied)

    dk = jnp.linalg.norm(k_vbm - k_cbm)
    return dk < tolerance


# =============================================================================
# Effective mass
# =============================================================================


def effective_mass(
    bands: BandStructure, band_index: int, k_index: int, direction: Array
) -> float:
    """
    Compute effective mass at a k-point using finite differences.

    m* = ℏ² / (d²E/dk²)

    Parameters
    ----------
    bands : BandStructure
        Band structure.
    band_index : int
        Which band.
    k_index : int
        Which k-point (should be extremum).
    direction : Array
        Direction for mass calculation (normalized).

    Returns
    -------
    m_eff : float
        Effective mass in units of electron mass.
    """
    # Use neighboring points for second derivative
    # This is approximate - proper implementation would use
    # finer k-grid around the point

    E = bands.eigenvalues[:, band_index]
    kdist = bands.kdist

    # Second derivative by finite difference
    if k_index == 0:
        k_index = 1
    elif k_index == len(kdist) - 1:
        k_index = len(kdist) - 2

    dk1 = kdist[k_index] - kdist[k_index - 1]
    dk2 = kdist[k_index + 1] - kdist[k_index]

    d2E_dk2 = (
        2
        * (
            (E[k_index + 1] - E[k_index]) / dk2
            - (E[k_index] - E[k_index - 1]) / dk1
        )
        / (dk1 + dk2)
    )

    # Convert to effective mass
    # m* = ℏ² / d²E/dk² with proper units
    # ℏ² = 7.62 eV·Å² / m_e
    hbar_sq = 7.62  # eV·Å²·m_e

    m_eff = hbar_sq / (d2E_dk2 + 1e-10)

    return float(m_eff)


# =============================================================================
# Fermi velocity (for Dirac materials)
# =============================================================================


def fermi_velocity(
    bands: BandStructure, band_index: int, k_index: int
) -> float:
    """
    Compute Fermi velocity v_F = (1/ℏ) dE/dk.

    Relevant for Dirac/Weyl materials near the Dirac point.

    Parameters
    ----------
    bands : BandStructure
        Band structure.
    band_index : int
        Which band.
    k_index : int
        Which k-point.

    Returns
    -------
    v_F : float
        Fermi velocity in m/s.
    """
    E = bands.eigenvalues[:, band_index]
    kdist = bands.kdist

    # First derivative
    if k_index == 0:
        dE_dk = (E[1] - E[0]) / (kdist[1] - kdist[0])
    elif k_index == len(kdist) - 1:
        dE_dk = (E[-1] - E[-2]) / (kdist[-1] - kdist[-2])
    else:
        dE_dk = (E[k_index + 1] - E[k_index - 1]) / (
            kdist[k_index + 1] - kdist[k_index - 1]
        )

    # Convert to m/s
    # v_F = dE/dk / ℏ
    # dE in eV, dk in 1/Å, ℏ = 6.582e-16 eV·s
    # 1 Å = 1e-10 m
    # v_F [m/s] = dE/dk [eV·Å] / ℏ [eV·s] * 1e-10 [m/Å]
    #           = dE/dk * 1e-10 / 6.582e-16
    #           = dE/dk * 1.52e5

    v_F = jnp.abs(dE_dk) * 1.52e5

    return float(v_F)
