build_graph.py

"""
build_graph.py - Convert a single CIF file to a PoreGCN graph dict.

Single public entry point:
    graph = cif_to_graph(cif_path, dataset='hmof_gas')

The returned dict has the following keys consumed by xai_engine.py:
    atom_fea         np.ndarray  [N_atoms, 120]
    bond_types       np.ndarray  [N_atoms, M_neighbors]
    nbr_fea_idx      np.ndarray  [N_atoms, M_neighbors]
    crystal_atom_idx list[torch.LongTensor]  (length 1, single-graph)
    pore_fea         np.ndarray  [N_pores, 8]
    atom_pore_edges  np.ndarray  [2, N_ap_edges]
    crystal_pore_idx list[torch.LongTensor]  (length 1)
    n_atoms          int
    n_pores          int
    structure        pymatgen.core.Structure
    pore_positions   np.ndarray  [N_pores, 3]  (Cartesian coords, Angstroms)
    pore_radii       np.ndarray  [N_pores]     (Angstroms)
    voronoi_used     bool        True if Zeo++ ran successfully

Zeo++ fallback: if ZEOPP_BINARY does not exist or the subprocess raises any
exception, a warning is logged and the graph is built in atom-only mode
(n_pores=0, pore_fea is empty). The PoreGCN model handles this gracefully
because AtomPoreConvLayer short-circuits when atom_pore_edges.numel() == 0.

Frontend contract: app.py calls
    from build_graph import cif_to_graph
"""

from __future__ import annotations

import json
import logging
import os
import subprocess
import tempfile
import warnings
from typing import Dict, List, Optional, Tuple

import numpy as np
import torch

try:
    from pymatgen.core import Structure
    from pymatgen.io.cif import CifParser
except ImportError:
    raise ImportError('pymatgen is required: pip install pymatgen')

from config import (
    ATOM_INIT_PATH,
    ATOM_PORE_CUTOFF,
    GRAPH_RADIUS,
    MAX_NUM_NEIGHBORS,
    MIN_PORE_RADIUS,
    ZEOPP_BINARY,
)

logger = logging.getLogger(__name__)

# =============================================================================
# Edge type indices (must match PoreGCN_unified/config.py EDGE_TYPES)
# =============================================================================

EDGE_NONE = 0
EDGE_SINGLE = 1
EDGE_DOUBLE = 2
EDGE_TRIPLE = 3
EDGE_ATOM_PORE = 4

BOND_THRESHOLDS = {
    'triple_max': 1.25,
    'double_max': 1.45,
    'single_max': 2.5,
}


def _classify_bond(distance: float) -> int:
    if distance < BOND_THRESHOLDS['triple_max']:
        return EDGE_TRIPLE
    if distance < BOND_THRESHOLDS['double_max']:
        return EDGE_DOUBLE
    if distance < BOND_THRESHOLDS['single_max']:
        return EDGE_SINGLE
    return EDGE_NONE


# =============================================================================
# Atom feature generation
# Mirrors build_data.py: 92-dim one-hot from atom_init.json + 28 MOF features
# =============================================================================

_ATOM_INIT: Optional[Dict[int, List[float]]] = None


def _load_atom_init() -> Dict[int, List[float]]:
    """Load atom_init.json once and cache it."""
    global _ATOM_INIT
    if _ATOM_INIT is not None:
        return _ATOM_INIT

    if not os.path.exists(ATOM_INIT_PATH):
        raise FileNotFoundError(
            f'atom_init.json not found at {ATOM_INIT_PATH}. '
            'Copy PoreGCN_unified/data/atom_init.json to hf_space/atom_init.json.'
        )

    with open(ATOM_INIT_PATH) as f:
        raw = json.load(f)

    # Keys are string atomic numbers; convert to int
    _ATOM_INIT = {int(k): v for k, v in raw.items()}
    return _ATOM_INIT


def _get_mof_atom_features(site, structure) -> np.ndarray:
    """
    28-dim MOF-specific real-valued features for a single atom site.
    Mirrors build_data.py:get_mof_atom_features.
    """
    features = []

    # Coordination number (simplified, r=3.0 A)
    try:
        neighbors = structure.get_neighbors(site, r=3.0)
        coord_num = len(neighbors)
    except Exception:
        coord_num = 0
    features.append(min(coord_num / 12.0, 1.0))

    # Is metal
    features.append(1.0 if site.specie.is_metal else 0.0)

    # Electronegativity
    try:
        en = site.specie.X / 4.0
    except Exception:
        en = 0.5
    features.append(en)

    # Atomic radius
    try:
        r = site.specie.atomic_radius
        features.append((r.real if hasattr(r, 'real') else float(r)) / 3.0)
    except Exception:
        features.append(1.0 / 3.0)

    # Pad to 28 features
    while len(features) < 28:
        features.append(0.0)

    return np.array(features[:28], dtype=np.float32)


def _build_atom_features(structure: Structure) -> np.ndarray:
    """Build [N_atoms, 120] atom feature matrix (92 one-hot + 28 MOF features)."""
    atom_init = _load_atom_init()
    rows = []
    for site in structure:
        atomic_num = site.specie.Z
        base_fea = atom_init.get(atomic_num, atom_init.get(1, [0.0] * 92))
        mof_fea = _get_mof_atom_features(site, structure)
        combined = np.concatenate([np.array(base_fea, dtype=np.float32), mof_fea])
        rows.append(combined)
    return np.array(rows, dtype=np.float32)


# =============================================================================
# Neighbor list and bond types
# =============================================================================

def _build_neighbor_list(
    structure: Structure,
    radius: float = GRAPH_RADIUS,
    max_neighbors: int = MAX_NUM_NEIGHBORS,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Returns:
        nbr_fea_idx  [N_atoms, max_neighbors]
        bond_types   [N_atoms, max_neighbors]
    """
    all_nbrs = structure.get_all_neighbors(radius, include_index=True)
    n_atoms = len(structure)
    nbr_fea_idx = np.zeros((n_atoms, max_neighbors), dtype=np.int64)
    bond_types = np.zeros((n_atoms, max_neighbors), dtype=np.int64)

    for i, nbrs in enumerate(all_nbrs):
        if not nbrs:
            # Isolated atom: self-loop padding
            nbr_fea_idx[i] = i
            bond_types[i] = EDGE_NONE
        else:
            nbrs_sorted = sorted(nbrs, key=lambda x: x[1])[:max_neighbors]
            for k, nbr in enumerate(nbrs_sorted):
                nbr_fea_idx[i, k] = nbr[2]
                bond_types[i, k] = _classify_bond(nbr[1])
            # Pad remaining slots with self-loops (NONE type)
            for k in range(len(nbrs_sorted), max_neighbors):
                nbr_fea_idx[i, k] = i
                bond_types[i, k] = EDGE_NONE

    return nbr_fea_idx, bond_types


# =============================================================================
# Voronoi / Zeo++ pore node extraction
# =============================================================================

def _run_zeopp(cif_path: str, tmp_dir: str) -> Optional[str]:
    """
    Run Zeo++ to generate a .nt2 Voronoi network file.

    Returns the path to the .nt2 file on success, or None on any failure.
    Falls back gracefully so the graph can still be built without pore nodes.
    """
    if not os.path.exists(ZEOPP_BINARY):
        logger.warning(
            'Zeo++ binary not found at %s. Building graph in atom-only mode.',
            ZEOPP_BINARY,
        )
        return None

    import shutil as _shutil
    basename = os.path.splitext(os.path.basename(cif_path))[0]

    # Zeo++ writes its output next to the input CIF (e.g. <basename>.nt2).
    # Copy the CIF into OUR tmp_dir first so the output location is predictable
    # regardless of where Gradio dropped the upload.
    local_cif = os.path.join(tmp_dir, f'{basename}.cif')
    try:
        _shutil.copy2(cif_path, local_cif)
    except Exception as exc:
        logger.warning('Could not copy CIF into tmp_dir: %s. Atom-only mode.', exc)
        return None

    expected_paths = [
        os.path.join(tmp_dir, f'{basename}.nt2'),
        os.path.join(tmp_dir, f'{basename}_vornet.txt'),
        os.path.join(os.path.dirname(cif_path), f'{basename}.nt2'),
        os.path.join(os.path.dirname(cif_path), f'{basename}_vornet.txt'),
    ]

    # Try -nt2 first (Zeo++ 0.3), then -vornet (Zeo++ 0.4+).
    # First successful invocation wins.
    invocations = [
        [str(ZEOPP_BINARY), '-ha', '-nt2', local_cif],
        [str(ZEOPP_BINARY), '-ha', '-vornet',
         os.path.join(tmp_dir, f'{basename}_vornet.txt'), local_cif],
    ]

    for cmd in invocations:
        try:
            subprocess.run(cmd, cwd=tmp_dir, check=True, capture_output=True, timeout=60)
        except subprocess.TimeoutExpired:
            logger.warning('Zeo++ timed out for %s.', cmd[2])
            continue
        except subprocess.CalledProcessError as exc:
            stderr = (exc.stderr or b'').decode('utf-8', errors='replace')[:200]
            logger.info('Zeo++ %s not supported (%s); trying next invocation.', cmd[2], stderr.strip())
            continue
        except Exception as exc:
            logger.warning('Zeo++ unexpected error: %s', exc)
            continue
        # Check all known output locations
        for cand in expected_paths:
            if os.path.exists(cand) and os.path.getsize(cand) > 0:
                return cand

    logger.warning('Zeo++ ran but no Voronoi output file produced. Atom-only mode.')
    return None


def _parse_nt2(nt2_path: str, min_radius: float = MIN_PORE_RADIUS) -> List[Dict]:
    """
    Parse a Zeo++ Voronoi output file. Supports two formats:

    1. -vornet format (default in PoreGCN_unified):
         lines beginning with 'v' (vertices) or 'e' (edges)
         vertex line: v <id> <x> <y> <z> <radius> [accessible_flag]

    2. -nt2 format (legacy):
         "Vertex table:" header followed by lines starting with vertex_id
         vertex line: <id> <x> <y> <z> <radius> [connected_ids...]

    Only vertices with radius >= min_radius are retained.
    Returns list of dicts with keys: x, y, z, radius, accessible.
    """
    vertices: List[Dict] = []

    try:
        with open(nt2_path) as f:
            content = f.read()
    except OSError:
        return vertices

    is_nt2_format = 'Vertex table:' in content

    if is_nt2_format:
        in_vertex_section = False
        for line in content.splitlines():
            stripped = line.strip()
            if stripped == 'Vertex table:':
                in_vertex_section = True
                continue
            if not in_vertex_section or not stripped:
                continue
            parts = stripped.split()
            if len(parts) < 5:
                continue
            try:
                _ = int(parts[0])
                x, y, z = float(parts[1]), float(parts[2]), float(parts[3])
                radius = float(parts[4])
            except (ValueError, IndexError):
                continue
            if radius >= min_radius:
                vertices.append({'x': x, 'y': y, 'z': z, 'radius': radius, 'accessible': 1})
    else:
        # -vornet format: 'v <id> <x> <y> <z> <radius> [accessible]'
        for line in content.splitlines():
            stripped = line.strip()
            if not stripped or stripped.startswith('#'):
                continue
            parts = stripped.split()
            # Vertex lines start with 'v' OR with a numeric id
            if parts[0] == 'v' and len(parts) >= 6:
                try:
                    x, y, z = float(parts[2]), float(parts[3]), float(parts[4])
                    radius = float(parts[5])
                    accessible = int(parts[6]) if len(parts) >= 7 else 1
                except (ValueError, IndexError):
                    continue
                if radius >= min_radius:
                    vertices.append({'x': x, 'y': y, 'z': z,
                                     'radius': radius, 'accessible': accessible})
            elif parts[0].lstrip('-').isdigit() and len(parts) >= 5:
                # Plain vertex line without 'v' prefix (alternative format)
                try:
                    x, y, z = float(parts[1]), float(parts[2]), float(parts[3])
                    radius = float(parts[4])
                except (ValueError, IndexError):
                    continue
                if radius >= min_radius:
                    vertices.append({'x': x, 'y': y, 'z': z,
                                     'radius': radius, 'accessible': 1})

    return vertices


def _create_pore_features(
    vertices: List[Dict],
    structure: Structure,
    atom_coords: np.ndarray,
    cutoff: float = ATOM_PORE_CUTOFF,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Build pore feature matrix and atom-pore edges.

    Pore feature vector (8-dim, mirrors PORE_NODE_FEATURES in config.py):
        [0:3] fractional x, y, z (wrapped to [0,1])
        [3]   radius / 10.0
        [4]   accessible flag
        [5]   n_atom_neighbors / 10.0 (normalized)
        [6]   max_neighbor_dist / 10.0
        [7]   min_neighbor_dist / 10.0

    Returns:
        pore_fea      [N_pores, 8]
        pore_positions [N_pores, 3] Cartesian (Angstroms)
        pore_radii    [N_pores]
    """
    if not vertices:
        empty_fea = np.zeros((0, 8), dtype=np.float32)
        empty_pos = np.zeros((0, 3), dtype=np.float32)
        empty_rad = np.zeros(0, dtype=np.float32)
        return empty_fea, empty_pos, empty_rad

    lattice = structure.lattice
    pore_features: List[List[float]] = []
    pore_positions = []
    pore_radii = []

    for v in vertices:
        cart = np.array([v['x'], v['y'], v['z']])
        frac = lattice.get_fractional_coords(cart) % 1.0

        # Compute neighbor statistics
        dists = np.linalg.norm(atom_coords - cart, axis=1)
        nearby_mask = dists < cutoff
        nearby_dists = dists[nearby_mask]
        n_nbr = len(nearby_dists)

        fea = [
            float(frac[0]), float(frac[1]), float(frac[2]),
            v['radius'] / 10.0,
            float(v.get('accessible', 1)),
            n_nbr / 10.0,
            float(nearby_dists.max() / 10.0) if n_nbr > 0 else 0.0,
            float(nearby_dists.min() / 10.0) if n_nbr > 0 else 0.0,
        ]
        pore_features.append(fea)
        pore_positions.append(cart)
        pore_radii.append(v['radius'])

    return (
        np.array(pore_features, dtype=np.float32),
        np.array(pore_positions, dtype=np.float32),
        np.array(pore_radii, dtype=np.float32),
    )


def _create_atom_pore_edges(
    atom_coords: np.ndarray,
    pore_positions: np.ndarray,
    cutoff: float = ATOM_PORE_CUTOFF,
) -> np.ndarray:
    """
    Build atom-pore edge index [2, N_edges].
    Row 0 = atom indices, row 1 = pore indices.
    """
    if len(pore_positions) == 0 or len(atom_coords) == 0:
        return np.zeros((2, 0), dtype=np.int64)

    edges = []
    for pore_idx, pore_pos in enumerate(pore_positions):
        dists = np.linalg.norm(atom_coords - pore_pos, axis=1)
        nearby_atoms = np.where(dists < cutoff)[0]
        for atom_idx in nearby_atoms:
            edges.append([atom_idx, pore_idx])

    if not edges:
        return np.zeros((2, 0), dtype=np.int64)

    return np.array(edges, dtype=np.int64).T


# =============================================================================
# CIF loader with fallbacks
# =============================================================================

def _load_structure(cif_path: str) -> Structure:
    """Load a CIF file with progressive fallbacks for non-standard formats.

    Handles CIFs with duplicate atom labels (P1 simulation outputs), partial
    occupancy > 1.0, powder-diffraction reflection blocks, and non-integer
    formulas by trying three strategies in order.
    """
    # Attempt 1: standard pymatgen path
    try:
        return Structure.from_file(cif_path)
    except Exception:
        pass

    # Attempt 2: pymatgen CifParser with relaxed occupancy tolerance and no
    # primitive reduction (avoids failures on partial-occupancy disorder sites)
    try:
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            parser = CifParser(cif_path, occupancy_tolerance=2.0)
            structs = parser.get_structures(primitive=False)
        if structs:
            return structs[0]
    except Exception:
        pass

    # Attempt 3: ASE CIF reader → pymatgen Structure (handles duplicate labels
    # and non-integer formulas that trip up pymatgen's parser)
    try:
        from ase.io import read as ase_read
        from pymatgen.io.ase import AseAtomsAdaptor
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            atoms = ase_read(cif_path)
        return AseAtomsAdaptor.get_structure(atoms)
    except Exception as exc:
        raise ValueError(
            f'pymatgen failed to parse {cif_path}: {exc}'
        ) from exc


# =============================================================================
# Public entry point
# =============================================================================

def cif_to_graph(cif_path: str, dataset: str = 'hmof_gas') -> Dict:
    """
    Convert a single CIF file to a PoreGCN graph dict.

    Steps:
        1. Load structure with pymatgen.
        2. Run Zeo++ (subprocess) in a temp dir to get .nt2 Voronoi network.
           If Zeo++ binary is absent or the subprocess fails, log a warning
           and continue in atom-only mode (n_pores=0).
        3. Parse .nt2 to extract pore vertices and radii.
        4. Build atom features (92 one-hot from atom_init.json + 28 MOF features).
        5. Build neighbor list and bond-type matrix with PBC.
        6. Build pore feature matrix [N_pores, 8].
        7. Build atom-pore edges with KD-tree-style distance search.
        8. Return graph dict.

    Args:
        cif_path: Absolute or relative path to a CIF file.
        dataset:  One of 'hmof_gas', 'core_mof', 'hmof_geometric'.
                  Currently used only for logging; graph construction is
                  identical across datasets.

    Returns:
        Dict with keys:
            atom_fea         np.ndarray  [N_atoms, 120]
            bond_types       np.ndarray  [N_atoms, M]  int64
            nbr_fea_idx      np.ndarray  [N_atoms, M]  int64
            crystal_atom_idx list of 1 LongTensor  (for single-graph batch)
            pore_fea         np.ndarray  [N_pores, 8]  float32
            atom_pore_edges  np.ndarray  [2, N_ap_edges]  int64
            crystal_pore_idx list of 1 LongTensor
            n_atoms          int
            n_pores          int
            structure        pymatgen.core.Structure
            pore_positions   np.ndarray  [N_pores, 3]  Cartesian Angstroms
            pore_radii       np.ndarray  [N_pores]     Angstroms
            voronoi_used     bool

    Raises:
        FileNotFoundError: if cif_path does not exist.
        ValueError: if structure has 0 atoms.
    """
    if not os.path.exists(cif_path):
        raise FileNotFoundError(f'CIF not found: {cif_path}')

    # 1. Load structure (with fallbacks for non-standard CIFs)
    try:
        structure = _load_structure(cif_path)
    except ValueError:
        raise
    except Exception as exc:
        raise ValueError(f'pymatgen failed to parse {cif_path}: {exc}') from exc

    n_atoms = len(structure)
    if n_atoms == 0:
        raise ValueError(f'Structure has 0 atoms: {cif_path}')

    atom_coords = structure.cart_coords  # [N_atoms, 3]

    # 2. Run Zeo++
    voronoi_used = False
    vertices: List[Dict] = []

    with tempfile.TemporaryDirectory() as tmp_dir:
        nt2_path = _run_zeopp(cif_path, tmp_dir)
        if nt2_path is not None:
            vertices = _parse_nt2(nt2_path, min_radius=MIN_PORE_RADIUS)
            if vertices:
                voronoi_used = True
                logger.info(
                    'Zeo++ succeeded for %s: %d pore vertices (radius >= %.1f A)',
                    os.path.basename(cif_path), len(vertices), MIN_PORE_RADIUS,
                )
            else:
                logger.info(
                    'Zeo++ ran but found no vertices >= %.1f A for %s. Atom-only mode.',
                    MIN_PORE_RADIUS, os.path.basename(cif_path),
                )

    # 3 & 4. Atom features
    atom_fea = _build_atom_features(structure)  # [N_atoms, 120]

    # 5. Neighbor list
    nbr_fea_idx, bond_types = _build_neighbor_list(structure)

    # 6 & 7. Pore features and atom-pore edges
    pore_fea, pore_positions, pore_radii = _create_pore_features(
        vertices, structure, atom_coords
    )
    atom_pore_edges = _create_atom_pore_edges(atom_coords, pore_positions)

    n_pores = len(pore_fea)

    # Wrap as single-graph "batch" (crystal_atom_idx and crystal_pore_idx are
    # lists of one tensor each, matching the PoreGCN.forward() signature)
    crystal_atom_idx = [torch.arange(n_atoms, dtype=torch.long)]
    crystal_pore_idx = [torch.arange(n_pores, dtype=torch.long)]

    return {
        'atom_fea': atom_fea,
        'bond_types': bond_types,
        'nbr_fea_idx': nbr_fea_idx,
        'crystal_atom_idx': crystal_atom_idx,
        'pore_fea': pore_fea,
        'atom_pore_edges': atom_pore_edges,
        'crystal_pore_idx': crystal_pore_idx,
        'n_atoms': n_atoms,
        'n_pores': n_pores,
        'structure': structure,
        'pore_positions': pore_positions,
        'pore_radii': pore_radii,
        'voronoi_used': voronoi_used,
    }


def graph_to_tensors(graph: Dict, device: str = 'cpu') -> Dict:
    """
    Convert numpy arrays in a graph dict to torch tensors on device.

    Called by xai_engine.py immediately before model forward pass.
    """
    return {
        'atom_fea': torch.tensor(graph['atom_fea'], dtype=torch.float32, device=device),
        'bond_types': torch.tensor(graph['bond_types'], dtype=torch.long, device=device),
        'nbr_fea_idx': torch.tensor(graph['nbr_fea_idx'], dtype=torch.long, device=device),
        'crystal_atom_idx': [idx.to(device) for idx in graph['crystal_atom_idx']],
        'pore_fea': torch.tensor(graph['pore_fea'], dtype=torch.float32, device=device),
        'atom_pore_edges': torch.tensor(graph['atom_pore_edges'], dtype=torch.long, device=device),
        'crystal_pore_idx': [idx.to(device) for idx in graph['crystal_pore_idx']],
    }