"""
Slater-Koster table reading and interpolation.

This module provides functionality for:

- Reading SK parameter files (DFTB+ format)
- Differentiable interpolation of SK integrals
- Rotation of SK integrals for arbitrary bond orientations
"""

from __future__ import annotations

from functools import partial

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

# Slater-Koster integral types
# For s, p orbitals: ssσ, spσ, ppσ, ppπ
# For s, p, d orbitals: add sdσ, pdσ, pdπ, ddσ, ddπ, ddδ
SK_INTEGRALS_SP = ["ss_sigma", "sp_sigma", "pp_sigma", "pp_pi"]
SK_INTEGRALS_SPD = SK_INTEGRALS_SP + [
    "sd_sigma",
    "pd_sigma",
    "pd_pi",
    "dd_sigma",
    "dd_pi",
    "dd_delta",
]


class OrbitalBasis(NamedTuple):
    """
    Orbital basis specification for an element.

    Attributes
    ----------
        element: Element symbol
        orbitals: List of orbital types ('s', 'p', 'd', etc.)
        n_orbitals: Total number of orbitals
        onsite_energies: On-site energies for each orbital, shape (n_orbitals,)
        hubbard_u: Hubbard U parameters for SCC-DFTB, shape (n_orbitals,)
        spin_w: Spin coupling constants for spin-DFTB, shape (n_orbitals,)
        soc_xi: Spin-orbit coupling constants, shape (n_shells,)
    """

    element: str
    orbitals: List[str]
    n_orbitals: int
    onsite_energies: Array
    hubbard_u: Optional[Array] = None
    spin_w: Optional[Array] = None
    soc_xi: Optional[Array] = None


class SlaterKosterTable(NamedTuple):
    """
    Slater-Koster parameter table for a pair of elements.

    Attributes
    ----------
        element1: First element symbol
        element2: Second element symbol
        distances: Tabulated distances in Angstroms, shape (n_points,)
        H_integrals: Hamiltonian SK integrals, shape (n_points, n_integrals)
        S_integrals: Overlap SK integrals, shape (n_points, n_integrals)
        repulsive: Repulsive potential values, shape (n_points,) or None
        integral_names: Names of the SK integrals
    """

    element1: str
    element2: str
    distances: Array
    H_integrals: Array
    S_integrals: Array
    repulsive: Optional[Array]
    integral_names: List[str]

    @property
    def r_min(self) -> float:
        """Minimum tabulated distance."""
        return float(self.distances[0])

    @property
    def r_max(self) -> float:
        """Maximum tabulated distance."""
        return float(self.distances[-1])

    @property
    def n_integrals(self) -> int:
        """Number of SK integral types."""
        return len(self.integral_names)


class SlaterKosterParams(NamedTuple):
    """
    Complete set of Slater-Koster parameters for a system.

    Attributes
    ----------
        elements: Dictionary mapping element symbols to OrbitalBasis
        tables: Dictionary mapping element pairs to SlaterKosterTable
    """

    elements: Dict[str, OrbitalBasis]
    tables: Dict[Tuple[str, str], SlaterKosterTable]

    def get_table(self, elem1: str, elem2: str) -> SlaterKosterTable:
        """Get SK table for element pair (order-independent)."""
        if (elem1, elem2) in self.tables:
            return self.tables[(elem1, elem2)]
        if (elem2, elem1) in self.tables:
            return self.tables[(elem2, elem1)]
        raise KeyError(f"No SK table for pair ({elem1}, {elem2})")


# =============================================================================
# Differentiable interpolation
# =============================================================================


@partial(jax.jit, static_argnums=(2,))
def interpolate_sk_integrals(
    r: Array, table: SlaterKosterTable, kind: str = "cubic"
) -> Tuple[Array, Array]:
    """
    Interpolate SK integrals at given distances.

    Uses smooth interpolation that is differentiable with respect to r.

    Parameters
    ----------
    r : Array
        Distances at which to evaluate, shape (...).
    table : SlaterKosterTable
        SlaterKosterTable with tabulated values.
    kind : str, optional
        Interpolation method ('linear' or 'cubic'). Default is 'cubic'.

    Returns
    -------
    H_sk : Array
        Hamiltonian integrals, shape (..., n_integrals).
    S_sk : Array
        Overlap integrals, shape (..., n_integrals).
    """
    # Flatten r for interpolation
    r_flat = r.ravel()

    # Cubic spline interpolation for smoothness
    H_interp = cubic_spline_interpolate(
        table.distances, table.H_integrals, r_flat
    )
    S_interp = cubic_spline_interpolate(
        table.distances, table.S_integrals, r_flat
    )

    # Reshape back
    output_shape = r.shape + (table.n_integrals,)
    H_sk = H_interp.reshape(output_shape)
    S_sk = S_interp.reshape(output_shape)

    # Apply smooth cutoff near r_max
    cutoff = smooth_cutoff(r, table.r_max - 0.5, table.r_max)
    H_sk = H_sk * cutoff[..., None]
    S_sk = S_sk * cutoff[..., None]

    return H_sk, S_sk


@jax.jit
def cubic_spline_interpolate(
    x_tab: Array, y_tab: Array, x_new: Array
) -> Array:
    """
    Cubic spline interpolation.

    Parameters
    ----------
    x_tab : Array
        Tabulated x values, shape (n,).
    y_tab : Array
        Tabulated y values, shape (n, m).
    x_new : Array
        New x values, shape (k,).

    Returns
    -------
    y_interp : Array
        Interpolated y values, shape (k, m).
    """
    # Find indices for interpolation
    idx = jnp.searchsorted(x_tab, x_new) - 1
    idx = jnp.clip(idx, 0, len(x_tab) - 2)

    # Get surrounding points
    x0 = x_tab[idx]
    x1 = x_tab[idx + 1]
    y0 = y_tab[idx]
    y1 = y_tab[idx + 1]

    # Normalized coordinate
    t = (x_new - x0) / (x1 - x0)
    t = t[:, None]  # Shape (k, 1) for broadcasting

    # For simplicity, use linear interpolation here
    # TODO: Implement proper cubic spline with continuous derivatives
    return y0 + t * (y1 - y0)


@jax.jit
def smooth_cutoff(r: Array, r_inner: float, r_outer: float) -> Array:
    """
    Smooth cutoff function.

    Goes from 1 at r_inner to 0 at r_outer using a cosine.

    Parameters
    ----------
    r : Array
        Distances, shape (...).
    r_inner : float
        Inner cutoff radius.
    r_outer : float
        Outer cutoff radius.

    Returns
    -------
    cutoff : Array
        Cutoff values, shape (...).
    """
    x = (r - r_inner) / (r_outer - r_inner)
    x = jnp.clip(x, 0.0, 1.0)
    return 0.5 * (1.0 + jnp.cos(jnp.pi * x))


# =============================================================================
# Slater-Koster rotation matrices
# =============================================================================


@jax.jit
def direction_cosines(r_vec: Array) -> Tuple[Array, Array, Array]:
    """
    Compute direction cosines from a displacement vector.

    Parameters
    ----------
    r_vec : Array
        Displacement vector, shape (..., 3).

    Returns
    -------
    l : Array
        Direction cosine along x, shape (...).
    m : Array
        Direction cosine along y, shape (...).
    n : Array
        Direction cosine along z, shape (...).
    """
    r = jnp.linalg.norm(r_vec, axis=-1, keepdims=True)
    r = jnp.maximum(r, 1e-10)  # Avoid division by zero

    cosines = r_vec / r
    l = cosines[..., 0]
    m = cosines[..., 1]
    n = cosines[..., 2]

    return l, m, n


@jax.jit
def sk_rotation_ss(l: Array, m: Array, n: Array, ss_sigma: Array) -> Array:
    """s-s interaction (trivial)."""
    return ss_sigma


@jax.jit
def sk_rotation_sp(l: Array, m: Array, n: Array, sp_sigma: Array) -> Array:
    """
    s-p interaction rotation.

    Returns matrix elements: (s|H|px), (s|H|py), (s|H|pz).

    Parameters
    ----------
    l : Array
        Direction cosine along x.
    m : Array
        Direction cosine along y.
    n : Array
        Direction cosine along z.
    sp_sigma : Array
        s-p sigma integral.

    Returns
    -------
    sp_elements : Array
        Matrix elements, shape (..., 3).
    """
    return sp_sigma[..., None] * jnp.stack([l, m, n], axis=-1)


@jax.jit
def sk_rotation_pp(
    l: Array, m: Array, n: Array, pp_sigma: Array, pp_pi: Array
) -> Array:
    """
    p-p interaction rotation.

    Returns 3x3 matrix of (pi|H|pj) elements.

    Parameters
    ----------
    l : Array
        Direction cosine along x.
    m : Array
        Direction cosine along y.
    n : Array
        Direction cosine along z.
    pp_sigma : Array
        p-p sigma integral.
    pp_pi : Array
        p-p pi integral.

    Returns
    -------
    pp_matrix : Array
        Matrix of (pi|H|pj) elements, shape (..., 3, 3).
    """
    # Diagonal elements
    ll = l * l
    mm = m * m
    nn = n * n

    H_xx = ll * pp_sigma + (1 - ll) * pp_pi
    H_yy = mm * pp_sigma + (1 - mm) * pp_pi
    H_zz = nn * pp_sigma + (1 - nn) * pp_pi

    # Off-diagonal elements
    lm = l * m
    ln = l * n
    mn = m * n

    H_xy = lm * (pp_sigma - pp_pi)
    H_xz = ln * (pp_sigma - pp_pi)
    H_yz = mn * (pp_sigma - pp_pi)

    # Construct matrix
    row1 = jnp.stack([H_xx, H_xy, H_xz], axis=-1)
    row2 = jnp.stack([H_xy, H_yy, H_yz], axis=-1)
    row3 = jnp.stack([H_xz, H_yz, H_zz], axis=-1)

    return jnp.stack([row1, row2, row3], axis=-2)


@jax.jit
def sk_rotation_sd(l: Array, m: Array, n: Array, sd_sigma: Array) -> Array:
    """
    s-d interaction rotation.

    Returns (s|H|d) for 5 d-orbitals in order: dxy, dyz, dzx, dx2-y2, dz2.

    Parameters
    ----------
    l : Array
        Direction cosine along x.
    m : Array
        Direction cosine along y.
    n : Array
        Direction cosine along z.
    sd_sigma : Array
        s-d sigma integral.

    Returns
    -------
    sd_elements : Array
        Matrix elements for 5 d-orbitals, shape (..., 5).
    """
    sqrt3 = jnp.sqrt(3.0)

    # d orbital order: xy, yz, zx, x2-y2, 3z2-r2
    d_xy = sqrt3 * l * m * sd_sigma
    d_yz = sqrt3 * m * n * sd_sigma
    d_zx = sqrt3 * n * l * sd_sigma
    d_x2y2 = 0.5 * sqrt3 * (l * l - m * m) * sd_sigma
    d_z2 = (n * n - 0.5 * (l * l + m * m)) * sd_sigma

    return jnp.stack([d_xy, d_yz, d_zx, d_x2y2, d_z2], axis=-1)


@jax.jit
def sk_rotation_pd(
    l: Array, m: Array, n: Array, pd_sigma: Array, pd_pi: Array
) -> Array:
    """
    p-d interaction rotation.

    Returns 3x5 matrix of (pi|H|dj) elements.

    Parameters
    ----------
    l : Array
        Direction cosine along x.
    m : Array
        Direction cosine along y.
    n : Array
        Direction cosine along z.
    pd_sigma : Array
        p-d sigma integral.
    pd_pi : Array
        p-d pi integral.

    Returns
    -------
    pd_matrix : Array
        Matrix of (pi|H|dj) elements, shape (..., 3, 5).
    """
    sqrt3 = jnp.sqrt(3.0)
    ll, mm, nn = l * l, m * m, n * n
    lm, mn, nl = l * m, m * n, n * l

    # px row
    px_dxy = sqrt3 * ll * m * pd_sigma + m * (1 - 2 * ll) * pd_pi
    px_dyz = sqrt3 * l * mn * pd_sigma - 2 * l * mn * pd_pi
    px_dzx = sqrt3 * ll * n * pd_sigma + n * (1 - 2 * ll) * pd_pi
    px_dx2y2 = (
        0.5 * sqrt3 * l * (ll - mm) * pd_sigma + l * (1 - ll + mm) * pd_pi
    )
    px_dz2 = l * (nn - 0.5 * (ll + mm)) * pd_sigma - sqrt3 * l * nn * pd_pi

    # py row
    py_dxy = sqrt3 * mm * l * pd_sigma + l * (1 - 2 * mm) * pd_pi
    py_dyz = sqrt3 * mm * n * pd_sigma + n * (1 - 2 * mm) * pd_pi
    py_dzx = sqrt3 * m * nl * pd_sigma - 2 * m * nl * pd_pi
    py_dx2y2 = (
        0.5 * sqrt3 * m * (ll - mm) * pd_sigma - m * (1 + ll - mm) * pd_pi
    )
    py_dz2 = m * (nn - 0.5 * (ll + mm)) * pd_sigma - sqrt3 * m * nn * pd_pi

    # pz row
    pz_dxy = sqrt3 * lm * n * pd_sigma - 2 * lm * n * pd_pi
    pz_dyz = sqrt3 * nn * m * pd_sigma + m * (1 - 2 * nn) * pd_pi
    pz_dzx = sqrt3 * nn * l * pd_sigma + l * (1 - 2 * nn) * pd_pi
    pz_dx2y2 = 0.5 * sqrt3 * n * (ll - mm) * pd_sigma - n * (ll - mm) * pd_pi
    pz_dz2 = (
        n * (nn - 0.5 * (ll + mm)) * pd_sigma + sqrt3 * n * (ll + mm) * pd_pi
    )

    px_row = jnp.stack([px_dxy, px_dyz, px_dzx, px_dx2y2, px_dz2], axis=-1)
    py_row = jnp.stack([py_dxy, py_dyz, py_dzx, py_dx2y2, py_dz2], axis=-1)
    pz_row = jnp.stack([pz_dxy, pz_dyz, pz_dzx, pz_dx2y2, pz_dz2], axis=-1)

    return jnp.stack([px_row, py_row, pz_row], axis=-2)


# =============================================================================
# File I/O for SK parameters
# =============================================================================


def load_sk_tables(
    directory: str, elements: List[str], format: str = "dftb+"
) -> SlaterKosterParams:
    """
    Load Slater-Koster parameter files.

    Parameters
    ----------
    directory : str
        Path to directory containing SK files.
    elements : List[str]
        List of element symbols to load.
    format : str, optional
        File format ('dftb+' or 'hotbit'). Default is 'dftb+'.

    Returns
    -------
    params : SlaterKosterParams
        SlaterKosterParams with loaded tables.
    """
    # TODO: Implement actual file parsing
    # This is a placeholder that should be replaced with real implementation
    raise NotImplementedError(
        "SK file loading not yet implemented. "
        "Use create_test_sk_params() for testing."
    )


def create_test_sk_params(elements: List[str] = None) -> SlaterKosterParams:
    """
    Create test SK parameters for development.

    This generates placeholder parameters that are differentiable
    but not physically accurate. For testing only.

    Parameters
    ----------
    elements : List[str], optional
        List of element symbols. Default is ['Bi', 'Te'].

    Returns
    -------
    params : SlaterKosterParams
        SlaterKosterParams with placeholder values.
    """
    # Create orbital bases
    if elements is None:
        elements = ["Bi", "Te"]
    element_bases = {}
    for elem in elements:
        if elem in ["Bi", "Sb"]:
            # Heavy p-block: s, p orbitals
            element_bases[elem] = OrbitalBasis(
                element=elem,
                orbitals=["s", "px", "py", "pz"],
                n_orbitals=4,
                onsite_energies=jnp.array([-10.0, -1.0, -1.0, -1.0]),
                soc_xi=jnp.array([0.0, 1.5]),  # s, p SOC constants (eV)
            )
        elif elem in ["Te", "Se"]:
            element_bases[elem] = OrbitalBasis(
                element=elem,
                orbitals=["s", "px", "py", "pz"],
                n_orbitals=4,
                onsite_energies=jnp.array([-12.0, -2.0, -2.0, -2.0]),
                soc_xi=jnp.array([0.0, 0.4]),  # s, p SOC constants (eV)
            )
        elif elem == "Mn":
            # Transition metal: s, p, d orbitals
            element_bases[elem] = OrbitalBasis(
                element=elem,
                orbitals=[
                    "s",
                    "px",
                    "py",
                    "pz",
                    "dxy",
                    "dyz",
                    "dzx",
                    "dx2y2",
                    "dz2",
                ],
                n_orbitals=9,
                onsite_energies=jnp.array(
                    [
                        -8.0,  # s
                        -3.0,
                        -3.0,
                        -3.0,  # p
                        -5.0,
                        -5.0,
                        -5.0,
                        -5.0,
                        -5.0,  # d
                    ]
                ),
                spin_w=jnp.array(
                    [
                        -0.02,  # s
                        -0.02,
                        -0.02,
                        -0.02,  # p
                        -0.04,
                        -0.04,
                        -0.04,
                        -0.04,
                        -0.04,  # d (Stoner parameter)
                    ]
                ),
                soc_xi=jnp.array([0.0, 0.05, 0.05]),  # s, p, d SOC
            )
        else:
            # Default: s, p orbitals
            element_bases[elem] = OrbitalBasis(
                element=elem,
                orbitals=["s", "px", "py", "pz"],
                n_orbitals=4,
                onsite_energies=jnp.array([-10.0, -2.0, -2.0, -2.0]),
            )

    # Create SK tables for all pairs
    tables = {}
    distances = jnp.linspace(1.0, 10.0, 100)  # Angstroms

    for i, elem1 in enumerate(elements):
        for elem2 in elements[i:]:
            # Generate placeholder SK integrals (exponential decay)
            n_integrals = len(SK_INTEGRALS_SP)
            H_integrals = jnp.zeros((len(distances), n_integrals))
            S_integrals = jnp.zeros((len(distances), n_integrals))

            # Typical decay lengths
            for j, name in enumerate(SK_INTEGRALS_SP):
                amplitude_H = -2.0 if "sigma" in name else -0.5
                amplitude_S = 0.5 if j == 0 else 0.3
                decay = 1.0

                H_integrals = H_integrals.at[:, j].set(
                    amplitude_H * jnp.exp(-distances / decay)
                )
                S_integrals = S_integrals.at[:, j].set(
                    amplitude_S * jnp.exp(-distances / decay)
                )

            tables[(elem1, elem2)] = SlaterKosterTable(
                element1=elem1,
                element2=elem2,
                distances=distances,
                H_integrals=H_integrals,
                S_integrals=S_integrals,
                repulsive=None,
                integral_names=SK_INTEGRALS_SP,
            )

    return SlaterKosterParams(elements=element_bases, tables=tables)
