"""
Differentiable eigensolver for tight-binding calculations.

This module provides:

- Generalized eigenvalue problem solver (H·c = E·S·c)
- Stable gradients near degeneracies via Lorentzian broadening
- Batched eigensolvers for multiple k-points
"""

from __future__ import annotations

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


class EigenResult(NamedTuple):
    """
    Result of eigenvalue computation.

    Attributes
    ----------
    eigenvalues : Array
        Eigenvalues (energies), shape (..., n_orbitals).
    eigenvectors : Array
        Eigenvectors (coefficients), shape (..., n_orbitals, n_orbitals).
        eigenvectors[..., :, i] is the i-th eigenvector.
    """

    eigenvalues: Array
    eigenvectors: Array


# =============================================================================
# Core eigensolvers
# =============================================================================


@jax.jit
def solve_eigenproblem(
    H: Array, S: Optional[Array] = None, hermitian: bool = True
) -> EigenResult:
    """
    Solve the (generalized) eigenvalue problem.

    For standard eigenvalue problem (S=None): H·c = E·c
    For generalized eigenvalue problem: H·c = E·S·c

    Parameters
    ----------
    H : Array
        Hamiltonian matrix, shape (..., n, n).
    S : Array, optional
        Overlap matrix, shape (..., n, n) or None for orthogonal basis.
    hermitian : bool, optional
        If True, use Hermitian eigensolver (faster, more stable).
        Default is True.

    Returns
    -------
    result : EigenResult
        EigenResult with eigenvalues and eigenvectors.
    """
    if S is None:
        # Standard eigenvalue problem
        if hermitian:
            eigenvalues, eigenvectors = jnp.linalg.eigh(H)
        else:
            eigenvalues, eigenvectors = jnp.linalg.eig(H)
    else:
        # Generalized eigenvalue problem: H·c = E·S·c
        # Transform to standard form via Löwdin orthogonalization
        eigenvalues, eigenvectors = solve_generalized_eigh(H, S)

    return EigenResult(eigenvalues=eigenvalues, eigenvectors=eigenvectors)


@jax.jit
def solve_generalized_eigh(H: Array, S: Array) -> Tuple[Array, Array]:
    """
    Solve generalized Hermitian eigenvalue problem H·c = E·S·c.

    Uses Löwdin symmetric orthogonalization:

    1. Diagonalize S = U·s·U†
    2. Form S^{-1/2} = U·s^{-1/2}·U†
    3. Transform: H' = S^{-1/2}·H·S^{-1/2}
    4. Solve standard problem: H'·c' = E·c'
    5. Back-transform: c = S^{-1/2}·c'

    Parameters
    ----------
    H : Array
        Hamiltonian matrix, shape (..., n, n).
    S : Array
        Overlap matrix, shape (..., n, n).

    Returns
    -------
    eigenvalues : Array
        Shape (..., n).
    eigenvectors : Array
        Shape (..., n, n).
    """
    # Diagonalize overlap matrix
    s_eig, U = jnp.linalg.eigh(S)

    # Regularize small eigenvalues to avoid numerical issues
    s_eig = jnp.maximum(s_eig, 1e-10)

    # S^{-1/2} in eigenbasis
    s_inv_sqrt = 1.0 / jnp.sqrt(s_eig)

    # Transform Hamiltonian: H' = S^{-1/2} H S^{-1/2}
    # S^{-1/2} = U @ diag(s^{-1/2}) @ U†
    S_inv_sqrt = U * s_inv_sqrt[..., None, :] @ jnp.conj(U).swapaxes(-2, -1)

    H_prime = S_inv_sqrt @ H @ S_inv_sqrt

    # Solve standard eigenvalue problem
    eigenvalues, c_prime = jnp.linalg.eigh(H_prime)

    # Back-transform eigenvectors
    eigenvectors = S_inv_sqrt @ c_prime

    return eigenvalues, eigenvectors


# =============================================================================
# Gradient-stable spectral functions
# =============================================================================


@jax.jit
def lorentzian_broadening(
    E: Array, eigenvalues: Array, eta: float = 0.01
) -> Array:
    """
    Lorentzian broadening for spectral function.

    A(E) = -1/π Im sum_n 1/(E - E_n + iη)
         = 1/π sum_n η / [(E - E_n)² + η²]

    This provides smooth, differentiable spectral functions even near
    degenerate eigenvalues.

    Parameters
    ----------
    E : Array
        Energy points, shape (n_E,).
    eigenvalues : Array
        Eigenvalues, shape (..., n_orbitals).
    eta : float, optional
        Broadening parameter in eV. Default is 0.01.

    Returns
    -------
    spectral : Array
        Spectral function, shape (..., n_E).
    """
    # E: (n_E,), eigenvalues: (..., n_orb) -> diff: (..., n_orb, n_E)
    diff = E[None, :] - eigenvalues[..., :, None]

    # Lorentzian for each eigenvalue
    lorentzians = eta / (diff**2 + eta**2) / jnp.pi

    # Sum over eigenvalues
    return jnp.sum(lorentzians, axis=-2)


@jax.jit
def gaussian_broadening(
    E: Array, eigenvalues: Array, sigma: float = 0.01
) -> Array:
    """
    Gaussian broadening for spectral function.

    A(E) = sum_n 1/(σ√(2π)) exp[-(E - E_n)²/(2σ²)]

    Parameters
    ----------
    E : Array
        Energy points, shape (n_E,).
    eigenvalues : Array
        Eigenvalues, shape (..., n_orbitals).
    sigma : float, optional
        Broadening parameter in eV. Default is 0.01.

    Returns
    -------
    spectral : Array
        Spectral function, shape (..., n_E).
    """
    diff = E[None, :] - eigenvalues[..., :, None]

    gaussians = jnp.exp(-0.5 * (diff / sigma) ** 2) / (
        sigma * jnp.sqrt(2 * jnp.pi)
    )

    return jnp.sum(gaussians, axis=-2)


# =============================================================================
# Orbital-resolved spectral functions
# =============================================================================


@jax.jit
def orbital_spectral_function(
    E: Array,
    eigenvalues: Array,
    eigenvectors: Array,
    orbital_indices: Array,
    eta: float = 0.01,
) -> Array:
    """
    Compute orbital-projected spectral function.

    A_μ(E) = sum_n |c_{n,μ}|² δ(E - E_n)

    where c_{n,μ} is the coefficient of orbital μ in eigenstate n.

    Parameters
    ----------
    E : Array
        Energy points, shape (n_E,).
    eigenvalues : Array
        Eigenvalues, shape (..., n_orbitals).
    eigenvectors : Array
        Eigenvectors, shape (..., n_orbitals, n_orbitals).
    orbital_indices : Array
        Which orbitals to project onto, shape (n_proj,).
    eta : float, optional
        Broadening parameter. Default is 0.01.

    Returns
    -------
    spectral : Array
        Orbital-projected spectral function, shape (..., n_E).
    """
    # Get weights for specified orbitals
    # eigenvectors[..., μ, n] = coefficient of orbital μ in state n
    weights = jnp.sum(
        jnp.abs(eigenvectors[..., orbital_indices, :]) ** 2, axis=-2
    )

    # weights shape: (..., n_orbitals) - weight of each eigenstate

    # Compute weighted spectral function
    diff = E[None, :] - eigenvalues[..., :, None]
    lorentzians = eta / (diff**2 + eta**2) / jnp.pi

    # Weight each Lorentzian by orbital character
    weighted = lorentzians * weights[..., :, None]

    return jnp.sum(weighted, axis=-2)


@jax.jit
def spin_spectral_function(
    E: Array,
    eigenvalues: Array,
    eigenvectors: Array,
    spin_operator: Array,
    eta: float = 0.01,
) -> Array:
    """
    Compute spin-projected spectral function.

    A_s(E) = sum_n <n|S_z|n> δ(E - E_n)

    Parameters
    ----------
    E : Array
        Energy points, shape (n_E,).
    eigenvalues : Array
        Eigenvalues, shape (..., n_orbitals).
    eigenvectors : Array
        Eigenvectors, shape (..., n_orbitals, n_orbitals).
    spin_operator : Array
        Spin operator matrix, shape (n_orbitals, n_orbitals).
    eta : float, optional
        Broadening parameter. Default is 0.01.

    Returns
    -------
    spectral : Array
        Spin-projected spectral function, shape (..., n_E).
    """
    # Compute <n|S|n> for each eigenstate
    # spin_expectation[..., n] = <ψ_n|S|ψ_n>
    psi = eigenvectors  # (..., n_orb, n_states)
    spin_expectation = jnp.einsum(
        "...in,...ij,...jn->...n", jnp.conj(psi), spin_operator, psi
    ).real

    # Weighted spectral function
    diff = E[None, :] - eigenvalues[..., :, None]
    lorentzians = eta / (diff**2 + eta**2) / jnp.pi

    weighted = lorentzians * spin_expectation[..., :, None]

    return jnp.sum(weighted, axis=-2)


# =============================================================================
# Density of states
# =============================================================================


@jax.jit
def density_of_states(
    E: Array,
    eigenvalues: Array,
    weights: Optional[Array] = None,
    eta: float = 0.01,
) -> Array:
    """
    Compute density of states from eigenvalues.

    DOS(E) = sum_k w_k sum_n δ(E - E_n(k))

    Parameters
    ----------
    E : Array
        Energy points, shape (n_E,).
    eigenvalues : Array
        Eigenvalues, shape (n_k, n_orbitals).
    weights : Array, optional
        k-point weights, shape (n_k,) or None for uniform.
    eta : float, optional
        Broadening parameter. Default is 0.01.

    Returns
    -------
    dos : Array
        DOS, shape (n_E,).
    """
    n_k = eigenvalues.shape[0]

    if weights is None:
        weights = jnp.ones(n_k) / n_k

    # Spectral function at each k-point
    dos_k = lorentzian_broadening(E, eigenvalues, eta)  # (n_k, n_E)

    # Weight and sum
    return jnp.sum(dos_k * weights[:, None], axis=0)


# =============================================================================
# Fermi level finding
# =============================================================================


@jax.jit
def find_fermi_level(
    eigenvalues: Array,
    n_electrons: float,
    weights: Optional[Array] = None,
    temperature: float = 0.001,
) -> float:
    """
    Find Fermi level for given number of electrons.

    Uses bisection with Fermi-Dirac occupations.

    Parameters
    ----------
    eigenvalues : Array
        Eigenvalues, shape (n_k, n_orbitals).
    n_electrons : float
        Target number of electrons.
    weights : Array, optional
        k-point weights.
    temperature : float, optional
        Electronic temperature in eV (for smearing). Default is 0.001.

    Returns
    -------
    E_F : float
        Fermi level in eV.
    """
    n_k = eigenvalues.shape[0]

    if weights is None:
        weights = jnp.ones(n_k) / n_k

    # Bounds for bisection
    E_min = jnp.min(eigenvalues) - 1.0
    E_max = jnp.max(eigenvalues) + 1.0

    def count_electrons(E_F):
        # Fermi-Dirac occupation
        x = (eigenvalues - E_F) / temperature
        x = jnp.clip(x, -50, 50)  # Avoid overflow
        f = 1.0 / (1.0 + jnp.exp(x))

        # Factor of 2 for spin (if not spin-polarized)
        return 2.0 * jnp.sum(f * weights[:, None])

    # Bisection
    for _ in range(50):
        E_mid = 0.5 * (E_min + E_max)
        n_e = count_electrons(E_mid)

        if n_e < n_electrons:
            E_min = E_mid
        else:
            E_max = E_mid

    return 0.5 * (E_min + E_max)


# =============================================================================
# Utility functions
# =============================================================================


def get_band_gap(
    eigenvalues: Array, n_electrons: float, weights: Optional[Array] = None
) -> Tuple[float, float, float]:
    """
    Compute band gap from eigenvalues.

    Parameters
    ----------
    eigenvalues : Array
        Eigenvalues, shape (n_k, n_orbitals).
    n_electrons : float
        Number of electrons.
    weights : Array, optional
        k-point weights.

    Returns
    -------
    gap : float
        Band gap in eV.
    vbm : float
        Valence band maximum.
    cbm : float
        Conduction band minimum.
    """
    E_F = find_fermi_level(eigenvalues, n_electrons, weights)

    # Find VBM (highest occupied) and CBM (lowest unoccupied)
    occupied = eigenvalues < E_F

    vbm = jnp.max(jnp.where(occupied, eigenvalues, -jnp.inf))
    cbm = jnp.min(jnp.where(~occupied, eigenvalues, jnp.inf))

    gap = cbm - vbm

    return float(gap), float(vbm), float(cbm)
