"""
Spin polarization for magnetic DFTB calculations.

This module provides:

- Spin-polarized Hamiltonian construction
- Stoner-like exchange splitting
- Spin configuration management
"""

from __future__ import annotations

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


class SpinConfig(NamedTuple):
    """
    Spin configuration for magnetic calculations.

    Attributes
    ----------
    collinear : bool
        If True, use collinear spin (separate up/down channels).
    initial_moments : Array, optional
        Initial magnetic moments per atom, shape (n_atoms,).
    spin_coupling : dict, optional
        Spin coupling constants W per orbital type.
    mixing : float
        Spin mixing parameter for self-consistency.
    """

    collinear: bool = True
    initial_moments: Optional[Array] = None
    spin_coupling: Optional[dict] = None
    mixing: float = 0.3


class SpinHamiltonian(NamedTuple):
    """
    Spin-polarized Hamiltonian.

    For collinear spin:
        H_up, H_down are separate Hamiltonians

    For non-collinear spin (with SOC):
        H is a 2N × 2N matrix in spinor basis

    Attributes
    ----------
    H_up : Array
        Spin-up Hamiltonian, shape (n_orb, n_orb).
    H_down : Array
        Spin-down Hamiltonian, shape (n_orb, n_orb).
    S : Array
        Overlap matrix (spin-independent), shape (n_orb, n_orb).
    magnetization : Array, optional
        Local magnetization per atom, shape (n_atoms,).
    """

    H_up: Array
    H_down: Array
    S: Array
    magnetization: Optional[Array] = None


# =============================================================================
# Collinear spin polarization
# =============================================================================


@jax.jit
def add_spin_polarization(
    H: Array,
    S: Array,
    magnetic_moments: Array,
    spin_coupling: Array,
    atom_orbital_map: Array,
) -> SpinHamiltonian:
    """
    Add spin polarization to a spin-unpolarized Hamiltonian.

    In collinear DFTB, the spin splitting is:
        H_σ = H_0 + σ * W * m

    where σ = ±1/2 for up/down, W is the spin coupling constant,
    and m is the local magnetic moment.

    Parameters
    ----------
    H : Array
        Spin-unpolarized Hamiltonian, shape (n_orb, n_orb).
    S : Array
        Overlap matrix, shape (n_orb, n_orb).
    magnetic_moments : Array
        Magnetic moment per atom, shape (n_atoms,).
    spin_coupling : Array
        W parameters per orbital, shape (n_orb,).
    atom_orbital_map : Array
        Maps orbital index to atom index, shape (n_orb,).

    Returns
    -------
    spin_hamiltonian : SpinHamiltonian
        SpinHamiltonian with H_up and H_down.
    """
    H.shape[0]

    # Build the spin splitting for each orbital
    # Δ_i = W_i * m_atom(i)
    atom_moments = magnetic_moments[atom_orbital_map]  # (n_orb,)
    spin_splitting = spin_coupling * atom_moments  # (n_orb,)

    # Add to diagonal
    # H_up = H_0 - 0.5 * W * m (spin up = -1/2 in convention)
    # H_down = H_0 + 0.5 * W * m
    # Note: convention varies, this gives lower energy for majority spin

    H_up = H - 0.5 * jnp.diag(spin_splitting)
    H_down = H + 0.5 * jnp.diag(spin_splitting)

    return SpinHamiltonian(
        H_up=H_up, H_down=H_down, S=S, magnetization=magnetic_moments
    )


@jax.jit
def compute_magnetization(
    eigenvalues_up: Array,
    eigenvalues_down: Array,
    eigenvectors_up: Array,
    eigenvectors_down: Array,
    n_electrons: float,
    atom_orbital_map: Array,
    temperature: float = 0.001,
) -> Array:
    """
    Compute local magnetic moments from spin-polarized calculation.

    m_atom = sum_{orb in atom} (n_up - n_down)

    Parameters
    ----------
    eigenvalues_up : Array
        Eigenvalues for spin-up, shape (n_k, n_orb).
    eigenvalues_down : Array
        Eigenvalues for spin-down, shape (n_k, n_orb).
    eigenvectors_up : Array
        Eigenvectors for spin-up, shape (n_k, n_orb, n_orb).
    eigenvectors_down : Array
        Eigenvectors for spin-down, shape (n_k, n_orb, n_orb).
    n_electrons : float
        Total number of electrons.
    atom_orbital_map : Array
        Maps orbital to atom, shape (n_orb,).
    temperature : float, optional
        Electronic temperature for smearing. Default is 0.001.

    Returns
    -------
    magnetization : Array
        Magnetic moments per atom, shape (n_atoms,).
    """
    n_atoms = int(jnp.max(atom_orbital_map)) + 1
    n_k, n_orb = eigenvalues_up.shape

    # Find Fermi levels for each spin channel
    # For simplicity, assume same Fermi level (charge neutral)
    all_eigenvalues = jnp.concatenate(
        [eigenvalues_up, eigenvalues_down], axis=-1
    )
    E_F = find_fermi_level_spin(all_eigenvalues, n_electrons, temperature)

    # Compute occupations
    f_up = fermi_dirac(eigenvalues_up, E_F, temperature)
    f_down = fermi_dirac(eigenvalues_down, E_F, temperature)

    # Compute orbital populations
    # n_orb_up[k, i] = sum_n f_up[k,n] * |c_up[k,i,n]|^2
    pop_up = jnp.einsum("kn,kin->ki", f_up, jnp.abs(eigenvectors_up) ** 2)
    pop_down = jnp.einsum(
        "kn,kin->ki", f_down, jnp.abs(eigenvectors_down) ** 2
    )

    # Sum over k-points (assuming uniform weights)
    pop_up = jnp.mean(pop_up, axis=0)  # (n_orb,)
    pop_down = jnp.mean(pop_down, axis=0)

    # Sum over orbitals per atom
    magnetization = jnp.zeros(n_atoms)
    for i_orb in range(n_orb):
        i_atom = atom_orbital_map[i_orb]
        magnetization = magnetization.at[i_atom].add(
            pop_up[i_orb] - pop_down[i_orb]
        )

    return magnetization


@jax.jit
def fermi_dirac(E: Array, E_F: float, T: float) -> Array:
    """Fermi-Dirac distribution."""
    x = (E - E_F) / T
    x = jnp.clip(x, -50, 50)
    return 1.0 / (1.0 + jnp.exp(x))


def find_fermi_level_spin(
    eigenvalues: Array, n_electrons: float, temperature: float
) -> float:
    """Find Fermi level for spin-polarized system."""
    E_min = jnp.min(eigenvalues) - 1.0
    E_max = jnp.max(eigenvalues) + 1.0

    for _ in range(50):
        E_mid = 0.5 * (E_min + E_max)
        n_e = jnp.sum(fermi_dirac(eigenvalues, E_mid, temperature))

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

    return 0.5 * (E_min + E_max)


# =============================================================================
# Self-consistent spin polarization
# =============================================================================


def self_consistent_spin(
    H_0: Array,
    S: Array,
    spin_coupling: Array,
    atom_orbital_map: Array,
    n_electrons: float,
    initial_moments: Optional[Array] = None,
    max_iter: int = 100,
    tolerance: float = 1e-5,
    mixing: float = 0.3,
    temperature: float = 0.001,
) -> Tuple[SpinHamiltonian, Array]:
    """
    Self-consistent calculation of spin polarization.

    Iterates until magnetic moments converge:

    1. Build H_up, H_down from current moments
    2. Diagonalize each spin channel
    3. Compute new moments from eigenstates
    4. Mix with old moments
    5. Check convergence

    Parameters
    ----------
    H_0 : Array
        Spin-unpolarized Hamiltonian.
    S : Array
        Overlap matrix.
    spin_coupling : Array
        W parameters per orbital.
    atom_orbital_map : Array
        Orbital to atom mapping.
    n_electrons : float
        Number of electrons.
    initial_moments : Array, optional
        Starting magnetic moments.
    max_iter : int, optional
        Maximum iterations. Default is 100.
    tolerance : float, optional
        Convergence threshold for moments. Default is 1e-5.
    mixing : float, optional
        Linear mixing parameter. Default is 0.3.
    temperature : float, optional
        Electronic temperature. Default is 0.001.

    Returns
    -------
    spin_hamiltonian : SpinHamiltonian
        Converged SpinHamiltonian.
    moments : Array
        Converged magnetic moments.
    """
    H_0.shape[0]
    n_atoms = int(jnp.max(atom_orbital_map)) + 1

    # Initialize moments
    if initial_moments is None:
        moments = jnp.ones(n_atoms) * 0.5  # Small initial polarization
    else:
        moments = initial_moments

    for _iteration in range(max_iter):
        # Build spin-polarized Hamiltonian
        spin_H = add_spin_polarization(
            H_0, S, moments, spin_coupling, atom_orbital_map
        )

        # Diagonalize (at Gamma point for simplicity)
        E_up, V_up = jnp.linalg.eigh(jnp.linalg.solve(spin_H.S, spin_H.H_up))
        E_down, V_down = jnp.linalg.eigh(
            jnp.linalg.solve(spin_H.S, spin_H.H_down)
        )

        # Compute new moments
        new_moments = compute_magnetization(
            E_up[None, :],
            E_down[None, :],
            V_up[None, :, :],
            V_down[None, :, :],
            n_electrons,
            atom_orbital_map,
            temperature,
        )

        # Check convergence
        diff = jnp.max(jnp.abs(new_moments - moments))
        if diff < tolerance:
            break

        # Mix
        moments = mixing * new_moments + (1 - mixing) * moments

    # Final Hamiltonian
    final_H = add_spin_polarization(
        H_0, S, moments, spin_coupling, atom_orbital_map
    )

    return final_H, moments


# =============================================================================
# Spin operators
# =============================================================================


def pauli_matrices(n_spatial_orb: int) -> Array:
    """
    Construct Pauli matrices in the spinor orbital basis.

    For a basis |orb, σ> with n_spatial_orb spatial orbitals,
    the total dimension is 2 * n_spatial_orb.

    Ordering: first all spin-up orbitals, then all spin-down
    |1↑>, |2↑>, ..., |n↑>, |1↓>, |2↓>, ..., |n↓>.

    Parameters
    ----------
    n_spatial_orb : int
        Number of spatial orbitals.

    Returns
    -------
    spin_matrices : Array
        Pauli matrices σ_x, σ_y, σ_z, shape (3, 2n, 2n).
    """
    n = n_spatial_orb
    2 * n

    # Identity for orbital space
    I_orb = jnp.eye(n)

    # Pauli matrices in spin space
    sigma_x = jnp.array([[0, 1], [1, 0]], dtype=complex)
    sigma_y = jnp.array([[0, -1j], [1j, 0]], dtype=complex)
    sigma_z = jnp.array([[1, 0], [0, -1]], dtype=complex)

    # Kronecker product: σ ⊗ I_orb
    # Result is block structure: [[0, I], [I, 0]] for σ_x, etc.
    S_x = jnp.kron(sigma_x, I_orb)
    S_y = jnp.kron(sigma_y, I_orb)
    S_z = jnp.kron(sigma_z, I_orb)

    # Factor of 1/2 for spin-1/2
    return 0.5 * jnp.stack([S_x, S_y, S_z], axis=0)


def spin_expectation(eigenvectors: Array, spin_matrices: Array) -> Array:
    """
    Compute spin expectation values <ψ|S|ψ> for each eigenstate.

    Parameters
    ----------
    eigenvectors : Array
        Spinor eigenvectors, shape (..., 2n_orb, n_bands).
    spin_matrices : Array
        Pauli matrices, shape (3, 2n_orb, 2n_orb).

    Returns
    -------
    spin_expect : Array
        Spin expectation values, shape (..., n_bands, 3).
    """
    # <ψ_n|S_i|ψ_n> = sum_ab ψ*_a,n S_i,ab ψ_b,n
    return jnp.einsum(
        "iab,...an,...bn->...ni",
        spin_matrices,
        jnp.conj(eigenvectors),
        eigenvectors,
    ).real
