from __future__ import annotations
import numpy as np
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional, Any
from enum import Enum

# ─────────────────────────────────────────────────────────────────────────────
# 1. Enums & Dataclasses
# ─────────────────────────────────────────────────────────────────────────────

class OrbitalType(int, Enum):
    S = 0
    PX = 1; PY = 2; PZ = 3
    DXY = 4; DYZ = 5; DZX = 6; DX2Y2 = 7; DZ2 = 8

class InteractionType(str, Enum):
    SS_SIGMA = "ss_sigma"
    SP_SIGMA = "sp_sigma"
    PP_SIGMA = "pp_sigma"
    PP_PI    = "pp_pi"
    SD_SIGMA = "sd_sigma"
    PD_SIGMA = "pd_sigma"
    PD_PI    = "pd_pi"
    DD_SIGMA = "dd_sigma"
    DD_PI    = "dd_pi"
    DD_DELTA = "dd_delta"

class BasisMode(str, Enum):
    HEURISTIC = "heuristic"
    S = "s"
    SP = "sp"
    SPD = "spd"

class EnergyReference(str, Enum):
    FERMI = "fermi"
    HOMO = "homo"
    ZERO = "zero"

@dataclass
class RadialParam:
    v0: float
    beta: float
    r0: float
    rc: float
    rs: float # r_smooth

@dataclass
class TBParameters:
    # Hamiltonian Controls
    basis:          BasisMode = BasisMode.HEURISTIC
    hopping_scale:  float = 1.0
    cutoff_radius:  float = 0.0 # 0.0 means use DB default
    beta_scale:     float = 1.0
    smooth_onset:   float = 0.8
    use_cf:         bool  = False
    
    # Solver Controls
    n_electrons:    int   = -1  # -1 means auto-detect from valence
    smearing_kT:    float = 0.01
    energy_ref:     EnergyReference = EnergyReference.FERMI
    
    # Analysis Controls
    dos_sigma:      float = 0.1
    dos_n_bins:     int   = 400
    
    # Visualization Controls (Phase 4)
    grid_spacing:   float = 0.20
    grid_padding:   float = 3.5
    iso_value:      float = 0.05
    
    # Store options
    store_hamiltonian: bool = False
    store_neighbors:   bool = False
    store_orb_info:    bool = False

@dataclass
class TBResults:
    eigenvalues:  np.ndarray = field(default_factory=lambda: np.array([], dtype=np.float32))
    eigenvectors: np.ndarray = field(default_factory=lambda: np.array([[]], dtype=np.float32))
    occupations:  np.ndarray = field(default_factory=lambda: np.array([], dtype=np.float32))
    fermi_level:  float = 0.0
    homo:         float = 0.0
    lumo:         float = 0.0
    gap:          float = 0.0
    n_electrons:  int = 0
    n_occ_states: int = 0
    basis:        BasisMode = BasisMode.HEURISTIC
    
    local_charge:     np.ndarray = field(default_factory=lambda: np.array([], dtype=np.float32))
    approx_charge:    np.ndarray = field(default_factory=lambda: np.array([], dtype=np.float32))
    state_atom_weights: np.ndarray = field(default_factory=lambda: np.array([[]], dtype=np.float32))
    ipr:              np.ndarray = field(default_factory=lambda: np.array([], dtype=np.float32))
    bond_order:       Dict[Tuple[int,int], float] = field(default_factory=dict)
    
    dos_x:        np.ndarray = field(default_factory=lambda: np.array([], dtype=np.float32))
    dos_y:        np.ndarray = field(default_factory=lambda: np.array([], dtype=np.float32))
    pdos_y:       Dict[str, np.ndarray] = field(default_factory=dict)
    pdos_orb:     Dict[str, np.ndarray] = field(default_factory=dict)
    
    # Field Reconstruction (Phase 4)
    orbital_fields: Dict[int, Any] = field(default_factory=dict) # state_idx -> ScalarField3D
    density_fields: Dict[str, Any] = field(default_factory=dict) # name -> ScalarField3D
    isosurfaces:    Dict[str, Any] = field(default_factory=dict) # name -> IsosurfaceMesh
    
    warnings:     List[str] = field(default_factory=list)
    atom_to_orbs: List[List[int]] = field(default_factory=list)
    
    # Optional Storage for Debug/Export
    hamiltonian:  Optional[np.ndarray] = None
    neighbors:    Optional[List[Tuple[int, int, float, np.ndarray]]] = None
    orb_info:     Optional[List[Tuple[int, OrbitalType]]] = None

# ─────────────────────────────────────────────────────────────────────────────
# 2. Universal Fallbacks (Periodic Trends)
# ─────────────────────────────────────────────────────────────────────────────

# (Atomic Number, Valence Electrons, Group)
PERIODIC_DATA = {
    "H": (1, 1, 1), "He": (2, 2, 18),
    "Li": (3, 1, 1), "Be": (4, 2, 2), "B": (5, 3, 13), "C": (6, 4, 14), "N": (7, 5, 15), "O": (8, 6, 16), "F": (9, 7, 17), "Ne": (10, 8, 18),
    "Na": (11, 1, 1), "Mg": (12, 2, 2), "Al": (13, 3, 13), "Si": (14, 4, 14), "P": (15, 5, 15), "S": (16, 6, 16), "Cl": (17, 7, 17), "Ar": (18, 8, 18),
    "K": (19, 1, 1), "Ca": (20, 2, 2), "Sc": (21, 3, 3), "Ti": (22, 4, 4), "V": (23, 5, 5), "Cr": (24, 6, 6), "Mn": (25, 7, 7), "Fe": (26, 8, 8), "Co": (27, 9, 9), "Ni": (28, 10, 10), "Cu": (29, 11, 11), "Zn": (30, 12, 12), "Ga": (31, 3, 13), "Ge": (32, 4, 14), "As": (33, 5, 15), "Se": (34, 6, 16), "Br": (35, 7, 17), "Kr": (36, 8, 18),
    "Pd": (46, 10, 10), "Ag": (47, 11, 11), "Au": (79, 11, 11)
}

def fallback_valence(el: str) -> int:
    return PERIODIC_DATA.get(el, (0, 4, 14))[1]

def fallback_basis(el: str, mode: BasisMode = BasisMode.HEURISTIC) -> List[OrbitalType]:
    z, _, group = PERIODIC_DATA.get(el, (0, 0, 14))
    if mode == BasisMode.S: return [OrbitalType.S]
    if mode == BasisMode.SP: return [OrbitalType.S, OrbitalType.PX, OrbitalType.PY, OrbitalType.PZ]
    if mode == BasisMode.SPD: return [OrbitalType.S, OrbitalType.PX, OrbitalType.PY, OrbitalType.PZ, OrbitalType.DXY, OrbitalType.DYZ, OrbitalType.DZX, OrbitalType.DX2Y2, OrbitalType.DZ2]
    
    # Heuristic Fallback
    res = [OrbitalType.S]
    if group >= 13 or z in [5,6,7,8,9]: res += [OrbitalType.PX, OrbitalType.PY, OrbitalType.PZ]
    if 3 <= group <= 12 or (z > 20 and group <= 12): res += [OrbitalType.DXY, OrbitalType.DYZ, OrbitalType.DZX, OrbitalType.DX2Y2, OrbitalType.DZ2]
    return res

def fallback_onsite(el: str, orb: OrbitalType) -> float:
    z, _, _ = PERIODIC_DATA.get(el, (0, 0, 14))
    if orb == OrbitalType.S: return -10.0 - 0.5 * (z % 10)
    if orb.value <= 3: return -6.0 - 0.3 * (z % 10)
    return -8.0 - 0.1 * (z % 10)

def fallback_hopping(elA: str, elB: str, itype: InteractionType) -> RadialParam:
    v0_map = {
        InteractionType.SS_SIGMA: -1.0, InteractionType.SP_SIGMA: 1.2,
        InteractionType.PP_SIGMA: 1.5,  InteractionType.PP_PI: -0.4,
        InteractionType.SD_SIGMA: -0.8, InteractionType.PD_SIGMA: -1.0, InteractionType.PD_PI: 0.5,
        InteractionType.DD_SIGMA: -0.6, InteractionType.DD_PI: 0.3, InteractionType.DD_DELTA: -0.1
    }
    return RadialParam(v0=v0_map.get(itype, 0.0), beta=1.8, r0=1.5, rc=4.0, rs=3.2)

def fallback_alpha(el: str, orb: OrbitalType) -> float:
    z, _, _ = PERIODIC_DATA.get(el, (0, 0, 14))
    if orb == OrbitalType.S: return 0.8 + 0.1 * (z // 10)
    if orb.value <= 3: return 0.6 + 0.1 * (z // 10) # P
    return 1.8 + 0.2 * (z // 10) # D: compact

# ─────────────────────────────────────────────────────────────────────────────
# 3. Parameter Database
# ─────────────────────────────────────────────────────────────────────────────

class ParameterDB:
    def __init__(self):
        self.onsite: Dict[Tuple[str, OrbitalType], float] = {}
        self.valence: Dict[str, int] = {}
        self.basis: Dict[str, List[OrbitalType]] = {}
        self.hopping: Dict[Tuple[str, str, OrbitalType, OrbitalType, InteractionType], RadialParam] = {}
        self.crystal_field: Dict[str, Dict[OrbitalType, float]] = {}
        self.orbital_alpha: Dict[Tuple[str, OrbitalType], float] = {}
        self.max_cutoff: float = 0.0
        self.warnings: List[str] = []

    def set_onsite(self, el: str, orb: OrbitalType, val: float):
        self.onsite[(el, orb)] = val

    def set_valence(self, el: str, val: int):
        self.valence[el] = val

    def set_basis(self, el: str, orbs: List[OrbitalType]):
        self.basis[el] = orbs

    def set_hopping(self, elA: str, elB: str, orbA: OrbitalType, orbB: OrbitalType, itype: InteractionType, rad: RadialParam):
        if elA > elB:
            elA, elB = elB, elA
            orbA, orbB = orbB, orbA
        elif elA == elB and orbA > orbB:
            orbA, orbB = orbB, orbA
        
        self.hopping[(elA, elB, orbA, orbB, itype)] = rad
        self.max_cutoff = max(self.max_cutoff, rad.rc)

    def set_crystal_field(self, el: str, shifts: Dict[OrbitalType, float]):
        self.crystal_field[el] = shifts

    def set_alpha(self, el: str, orb: OrbitalType, val: float):
        self.orbital_alpha[(el, orb)] = val

    def get_onsite(self, el: str, orb: OrbitalType) -> float:
        if (el, orb) in self.onsite: return self.onsite[(el, orb)]
        msg = f"Fallback onsite for {el} {orb.name}"
        if msg not in self.warnings: self.warnings.append(msg)
        return fallback_onsite(el, orb)

    def get_alpha(self, el: str, orb: OrbitalType) -> float:
        if (el, orb) in self.orbital_alpha: return self.orbital_alpha[(el, orb)]
        return fallback_alpha(el, orb)

    def get_hopping(self, elA: str, elB: str, orbA: OrbitalType, orbB: OrbitalType, itype: InteractionType) -> Optional[RadialParam]:
        if elA > elB:
            elA, elB = elB, elA
            orbA, orbB = orbB, orbA
        elif elA == elB and orbA > orbB:
            orbA, orbB = orbB, orbA
        return self.hopping.get((elA, elB, orbA, orbB, itype))

# Global DB Instance
DB = ParameterDB()

def init_db():
    # --- Elements (Level 0/1) ---
    for el, v in [("H",1), ("C",4), ("N",5), ("O",6), ("Ni",10), ("Cu",11)]:
        DB.set_valence(el, v)
    
    # --- Basis Sets (Source of Truth) ---
    DB.set_basis("H",  [OrbitalType.S])
    DB.set_basis("C",  [OrbitalType.S, OrbitalType.PX, OrbitalType.PY, OrbitalType.PZ])
    DB.set_basis("N",  [OrbitalType.S, OrbitalType.PX, OrbitalType.PY, OrbitalType.PZ])
    DB.set_basis("O",  [OrbitalType.S, OrbitalType.PX, OrbitalType.PY, OrbitalType.PZ])
    DB.set_basis("Ni", [OrbitalType.S, OrbitalType.DXY, OrbitalType.DYZ, OrbitalType.DZX, OrbitalType.DX2Y2, OrbitalType.DZ2])
    DB.set_basis("Cu", [OrbitalType.S, OrbitalType.DXY, OrbitalType.DYZ, OrbitalType.DZX, OrbitalType.DX2Y2, OrbitalType.DZ2])

    # --- On-sites ---
    data_onsite = [
        ("H", OrbitalType.S, -13.6),
        ("C", OrbitalType.S, -15.2), ("C", OrbitalType.PX, -8.5), ("C", OrbitalType.PY, -8.5), ("C", OrbitalType.PZ, -8.5),
        ("N", OrbitalType.S, -18.0), ("N", OrbitalType.PX, -10.8), ("N", OrbitalType.PY, -10.8), ("N", OrbitalType.PZ, -10.8),
        ("O", OrbitalType.S, -21.2), ("O", OrbitalType.PX, -12.3), ("O", OrbitalType.PY, -12.3), ("O", OrbitalType.PZ, -12.3),
        ("Ni", OrbitalType.S, -6.0), ("Ni", OrbitalType.DXY, -8.8), ("Ni", OrbitalType.DYZ, -8.8), ("Ni", OrbitalType.DZX, -8.8), ("Ni", OrbitalType.DX2Y2, -8.8), ("Ni", OrbitalType.DZ2, -8.8),
        ("Cu", OrbitalType.S, -5.5), ("Cu", OrbitalType.DXY, -9.5), ("Cu", OrbitalType.DYZ, -9.5), ("Cu", OrbitalType.DZX, -9.5), ("Cu", OrbitalType.DX2Y2, -9.5), ("Cu", OrbitalType.DZ2, -9.5),
    ]
    for el, orb, val in data_onsite:
        DB.set_onsite(el, orb, val)

    # --- Crystal Field (Level 2) ---
    DB.set_crystal_field("Ni", {OrbitalType.DZ2: 0.20, OrbitalType.DX2Y2: 0.20, OrbitalType.DXY: -0.10, OrbitalType.DYZ: -0.10, OrbitalType.DZX: -0.10})
    DB.set_crystal_field("Cu", {OrbitalType.DX2Y2: 0.25, OrbitalType.DZ2: 0.15})

    # --- Hoppings (Level 1/2) ---
    def add_sk(a, b, oA, oB, itype, v0, beta, r0, rc, rs):
        DB.set_hopping(a, b, oA, oB, itype, RadialParam(v0, beta, r0, rc, rs))

    # H-H
    add_sk("H","H", OrbitalType.S, OrbitalType.S, InteractionType.SS_SIGMA, -2.5, 1.5, 0.74, 3.0, 2.6)
    # C-C
    add_sk("C","C", OrbitalType.S, OrbitalType.S, InteractionType.SS_SIGMA, -2.8, 1.4, 1.54, 3.5, 3.0)
    add_sk("C","C", OrbitalType.S, OrbitalType.PX, InteractionType.SP_SIGMA, 3.0, 1.4, 1.54, 3.5, 3.0)
    add_sk("C","C", OrbitalType.PX, OrbitalType.PX, InteractionType.PP_SIGMA, 3.5, 1.4, 1.54, 3.5, 3.0)
    add_sk("C","C", OrbitalType.PX, OrbitalType.PY, InteractionType.PP_PI, -1.0, 1.4, 1.54, 3.5, 3.0)
    # N-N
    add_sk("N","N", OrbitalType.S, OrbitalType.S, InteractionType.SS_SIGMA, -2.6, 1.4, 1.45, 3.5, 3.0)
    add_sk("N","N", OrbitalType.S, OrbitalType.PX, InteractionType.SP_SIGMA, 3.0, 1.4, 1.45, 3.5, 3.0)
    add_sk("N","N", OrbitalType.PX, OrbitalType.PX, InteractionType.PP_SIGMA, 3.4, 1.4, 1.45, 3.5, 3.0)
    add_sk("N","N", OrbitalType.PX, OrbitalType.PY, InteractionType.PP_PI, -1.0, 1.4, 1.45, 3.5, 3.0)
    # O-O
    add_sk("O","O", OrbitalType.S, OrbitalType.S, InteractionType.SS_SIGMA, -2.4, 1.5, 1.48, 3.2, 2.8)
    add_sk("O","O", OrbitalType.S, OrbitalType.PX, InteractionType.SP_SIGMA, 2.9, 1.5, 1.48, 3.2, 2.8)
    add_sk("O","O", OrbitalType.PX, OrbitalType.PX, InteractionType.PP_SIGMA, 3.3, 1.5, 1.48, 3.2, 2.8)
    add_sk("O","O", OrbitalType.PX, OrbitalType.PY, InteractionType.PP_PI, -1.0, 1.5, 1.48, 3.2, 2.8)
    # Ni-Ni
    add_sk("Ni","Ni", OrbitalType.S, OrbitalType.S, InteractionType.SS_SIGMA, -1.5, 1.2, 2.49, 3.8, 3.3)
    add_sk("Ni","Ni", OrbitalType.S, OrbitalType.DXY, InteractionType.SD_SIGMA, -2.2, 1.2, 2.49, 3.8, 3.3)
    add_sk("Ni","Ni", OrbitalType.DXY, OrbitalType.DXY, InteractionType.DD_SIGMA, -1.6, 1.2, 2.49, 3.8, 3.3)
    add_sk("Ni","Ni", OrbitalType.DXY, OrbitalType.DYZ, InteractionType.DD_PI, 0.9, 1.2, 2.49, 3.8, 3.3)
    add_sk("Ni","Ni", OrbitalType.DXY, OrbitalType.DX2Y2, InteractionType.DD_DELTA, -0.25, 1.2, 2.49, 3.8, 3.3)
    # Cu-Cu
    add_sk("Cu","Cu", OrbitalType.S, OrbitalType.S, InteractionType.SS_SIGMA, -1.5, 1.2, 2.56, 3.8, 3.3)
    add_sk("Cu","Cu", OrbitalType.S, OrbitalType.DXY, InteractionType.SD_SIGMA, -2.1, 1.2, 2.56, 3.8, 3.3)
    add_sk("Cu","Cu", OrbitalType.DXY, OrbitalType.DXY, InteractionType.DD_SIGMA, -1.5, 1.2, 2.56, 3.8, 3.3)
    add_sk("Cu","Cu", OrbitalType.DXY, OrbitalType.DYZ, InteractionType.DD_PI, 0.9, 1.2, 2.56, 3.8, 3.3)
    add_sk("Cu","Cu", OrbitalType.DXY, OrbitalType.DX2Y2, InteractionType.DD_DELTA, -0.25, 1.2, 2.56, 3.8, 3.3)
    # Ni-O
    add_sk("Ni","O", OrbitalType.S, OrbitalType.S, InteractionType.SS_SIGMA, -1.8, 1.2, 2.00, 3.5, 3.0)
    add_sk("Ni","O", OrbitalType.S, OrbitalType.PX, InteractionType.SP_SIGMA, 2.8, 1.2, 2.00, 3.5, 3.0)
    add_sk("Ni","O", OrbitalType.DXY, OrbitalType.PX, InteractionType.PD_SIGMA, -2.8, 1.2, 2.00, 3.5, 3.0)
    add_sk("Ni","O", OrbitalType.DXY, OrbitalType.PY, InteractionType.PD_PI, 1.3, 1.2, 2.00, 3.5, 3.0)
    # Cu-O
    add_sk("Cu","O", OrbitalType.S, OrbitalType.S, InteractionType.SS_SIGMA, -1.7, 1.2, 2.00, 3.5, 3.0)
    add_sk("Cu","O", OrbitalType.S, OrbitalType.PX, InteractionType.SP_SIGMA, 2.7, 1.2, 2.00, 3.5, 3.0)
    add_sk("Cu","O", OrbitalType.DXY, OrbitalType.PX, InteractionType.PD_SIGMA, -2.6, 1.2, 2.00, 3.5, 3.0)
    add_sk("Cu","O", OrbitalType.DXY, OrbitalType.PY, InteractionType.PD_PI, 1.2, 1.2, 2.00, 3.5, 3.0)
    # C-O
    add_sk("C","O", OrbitalType.S, OrbitalType.S, InteractionType.SS_SIGMA, -2.7, 1.5, 1.43, 3.4, 2.9)
    add_sk("C","O", OrbitalType.S, OrbitalType.PX, InteractionType.SP_SIGMA, 3.0, 1.5, 1.43, 3.4, 2.9)
    add_sk("C","O", OrbitalType.PX, OrbitalType.PX, InteractionType.PP_SIGMA, 3.5, 1.5, 1.43, 3.4, 2.9)
    add_sk("C","O", OrbitalType.PX, OrbitalType.PY, InteractionType.PP_PI, -1.0, 1.5, 1.43, 3.4, 2.9)

init_db()


# ─────────────────────────────────────────────────────────────────────────────
# 4. Slater-Koster 9x9 Block
# ─────────────────────────────────────────────────────────────────────────────

def get_sk_block(l: float, m: float, n: float, p: Dict[InteractionType, float]) -> np.ndarray:
    H = np.zeros((9, 9))
    l2, m2, n2 = l*l, m*m, n*n
    root3 = np.sqrt(3.0)
    
    def get_v(it: InteractionType): return p.get(it, 0.0)

    # ss
    H[0, 0] = get_v(InteractionType.SS_SIGMA)
    # sp
    v_sp = get_v(InteractionType.SP_SIGMA)
    H[0, 1], H[1, 0] =  l * v_sp, -l * v_sp
    H[0, 2], H[2, 0] =  m * v_sp, -m * v_sp
    H[0, 3], H[3, 0] =  n * v_sp, -n * v_sp
    # pp
    vps, vpp = get_v(InteractionType.PP_SIGMA), get_v(InteractionType.PP_PI)
    H[1, 1] = l2 * vps + (1-l2) * vpp
    H[2, 2] = m2 * vps + (1-m2) * vpp
    H[3, 3] = n2 * vps + (1-n2) * vpp
    H[1, 2] = H[2, 1] = l*m * (vps - vpp)
    H[1, 3] = H[3, 1] = l*n * (vps - vpp)
    H[2, 3] = H[3, 2] = m*n * (vps - vpp)
    # sd
    vsd = get_v(InteractionType.SD_SIGMA)
    H[0, 4] = root3 * l * m * vsd
    H[0, 5] = root3 * m * n * vsd
    H[0, 6] = root3 * l * n * vsd
    H[0, 7] = 0.5 * root3 * (l2 - m2) * vsd
    H[0, 8] = (n2 - 0.5*(l2 + m2)) * vsd
    for i in range(4, 9): H[i, 0] = H[0, i]
    # pd
    vpds, vpdp = get_v(InteractionType.PD_SIGMA), get_v(InteractionType.PD_PI)
    H[1, 4] = root3 * l2 * m * vpds + m * (1.0 - 2.0 * l2) * vpdp
    H[1, 5] = root3 * l * m * n * vpds - 2.0 * l * m * n * vpdp
    H[1, 6] = root3 * l2 * n * vpds + n * (1.0 - 2.0 * l2) * vpdp
    H[1, 7] = 0.5 * root3 * l * (l2 - m2) * vpds + l * (1.0 - l2 + m2) * vpdp
    H[1, 8] = root3 * l * (n2 - 0.5 * (l2 + m2)) * vpds - root3 * l * n2 * vpdp
    H[2, 4] = root3 * l * m2 * vpds + l * (1.0 - 2.0 * m2) * vpdp
    H[2, 5] = root3 * m2 * n * vpds + n * (1.0 - 2.0 * m2) * vpdp
    H[2, 6] = H[1, 5]
    H[2, 7] = 0.5 * root3 * m * (l2 - m2) * vpds - m * (1.0 + l2 - m2) * vpdp
    H[2, 8] = root3 * m * (n2 - 0.5 * (l2 + m2)) * vpds - root3 * m * n2 * vpdp
    H[3, 4] = H[1, 5]
    H[3, 5] = root3 * m * n2 * vpds + m * (1.0 - 2.0 * n2) * vpdp
    H[3, 6] = root3 * l * n2 * vpds + l * (1.0 - 2.0 * n2) * vpdp
    H[3, 7] = 0.5 * root3 * n * (l2 - m2) * vpds - root3 * n * (l2 - m2) * vpdp
    H[3, 8] = n * (n2 - 0.5 * (l2 + m2)) * vpds + root3 * n * (l2 + m2) * vpdp
    for i in range(1, 4):
        for j in range(4, 9): H[j, i] = H[i, j]
    # dd
    vds, vdp, vdd = get_v(InteractionType.DD_SIGMA), get_v(InteractionType.DD_PI), get_v(InteractionType.DD_DELTA)
    H[4, 4] = 3.0 * l2 * m2 * vds + (l2 + m2 - 4.0 * l2 * m2) * vdp + (n2 + l2 * m2) * vdd
    H[5, 5] = 3.0 * m2 * n2 * vds + (m2 + n2 - 4.0 * m2 * n2) * vdp + (l2 + m2 * n2) * vdd
    H[6, 6] = 3.0 * l2 * n2 * vds + (l2 + n2 - 4.0 * l2 * n2) * vdp + (m2 + l2 * n2) * vdd
    H[7, 7] = 0.75 * (l2 - m2)**2 * vds + (l2 + m2 - (l2 - m2)**2) * vdp + (n2 + 0.25 * (l2 - m2)**2) * vdd
    H[8, 8] = (n2 - 0.5 * (l2 + m2))**2 * vds + 3.0 * n2 * (l2 + m2) * vdp + 0.75 * (l2 + m2)**2 * vdd
    H[4, 5] = H[5, 4] = 3.0 * l * m2 * n * vds + l * n * (1.0 - 4.0 * m2) * vdp + l * n * (m2 - 1.0) * vdd
    H[4, 6] = H[6, 4] = 3.0 * l2 * m * n * vds + m * n * (1.0 - 4.0 * l2) * vdp + m * n * (l2 - 1.0) * vdd
    H[4, 7] = H[7, 4] = 1.5 * l * m * (l2 - m2) * vds + 2.0 * l * m * (m2 - l2) * vdp + 0.5 * l * m * (l2 - m2) * vdd
    H[4, 8] = H[8, 4] = root3 * l * m * (n2 - 0.5 * (l2 + m2)) * vds - 2.0 * root3 * l * m * n2 * vdp + 0.5 * root3 * l * m * (1.0 + n2) * vdd
    H[5, 6] = H[6, 5] = 3.0 * l * m * n2 * vds + l * m * (1.0 - 4.0 * n2) * vdp + l * m * (n2 - 1.0) * vdd
    H[5, 7] = H[7, 5] = 1.5 * m * n * (l2 - m2) * vds - m * n * (1.0 + 2.0 * (l2 - m2)) * vdp + m * n * (1.0 + 0.5 * (l2 - m2)) * vdd
    H[5, 8] = H[8, 5] = m * n * (n2 - 0.5 * (l2 + m2)) * vds + root3 * m * n * (l2 + m2 - n2) * vdp - 0.5 * root3 * m * n * (l2 + m2) * vdd
    H[6, 7] = H[7, 6] = 1.5 * l * n * (l2 - m2) * vds + l * n * (1.0 - 2.0 * (l2 - m2)) * vdp - l * n * (1.0 - 0.5 * (l2 - m2)) * vdd
    H[6, 8] = H[8, 6] = l * n * (n2 - 0.5 * (l2 + m2)) * vds + root3 * l * n * (l2 + m2 - n2) * vdp - 0.5 * root3 * l * n * (l2 + m2) * vdd
    H[7, 8] = H[8, 7] = 0.5 * root3 * (l2 - m2) * (n2 - 0.5 * (l2 + m2)) * vds + root3 * n2 * (m2 - l2) * vdp + 0.25 * root3 * (1.0 + n2) * (l2 - m2) * vdd
    return H

# ─────────────────────────────────────────────────────────────────────────────
# 4. Assembler Utilities
# ─────────────────────────────────────────────────────────────────────────────

def build_orbital_basis(elements: List[str], params: TBParameters) -> Tuple[List[List[int]], List[Tuple[int, OrbitalType]]]:
    atom_orbs: List[List[int]] = []
    orb_info: List[Tuple[int, OrbitalType]] = []
    for i, el in enumerate(elements):
        start = len(orb_info)
        
        # Use DB basis if available, otherwise fallback
        if el in DB.basis:
            orbs = DB.basis[el]
        else:
            orbs = fallback_basis(el, params.basis)
        
        # Specific mode overrides (manual narrowing)
        if params.basis == BasisMode.S: orbs = [OrbitalType.S]
        elif params.basis == BasisMode.SP: orbs = [ot for ot in orbs if ot.value <= 3]
        elif params.basis == BasisMode.SPD: pass # Use all heuristic or DB ones
            
        for ot in orbs: orb_info.append((i, ot))
        atom_orbs.append(list(range(start, len(orb_info))))
    return atom_orbs, orb_info

def build_pair_template_cache(elements: List[str]) -> Dict[Tuple[str, str], Dict[InteractionType, Optional[RadialParam]]]:
    cache = {}
    unique_elements = sorted(set(elements))
    for i, elA in enumerate(unique_elements):
        for elB in unique_elements[i:]:
            key = (elA, elB)
            template = {}
            for itype in InteractionType:
                # 1. Try DB specific orbital pairs for this interaction
                rad = None
                if itype == InteractionType.SS_SIGMA: rad = DB.get_hopping(elA, elB, OrbitalType.S, OrbitalType.S, itype)
                elif itype == InteractionType.SP_SIGMA: rad = DB.get_hopping(elA, elB, OrbitalType.S, OrbitalType.PX, itype)
                elif itype == InteractionType.PP_SIGMA: rad = DB.get_hopping(elA, elB, OrbitalType.PX, OrbitalType.PX, InteractionType.PP_SIGMA)
                elif itype == InteractionType.PP_PI:    rad = DB.get_hopping(elA, elB, OrbitalType.PX, OrbitalType.PY, InteractionType.PP_PI)
                elif itype == InteractionType.SD_SIGMA: rad = DB.get_hopping(elA, elB, OrbitalType.S, OrbitalType.DXY, InteractionType.SD_SIGMA)
                elif itype == InteractionType.PD_SIGMA: rad = DB.get_hopping(elA, elB, OrbitalType.DXY, OrbitalType.PX, InteractionType.PD_SIGMA)
                elif itype == InteractionType.PD_PI:    rad = DB.get_hopping(elA, elB, OrbitalType.DXY, OrbitalType.PY, InteractionType.PD_PI)
                elif itype == InteractionType.DD_SIGMA: rad = DB.get_hopping(elA, elB, OrbitalType.DXY, OrbitalType.DXY, InteractionType.DD_SIGMA)
                elif itype == InteractionType.DD_PI:    rad = DB.get_hopping(elA, elB, OrbitalType.DXY, OrbitalType.DYZ, InteractionType.DD_PI)
                elif itype == InteractionType.DD_DELTA: rad = DB.get_hopping(elA, elB, OrbitalType.DXY, OrbitalType.DX2Y2, InteractionType.DD_DELTA)
                
                # 2. If missing, use universal fallback
                if rad is None:
                    rad = fallback_hopping(elA, elB, itype)
                    # Optionally log warning for missing pair
                
                template[itype] = rad
            cache[key] = template
    return cache

def get_radial_val(rad: RadialParam, r: float, params: TBParameters) -> float:
    # Scale parameters from UI
    rc = params.cutoff_radius if params.cutoff_radius > 1e-4 else rad.rc
    rs = params.smooth_onset * rc
    beta = rad.beta * params.beta_scale
    
    if r >= rc: return 0.0
    smooth = 1.0
    if r > rs:
        x = (r - rs) / (rc - rs)
        smooth = 0.5 * (np.cos(np.pi * x) + 1.0)
    return rad.v0 * np.exp(-beta * (r - rad.r0)) * smooth

def build_neighbor_list(pos: np.ndarray, cutoff: float, lat: Optional[np.ndarray] = None):
    N = len(pos)
    neighbors = []
    
    # Check if we have a valid lattice for MIC
    if lat is not None and np.linalg.norm(lat) > 1e-4:
        # Minimum Image Convention (MIC) - assumes orthorhombic cell for simplicity
        is_ortho = (np.abs(lat - np.diag(np.diag(lat))).sum() < 1e-6)
        L = np.diag(lat) if is_ortho else None
        
        for i in range(N):
            for j in range(i + 1, N):
                dv = pos[j] - pos[i]
                if L is not None:
                    # Optimized MIC for orthorhombic cell
                    dv -= np.round(dv / L) * L
                d = float(np.linalg.norm(dv))
                if 0.01 < d < cutoff:
                    neighbors.append((i, j, d, dv))
                    
        # Self-interaction images (only for PBC)
        if L is not None:
            for i in range(N):
                # Check nearest images of the SAME atom (excluding self at 0,0,0)
                for axes in [(1,0,0), (0,1,0), (0,0,1)]:
                    dv = np.array(axes) * L
                    d = float(np.linalg.norm(dv))
                    if 0.01 < d < cutoff:
                        neighbors.append((i, i, d, dv))
    else:
        # Use cKDTree for non-periodic systems (fast O(N log N))
        try:
            from scipy.spatial import cKDTree
            tree = cKDTree(pos)
            pairs = tree.query_pairs(cutoff)
            for i, j in pairs:
                dv = pos[j] - pos[i]
                d = float(np.linalg.norm(dv))
                if d > 0.01:
                    neighbors.append((i, j, d, dv))
        except ImportError:
            for i in range(N):
                for j in range(i+1, N):
                    dv = pos[j] - pos[i]
                    d = float(np.linalg.norm(dv))
                    if 0.01 < d < cutoff:
                        neighbors.append((i, j, d, dv))
    return neighbors

# ─────────────────────────────────────────────────────────────────────────────
# 5. Hamiltonian Assembler
# ─────────────────────────────────────────────────────────────────────────────

def build_hamiltonian(pos: np.ndarray, elements: List[str], params: TBParameters, lat: Optional[np.ndarray]=None):
    atom_orbs, orb_info = build_orbital_basis(elements, params)
    M = len(orb_info)
    H = np.zeros((M, M), dtype=np.float64)
    
    # Diagonal + Crystal Field
    cf_shifts = DB.crystal_field if params.use_cf else {}
    for mu, (at_idx, ot) in enumerate(orb_info):
        el = elements[at_idx]
        val = DB.get_onsite(el, ot)
        if params.use_cf and el in cf_shifts:
            val += cf_shifts[el].get(ot, 0.0)
        H[mu, mu] = val
        
    # Neighbors & Caching
    pair_cache = build_pair_template_cache(elements)
    rc_max = max(DB.max_cutoff, params.cutoff_radius)
    neighbors = build_neighbor_list(pos, rc_max, lat)
    
    for (i, j, r, dv) in neighbors:
        elA, elB = elements[i], elements[j]
        l, m, n = dv / r
        
        # Get pairwise radial lookups from cache
        key = tuple(sorted((elA, elB)))
        template = pair_cache.get(key, {})
        
        pair_hops = {}
        for itype, rad in template.items():
            if rad:
                val = get_radial_val(rad, r, params) * params.hopping_scale
                if abs(val) > 1e-10:
                    pair_hops[itype] = val
        
        if not pair_hops: continue
        
        # Evaluate 9x9 SK block (only for the orbitals present)
        blk = get_sk_block(l, m, n, pair_hops)
        
        # Block-wise assembly
        idx_i = atom_orbs[i]
        idx_j = atom_orbs[j]
        
        # We need to map the full 9x9 block into the local H slices
        # based on which orbital types are present in atom i and atom j
        orb_vals_i = [ot.value for _, ot in [orb_info[m] for m in idx_i]]
        orb_vals_j = [ot.value for _, ot in [orb_info[n] for n in idx_j]]
        
        sub = blk[np.ix_(orb_vals_i, orb_vals_j)]
        H[np.ix_(idx_i, idx_j)] += sub
        H[np.ix_(idx_j, idx_i)] += sub.T
                
    return H, neighbors, atom_orbs, orb_info

# ─────────────────────────────────────────────────────────────────────────────
# 6. Solver & Analysis
# ─────────────────────────────────────────────────────────────────────────────

def solve_tb(H: np.ndarray, elements: List[str], params: TBParameters, 
             atom_to_orbs: List[List[int]], orb_info: List[Tuple[int, OrbitalType]],
             neighbors: List[Any] = []) -> TBResults:
    M = H.shape[0]; N = len(elements)
    results = TBResults(basis=params.basis, atom_to_orbs=atom_to_orbs)
    results.warnings = list(DB.warnings)
    
    if params.store_hamiltonian: results.hamiltonian = H.copy()
    if params.store_neighbors: results.neighbors = list(neighbors)
    if params.store_orb_info: results.orb_info = list(orb_info)
    
    try: evals, evecs = np.linalg.eigh(H)
    except: 
        results.warnings.append("Diagonalization failed")
        return results
    
    n_e = params.n_electrons
    if n_e < 0: 
        n_e = sum(DB.valence.get(el, fallback_valence(el)) for el in elements)
    n_e = min(n_e, 2*M)
    results.n_electrons = n_e
    
    kT = params.smearing_kT
    if kT < 1.0e-4:
        # T=0 limit
        n_occ = (n_e + 1) // 2
        occ = np.zeros(M)
        occ[:n_occ] = 2.0
        if n_e % 2 == 1: occ[n_occ-1] = 1.0
        fermi = float(evals[n_occ-1]) if n_occ > 0 else 0.0
    else:
        # Robust Bisection Solver for Fermi Level
        mu_lo, mu_hi = evals.min() - 5*kT, evals.max() + 5*kT
        for _ in range(60):
            mu = 0.5 * (mu_lo + mu_hi)
            occ = 2.0 / (np.exp(np.clip((evals - mu)/kT, -50, 50)) + 1.0)
            if occ.sum() > n_e: mu_hi = mu
            else: mu_lo = mu
        fermi = 0.5 * (mu_lo + mu_hi)

    results.eigenvalues = evals.astype(np.float32)
    results.eigenvectors = evecs.astype(np.float32)
    results.occupations = occ.astype(np.float32)
    results.fermi_level = fermi
    
    # Gap analysis
    occ_idx = np.where(occ > 1e-4)[0]
    unocc_idx = np.where(occ < 2.0 - 1e-4)[0]
    results.homo = float(evals[occ_idx[-1]]) if len(occ_idx) > 0 else np.nan
    results.lumo = float(evals[unocc_idx[0]]) if len(unocc_idx) > 0 else np.nan
    results.gap = max(0.0, results.lumo - results.homo) if not np.isnan(results.homo) else 0.0
    results.n_occ_states = len(occ_idx)
    
    # State-Atom Weights (Contribution per atom to each eigenstate)
    state_atom_weights = np.zeros((M, N), dtype=np.float32)
    for i in range(N):
        orbs = atom_to_orbs[i]
        state_atom_weights[:, i] = np.sum(evecs[orbs, :]**2, axis=0)
    results.state_atom_weights = state_atom_weights
    
    # Population & Charge
    dens = (evecs * occ[np.newaxis, :]) @ evecs.T
    results.local_charge = np.array([sum(dens[mu,mu] for mu in atom_to_orbs[i]) for i in range(N)], dtype=np.float32)
    results.approx_charge = np.array([float(DB.valence.get(elements[i], fallback_valence(elements[i]))) - results.local_charge[i] for i in range(N)], dtype=np.float32)
    results.ipr = np.sum(evecs**4, axis=0).astype(np.float32)
    
    # Bond Order
    for (i, j, r, dv) in neighbors:
        if i >= j: continue
        key = (i, j)
        bo = 0.0
        for mu in atom_to_orbs[i]:
            for nu in atom_to_orbs[j]:
                # Weighted overlap population Re(P_munu * np.sign(H_munu))
                bo += 2.0 * dens[mu, nu] * np.sign(H[mu, nu])
        results.bond_order[key] = bo
    
    # DOS / PDOS
    dos_x = np.linspace(evals.min()-2, evals.max()+2, params.dos_n_bins)
    gauss = np.exp(-0.5*((dos_x[:,None]-evals[None,:])/params.dos_sigma)**2)/(params.dos_sigma*np.sqrt(2*np.pi))
    results.dos_y = gauss.sum(axis=1).astype(np.float32)
    
    shells = {"s": [OrbitalType.S], "p": [1,2,3], "d": [4,5,6,7,8]}
    for sp in sorted(set(elements)):
        for sname, stypes in shells.items():
            idx = [mu for mu, (ai, ot) in enumerate(orb_info) if elements[ai] == sp and int(ot) in stypes]
            if not idx: continue
            w = np.sum(evecs[idx, :]**2, axis=0)
            results.pdos_y[f"{sp}_{sname}"] = (gauss * w[None, :]).sum(axis=1).astype(np.float32)
            
    for sname, stypes in shells.items():
        idx = [mu for mu, (ai, ot) in enumerate(orb_info) if int(ot) in stypes]
        if idx:
            w = np.sum(evecs[idx, :]**2, axis=0)
            results.pdos_orb[sname] = (gauss * w[None, :]).sum(axis=1).astype(np.float32)

    # Reference Shifting
    ref = 0.0
    if params.energy_ref == EnergyReference.FERMI: ref = fermi
    elif params.energy_ref == EnergyReference.HOMO: ref = results.homo
    
    results.dos_x = (dos_x - ref).astype(np.float32)
    results.eigenvalues -= ref; results.fermi_level -= ref; results.homo -= ref; results.lumo -= ref
    return results

def run_tight_binding(pos: np.ndarray, elements: List[str], params: TBParameters, lat: Optional[np.ndarray]=None) -> TBResults:
    if len(elements) == 0: return TBResults()
    H, neighbors, atom_to_orbs, orb_info = build_hamiltonian(pos, elements, params, lat)
    return solve_tb(H, elements, params, atom_to_orbs, orb_info, neighbors)
