"""
Spectral function computation for ARPES simulation.

This module provides:

- A(k, E) spectral function from tight-binding
- Matrix element effects
- Interface to ARPESim
"""

from __future__ import annotations

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


class ARPESSpectrum(NamedTuple):
    """
    ARPES spectrum data.

    Attributes
    ----------
    kpoints : Array
        k-points, shape (nk, 3).
    energies : Array
        Energy values, shape (nE,).
    intensity : Array
        ARPES intensity I(k, E), shape (nk, nE).
    spin_polarization : Array, optional
        Spin polarization P(k, E), shape (nk, nE, 3).
    """

    kpoints: Array
    energies: Array
    intensity: Array
    spin_polarization: Optional[Array] = None


@jax.jit
def spectral_function(
    eigenvalues: Array, eigenvectors: Array, E: Array, eta: float = 0.01
) -> Array:
    """
    Compute the spectral function A(k, E).

    A(k, E) = -1/π Im G(k, E) = sum_n |ψ_n(k)|² δ(E - E_n(k))

    The delta function is broadened to a Lorentzian for smoothness.

    Parameters
    ----------
    eigenvalues : Array
        Band energies, shape (nk, n_bands).
    eigenvectors : Array
        Eigenvectors, shape (nk, n_orb, n_bands).
    E : Array
        Energy values, shape (nE,).
    eta : float, optional
        Lorentzian broadening in eV. Default is 0.01.

    Returns
    -------
    spectral : Array
        A(k, E), shape (nk, nE).
    """
    # eigenvalues: (nk, n_bands)
    # E: (nE,)
    # diff: (nk, n_bands, nE)
    diff = E[None, None, :] - eigenvalues[:, :, None]

    # Lorentzian broadening
    spectral = eta / (diff**2 + eta**2) / jnp.pi

    # Sum over bands (each band contributes equally for now)
    return jnp.sum(spectral, axis=1)


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

    A_orb(k, E) = sum_n |<orb|ψ_n(k)>|² δ(E - E_n(k))

    Parameters
    ----------
    eigenvalues : Array
        Band energies, shape (nk, n_bands).
    eigenvectors : Array
        Eigenvectors, shape (nk, n_orb, n_bands).
    E : Array
        Energy values, shape (nE,).
    orbital_weights : Array
        Projection onto orbitals, shape (n_orb,).
    eta : float, optional
        Lorentzian broadening. Default is 0.01.

    Returns
    -------
    spectral : Array
        Orbital-resolved A(k, E), shape (nk, nE).
    """
    # Compute orbital character of each band
    # |<orb|ψ_n>|² summed over specified orbitals
    # eigenvectors: (nk, n_orb, n_bands)
    # orbital_weights: (n_orb,)

    orbital_character = jnp.einsum(
        "o,kob->kb", orbital_weights, jnp.abs(eigenvectors) ** 2
    )  # (nk, n_bands)

    # Spectral function weighted by orbital character
    diff = E[None, None, :] - eigenvalues[:, :, None]
    lorentzians = eta / (diff**2 + eta**2) / jnp.pi  # (nk, n_bands, nE)

    weighted = lorentzians * orbital_character[:, :, None]

    return jnp.sum(weighted, axis=1)


@jax.jit
def spin_resolved_spectral(
    eigenvalues: Array,
    eigenvectors: Array,
    E: Array,
    spin_matrices: Array,
    eta: float = 0.01,
) -> Tuple[Array, Array]:
    """
    Compute spin-resolved spectral function.

    Returns both total intensity and spin polarization vector.

    Parameters
    ----------
    eigenvalues : Array
        Band energies, shape (nk, n_bands).
    eigenvectors : Array
        Eigenvectors in spinor basis, shape (nk, n_orb, n_bands).
        For spin-1/2: n_orb = 2 * n_spatial_orb.
    E : Array
        Energy values, shape (nE,).
    spin_matrices : Array
        Pauli matrices σx, σy, σz, shape (3, n_orb, n_orb).
    eta : float, optional
        Lorentzian broadening. Default is 0.01.

    Returns
    -------
    intensity : Array
        Total A(k, E), shape (nk, nE).
    polarization : Array
        Spin polarization P(k, E), shape (nk, nE, 3).
    """
    nk, n_orb, n_bands = eigenvectors.shape
    len(E)

    # Compute spin expectation values for each band
    # <ψ_n|σ_i|ψ_n> for i = x, y, z
    spin_expect = jnp.zeros((nk, n_bands, 3))

    for i in range(3):
        # <ψ|σ_i|ψ> = sum_ab ψ*_a σ_i_ab ψ_b
        spin_expect = spin_expect.at[:, :, i].set(
            jnp.einsum(
                "kab,kbn,kan->kn",
                spin_matrices[i][None, :, :],
                eigenvectors,
                jnp.conj(eigenvectors),
            ).real
        )

    # Lorentzians for each band
    diff = E[None, None, :] - eigenvalues[:, :, None]  # (nk, n_bands, nE)
    lorentzians = eta / (diff**2 + eta**2) / jnp.pi

    # Total intensity
    intensity = jnp.sum(lorentzians, axis=1)  # (nk, nE)

    # Spin-weighted intensity for polarization numerator
    spin_weighted = lorentzians[:, :, :, None] * spin_expect[:, :, None, :]
    spin_intensity = jnp.sum(spin_weighted, axis=1)  # (nk, nE, 3)

    # Polarization = spin_intensity / total_intensity
    polarization = spin_intensity / (intensity[:, :, None] + 1e-10)

    return intensity, polarization


@jax.jit
def matrix_element_correction(
    eigenvectors: Array, photon_polarization: Array, orbital_dipoles: Array
) -> Array:
    """
    Compute photoemission matrix element corrections.

    M_n(k) = |<f|ε·r|ψ_n(k)>|²

    where ε is the photon polarization and r is the dipole operator.

    Parameters
    ----------
    eigenvectors : Array
        Eigenvectors, shape (nk, n_orb, n_bands).
    photon_polarization : Array
        Polarization vector, shape (3,).
    orbital_dipoles : Array
        Dipole matrix elements <orb|r|f>, shape (n_orb, 3).

    Returns
    -------
    weights : Array
        Matrix element weights, shape (nk, n_bands).
    """
    # Project polarization onto orbital dipoles
    dipole_proj = orbital_dipoles @ photon_polarization  # (n_orb,)

    # Matrix element for each band
    # M_n = |sum_orb c_orb * dipole_proj_orb|²
    return jnp.abs(jnp.einsum("o,kob->kb", dipole_proj, eigenvectors)) ** 2


def compute_arpes_spectrum(
    eigenvalues: Array,
    eigenvectors: Array,
    kpoints: Array,
    E_range: Tuple[float, float] = (-2.0, 0.5),
    n_E: int = 200,
    eta: float = 0.02,
    spin_matrices: Optional[Array] = None,
) -> ARPESSpectrum:
    """
    Compute full ARPES spectrum I(k, E).

    This is the main interface function for generating ARPES-comparable
    spectra from tight-binding calculations.

    Parameters
    ----------
    eigenvalues : Array
        Band energies, shape (nk, n_bands).
    eigenvectors : Array
        Eigenvectors, shape (nk, n_orb, n_bands).
    kpoints : Array
        k-points, shape (nk, 3).
    E_range : Tuple[float, float], optional
        Energy range (E_min, E_max) in eV. Default is (-2.0, 0.5).
    n_E : int, optional
        Number of energy points. Default is 200.
    eta : float, optional
        Broadening in eV. Default is 0.02.
    spin_matrices : Array, optional
        Pauli matrices for spin resolution.

    Returns
    -------
    spectrum : ARPESSpectrum
        ARPESSpectrum object.
    """
    E = jnp.linspace(E_range[0], E_range[1], n_E)

    if spin_matrices is None:
        # Spin-integrated
        intensity = spectral_function(eigenvalues, eigenvectors, E, eta)
        spin_pol = None
    else:
        # Spin-resolved
        intensity, spin_pol = spin_resolved_spectral(
            eigenvalues, eigenvectors, E, spin_matrices, eta
        )

    return ARPESSpectrum(
        kpoints=kpoints,
        energies=E,
        intensity=intensity,
        spin_polarization=spin_pol,
    )


# =============================================================================
# Loss functions for fitting to experimental ARPES
# =============================================================================


@jax.jit
def arpes_mse_loss(
    predicted: Array, experimental: Array, mask: Optional[Array] = None
) -> float:
    """
    Mean squared error loss between predicted and experimental ARPES.

    Parameters
    ----------
    predicted : Array
        Predicted I(k, E), shape (nk, nE).
    experimental : Array
        Experimental I(k, E), same shape.
    mask : Array, optional
        Mask for valid data points.

    Returns
    -------
    loss : float
        MSE loss value.
    """
    diff = predicted - experimental

    if mask is not None:
        diff = diff * mask
        n_valid = jnp.sum(mask)
    else:
        n_valid = diff.size

    return jnp.sum(diff**2) / n_valid


@jax.jit
def arpes_correlation_loss(predicted: Array, experimental: Array) -> float:
    """
    Negative correlation loss (for maximizing correlation).

    Parameters
    ----------
    predicted : Array
        Predicted I(k, E), shape (nk, nE).
    experimental : Array
        Experimental I(k, E), same shape.

    Returns
    -------
    loss : float
        1 - correlation coefficient.
    """
    pred_flat = predicted.ravel()
    exp_flat = experimental.ravel()

    pred_centered = pred_flat - jnp.mean(pred_flat)
    exp_centered = exp_flat - jnp.mean(exp_flat)

    correlation = jnp.sum(pred_centered * exp_centered) / (
        jnp.sqrt(jnp.sum(pred_centered**2))
        * jnp.sqrt(jnp.sum(exp_centered**2))
        + 1e-10
    )

    return 1.0 - correlation


@jax.jit
def spin_arpes_loss(
    predicted_intensity: Array,
    predicted_polarization: Array,
    exp_intensity: Array,
    exp_polarization: Array,
    intensity_weight: float = 1.0,
    spin_weight: float = 1.0,
) -> float:
    """
    Combined loss for spin-resolved ARPES.

    Parameters
    ----------
    predicted_intensity : Array
        Predicted I(k, E).
    predicted_polarization : Array
        Predicted P(k, E, 3).
    exp_intensity : Array
        Experimental I(k, E).
    exp_polarization : Array
        Experimental P(k, E, 3).
    intensity_weight : float, optional
        Weight for intensity term. Default is 1.0.
    spin_weight : float, optional
        Weight for spin polarization term. Default is 1.0.

    Returns
    -------
    loss : float
        Total loss.
    """
    # Intensity loss
    loss_I = arpes_mse_loss(predicted_intensity, exp_intensity)

    # Spin polarization loss
    loss_P = jnp.mean((predicted_polarization - exp_polarization) ** 2)

    return intensity_weight * loss_I + spin_weight * loss_P
