# ============================
# hbonds.py — Hydrogen Bond Analysis
# ============================
from __future__ import annotations
import numpy as np
from typing import Tuple, List, Optional

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

def compute_hbonds(positions: np.ndarray, elements: List[str], lattice: np.ndarray, 
                   d_ha_cut: float = 2.5, angle_cut: float = 120.0) -> Tuple[np.ndarray, np.ndarray]:
    """
    Detects Hydrogen Bonds and returns (distances, pairs).
    Pairs are (H_index, A_index).
    Criteria: D, A in {N, O, F}; H in {H}; dist(D,H)<1.3; dist(H,A)<d_ha_cut; angle(D-H...A)>angle_cut.
    """
    if not elements or len(elements) != len(positions): return np.array([]), np.array([])
    
    # Indices
    donors_acceptors = [i for i, e in enumerate(elements) if e in ('N', 'O', 'F')]
    hydrogens = [i for i, e in enumerate(elements) if e == 'H']
    
    if not donors_acceptors or not hydrogens: return np.array([]), np.array([])
    
    hb_dists = []
    hb_pairs = []
    
    # Box size for PBC (assume orthorhombic for simplicity)
    box = np.diag(lattice) if lattice is not None else None
    
    # Helper for MIC vector
    def get_vec_mic(v, box):
        if box is None: return v
        return v - np.round(v / box) * box

    # Using cKDTree for speed if available
    if HAS_SCIPY:
        # Prepare data for Tree
        pos_data = positions[donors_acceptors]
        if box is not None:
            pos_data = pos_data % box
            
        try:
            tree_da = cKDTree(pos_data, boxsize=box)
        except TypeError:
            # Fallback for older scipy
            tree_da = cKDTree(positions[donors_acceptors])
            pass
        
        for h_idx in hydrogens:
            h_pos = positions[h_idx]
            h_pos_query = h_pos % box if box is not None else h_pos
            
            # 1. Find covalently bonded Donor (D)
            try:
                d_dists, d_indices = tree_da.query(h_pos_query, k=1, distance_upper_bound=1.3)
            except TypeError:
                 d_dists, d_indices = tree_da.query(h_pos, k=1, distance_upper_bound=1.3)

            if d_dists == float('inf'): continue
            
            d_real_idx = donors_acceptors[d_indices]
            d_pos = positions[d_real_idx]
            
            # 2. Find Acceptors (A) within d_ha_cut
            try:
                a_indices_local = tree_da.query_ball_point(h_pos_query, d_ha_cut)
            except TypeError:
                a_indices_local = tree_da.query_ball_point(h_pos, d_ha_cut)
            
            for a_idx_local in a_indices_local:
                a_real_idx = donors_acceptors[a_idx_local]
                if a_real_idx == d_real_idx: continue # D and A cannot be same atom
                
                a_pos = positions[a_real_idx]
                
                # Check Angle D-H...A using MIC vectors
                v_hd = get_vec_mic(d_pos - h_pos, box)
                v_ha = get_vec_mic(a_pos - h_pos, box)
                
                l_hd = np.linalg.norm(v_hd)
                l_ha = np.linalg.norm(v_ha)
                
                if l_hd < 1e-3 or l_ha < 1e-3: continue
                
                dot = np.dot(v_hd, v_ha)
                cos_theta = dot / (l_hd * l_ha)
                
                angle_rad = np.arccos(np.clip(cos_theta, -1.0, 1.0))
                angle_deg = np.degrees(angle_rad)
                
                if angle_deg > angle_cut:
                    hb_dists.append(l_ha)
                    hb_pairs.append((h_idx, a_real_idx))
                    
    else:
        # Fallback (slow)
        for h_idx in hydrogens:
            h_pos = positions[h_idx]
            # Find D
            best_d = -1
            best_d_dist = 1.3
            for d_idx in donors_acceptors:
                v = get_vec_mic(positions[d_idx] - h_pos, box)
                dist = np.linalg.norm(v)
                if dist < best_d_dist:
                    best_d_dist = dist
                    best_d = d_idx
            
            if best_d == -1: continue
            d_pos = positions[best_d]
            
            # Find A
            for a_idx in donors_acceptors:
                if a_idx == best_d: continue
                a_pos = positions[a_idx]
                v_ha = get_vec_mic(a_pos - h_pos, box)
                dist_ha = np.linalg.norm(v_ha)
                
                if dist_ha < d_ha_cut:
                    # Angle
                    v_hd = get_vec_mic(d_pos - h_pos, box)
                    
                    l_hd = np.linalg.norm(v_hd)
                    l_ha = np.linalg.norm(v_ha)
                    
                    angle = np.degrees(np.arccos(np.clip(np.dot(v_hd, v_ha) / (l_hd * l_ha), -1, 1)))
                    if angle > angle_cut:
                        hb_dists.append(dist_ha)
                        hb_pairs.append((h_idx, a_idx))

    return np.array(hb_dists), np.array(hb_pairs)
