import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Tuple, Optional, Any
from .tight_binding import OrbitalType, DB, TBResults, TBParameters

@dataclass
class GridSpec:
    origin: np.ndarray        # (3,)
    shape:  Tuple[int, int, int]
    spacing: float
    
    @property
    def rmin(self): return self.origin
    @property
    def rmax(self): return self.origin + np.array(self.shape) * self.spacing

@dataclass
class ScalarField3D:
    grid: GridSpec
    values: np.ndarray        # shape (nx, ny, nz)
    name: str = ""
    is_signed: bool = False

@dataclass
class IsosurfaceMesh:
    vertices: np.ndarray      # (V, 3)
    faces: np.ndarray         # (F, 3)
    normals: Optional[np.ndarray] = None
    iso_value: float = 0.0
    name: str = ""

# ─────────────────────────────────────────────────────────────────────────────
# 1. Real-Space Basis Evaluator
# ─────────────────────────────────────────────────────────────────────────────

def eval_basis_function(orb: OrbitalType, dr: np.ndarray, alpha: float) -> float:
    """
    Evaluates a Normalized Gaussian-type orbital at relative vector dr.
    """
    r2 = np.sum(dr**2)
    x, y, z = dr
    
    # Normalization N: integral[ |N * Poly * exp(-alpha*r2)|^2 ] = 1
    # For S: N = (2*alpha/pi)^(3/4)
    n_s = (2.0 * alpha / np.pi)**0.75
    g = n_s * np.exp(-alpha * r2)
    
    if orb == OrbitalType.S:
        return g
    
    # For P: N = (2*alpha/pi)^(3/4) * sqrt(4*alpha)
    n_p = n_s * np.sqrt(4.0 * alpha)
    if orb == OrbitalType.PX: return n_p * x * g
    elif orb == OrbitalType.PY: return n_p * y * g
    elif orb == OrbitalType.PZ: return n_p * z * g
    
    # For D: N = (2*alpha/pi)^(3/4) * 4*alpha / sqrt(3) (simplified)
    n_d = n_s * (4.0 * alpha) / np.sqrt(3.0) 
    if orb == OrbitalType.DXY: return n_d * x * y * g
    elif orb == OrbitalType.DYZ: return n_d * y * z * g
    elif orb == OrbitalType.DZX: return n_d * z * x * g
    elif orb == OrbitalType.DX2Y2: return 0.5 * n_d * (x*x - y*y) * g
    elif orb == OrbitalType.DZ2:   return (1.0/np.sqrt(12.0)) * n_d * (2.0*z*z - x*x - y*y) * g
    
    return 0.0

# ─────────────────────────────────────────────────────────────────────────────
# 2. Grid & Field Builders
# ─────────────────────────────────────────────────────────────────────────────

def make_grid_from_atoms(pos: np.ndarray, params: TBParameters) -> GridSpec:
    pad = params.grid_padding
    spacing = params.grid_spacing
    
    rmin = np.min(pos, axis=0) - pad
    rmax = np.max(pos, axis=0) + pad
    
    shape = tuple(np.ceil((rmax - rmin) / spacing).astype(int) + 1)
    return GridSpec(origin=rmin, shape=shape, spacing=spacing)

def build_state_field(pos: np.ndarray, elements: List[str], orb_info: List[Tuple[int, OrbitalType]], 
                      evecs: np.ndarray, state_idx: int, grid: GridSpec) -> ScalarField3D:
    """Computes signed orbital amplitude psi_n(r) on the grid."""
    nx, ny, nz = grid.shape
    field_vals = np.zeros(grid.shape, dtype=np.float32)
    
    # 1. Prepare grid coordinates
    x = grid.origin[0] + np.arange(nx) * grid.spacing
    y = grid.origin[1] + np.arange(ny) * grid.spacing
    z = grid.origin[2] + np.arange(nz) * grid.spacing
    
    # 2. Reconstruct state (ψ_n = Σ c_μn χ_μ)
    for mu, (ai, ot) in enumerate(orb_info):
        coeff = evecs[mu, state_idx]
        if abs(coeff) < 1e-4: continue
        
        atom_pos = pos[ai]
        alpha = DB.get_alpha(elements[ai], ot)
        
        # Orbital Cutoff: skip points where Gaussian < 1e-4
        # exp(-alpha * r2) > 1e-4 => alpha * r2 < ln(1e4) => r < sqrt(9.2 / alpha)
        r_cut = np.sqrt(9.2 / alpha)
        
        for ix in range(nx):
            dx = x[ix] - atom_pos[0]
            if abs(dx) > r_cut: continue
            for iy in range(ny):
                dy = y[iy] - atom_pos[1]
                if abs(dy) > r_cut: continue
                
                # Vectorized Z
                dz = z - atom_pos[2]
                mask = np.abs(dz) <= r_cut
                if not np.any(mask): continue
                
                # dr is (N_mask, 3) implicitly
                for iz in np.where(mask)[0]:
                    dr = np.array([dx, dy, dz[iz]])
                    field_vals[ix, iy, iz] += coeff * eval_basis_function(ot, dr, alpha)
    
    return ScalarField3D(grid, field_vals, name=f"State {state_idx}", is_signed=True)

def build_total_density_field(pos: np.ndarray, elements: List[str], orb_info: List[Tuple[int, OrbitalType]],
                             evecs: np.ndarray, occupations: np.ndarray, grid: GridSpec) -> ScalarField3D:
    """Computes total occupied electron density rho(r)."""
    density = np.zeros(grid.shape, dtype=np.float32)
    for n in range(len(occupations)):
        occ = occupations[n]
        if occ < 1e-4: continue
        state_field = build_state_field(pos, elements, orb_info, evecs, n, grid)
        density += occ * (state_field.values**2)
    return ScalarField3D(grid, density, name="Total Density", is_signed=False)

def build_window_ldos(pos: np.ndarray, elements: List[str], orb_info: List[Tuple[int, OrbitalType]],
                      evals: np.ndarray, evecs: np.ndarray, e_min: float, e_max: float, grid: GridSpec) -> ScalarField3D:
    """Computes local density of states (LDOS) in an energy window [e_min, e_max]."""
    ldos = np.zeros(grid.shape, dtype=np.float32)
    for n, e in enumerate(evals):
        if e_min <= e <= e_max:
            # For LDOS we treat each state within window with weight 1.0 (delta-like)
            # or we could use Gaussian smearing here if needed.
            state_field = build_state_field(pos, elements, orb_info, evecs, n, grid)
            ldos += (state_field.values**2)
    return ScalarField3D(grid, ldos, name=f"LDOS [{e_min:.2f}, {e_max:.2f}]", is_signed=False)

# ─────────────────────────────────────────────────────────────────────────────
# 3. Surface Extraction
# ─────────────────────────────────────────────────────────────────────────────

def extract_isosurface(field: ScalarField3D, iso_val: float) -> Optional[IsosurfaceMesh]:
    try:
        from skimage import measure
    except ImportError:
        return None
        
    vals = field.values
    v_min, v_max = vals.min(), vals.max()
    
    # Range check
    if iso_val < v_min or iso_val > v_max:
        return None
        
    # Marching Cubes
    try:
        # spacing=(dx, dy, dz)
        verts, faces, normals, _ = measure.marching_cubes(
            vals, level=iso_val, spacing=(field.grid.spacing,)*3
        )
        # Shift to origin
        verts += field.grid.origin
        return IsosurfaceMesh(vertices=verts.astype(np.float32), 
                              faces=faces.astype(np.int32), 
                              normals=normals.astype(np.float32),
                              iso_value=iso_val,
                              name=field.name)
    except Exception as e:
        print(f"Marching Cubes Error: {e}")
        return None
