# ============================
# coordination.py — Coordination Analysis
# ============================
from __future__ import annotations
import numpy as np
from typing import Tuple, List, Optional
from .geometry import compute_angle
from ..constants import ATOMIC_RADII

try:
    from scipy.spatial import cKDTree
    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False

def compute_cn(positions: np.ndarray, lattice: np.ndarray, r_cut: float = 3.0) -> Tuple[float, np.ndarray]:
    """Calculates average Coordination Number (CN) with PBC awareness."""
    if positions.shape[0] < 2: return 0.0, np.array([])
    
    box = np.diag(lattice) if lattice is not None else None
    
    if HAS_SCIPY:
        # Wrap positions to primary cell for tree
        pos_data = positions % box if box is not None else positions
        tree = cKDTree(pos_data, boxsize=box)
        # query_ball_point returns list of neighbors for each point
        indices = tree.query_ball_point(pos_data, r_cut)
        # Subtract 1 because query includes self
        counts = np.array([len(idx_list) - 1 for idx_list in indices], dtype=np.int32)
        return float(np.mean(counts)), counts
    else:
        # Fallback (Slow) with Minimum Image Convention
        counts = []
        for i in range(len(positions)):
            v_all = positions - positions[i]
            if box is not None:
                v_all = v_all - box * np.round(v_all / box)
            d = np.linalg.norm(v_all, axis=1)
            cnt = np.sum((d > 0) & (d < r_cut))
            counts.append(cnt)
        counts = np.array(counts, dtype=np.int32)
        return float(np.mean(counts)), counts

def compute_coordination_sphere(positions: np.ndarray, elements: List[str], lattice: np.ndarray, factor: float = 1.2) -> Tuple[float, np.ndarray, float]:
    """
    Calculates Coordination Number based on overlapping coordination spheres.
    Criterion: d_ij < factor * (r_i + r_j)
    Returns: (avg_cn, counts, avg_bond_len)
    """
    if not elements or len(elements) != len(positions): return 0.0, np.array([]), 0.0
    
    N = len(positions)
    # Get radii, default to 1.5 if unknown
    radii = np.array([ATOMIC_RADII.get(e, 1.5) for e in elements])
    
    # Max possible cutoff to limit search
    max_r = np.max(radii)
    max_cutoff = factor * (max_r + max_r)
    
    counts = np.zeros(N, dtype=np.int32)
    total_bond_len = 0.0
    n_bonds = 0
    
    # Box for PBC
    box = np.diag(lattice) if lattice is not None else None
    
    bonds = []
    if HAS_SCIPY:
        pos_data = positions
        if box is not None:
             pos_data = positions % box
        try: tree = cKDTree(pos_data, boxsize=box)
        except: tree = cKDTree(positions)
        
        try: pairs = tree.query_pairs(max_cutoff)
        except: pairs = []
        
        for i, j in pairs:
            v_orig = positions[j] - positions[i]
            shift = np.zeros(3)
            if box is not None:
                shift = -np.round(v_orig / box)
                v = v_orig + shift * box
            else:
                v = v_orig
            dist = np.linalg.norm(v)
            if dist < factor * (radii[i] + radii[j]):
                counts[i] += 1; counts[j] += 1
                total_bond_len += dist; n_bonds += 1
                bonds.append((i, j, shift[0], shift[1], shift[2]))
    else:
        for i in range(N):
            for j in range(i + 1, N):
                v_orig = positions[j] - positions[i]
                shift = np.zeros(3)
                if box is not None:
                    shift = -np.round(v_orig / box)
                    v = v_orig + shift * box
                else:
                    v = v_orig
                dist = np.linalg.norm(v)
                if dist < factor * (radii[i] + radii[j]):
                    counts[i] += 1; counts[j] += 1
                    total_bond_len += dist; n_bonds += 1
                    bonds.append((i, j, shift[0], shift[1], shift[2]))
                    
    avg_bond_len = total_bond_len / n_bonds if n_bonds > 0 else 0.0
    return float(np.mean(counts)), counts, avg_bond_len, bonds

def compute_bad(positions: np.ndarray, lattice: np.ndarray, r_cut: float = 3.0, bins: int = 180) -> Tuple[np.ndarray, np.ndarray]:
    """Calculates Bond Angle Distribution (BAD) histogram."""
    angles = []
    if HAS_SCIPY:
        tree = cKDTree(positions)
        # Find all pairs within r_cut
        pairs = tree.query_pairs(r_cut)
        
        # Build adjacency list
        adj = {i: [] for i in range(len(positions))}
        for i, j in pairs:
            adj[i].append(j)
            adj[j].append(i)
            
        # Iterate over central atoms (j)
        for j, neighbors in adj.items():
            if len(neighbors) < 2: continue
            # Iterate over unique pairs of neighbors (i, k)
            for i_idx in range(len(neighbors)):
                for k_idx in range(i_idx + 1, len(neighbors)):
                    i, k = neighbors[i_idx], neighbors[k_idx]
                    p1, p2, p3 = positions[i], positions[j], positions[k]
                    angles.append(compute_angle(p1, p2, p3))
    else:
        # Fallback: Brute Force (Slow but functional for small N)
        # Build adjacency list by checking all pairs
        adj = {i: [] for i in range(len(positions))}
        for i in range(len(positions)):
            for j in range(i + 1, len(positions)):
                d = np.linalg.norm(positions[i] - positions[j])
                if d < r_cut:
                    adj[i].append(j)
                    adj[j].append(i)
        
        for j, neighbors in adj.items():
            if len(neighbors) < 2: continue
            for i_idx in range(len(neighbors)):
                for k_idx in range(i_idx + 1, len(neighbors)):
                    i, k = neighbors[i_idx], neighbors[k_idx]
                    p1, p2, p3 = positions[i], positions[j], positions[k]
                    angles.append(compute_angle(p1, p2, p3))
    
    if not angles:
        return np.linspace(0, 180, bins), np.zeros(bins)
        
    hist, edges = np.histogram(angles, bins=bins, range=(0, 180))
    centers = (edges[:-1] + edges[1:]) / 2
    # Normalize PDF
    hist = hist / np.sum(hist) if np.sum(hist) > 0 else hist
    return centers, hist
def compute_layer_cn(positions: np.ndarray, elements: List[str], lattice: np.ndarray, 
                     z_center: float, thickness: float, 
                     method: str = "sphere", **kwargs) -> Tuple[float, np.ndarray]:
    """
    Coordination analysis for a specific Z-layer.
    Filter atoms in layer and compute their coordination.
    """
    z_min = z_center - thickness/2
    z_max = z_center + thickness/2
    mask = (positions[:, 2] >= z_min) & (positions[:, 2] <= z_max)
    
    if not np.any(mask):
        return 0.0, np.array([])
        
    layer_pos = positions[mask]
    layer_elements = [elements[i] for i, m in enumerate(mask) if m]
    
    if method == "sphere":
        avg, counts, _ = compute_coordination_sphere(layer_pos, layer_elements, lattice, **kwargs)
        return avg, counts
    else:
        # Simple r_cut method
        r_cut = kwargs.get("r_cut", 3.0)
        return compute_cn(layer_pos, lattice, r_cut=r_cut)
