"""
Hamiltonian construction for DFTB.

This module provides functionality for:

- Building the tight-binding Hamiltonian matrix from SK parameters
- k-dependent Hamiltonian for periodic systems
- Spin-polarized Hamiltonian construction
"""

from __future__ import annotations

from functools import partial

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

from difftb.io.slater_koster import (
    SlaterKosterParams,
    direction_cosines,
    interpolate_sk_integrals,
    sk_rotation_pp,
    sk_rotation_sp,
)


class Hamiltonian(NamedTuple):
    """
    Tight-binding Hamiltonian and overlap matrices.

    Attributes
    ----------
    H : Array
        Hamiltonian matrix, shape (n_orbitals, n_orbitals) or
        (nk, n_orbitals, n_orbitals).
    S : Array
        Overlap matrix, same shape as H.
    n_orbitals : int
        Total number of orbitals in the basis.
    orbital_info : List[Tuple[int, str]], optional
        List of (atom_index, orbital_type) for each orbital.
    """

    H: Array
    S: Array
    n_orbitals: int
    orbital_info: Optional[List[Tuple[int, str]]] = None


class AtomicStructure(NamedTuple):
    """
    Atomic structure specification.

    Attributes
    ----------
    positions : Array
        Atomic positions in fractional coordinates, shape (n_atoms, 3).
    elements : List[str]
        Element symbol for each atom.
    lattice_vectors : Array
        Lattice vectors as rows, shape (3, 3).
    """

    positions: Array
    elements: List[str]
    lattice_vectors: Array

    @property
    def n_atoms(self) -> int:
        return len(self.elements)

    @property
    def positions_cartesian(self) -> Array:
        """Atomic positions in Cartesian coordinates (Angstroms)."""
        return self.positions @ self.lattice_vectors


# =============================================================================
# Neighbor list construction
# =============================================================================


def get_neighbor_pairs(
    structure: AtomicStructure, cutoff: float, max_cells: int = 3
) -> Tuple[Array, Array, Array, Array]:
    """
    Find all atom pairs within cutoff distance, including periodic images.

    Parameters
    ----------
    structure : AtomicStructure
        Atomic structure.
    cutoff : float
        Cutoff distance in Angstroms.
    max_cells : int, optional
        Maximum number of periodic cells to search in each direction.
        Default is 3.

    Returns
    -------
    atom_i : Array
        First atom index for each pair, shape (n_pairs,).
    atom_j : Array
        Second atom index for each pair, shape (n_pairs,).
    r_vec : Array
        Displacement vectors from i to j, shape (n_pairs, 3).
    cell_shift : Array
        Periodic cell shift for each pair, shape (n_pairs, 3).
    """
    n_atoms = structure.n_atoms
    positions = structure.positions_cartesian
    lattice = structure.lattice_vectors

    # Generate all cell shifts to consider
    shifts = jnp.array(
        [
            [i, j, k]
            for i in range(-max_cells, max_cells + 1)
            for j in range(-max_cells, max_cells + 1)
            for k in range(-max_cells, max_cells + 1)
        ]
    )

    # For each pair of atoms and each cell shift, compute distance
    pairs_i = []
    pairs_j = []
    r_vecs = []
    cell_shifts = []

    for i in range(n_atoms):
        for j in range(n_atoms):
            for shift in shifts:
                # Skip self-interaction in home cell
                if i == j and jnp.all(shift == 0):
                    continue

                # Compute displacement
                r_j = positions[j] + shift @ lattice
                r_vec = r_j - positions[i]
                r = jnp.linalg.norm(r_vec)

                if r < cutoff:
                    pairs_i.append(i)
                    pairs_j.append(j)
                    r_vecs.append(r_vec)
                    cell_shifts.append(shift)

    return (
        jnp.array(pairs_i),
        jnp.array(pairs_j),
        jnp.stack(r_vecs, axis=0) if r_vecs else jnp.zeros((0, 3)),
        (
            jnp.stack(cell_shifts, axis=0)
            if cell_shifts
            else jnp.zeros((0, 3), dtype=int)
        ),
    )


# =============================================================================
# Hamiltonian matrix construction
# =============================================================================


def build_hamiltonian_k(
    k: Array,
    structure: AtomicStructure,
    sk_params: SlaterKosterParams,
    cutoff: float = 10.0,
) -> Hamiltonian:
    """
    Build the k-dependent Hamiltonian matrix.

    This is the core function that constructs H(k) and S(k) for a given k-point.

    Parameters
    ----------
    k : Array
        k-point in fractional coordinates, shape (3,).
    structure : AtomicStructure
        Atomic structure.
    sk_params : SlaterKosterParams
        Slater-Koster parameters.
    cutoff : float, optional
        Neighbor cutoff distance. Default is 10.0.

    Returns
    -------
    hamiltonian : Hamiltonian
        Hamiltonian object with H(k) and S(k).
    """
    # Determine total number of orbitals
    n_orbitals = 0
    orbital_info = []
    atom_orbital_start = []

    for i, elem in enumerate(structure.elements):
        basis = sk_params.elements[elem]
        atom_orbital_start.append(n_orbitals)
        for orb in basis.orbitals:
            orbital_info.append((i, orb))
            n_orbitals += 1

    # Initialize matrices
    H_k = jnp.zeros((n_orbitals, n_orbitals), dtype=complex)
    S_k = jnp.zeros((n_orbitals, n_orbitals), dtype=complex)

    # Add on-site energies
    idx = 0
    for _i, elem in enumerate(structure.elements):
        basis = sk_params.elements[elem]
        for _j, E_onsite in enumerate(basis.onsite_energies):
            H_k = H_k.at[idx, idx].set(E_onsite)
            S_k = S_k.at[idx, idx].set(1.0)
            idx += 1

    # Get neighbor pairs
    atom_i, atom_j, r_vecs, cell_shifts = get_neighbor_pairs(structure, cutoff)

    # Add hopping terms
    for pair_idx in range(len(atom_i)):
        i = int(atom_i[pair_idx])
        j = int(atom_j[pair_idx])
        r_vec = r_vecs[pair_idx]
        shift = cell_shifts[pair_idx]

        elem_i = structure.elements[i]
        elem_j = structure.elements[j]

        # Get SK table for this element pair
        table = sk_params.get_table(elem_i, elem_j)

        # Compute distance and direction cosines
        r = jnp.linalg.norm(r_vec)
        l, m, n = direction_cosines(r_vec[None, :])
        l, m, n = l[0], m[0], n[0]

        # Interpolate SK integrals
        H_sk, S_sk = interpolate_sk_integrals(r[None], table)
        H_sk, S_sk = H_sk[0], S_sk[0]  # Remove batch dimension

        # Phase factor for periodic boundary conditions
        # k is in fractional coordinates, shift is in lattice vectors
        phase = jnp.exp(2j * jnp.pi * jnp.dot(k, shift))

        # Build the orbital block
        basis_i = sk_params.elements[elem_i]
        basis_j = sk_params.elements[elem_j]

        # Get orbital block matrices from SK rotation
        H_block, S_block = build_orbital_block(
            l, m, n, H_sk, S_sk, basis_i, basis_j, table.integral_names
        )

        # Add to full matrices with phase
        i_start = atom_orbital_start[i]
        j_start = atom_orbital_start[j]
        n_i = basis_i.n_orbitals
        n_j = basis_j.n_orbitals

        H_k = H_k.at[i_start : i_start + n_i, j_start : j_start + n_j].add(
            phase * H_block
        )
        S_k = S_k.at[i_start : i_start + n_i, j_start : j_start + n_j].add(
            phase * S_block
        )

    # Hermitianize (in case of numerical asymmetry)
    H_k = 0.5 * (H_k + H_k.conj().T)
    S_k = 0.5 * (S_k + S_k.conj().T)

    return Hamiltonian(
        H=H_k, S=S_k, n_orbitals=n_orbitals, orbital_info=orbital_info
    )


def build_orbital_block(
    l: Array,
    m: Array,
    n: Array,
    H_sk: Array,
    S_sk: Array,
    basis_i,
    basis_j,
    integral_names: List[str],
) -> Tuple[Array, Array]:
    """
    Build the orbital block matrix from SK integrals.

    This applies the Slater-Koster rotation rules to transform the
    two-center integrals into the full orbital matrix.

    Parameters
    ----------
    l : Array
        Direction cosine along x.
    m : Array
        Direction cosine along y.
    n : Array
        Direction cosine along z.
    H_sk : Array
        Interpolated Hamiltonian SK integrals.
    S_sk : Array
        Interpolated overlap SK integrals.
    basis_i : OrbitalBasis
        Orbital basis for atom i.
    basis_j : OrbitalBasis
        Orbital basis for atom j.
    integral_names : List[str]
        Names of SK integral types.

    Returns
    -------
    H_block : Array
        Hamiltonian block, shape (n_orbitals_i, n_orbitals_j).
    S_block : Array
        Overlap block, same shape.
    """
    n_i = basis_i.n_orbitals
    n_j = basis_j.n_orbitals

    H_block = jnp.zeros((n_i, n_j))
    S_block = jnp.zeros((n_i, n_j))

    # Extract SK integrals by name
    sk_dict_H = {name: H_sk[i] for i, name in enumerate(integral_names)}
    sk_dict_S = {name: S_sk[i] for i, name in enumerate(integral_names)}

    # For sp basis, we have:
    # orbitals = ['s', 'px', 'py', 'pz'] -> indices 0, 1, 2, 3

    # s-s interaction (0,0)
    if "ss_sigma" in sk_dict_H:
        H_block = H_block.at[0, 0].set(sk_dict_H["ss_sigma"])
        S_block = S_block.at[0, 0].set(sk_dict_S["ss_sigma"])

    # s-p interactions (0, 1:4) and (1:4, 0)
    if "sp_sigma" in sk_dict_H:
        sp_vec_H = sk_rotation_sp(l, m, n, sk_dict_H["sp_sigma"])
        sp_vec_S = sk_rotation_sp(l, m, n, sk_dict_S["sp_sigma"])

        # s_i -> p_j
        H_block = H_block.at[0, 1:4].set(sp_vec_H)
        S_block = S_block.at[0, 1:4].set(sp_vec_S)

        # p_i -> s_j (sign flip for sp)
        H_block = H_block.at[1:4, 0].set(-sp_vec_H)
        S_block = S_block.at[1:4, 0].set(sp_vec_S)

    # p-p interactions (1:4, 1:4)
    if "pp_sigma" in sk_dict_H and "pp_pi" in sk_dict_H:
        pp_mat_H = sk_rotation_pp(
            l, m, n, sk_dict_H["pp_sigma"], sk_dict_H["pp_pi"]
        )
        pp_mat_S = sk_rotation_pp(
            l, m, n, sk_dict_S["pp_sigma"], sk_dict_S["pp_pi"]
        )

        H_block = H_block.at[1:4, 1:4].set(pp_mat_H)
        S_block = S_block.at[1:4, 1:4].set(pp_mat_S)

    return H_block, S_block


# =============================================================================
# Vectorized Hamiltonian construction for multiple k-points
# =============================================================================


@partial(jax.jit, static_argnums=(2, 3))
def build_hamiltonian_kpoints(
    kpoints: Array,
    structure: AtomicStructure,
    sk_params: SlaterKosterParams,
    cutoff: float = 10.0,
) -> Hamiltonian:
    """
    Build Hamiltonian matrices for multiple k-points.

    Uses vmap for efficient parallelization.

    Parameters
    ----------
    kpoints : Array
        Array of k-points, shape (nk, 3).
    structure : AtomicStructure
        Atomic structure.
    sk_params : SlaterKosterParams
        Slater-Koster parameters.
    cutoff : float, optional
        Neighbor cutoff. Default is 10.0.

    Returns
    -------
    hamiltonian : Hamiltonian
        Hamiltonian with H and S of shape (nk, n_orbitals, n_orbitals).
    """

    # vmap over k-points
    def build_single(k):
        return build_hamiltonian_k(k, structure, sk_params, cutoff)

    # This would need to be vmapped properly - placeholder for now
    # In practice, the neighbor list should be computed once, then
    # only the phase factors vary with k
    return jax.vmap(build_single)(kpoints)


def build_hamiltonian_realspace(
    structure: AtomicStructure,
    sk_params: SlaterKosterParams,
    cutoff: float = 10.0,
    max_cells: int = 3,
) -> Tuple[Array, Array, Array]:
    """
    Build real-space Hamiltonian matrices H(R) for Fourier transform.

    This is more efficient for computing H(k) at many k-points.
    H(k) = sum_R H(R) * exp(i k·R).

    Parameters
    ----------
    structure : AtomicStructure
        Atomic structure.
    sk_params : SlaterKosterParams
        Slater-Koster parameters.
    cutoff : float, optional
        Neighbor cutoff. Default is 10.0.
    max_cells : int, optional
        Maximum cell index. Default is 3.

    Returns
    -------
    H_R : Array
        Real-space Hamiltonian blocks, shape (n_cells, n_orb, n_orb).
    S_R : Array
        Real-space overlap blocks, same shape.
    R_vectors : Array
        Cell vectors, shape (n_cells, 3).
    """
    # This is a more efficient implementation for band structure calculations
    # Build H(R) once, then H(k) = sum_R H(R) exp(ik·R)
    raise NotImplementedError(
        "Real-space Hamiltonian construction not yet implemented"
    )
