"""
Spin-orbit coupling for DFTB calculations.

This module provides:

- On-site L·S spin-orbit coupling
- SOC matrix construction for p, d orbitals
- Full spinor Hamiltonian with SOC
"""

from __future__ import annotations

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


class SOCParams(NamedTuple):
    """
    Spin-orbit coupling parameters.

    Attributes
    ----------
    xi : Array
        SOC coupling constants per shell (s, p, d, ...), shape varies.
    orbital_l : Array
        Angular momentum quantum number for each orbital.
    """

    xi: Array  # SOC strength in eV
    orbital_l: Array  # l quantum number for each orbital


# =============================================================================
# L·S matrix elements for atomic orbitals
# =============================================================================

# Angular momentum matrices for l=1 (p orbitals)
# Basis order: px, py, pz (real spherical harmonics)


def Lx_p() -> Array:
    """L_x matrix for p orbitals in real basis (px, py, pz)."""
    return jnp.array([[0, 0, 0], [0, 0, -1j], [0, 1j, 0]], dtype=complex)


def Ly_p() -> Array:
    """L_y matrix for p orbitals."""
    return jnp.array([[0, 0, 1j], [0, 0, 0], [-1j, 0, 0]], dtype=complex)


def Lz_p() -> Array:
    """L_z matrix for p orbitals."""
    return jnp.array([[0, -1j, 0], [1j, 0, 0], [0, 0, 0]], dtype=complex)


# Angular momentum matrices for l=2 (d orbitals)
# Basis order: dxy, dyz, dzx, dx2-y2, dz2


def Lx_d() -> Array:
    """L_x matrix for d orbitals in real basis."""
    sqrt3 = jnp.sqrt(3.0)
    return jnp.array(
        [
            [0, 0, 1j, 0, 0],
            [0, 0, 0, -1j, 1j * sqrt3],
            [-1j, 0, 0, 0, 0],
            [0, 1j, 0, 0, 0],
            [0, -1j * sqrt3, 0, 0, 0],
        ],
        dtype=complex,
    )


def Ly_d() -> Array:
    """L_y matrix for d orbitals."""
    sqrt3 = jnp.sqrt(3.0)
    return jnp.array(
        [
            [0, -1j, 0, 0, 0],
            [1j, 0, 0, 0, 0],
            [0, 0, 0, 1j, 1j * sqrt3],
            [0, 0, -1j, 0, 0],
            [0, 0, -1j * sqrt3, 0, 0],
        ],
        dtype=complex,
    )


def Lz_d() -> Array:
    """L_z matrix for d orbitals."""
    return jnp.array(
        [
            [0, 0, 0, 2j, 0],
            [0, 0, 1j, 0, 0],
            [0, -1j, 0, 0, 0],
            [-2j, 0, 0, 0, 0],
            [0, 0, 0, 0, 0],
        ],
        dtype=complex,
    )


def angular_momentum_matrices(l: int) -> Tuple[Array, Array, Array]:
    """
    Get angular momentum matrices L_x, L_y, L_z for given l.

    Parameters
    ----------
    l : int
        Angular momentum quantum number (0, 1, 2, ...).

    Returns
    -------
    L_x : Array
        L_x matrix, shape (2l+1, 2l+1).
    L_y : Array
        L_y matrix, shape (2l+1, 2l+1).
    L_z : Array
        L_z matrix, shape (2l+1, 2l+1).
    """
    if l == 0:
        # s orbitals: L = 0
        return jnp.zeros((1, 1)), jnp.zeros((1, 1)), jnp.zeros((1, 1))
    if l == 1:
        return Lx_p(), Ly_p(), Lz_p()
    if l == 2:
        return Lx_d(), Ly_d(), Lz_d()
    raise NotImplementedError(f"L matrices not implemented for l={l}")


# =============================================================================
# SOC Hamiltonian construction
# =============================================================================


def soc_matrix_atom(xi_shells: Array, orbital_ls: List[int]) -> Array:
    """
    Build SOC matrix for a single atom.

    H_SOC = ξ L·S = ξ (L_x S_x + L_y S_y + L_z S_z)

    In the spinor basis |orb, σ>, this becomes a 2N × 2N matrix.

    Parameters
    ----------
    xi_shells : Array
        SOC constants for each shell, shape (n_shells,).
    orbital_ls : List[int]
        l quantum number for each spatial orbital.

    Returns
    -------
    H_SOC : Array
        SOC matrix, shape (2*n_orb, 2*n_orb).
    """
    n_orb = len(orbital_ls)

    # Pauli spin matrices (spin-1/2)
    Sx = 0.5 * jnp.array([[0, 1], [1, 0]], dtype=complex)
    Sy = 0.5 * jnp.array([[0, -1j], [1j, 0]], dtype=complex)
    Sz = 0.5 * jnp.array([[1, 0], [0, -1]], dtype=complex)

    # Build L matrices in the full orbital space
    L_x = jnp.zeros((n_orb, n_orb), dtype=complex)
    L_y = jnp.zeros((n_orb, n_orb), dtype=complex)
    L_z = jnp.zeros((n_orb, n_orb), dtype=complex)

    # Build ξ diagonal
    xi_diag = jnp.zeros(n_orb)

    # Fill in blocks for each shell
    idx = 0
    current_l = -1

    for i, l in enumerate(orbital_ls):
        if l != current_l:
            current_l = l
            if l > 0:
                Lx_l, Ly_l, Lz_l = angular_momentum_matrices(l)
                block_size = 2 * l + 1
                start = idx
                end = idx + block_size

                L_x = L_x.at[start:end, start:end].set(Lx_l)
                L_y = L_y.at[start:end, start:end].set(Ly_l)
                L_z = L_z.at[start:end, start:end].set(Lz_l)

        # Set ξ value
        xi_diag = xi_diag.at[i].set(xi_shells[min(l, len(xi_shells) - 1)])
        idx += 1

    # Weight L matrices by ξ
    xi_mat = jnp.diag(xi_diag)
    L_x = xi_mat @ L_x
    L_y = xi_mat @ L_y
    L_z = xi_mat @ L_z

    # Build H_SOC = L_x ⊗ S_x + L_y ⊗ S_y + L_z ⊗ S_z
    # Kronecker product convention: spin is fast index
    return jnp.kron(L_x, Sx) + jnp.kron(L_y, Sy) + jnp.kron(L_z, Sz)


@jax.jit
def add_soc(
    H: Array, S: Array, soc_matrices: Array, atom_blocks: List[Tuple[int, int]]
) -> Tuple[Array, Array]:
    """
    Add spin-orbit coupling to a spin-unpolarized Hamiltonian.

    Transforms from scalar basis to spinor basis and adds SOC.

    Parameters
    ----------
    H : Array
        Hamiltonian in scalar basis, shape (n_orb, n_orb).
    S : Array
        Overlap in scalar basis, shape (n_orb, n_orb).
    soc_matrices : Array
        SOC matrices for each atom, shape (n_atoms, 2*n_atom_orb, 2*n_atom_orb).
    atom_blocks : List[Tuple[int, int]]
        Start and end indices for each atom's orbitals.

    Returns
    -------
    H_spinor : Array
        Hamiltonian in spinor basis, shape (2*n_orb, 2*n_orb).
    S_spinor : Array
        Overlap in spinor basis.
    """
    H.shape[0]

    # Expand to spinor basis: |orb⟩ → |orb,↑⟩, |orb,↓⟩
    # Use block diagonal structure: H_spinor = H ⊗ I_2
    I_spin = jnp.eye(2)
    H_spinor = jnp.kron(H, I_spin)
    S_spinor = jnp.kron(S, I_spin)

    # Add SOC for each atom
    for i_atom, (start, end) in enumerate(atom_blocks):
        end - start

        # Indices in spinor basis
        start_spin = 2 * start
        end_spin = 2 * end

        # Add SOC block
        H_spinor = H_spinor.at[start_spin:end_spin, start_spin:end_spin].add(
            soc_matrices[i_atom]
        )

    return H_spinor, S_spinor


def build_spinor_hamiltonian(
    H_scalar: Array,
    S_scalar: Array,
    elements: List[str],
    soc_params: dict,
    atom_orbital_info: List[Tuple[int, int, List[int]]],
) -> Tuple[Array, Array]:
    """
    Build full spinor Hamiltonian with SOC.

    This is the main interface for adding SOC to a calculation.

    Args:
        H_scalar: Hamiltonian without SOC, shape (n_orb, n_orb)
        S_scalar: Overlap matrix, shape (n_orb, n_orb)
        elements: Element symbols for each atom
        soc_params: Dictionary mapping element to ξ values
        atom_orbital_info: List of (atom_idx, start_orb, orbital_ls) for each atom

    Returns
    -------
        H_spinor: Full Hamiltonian with SOC, shape (2*n_orb, 2*n_orb)
        S_spinor: Overlap in spinor basis
    """
    H_scalar.shape[0]
    len(elements)

    # Build SOC matrix for each atom
    soc_matrices = []
    atom_blocks = []

    for i_atom, (_, start_orb, orbital_ls) in enumerate(atom_orbital_info):
        elem = elements[i_atom]
        n_atom_orb = len(orbital_ls)

        # Get ξ values for this element
        if elem in soc_params:
            xi_shells = jnp.array(soc_params[elem])
        else:
            xi_shells = jnp.zeros(3)  # No SOC

        # Build atomic SOC matrix
        soc_mat = soc_matrix_atom(xi_shells, orbital_ls)
        soc_matrices.append(soc_mat)
        atom_blocks.append((start_orb, start_orb + n_atom_orb))

    # Stack SOC matrices
    # Note: This assumes all atoms have same orbital structure
    # For variable orbital counts, need different handling

    # Add SOC to Hamiltonian
    H_spinor, S_spinor = add_soc(
        H_scalar, S_scalar, jnp.stack(soc_matrices, axis=0), atom_blocks
    )

    return H_spinor, S_spinor


# =============================================================================
# SOC constants for common elements
# =============================================================================

# Approximate ξ values in eV from literature
# Format: {element: [ξ_s, ξ_p, ξ_d]}
# Note: s orbitals have no SOC (l=0), but we include for indexing

DEFAULT_SOC_PARAMS = {
    # Light elements (weak SOC)
    "C": [0.0, 0.006, 0.0],
    "N": [0.0, 0.010, 0.0],
    "O": [0.0, 0.014, 0.0],
    # 3d transition metals
    "Mn": [0.0, 0.05, 0.05],
    "Fe": [0.0, 0.06, 0.06],
    "Co": [0.0, 0.07, 0.07],
    "Ni": [0.0, 0.08, 0.08],
    # 4d/5d elements
    "Mo": [0.0, 0.15, 0.15],
    "W": [0.0, 0.40, 0.40],
    # Heavy p-block (strong SOC)
    "Se": [0.0, 0.4, 0.0],
    "Te": [0.0, 0.5, 0.0],  # ~0.4-0.5 eV
    "Sb": [0.0, 1.0, 0.0],
    "Bi": [0.0, 1.5, 0.0],  # ~1.5-2.0 eV - very strong!
    # Topological insulator elements
    "Pb": [0.0, 0.9, 0.0],
}


def get_soc_params(elements: List[str]) -> dict:
    """
    Get SOC parameters for a list of elements.

    Args:
        elements: List of element symbols

    Returns
    -------
        Dictionary of ξ values for each unique element
    """
    params = {}
    for elem in set(elements):
        if elem in DEFAULT_SOC_PARAMS:
            params[elem] = DEFAULT_SOC_PARAMS[elem]
        else:
            print(f"Warning: No SOC parameters for {elem}, using zero")
            params[elem] = [0.0, 0.0, 0.0]
    return params


# =============================================================================
# Combined spin + SOC Hamiltonian
# =============================================================================


def build_magnetic_soc_hamiltonian(
    H_0: Array,
    S: Array,
    magnetic_moments: Array,
    spin_coupling: Array,
    atom_orbital_map: Array,
    elements: List[str],
    soc_params: dict,
    atom_orbital_info: List[Tuple[int, int, List[int]]],
) -> Tuple[Array, Array]:
    """
    Build Hamiltonian with both spin polarization and SOC.

    For magnetic systems with SOC, we need the full spinor basis
    with exchange splitting and L·S coupling.

    H = H_0 ⊗ I_spin + Σ_atom (ξ L·S + W m σ_z)

    Args:
        H_0: Spin-unpolarized Hamiltonian
        S: Overlap matrix
        magnetic_moments: Moments per atom
        spin_coupling: W per orbital
        atom_orbital_map: Orbital to atom mapping
        elements: Element symbols
        soc_params: SOC ξ values
        atom_orbital_info: Orbital structure info

    Returns
    -------
        H_full: Hamiltonian with spin + SOC
        S_full: Overlap in spinor basis
    """
    n_orb = H_0.shape[0]

    # First, expand to spinor basis
    I_spin = jnp.eye(2)
    H_spinor = jnp.kron(H_0, I_spin)
    S_spinor = jnp.kron(S, I_spin)

    # Add exchange splitting (diagonal in spinor basis)
    # H_exchange = diag(m_i W_i σ_z) for each orbital
    jnp.array([[1, 0], [0, -1]], dtype=complex)

    for i_orb in range(n_orb):
        i_atom = int(atom_orbital_map[i_orb])
        m = magnetic_moments[i_atom]
        W = spin_coupling[i_orb]

        # Indices in spinor basis
        i_up = 2 * i_orb
        i_dn = 2 * i_orb + 1

        # Add exchange: -W*m/2 for up, +W*m/2 for down
        H_spinor = H_spinor.at[i_up, i_up].add(-0.5 * W * m)
        H_spinor = H_spinor.at[i_dn, i_dn].add(+0.5 * W * m)

    # Add SOC for each atom
    for i_atom, (_, start_orb, orbital_ls) in enumerate(atom_orbital_info):
        elem = elements[i_atom]
        n_atom_orb = len(orbital_ls)

        if elem in soc_params:
            xi_shells = jnp.array(soc_params[elem])
        else:
            continue

        # Build atomic SOC matrix
        soc_mat = soc_matrix_atom(xi_shells, orbital_ls)

        # Add to Hamiltonian
        start_spin = 2 * start_orb
        end_spin = 2 * (start_orb + n_atom_orb)

        H_spinor = H_spinor.at[start_spin:end_spin, start_spin:end_spin].add(
            soc_mat
        )

    return H_spinor, S_spinor
