# ============================
# rdf.py — Advanced Radial Distribution Function
# ============================
from __future__ import annotations
import numpy as np
from typing import Tuple, Optional, List, Union

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

def compute_rdf(pos_ref: np.ndarray, 
                pos_target: np.ndarray, 
                lattice: Optional[np.ndarray] = None, 
                r_max: float = 10.0, 
                bins: int = 200, 
                mode: str = "bulk") -> Tuple[np.ndarray, np.ndarray]:
    """
    Robust RDF calculation with PBC (Periodic Boundary Conditions) awareness.
    """
    if pos_ref.size == 0 or pos_target.size == 0:
        return np.linspace(0, r_max, bins), np.zeros(bins)

    if not HAS_SCIPY:
        return _compute_rdf_simple(pos_ref, pos_target, r_max, bins, mode)

    # 1. Setup PBC Box
    box = None
    if lattice is not None and np.linalg.norm(lattice) > 1e-3:
        box = np.abs(np.diag(lattice))
        if np.any(box < 1e-3): box = None

    # Handle self-interaction mask
    is_self = (pos_ref is pos_target) or (pos_ref.shape == pos_target.shape and np.allclose(pos_ref, pos_target))
    
    # Wrap to primary cell for KDTree efficiency
    pos_target_wrapped = pos_target % box if box is not None else pos_target
    pos_ref_wrapped = pos_ref % box if box is not None else pos_ref
    
    tree_target = cKDTree(pos_target_wrapped, boxsize=box)
    
    # Query in batches or one-shot for speed
    indices = tree_target.query_ball_point(pos_ref_wrapped, r_max)
    
    all_dists = []
    for i, idx_list in enumerate(indices):
        if not idx_list: continue
        
        # Filter self
        if is_self:
            idx_list = [idx for idx in idx_list if idx != i]
        
        if not idx_list: continue
        
        # Calculate periodic distances (Minimum Image Convention)
        target_coords = pos_target_wrapped[idx_list]
        diffs = target_coords - pos_ref_wrapped[i]
        
        if box is not None:
            diffs = diffs - box * np.round(diffs / box)
            
        d = np.linalg.norm(diffs, axis=1)
        # Final filter for numerics
        d = d[d > 1e-6]
        all_dists.extend(d)
        
    distances = np.array(all_dists)
    if distances.size == 0:
        return np.linspace(0, r_max, bins), np.zeros(bins)

    # 2. Histogramging
    hist, edges = np.histogram(distances, bins=bins, range=(0, r_max))
    radii = (edges[:-1] + edges[1:]) / 2
    dr = edges[1] - edges[0]
    
    # 3. Normalization (Density Aware)
    n_ref = len(pos_ref)
    n_target = len(pos_target)
    
    if mode == "2d":
        # Area = |a x b|
        area = np.linalg.norm(np.cross(lattice[0], lattice[1])) if lattice is not None else 1.0
        rho = n_target / area
        shell_vol = 2.0 * np.pi * radii * dr
    elif mode == "cluster":
        # No density normalization for finite clusters
        g_r = hist / (n_ref + 1e-12)
        return radii, g_r
    else: # bulk 3D
        vol = np.abs(np.linalg.det(lattice)) if lattice is not None else 1.0
        rho = n_target / vol
        shell_vol = 4.0 * np.pi * (radii**2) * dr
        
    g_r = hist / (shell_vol * rho * n_ref + 1e-15)
    return radii, np.nan_to_num(g_r)

def compute_z_profile(positions: np.ndarray, bins: int = 100, z_range: Optional[Tuple[float, float]] = None) -> Tuple[np.ndarray, np.ndarray]:
    """Calculates density profile along Z axis."""
    z_coords = positions[:, 2]
    if z_range is None:
        z_range = (z_coords.min() - 0.5, z_coords.max() + 0.5)
    
    hist, edges = np.histogram(z_coords, bins=bins, range=z_range)
    centers = (edges[:-1] + edges[1:]) / 2
    # Normalize density (count / dz)
    dz = edges[1] - edges[0]
    return centers, hist / (dz + 1e-12)

def compute_layer_rdf(positions: np.ndarray, elements: List[str], lattice: np.ndarray, 
                      z_center: float, thickness: float, r_max: float = 10.0) -> Tuple[np.ndarray, np.ndarray]:
    """
    Calculates RDF for a specific layer.
    Filters atoms within z_center +/- thickness/2.
    Useful for interface and layered material analysis.
    """
    z_min = z_center - thickness/2
    z_max = z_center + thickness/2
    
    mask = (positions[:, 2] >= z_min) & (positions[:, 2] <= z_max)
    layer_pos = positions[mask]
    
    if layer_pos.size == 0:
        return np.linspace(0, r_max, 50), np.zeros(50)
        
    # We use 2D mode for layer-resolved RDF
    return compute_rdf(layer_pos, layer_pos, lattice, r_max=r_max, mode="2d")

def _compute_rdf_simple(pos_ref: np.ndarray, pos_target: np.ndarray, r_max: float, bins: int, mode: str):
    """Fallback O(N*M) calculation when Scipy is missing."""
    all_dists = []
    for p in pos_ref:
        d = np.linalg.norm(pos_target - p, axis=1)
        # Filter self
        d = d[d > 1e-3]
        all_dists.extend(d[d <= r_max])
    
    distances = np.array(all_dists)
    hist, edges = np.histogram(distances, bins=bins, range=(0, r_max))
    radii = (edges[:-1] + edges[1:]) / 2
    dr = edges[1] - edges[0]
    
    # Approximate normalization (Uniform)
    normalization = 1.0
    if mode == "2d": normalization = 2 * np.pi * radii * dr
    else: normalization = 4 * np.pi * (radii**2) * dr
    
    return radii, hist / (normalization + 1e-12)
