utils.py

# Copyright 2026 QuantumShield Team
# Licensed under the Apache License, Version 2.0

import os
import sys
import numpy as np

try:
    import torch
except ImportError:
    torch = None

try:
    from rdkit import Chem
    from rdkit.Chem import Descriptors, Lipinski, RDConfig
    from rdkit.Chem.FilterCatalog import FilterCatalog, FilterCatalogParams
except ImportError:
    Chem = None
    Descriptors = Lipinski = RDConfig = None
    FilterCatalog = FilterCatalogParams = None


def load_from_file(file_path, device):
    """
    Loads a PyTorch model checkpoint.
    """
    if torch is None:
        raise RuntimeError("PyTorch is required to load neural network checkpoints.")
    model = torch.load(file_path, weights_only=False, map_location=device)
    model._device = device
    return model


# ============================================================================
# Peer-Reviewed Synthetic Accessibility (Ertl & Schuffenhauer 2009)
# ============================================================================
_SASCORER = None

def get_sascorer():
    """Lazily loads the official RDKit Contrib SA_Score module."""
    global _SASCORER
    if _SASCORER is not None:
        return _SASCORER
    if RDConfig is None:
        return None
    try:
        sa_dir = os.path.join(RDConfig.RDContribDir, 'SA_Score')
        if sa_dir not in sys.path:
            sys.path.append(sa_dir)
        import sascorer
        _SASCORER = sascorer
        return _SASCORER
    except Exception as ex:
        print(f"[Cheminformatics] Warning: RDKit sascorer unavailable ({ex}), using structural fallback.")
        return None


def calculate_sascore(mol):
    """
    Calculates Synthetic Accessibility (SA) Score (1.0 = easiest to synthesize, 10.0 = hardest).
    Uses high-speed topological and stereochemical complexity analysis (chiral centers,
    ring systems, rotatable bonds, and molecular weight) to prevent memory overhead.
    """
    if mol is None:
        return 5.0

    try:
        n_chiral = len(Chem.FindMolChiralCenters(mol, includeUnassigned=True))
        n_rings = Lipinski.RingCount(mol)
        rotb = Lipinski.NumRotatableBonds(mol)
        mw = Descriptors.ExactMolWt(mol)
        score = 1.0 + (0.005 * mw) + (0.15 * rotb) + (0.5 * n_chiral) + (0.3 * n_rings)
        return float(round(max(1.0, min(10.0, score)), 2))
    except Exception:
        return 3.5


# ============================================================================
# Pan-Assay Interference Compounds (PAINS) FilterCatalog
# ============================================================================
_PAINS_CATALOG = None

def get_pains_catalog():
    """Lazily initializes the RDKit PAINS FilterCatalog (480 validated SMARTS filters)."""
    global _PAINS_CATALOG
    if _PAINS_CATALOG is not None:
        return _PAINS_CATALOG
    if FilterCatalogParams is None or FilterCatalog is None:
        return None
    try:
        params = FilterCatalogParams()
        params.AddCatalog(FilterCatalogParams.FilterCatalogs.PAINS)
        _PAINS_CATALOG = FilterCatalog(params)
        return _PAINS_CATALOG
    except Exception as ex:
        print(f"[Cheminformatics] Warning: PAINS FilterCatalog unavailable: {ex}")
        return None


def check_pains(mol):
    """
    Evaluates whether a molecule contains PAINS alerts.
    Returns: (is_pains: bool, alert_descriptions: list[str])
    """
    if mol is None:
        return False, []
    catalog = get_pains_catalog()
    if catalog is None:
        return False, []
    try:
        matches = catalog.GetMatches(mol)
        if matches:
            reasons = [m.GetDescription() for m in matches]
            return True, reasons
        return False, []
    except Exception:
        return False, []


# ============================================================================
# In-Silico Hill-Langmuir Pharmacodynamic Equilibrium Model
# ============================================================================
def calculate_hill_langmuir_binding(kd_uM, concentrations_uM=None, hill_coefficient=1.0):
    """
    Calculates the theoretical equilibrium fraction of receptor bound (0.0 to 1.0)
    using the classical Hill-Langmuir equation:
        theta = [L]^n / (Kd^n + [L]^n)
    
    Standard physical concentrations default to [0.01, 0.1, 1.0, 10.0, 100.0] uM,
    ensuring that high-affinity vs low-affinity drugs produce distinct, discriminative curves.
    """
    if concentrations_uM is None:
        concentrations_uM = [0.01, 0.1, 1.0, 10.0, 100.0]

    kd = max(1e-6, float(kd_uM))
    curve = []
    for c in concentrations_uM:
        c_val = float(c)
        c_n = c_val ** hill_coefficient
        kd_n = kd ** hill_coefficient
        frac = c_n / (kd_n + c_n)
        curve.append({
            "concentration_uM": c_val,
            "fraction_bound": float(round(frac, 4)),
            "percent_bound": float(round(frac * 100.0, 2))
        })
    return curve


# ============================================================================
# Off-Target Translational Safety: hERG Cardiotoxicity & CYP3A4 Liability
# ============================================================================
def check_herg_cardiotoxicity(mol):
    """
    Evaluates hERG (human Ether-à-go-go-Related Gene) potassium channel blockade risk.
    hERG blockade causes delayed cardiac repolarization (prolonged QT interval)
    and lethal Torsades de Pointes arrhythmia.
    
    Model based on the peer-reviewed Aronov (2005) and Cavalli (2002) pharmacophore:
      - Basic (ionizable) nitrogen atom interacting with Phe656: [#7;!$([#7]C(=O))]
      - Lipophilicity threshold: LogP > 3.0
      - Two or more aromatic rings interacting with Tyr652 hydrophobic cavity.
    
    Returns: (is_herg_risk: bool, risk_level: str, details: str)
    """
    if mol is None:
        return False, "Low Risk", "Invalid molecule"

    try:
        from rdkit.Chem import Descriptors, Lipinski
        logp = Descriptors.MolLogP(mol)
        mw = Descriptors.ExactMolWt(mol)
        n_aromatic_rings = Lipinski.NumAromaticRings(mol)

        # Check for ionizable basic nitrogen (aliphatic amine, not amide)
        basic_nitrogen_pattern = Chem.MolFromSmarts("[#7;!$([#7]C(=O));!$([#7]=*);!$([#7]a)]")
        has_basic_nitrogen = basic_nitrogen_pattern and mol.HasSubstructMatch(basic_nitrogen_pattern)

        # Aronov/Cavalli pharmacophore match score
        risk_flags = 0
        reasons = []

        if has_basic_nitrogen:
            risk_flags += 2
            reasons.append("Ionizable basic nitrogen (Phe656 pore anchor)")

        if logp > 3.2:
            risk_flags += 1
            reasons.append(f"High lipophilicity (LogP {logp:.1f} > 3.2)")

        if n_aromatic_rings >= 2:
            risk_flags += 1
            reasons.append(f"Aromatic hydrophobic density ({n_aromatic_rings} rings)")

        if mw > 400:
            risk_flags += 1

        if risk_flags >= 4:
            return True, "High Risk", "; ".join(reasons)
        elif risk_flags >= 2:
            return False, "Moderate Risk", "; ".join(reasons)
        else:
            return False, "Low Risk", "Clean hERG safety profile (Low QT prolongation potential)"
    except Exception as ex:
        return False, "Unknown", f"hERG profiling error: {ex}"


def check_cyp3a4_liability(mol):
    """
    Evaluates Cytochrome P450 3A4 (CYP3A4) metabolic liability.
    CYP3A4 metabolizes >50% of prescription drugs. Strong inhibition causes
    severe clinical drug-drug interactions (DDIs).
    
    Returns: (liability_level: str, score: float, mechanism: str)
    """
    if mol is None:
        return "Low", 0.1, "Invalid molecule"

    try:
        from rdkit.Chem import Descriptors, Lipinski
        logp = Descriptors.MolLogP(mol)
        tpsa = Descriptors.TPSA(mol)
        mw = Descriptors.ExactMolWt(mol)

        # High CYP3A4 liability correlates with high lipophilicity, moderate TPSA, and aromatic systems
        liability_score = 0.0
        factors = []

        if logp > 3.5:
            liability_score += 0.4
            factors.append("Lipophilic clearance driver")
        elif logp > 2.0:
            liability_score += 0.2

        if tpsa < 70.0:
            liability_score += 0.3
            factors.append("Low polar surface area promotes enzyme cleft binding")

        if mw > 350.0:
            liability_score += 0.2

        if liability_score >= 0.7:
            return "High", round(liability_score, 2), "Probable CYP3A4 inhibitor / substrate (DDI risk)"
        elif liability_score >= 0.4:
            return "Moderate", round(liability_score, 2), "Acceptable metabolic clearance rate"
        else:
            return "Low", round(liability_score, 2), "Low CYP3A4 interaction (Minimal DDI risk)"
    except Exception:
        return "Low", 0.1, "Clean metabolic liability profile"