import glob, pandas as pd, numpy as np
import matplotlib.pyplot as plt
from ase.io import read
from ase.neighborlist import natural_cutoffs, NeighborList
from matscipy.rings import ring_statistics
import re
from sklearn.linear_model import LinearRegression
import warnings
warnings.filterwarnings('ignore')

# ---------------------------------------------------------------
# BASIC GEOMETRY UTILITIES
# ---------------------------------------------------------------
def angle_between(v1, v2):
    cos_theta = np.dot(v1, v2)
    cos_theta = np.clip(cos_theta, -1.0, 1.0)
    return np.degrees(np.arccos(cos_theta))

def compute_angles(atoms, center_idx, neighbors):
    if len(neighbors) < 2:
        return 0.0, 0.0
    pos = atoms.positions
    center_pos = pos[center_idx]
    angles = []
    for i, j1 in enumerate(neighbors):
        for j2 in neighbors[i+1:]:
            vec1 = pos[j1] - center_pos
            vec2 = pos[j2] - center_pos
            if np.linalg.norm(vec1) > 0 and np.linalg.norm(vec2) > 0:
                angles.append(angle_between(vec1/np.linalg.norm(vec1),
                                            vec2/np.linalg.norm(vec2)))
    return np.mean(angles) if angles else 0.0, np.std(angles) if angles else 0.0


# ---------------------------------------------------------------
# GLOBAL RING CACHE
# ---------------------------------------------------------------
global_rings_cache = {}

def get_global_rings(atoms):
    key = f"{atoms.get_chemical_formula()}_{len(atoms)}"
    if key not in global_rings_cache:
        try:
            all_rings = ring_statistics(atoms, cutoff=1.6, maxlength=8)
            global_rings_cache[key] = {
                'total_rings_3_8': float(np.sum(all_rings[2:8])),
                'rings_4': float(all_rings[4]),
                'rings_5': float(all_rings[5]),
                'rings_6': float(all_rings[6]),
                'rings_7': float(all_rings[7])
            }
            print(f"Cached rings for {key}")
        except:
            global_rings_cache[key] = {
                'total_rings_3_8': 0.0,
                'rings_4': 0.0,
                'rings_5': 0.0,
                'rings_6': 0.0,
                'rings_7': 0.0
            }
    return global_rings_cache[key]


# ---------------------------------------------------------------
# MIC SAFE LOCAL RING FEATURES
# ---------------------------------------------------------------
def matscipy_ring_features(atoms, nearest_c, cutoff=1.6):
    try:
        global_rings = get_global_rings(atoms)

        idxs = np.arange(len(atoms))
        dists = atoms.get_distances(nearest_c, idxs, mic=True)

        mask = (dists < 5.0) & (dists > 1e-3)
        local_atoms = atoms[mask]

        if len(local_atoms) == 0:
            raise RuntimeError("Empty local selection")

        local_rings = ring_statistics(local_atoms, cutoff=cutoff, maxlength=8)

        local_rings_3 = local_rings[3] if len(local_rings) > 3 else 0
        local_rings_4 = local_rings[4] if len(local_rings) > 4 else 0
        local_rings_5 = local_rings[5] if len(local_rings) > 5 else 0
        local_rings_6 = local_rings[6] if len(local_rings) > 6 else 0
        local_rings_7 = local_rings[7] if len(local_rings) > 7 else 0

        local_rings_3_8 = np.sum(local_rings[2:8])
        rings_per_atom = local_rings_3_8 / len(local_atoms)

        ring_counts_3_8 = local_rings[2:8]
        valid_rings = ring_counts_3_8 > 0
        ring_sizes = [size for size, exists in zip(range(3, 9), valid_rings) if exists]
        smallest_ring = min(ring_sizes) if ring_sizes else 0

        return {
            'total_rings_3_8': global_rings['total_rings_3_8'],
            'global_rings_4': global_rings['rings_4'],
            'global_rings_5': global_rings['rings_5'],
            'global_rings_6': global_rings['rings_6'],
            'global_rings_7': global_rings['rings_7'],

            'local_rings_3_8': float(local_rings_3_8),
            'local_rings_3': float(local_rings_3),
            'local_rings_4': float(local_rings_4),
            'local_rings_5': float(local_rings_5),
            'local_rings_6': float(local_rings_6),
            'local_rings_7': float(local_rings_7),
            'local_hexagons': float(local_rings_6),
            'rings_per_atom': float(rings_per_atom),
            'smallest_ring': float(smallest_ring),
        }

    except Exception as e:
        print(f"Ring stats error: {e}")
        zeros = {k: 0.0 for k in [
            'total_rings_3_8','global_rings_4','global_rings_5',
            'global_rings_6','global_rings_7','local_rings_3_8',
            'local_rings_3','local_rings_4','local_rings_5',
            'local_rings_6','local_rings_7','local_hexagons',
            'rings_per_atom','smallest_ring'
        ]}
        return zeros


# ---------------------------------------------------------------
# ADVANCED LOCAL GEOMETRY
# ---------------------------------------------------------------
def advanced_features(atoms, nearest_c, c_neighbors, nl):
    pos = atoms.positions
    center_pos = pos[nearest_c]

    neighbor_dists_raw = np.linalg.norm(pos[c_neighbors] - center_pos, axis=1)
    valid_mask = (neighbor_dists_raw > 0.8) & (neighbor_dists_raw < 2.5)
    neighbor_dists = neighbor_dists_raw[valid_mask]
    valid_neighbors = [c_neighbors[i] for i in np.where(valid_mask)[0]]

    if len(neighbor_dists) < 2:
        neighbor_dists = np.array([1.42])

    curvature = np.std(neighbor_dists) / max(np.mean(neighbor_dists), 1.2)

    planarity = 0.0
    if len(valid_neighbors) >= 3:
        try:
            X = pos[valid_neighbors, :2]
            y = pos[valid_neighbors, 2]
            reg = LinearRegression().fit(X, y)
            planarity = np.mean((y - reg.predict(X))**2)
        except:
            pass

    bond_var = np.var(np.clip(neighbor_dists, 1.0, 2.0))

    angles_local = []
    for i in range(len(valid_neighbors)):
        for j in range(i+1, len(valid_neighbors)):
            vec1 = pos[valid_neighbors[i]] - center_pos
            vec2 = pos[valid_neighbors[j]] - center_pos
            if np.linalg.norm(vec1) > 0 and np.linalg.norm(vec2) > 0:
                angle = angle_between(vec1/np.linalg.norm(vec1),
                                      vec2/np.linalg.norm(vec2))
                angles_local.append(angle)

    q6 = np.abs(np.mean(np.exp(1j * np.radians(angles_local) * 6))) if angles_local else 0.0
    q6 = min(q6, 0.95)

    z_mean = np.mean(pos[:,2])
    surface_cn = sum(1 for nb in valid_neighbors if abs(pos[nb,2] - z_mean) < 1.0)

    neighbor_cn_var = 0.0
    try:
        neighbor_cns = [len(nl.get_neighbors(nb)[0]) for nb in valid_neighbors]
        neighbor_cn_var = np.var(neighbor_cns)
    except:
        pass

    h_idx = next((i for i,s in enumerate(atoms) if s.symbol=='H'), None)
    h_bond_dist = atoms.get_distance(nearest_c, h_idx) if h_idx is not None else 0.0

    return {
        'curvature': float(curvature),
        'planarity': float(planarity),
        'bond_var': float(bond_var),
        'q6_order': float(q6),
        'surface_cn': float(surface_cn),
        'neighbor_cn_var': float(neighbor_cn_var),
        'h_bond_dist': float(h_bond_dist)
    }


# ---------------------------------------------------------------
# HELPERS
# ---------------------------------------------------------------
def extract_id(filename):
    match = re.search(r'POSCAR_(\d+)', filename)
    return int(match.group(1)) if match else None


# ---------------------------------------------------------------
# NEW: MIC MULTISHELL DENSITIES
# ---------------------------------------------------------------
def compute_multishell_densities(atoms, center_idx):
    idxs = np.arange(len(atoms))
    d = atoms.get_distances(center_idx, idxs, mic=True)

    result = {}
    for R in (2.0, 3.0, 5.0):
        mask = (d < R) & (d > 1e-3)
        result[R] = int(np.sum(mask))

    return result


# ---------------------------------------------------------------
# MAIN LOCAL FEATURE BUILDER
# ---------------------------------------------------------------
def local_features(atoms):
    try:
        if len(atoms) < 10:
            return None

        h_indices = [i for i,s in enumerate(atoms) if s.symbol=='H']
        if not h_indices:
            return None
        h_idx = h_indices[0]

        c_indices = [i for i,s in enumerate(atoms) if s.symbol=='C']
        if not c_indices:
            return None

        dists = atoms.get_distances(h_idx, c_indices, mic=True)
        nearest_c = c_indices[np.argmin(dists)]

        cutoffs = natural_cutoffs(atoms)
        nl = NeighborList(cutoffs, self_interaction=False, bothways=True)
        nl.update(atoms)

        neighbors = nl.get_neighbors(nearest_c)[0]
        c_neighbors = [i for i in neighbors if atoms[i].symbol == 'C']

        cn = len(c_neighbors)
        theta_mean, theta_std = compute_angles(atoms, nearest_c, c_neighbors)

        idxs = np.arange(len(atoms))
        d_all = atoms.get_distances(nearest_c, idxs, mic=True)
        density_mask = (d_all < 5.0) & (d_all > 1e-3)
        local_density = np.sum(density_mask)

        # --- NEW MULTISHELL ---
        shell = compute_multishell_densities(atoms, nearest_c)

        height = abs(atoms[nearest_c].position[2] -
                     np.mean(atoms.positions[c_neighbors, 2])) if cn >= 3 else 0.0

        geo_feats = advanced_features(atoms, nearest_c, c_neighbors, nl)
        ring_feats = matscipy_ring_features(atoms, nearest_c)

        return {
            'cn': float(cn),
            'theta_mean': theta_mean,
            'theta_std': theta_std,
            'local_density': float(local_density),
            'density_2A': float(shell[2.0]),
            'density_3A': float(shell[3.0]),
            'density_5A': float(shell[5.0]),
            'height': height,
            **geo_feats,
            **ring_feats
        }

    except Exception as e:
        print(f"Feature error: {e}")
        return None


# ---------------------------------------------------------------
# EXECUTION
# ---------------------------------------------------------------
print("=== EXTENDED RING + GEOMETRY FEATURES ===")
features = []
poscar_files = sorted(glob.glob('POSCAR_*'))

for poscar in poscar_files:
    print(f"Processing {poscar}...", end=" ")
    try:
        atoms = read(poscar)
        feats = local_features(atoms)
        if feats:
            feats['filename'] = poscar
            feats['id'] = extract_id(poscar)
            features.append(feats)
            print("✓")
        else:
            print("✗")
    except Exception as e:
        print(f"✗ {e}")

df_final = pd.DataFrame(features)


# ---------------------------------------------------------------
# ENERGY MERGE
# ---------------------------------------------------------------
energy_data = []
try:
    with open('energy_pairs.txt', 'r') as f:
        for line in f:
            if line.startswith('#'):
                continue
            parts = line.split()
            if len(parts) >= 5:
                energy_data.append({'id': int(parts[0]), 'ΔG(eV)': float(parts[4])})
except:
    pass

if energy_data:
    df_final = df_final.merge(pd.DataFrame(energy_data), on='id', how='left')

df_final.to_csv('HER_sites_full.csv', index=False)

if 'ΔG(eV)' in df_final.columns:
    df_clean = df_final.dropna(subset=['ΔG(eV)'])
    df_clean.to_csv('HER_sites_cleaned.csv', index=False)
    print(f"✓ Full: {len(df_final)} → Cleaned: {len(df_clean)}")

else:
    print(f"Structural only: {len(df_final)} sites")


print("\nReady.")