app/engine.py

"""
Hybrid Enterprise Chemical Search Engine 
Includes: FAISS Retrieval, Multi-Fingerprint Fusion, Chemical-Aware Reranking, 
Z-Score Calibration, and MMR Diversity Control.
============================================================================
"""

import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

import os
import pickle
import numpy as np
import faiss

from rdkit import Chem
from rdkit.Chem import AllChem, MACCSkeys, rdMolDescriptors
from rdkit import DataStructs


class ChemicalSearchEngine:
    """
    Advanced Hybrid Chemical Similarity Search Engine (Drug-Discovery Grade).
    
    Pipeline Steps:
    1. FAISS Binary Retrieval: Rapid screen over Morgan fingerprints (Top 200).
    2. Multi-Fingerprint Fusion: Combines Morgan, MACCS, Atom Pairs, and Torsion.
    3. Chemical-Aware Reranking: Applies strict domain constraints (Aromaticity, Rings, Charge, Fragments).
    4. Similarity Calibration: Z-score normalization to probabilities.
    5. Diversity Control (MMR): Maximizes scaffold diversity.
    """

    def __init__(self, bit_size=2048):
        self.bit_size = bit_size
        self.index = None
        self.index_built = False
        self.total_compounds = 0
        
        self.metadata = []
        self.morgan_fps = []          
        self.maccs_fps = []           
        self.atom_pair_fps = []       
        self.torsion_fps = []         
        self.faiss_fingerprints = []  

    def _mol_to_all_fingerprints(self, mol):
        """Generates all 4 distinct tactical chemical fingerprints."""
        if mol is None:
            return None, None, None, None
        try:
            morgan = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=self.bit_size)
            maccs = MACCSkeys.GenMACCSKeys(mol)
            atom_pairs = rdMolDescriptors.GetHashedAtomPairFingerprintAsBitVect(mol, nBits=self.bit_size)
            torsions = rdMolDescriptors.GetHashedTopologicalTorsionFingerprintAsBitVect(mol, nBits=self.bit_size)
            return morgan, maccs, atom_pairs, torsions
        except Exception:
            return None, None, None, None

    def _extract_chemical_features(self, mol):
        """
        Extracts structural semantics for the Chemical-Aware Reranking Layer.
        Pre-computed during ingestion for zero overhead during search.
        """
        if mol is None:
            return 0, 0, 0, 1
        try:
            arom = sum(1 for atom in mol.GetAtoms() if atom.GetIsAromatic())
            rings = mol.GetRingInfo().NumRings()
            charge = sum(atom.GetFormalCharge() for atom in mol.GetAtoms())
            frags = len(Chem.GetMolFrags(mol))
            return arom, rings, charge, frags
        except Exception:
            return 0, 0, 0, 1

    def _bitvect_to_numpy(self, bv):
        """Converts an RDKit bit vector to a NumPy uint8 array."""
        arr = np.zeros((len(bv),), dtype=np.uint8)
        DataStructs.ConvertToNumpyArray(bv, arr)
        return arr

    def add_compounds(self, smiles_list, metadata_list=None):
        """Processes compounds, pre-computes chemical semantics, and builds engine."""
        print(f"[START] Processing {len(smiles_list)} compounds for Hybrid Engine...")
        
        for i, smiles in enumerate(smiles_list):
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                continue
                
            morgan, maccs, atom_pairs, torsions = self._mol_to_all_fingerprints(mol)
            if morgan is None:
                continue
                
            # Extract and cache chemical features (Aromaticity, Rings, Charge, Fragments)
            arom, rings, charge, frags = self._extract_chemical_features(mol)
                
            meta = metadata_list[i] if metadata_list else {"smiles": smiles}
            meta["smiles"] = smiles
            # Cache semantic features directly in metadata
            meta["chem_features"] = {"arom": arom, "rings": rings, "charge": charge, "frags": frags}
            self.metadata.append(meta)
            
            self.morgan_fps.append(morgan)
            self.maccs_fps.append(maccs)
            self.atom_pair_fps.append(atom_pairs)
            self.torsion_fps.append(torsions)
            self.faiss_fingerprints.append(self._bitvect_to_numpy(morgan))

        self.total_compounds = len(self.metadata)
        print(f"[SUCCESS] Feature extraction complete: {self.total_compounds} compounds")
        self._build_faiss_index()

    def _build_faiss_index(self):
        """Builds FAISS index with proper byte packing."""
        if len(self.faiss_fingerprints) == 0:
            return

        self.index = faiss.IndexBinaryFlat(self.bit_size)
        np_fps = np.array(self.faiss_fingerprints, dtype=np.uint8)
        packed_fps = np.packbits(np_fps, axis=1)
        
        self.index.add(packed_fps)
        self.index_built = True
        print(f"[SUCCESS] FAISS Binary index successfully built.")

    def search(self, query_smiles, k=3, lambda_param=0.6):
        """Executes the complete Drug-Discovery Grade Search Pipeline."""
        query_mol = Chem.MolFromSmiles(query_smiles)
        if query_mol is None or not self.index_built:
            return []

        # 1. Query fingerprints & Semantic features
        q_morgan, q_maccs, q_atom_pairs, q_torsions = self._mol_to_all_fingerprints(query_mol)
        q_arom, q_rings, q_charge, q_frags = self._extract_chemical_features(query_mol)
        
        if q_morgan is None:
            return []

        # 2. FAISS Fast Retrieval
        k_search = min(max(k * 20, 200), self.total_compounds)
        q_packed = np.packbits(self._bitvect_to_numpy(q_morgan).reshape(1, -1), axis=1)
        _, indices = self.index.search(q_packed, k_search)
        candidate_indices = [int(idx) for idx in indices[0] if idx >= 0]

        if not candidate_indices:
            return []

        # 3. Multi-FP & Chemical-Aware Scoring
        candidate_pool = []
        for idx in candidate_indices:
            # 3a. Multi-Fingerprint Fusion (Base Structural Score)
            score_morgan = DataStructs.FingerprintSimilarity(q_morgan, self.morgan_fps[idx])
            score_maccs = DataStructs.FingerprintSimilarity(q_maccs, self.maccs_fps[idx])
            score_atom_pairs = DataStructs.FingerprintSimilarity(q_atom_pairs, self.atom_pair_fps[idx])
            score_torsions = DataStructs.FingerprintSimilarity(q_torsions, self.torsion_fps[idx])
            
            base_score = (0.50 * score_morgan + 0.20 * score_maccs + 
                          0.20 * score_atom_pairs + 0.10 * score_torsions)

            # 3b. Chemical-Aware Penalties & Bonuses (O(1) lookup, ultra-fast)
            c_feats = self.metadata[idx]["chem_features"]
            c_arom, c_rings, c_charge, c_frags = c_feats["arom"], c_feats["rings"], c_feats["charge"], c_feats["frags"]

            # Aromaticity Bonus
            max_arom = max(q_arom, c_arom)
            arom_score = 1.0 if max_arom == 0 else 1.0 - (abs(q_arom - c_arom) / max_arom)

            # Ring System Bonus
            max_rings = max(q_rings, c_rings)
            ring_score = 1.0 if max_rings == 0 else 1.0 - (abs(q_rings - c_rings) / max_rings)

            # Strict Domain Penalties
            charge_pen = min(abs(c_charge) / 3.0, 1.0)
            frag_pen = 0.0 if c_frags <= 1 else min((c_frags - 1) * 0.3, 1.0)

            # 3c. Final Hybrid Score
            chem_aware_score = (
                0.70 * base_score +
                0.15 * arom_score +
                0.10 * ring_score -
                0.15 * charge_pen -
                0.10 * frag_pen
            )
            
            candidate_pool.append({
                "index": idx,
                "smiles": self.metadata[idx]["smiles"],
                "metadata": self.metadata[idx],
                "chem_aware_score": chem_aware_score,
                "base_structural_score": base_score
            })

        # 4. Calibration Layer (Z-Score)
        scores = [c["chem_aware_score"] for c in candidate_pool]
        mean_score = np.mean(scores)
        std_score = np.std(scores) if np.std(scores) > 0 else 1.0
        
        for c in candidate_pool:
            z = (c["chem_aware_score"] - mean_score) / std_score
            c["calibrated_score"] = 1.0 / (1.0 + np.exp(-z))

        # 5. Diversity Control (MMR)
        selected_results = []
        remaining_candidates = list(candidate_pool)
        
        remaining_candidates.sort(key=lambda x: x["calibrated_score"], reverse=True)
        selected_results.append(remaining_candidates.pop(0))
        
        while len(selected_results) < k and remaining_candidates:
            best_mmr = -float('inf')
            best_idx = -1
            
            for idx, cand in enumerate(remaining_candidates):
                max_sim = -float('inf')
                cand_fp = self.morgan_fps[cand["index"]]
                
                for sel in selected_results:
                    sim = DataStructs.FingerprintSimilarity(cand_fp, self.morgan_fps[sel["index"]])
                    if sim > max_sim: max_sim = sim
                
                mmr_val = (lambda_param * cand["calibrated_score"]) - ((1.0 - lambda_param) * max_sim)
                if mmr_val > best_mmr:
                    best_mmr = mmr_val
                    best_idx = idx
            
            if best_idx != -1:
                selected_results.append(remaining_candidates.pop(best_idx))
            else:
                break

        # Output formatting
        return [{
            "smiles": res["smiles"],
            "similarity_score": round(float(res["chem_aware_score"]), 4), 
            "calibrated_score": round(float(res["calibrated_score"]), 4),
            "metadata": res["metadata"],
            "index": res["index"]
        } for res in selected_results]

    def save_index(self, filepath):
        """Serializes the multi-fingerprint infrastructure and metadata safely."""
        if not self.index_built:
            return
        try:
            faiss_file = filepath.replace(".pkl", ".faiss")
            faiss.write_index_binary(self.index, faiss_file)
            
            data = {
                "metadata": self.metadata,
                "morgan_fps": self.morgan_fps,
                "maccs_fps": self.maccs_fps,
                "atom_pair_fps": self.atom_pair_fps,
                "torsion_fps": self.torsion_fps,
                "faiss_fingerprints": self.faiss_fingerprints,
                "bit_size": self.bit_size,
                "total_compounds": self.total_compounds
            }
            with open(filepath, "wb") as f:
                pickle.dump(data, f)
            print(f"[SAVE] Engine saved at: {filepath}")
        except Exception as e:
            print(f"[ERROR] Serialization failed: {e}")

    def load_index(self, filepath):
        """Loads and provisions the full hybrid matrix and FAISS architecture."""
        if not os.path.exists(filepath):
            return False
        try:
            with open(filepath, "rb") as f:
                data = pickle.load(f)
                
            self.metadata = data["metadata"]
            self.morgan_fps = data["morgan_fps"]
            self.maccs_fps = data["maccs_fps"]
            self.atom_pair_fps = data["atom_pair_fps"]
            self.torsion_fps = data["torsion_fps"]
            self.faiss_fingerprints = data["faiss_fingerprints"]
            self.bit_size = data.get("bit_size", 2048)
            self.total_compounds = data.get("total_compounds", len(self.metadata))
            
            faiss_file = filepath.replace(".pkl", ".faiss")
            self.index = faiss.read_index_binary(faiss_file)
            self.index_built = True
            
            print(f"[LOAD] Engine loaded completely: {self.total_compounds} compounds.")
            return True
        except Exception as e:
            print(f"[ERROR] Loading error: {e}")
            return False