src/util/utils.py

import os
import time
import math
import datetime
import warnings
import itertools
from copy import deepcopy
from functools import partial
from collections import Counter
from multiprocessing import Pool
from statistics import mean

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from scipy.spatial.distance import cosine as cos_distance

import torch
import wandb

from rdkit import Chem, DataStructs, RDLogger
from rdkit.Chem import (
    AllChem,
    Draw,
    Descriptors,
    Lipinski,
    Crippen,
    rdMolDescriptors,
    FilterCatalog,
)
from rdkit.Chem.Scaffolds import MurckoScaffold

# Disable RDKit warnings
RDLogger.DisableLog("rdApp.*")


class Metrics(object):
    """
    Collection of static methods to compute various metrics for molecules.
    """

    @staticmethod
    def valid(x):
        """
        Checks whether the molecule is valid.
        
        Args:
            x: RDKit molecule object.
        
        Returns:
            bool: True if molecule is valid and has a non-empty SMILES representation.
        """
        return x is not None and Chem.MolToSmiles(x) != ''

    @staticmethod
    def tanimoto_sim_1v2(data1, data2):
        """
        Computes the average Tanimoto similarity for paired fingerprints.
        
        Args:
            data1: Fingerprint data for first set.
            data2: Fingerprint data for second set.
        
        Returns:
            float: The average Tanimoto similarity between corresponding fingerprints.
        """
        # Determine the minimum size between two arrays for pairing
        min_len = data1.size if data1.size > data2.size else data2
        sims = []
        for i in range(min_len):
            sim = DataStructs.FingerprintSimilarity(data1[i], data2[i])
            sims.append(sim)
        # Use 'mean' from statistics; note that variable 'sim' was used, corrected to use sims list.
        mean_sim = mean(sims)
        return mean_sim

    @staticmethod
    def mol_length(x):
        """
        Computes the length of the largest fragment (by character count) in a SMILES string.
        
        Args:
            x (str): SMILES string.
        
        Returns:
            int: Number of alphabetic characters in the longest fragment of the SMILES.
        """
        if x is not None:
            # Split at dots (.) and take the fragment with maximum length, then count alphabetic characters.
            return len([char for char in max(x.split(sep="."), key=len).upper() if char.isalpha()])
        else:
            return 0

    @staticmethod
    def max_component(data, max_len):
        """
        Returns the average normalized length of molecules in the dataset.
        
        Each molecule's length is computed and divided by max_len, then averaged.
        
        Args:
            data (iterable): Collection of SMILES strings.
            max_len (int): Maximum possible length for normalization.
        
        Returns:
            float: Normalized average length.
        """
        lengths = np.array(list(map(Metrics.mol_length, data)), dtype=np.float32)
        return (lengths / max_len).mean()

    @staticmethod
    def mean_atom_type(data):
        """
        Computes the average number of unique atom types in the provided node data.
        
        Args:
            data (iterable): Iterable containing node data with unique atom types.
        
        Returns:
            float: The average count of unique atom types, subtracting one.
        """
        atom_types_used = []
        for i in data:
            # Assuming each element i has a .unique() method that returns unique atom types.
            atom_types_used.append(len(i.unique().tolist()))
        av_type = np.mean(atom_types_used) - 1
        return av_type


def mols2grid_image(mols, path):
    """
    Saves grid images for a list of molecules.
    
    For each molecule in the list, computes 2D coordinates and saves an image file.
    
    Args:
        mols (list): List of RDKit molecule objects.
        path (str): Directory where images will be saved.
    """
    # Replace None molecules with an empty molecule
    mols = [e if e is not None else Chem.RWMol() for e in mols]

    for i in range(len(mols)):
        if Metrics.valid(mols[i]):
            AllChem.Compute2DCoords(mols[i])
            file_path = os.path.join(path, "{}.png".format(i + 1))
            Draw.MolToFile(mols[i], file_path, size=(1200, 1200))
            # wandb.save(file_path)  # Optionally save to Weights & Biases
        else:
            continue


def save_smiles_matrices(mols, edges_hard, nodes_hard, path, data_source=None):
    """
    Saves the edge and node matrices along with SMILES strings to text files.
    
    Each file contains the edge matrix, node matrix, and SMILES representation for a molecule.
    
    Args:
        mols (list): List of RDKit molecule objects.
        edges_hard (torch.Tensor): Tensor of edge features.
        nodes_hard (torch.Tensor): Tensor of node features.
        path (str): Directory where files will be saved.
        data_source: Optional data source information (not used in function).
    """
    mols = [e if e is not None else Chem.RWMol() for e in mols]

    for i in range(len(mols)):
        if Metrics.valid(mols[i]):
            save_path = os.path.join(path, "{}.txt".format(i + 1))
            with open(save_path, "a") as f:
                np.savetxt(f, edges_hard[i].cpu().numpy(), header="edge matrix:\n", fmt='%1.2f')
                f.write("\n")
                np.savetxt(f, nodes_hard[i].cpu().numpy(), header="node matrix:\n", footer="\nsmiles:", fmt='%1.2f')
                f.write("\n")
            # Append the SMILES representation to the file
            with open(save_path, "a") as f:
                print(Chem.MolToSmiles(mols[i]), file=f)
            # wandb.save(save_path)  # Optionally save to Weights & Biases
        else:
            continue

def dense_to_sparse_with_attr(adj):
    """
    Converts a dense adjacency matrix to a sparse representation.
    
    Args:
        adj (torch.Tensor): Adjacency matrix tensor (2D or 3D) with square last two dimensions.
    
    Returns:
        tuple: A tuple containing indices and corresponding edge attributes.
    """
    assert adj.dim() >= 2 and adj.dim() <= 3
    assert adj.size(-1) == adj.size(-2)

    index = adj.nonzero(as_tuple=True)
    edge_attr = adj[index]

    if len(index) == 3:
        batch = index[0] * adj.size(-1)
        index = (batch + index[1], batch + index[2])
    return index, edge_attr


def mol_sample(sample_directory, edges, nodes, idx, i, matrices2mol, dataset_name):
    """
    Samples molecules from edge and node predictions, then saves grid images and text files.
    
    Args:
        sample_directory (str): Directory to save the samples.
        edges (torch.Tensor): Edge predictions tensor.
        nodes (torch.Tensor): Node predictions tensor.
        idx (int): Current index for naming the sample.
        i (int): Epoch/iteration index.
        matrices2mol (callable): Function to convert matrices to RDKit molecule.
        dataset_name (str): Name of the dataset for file naming.
    """
    sample_path = os.path.join(sample_directory, "{}_{}-epoch_iteration".format(idx + 1, i + 1))
    # Get the index of the maximum predicted feature along the last dimension
    g_edges_hat_sample = torch.max(edges, -1)[1]
    g_nodes_hat_sample = torch.max(nodes, -1)[1]
    # Convert matrices to molecule objects
    mol = [matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(),
                        strict=True, file_name=dataset_name)
           for e_, n_ in zip(g_edges_hat_sample, g_nodes_hat_sample)]

    if not os.path.exists(sample_path):
        os.makedirs(sample_path)

    mols2grid_image(mol, sample_path)
    save_smiles_matrices(mol, g_edges_hat_sample.detach(), g_nodes_hat_sample.detach(), sample_path)

    # Remove the directory if no files were saved
    if len(os.listdir(sample_path)) == 0:
        os.rmdir(sample_path)

    print("Valid molecules are saved.")
    print("Valid matrices and smiles are saved")


def logging(log_path, start_time, i, idx, loss, save_path, drug_smiles, edge, node, 
            matrices2mol, dataset_name, real_adj, real_annot, drug_vecs):
    """
    Logs training statistics and evaluation metrics.
    
    The function generates molecules from predictions, computes various metrics such as
    validity, uniqueness, novelty, and similarity scores, and logs them using wandb and a file.
    
    Args:
        log_path (str): Path to save the log file.
        start_time (float): Start time to compute elapsed time.
        i (int): Current iteration index.
        idx (int): Current epoch index.
        loss (dict): Dictionary to update with loss and metric values.
        save_path (str): Directory path to save sample outputs.
        drug_smiles (list): List of reference drug SMILES.
        edge (torch.Tensor): Edge prediction tensor.
        node (torch.Tensor): Node prediction tensor.
        matrices2mol (callable): Function to convert matrices to molecules.
        dataset_name (str): Dataset name.
        real_adj (torch.Tensor): Ground truth adjacency matrix tensor.
        real_annot (torch.Tensor): Ground truth annotation tensor.
        drug_vecs (list): List of drug vectors for similarity calculation.
    """
    g_edges_hat_sample = torch.max(edge, -1)[1]
    g_nodes_hat_sample = torch.max(node, -1)[1]

    a_tensor_sample = torch.max(real_adj, -1)[1].float()
    x_tensor_sample = torch.max(real_annot, -1)[1].float()

    # Generate molecules from predictions and real data
    mols = [matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(),
                         strict=True, file_name=dataset_name)
            for e_, n_ in zip(g_edges_hat_sample, g_nodes_hat_sample)]
    real_mol = [matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(),
                              strict=True, file_name=dataset_name)
                for e_, n_ in zip(a_tensor_sample, x_tensor_sample)]

    # Compute average number of atom types
    atom_types_average = Metrics.mean_atom_type(g_nodes_hat_sample)
    real_smiles = [Chem.MolToSmiles(x) for x in real_mol if x is not None]
    gen_smiles = []
    uniq_smiles = []
    for line in mols:
        if line is not None:
            gen_smiles.append(Chem.MolToSmiles(line))
            uniq_smiles.append(Chem.MolToSmiles(line))
        elif line is None:
            gen_smiles.append(None)

    # Process SMILES to take the longest fragment if multiple are present
    gen_smiles_saves = [None if x is None else max(x.split('.'), key=len) for x in gen_smiles]
    uniq_smiles_saves = [None if x is None else max(x.split('.'), key=len) for x in uniq_smiles]

    # Save the generated SMILES to a text file
    sample_save_dir = os.path.join(save_path, "samples.txt")
    with open(sample_save_dir, "a") as f:
        for s in gen_smiles_saves:
            if s is not None:
                f.write(s + "\n")

    k = len(set(uniq_smiles_saves) - {None})
    et = time.time() - start_time
    et = str(datetime.timedelta(seconds=et))[:-7]
    log_str = "Elapsed [{}], Epoch/Iteration [{}/{}]".format(et, idx, i + 1)
    
    # Generate molecular fingerprints for similarity computations
    gen_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in mols if x is not None]
    chembl_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in real_mol if x is not None]

    # Compute evaluation metrics: validity, uniqueness, novelty, similarity scores, and average maximum molecule length.
    valid = fraction_valid(gen_smiles_saves)
    unique = fraction_unique(uniq_smiles_saves, k)
    novel_starting_mol = novelty(gen_smiles_saves, real_smiles)
    novel_akt = novelty(gen_smiles_saves, drug_smiles)
    if len(uniq_smiles_saves) == 0:
        snn_chembl = 0
        snn_akt = 0
        maxlen = 0
    else:
        snn_chembl = average_agg_tanimoto(np.array(chembl_vecs), np.array(gen_vecs))
        snn_akt = average_agg_tanimoto(np.array(drug_vecs), np.array(gen_vecs))
        maxlen = Metrics.max_component(uniq_smiles_saves, 45)

    # Update loss dictionary with computed metrics
    loss.update({
        'Validity': valid,
        'Uniqueness': unique,
        'Novelty': novel_starting_mol,
        'Novelty_akt': novel_akt,
        'SNN_chembl': snn_chembl,
        'SNN_akt': snn_akt,
        'MaxLen': maxlen,
        'Atom_types': atom_types_average
    })

    # Log metrics using wandb
    wandb.log({
        "Validity": valid,
        "Uniqueness": unique,
        "Novelty": novel_starting_mol,
        "Novelty_akt": novel_akt,
        "SNN_chembl": snn_chembl,
        "SNN_akt": snn_akt,
        "MaxLen": maxlen,
        "Atom_types": atom_types_average
    })

    # Append each metric to the log string and write to the log file
    for tag, value in loss.items():
        log_str += ", {}: {:.4f}".format(tag, value)
    with open(log_path, "a") as f:
        f.write(log_str + "\n")
    print(log_str)
    print("\n")


def plot_grad_flow(named_parameters, model, itera, epoch, grad_flow_directory):
    """
    Plots the gradients flowing through different layers during training.
    
    This is useful to check for possible gradient vanishing or exploding problems.
    
    Args:
        named_parameters (iterable): Iterable of (name, parameter) tuples from the model.
        model (str): Name of the model (used for saving the plot).
        itera (int): Iteration index.
        epoch (int): Current epoch.
        grad_flow_directory (str): Directory to save the gradient flow plot.
    """
    ave_grads = []
    max_grads = []
    layers = []
    for n, p in named_parameters:
        if p.requires_grad and ("bias" not in n):
            layers.append(n)
            ave_grads.append(p.grad.abs().mean().cpu())
            max_grads.append(p.grad.abs().max().cpu())
    # Plot maximum gradients and average gradients for each layer
    plt.bar(np.arange(len(max_grads)), max_grads, alpha=0.1, lw=1, color="c")
    plt.bar(np.arange(len(max_grads)), ave_grads, alpha=0.1, lw=1, color="b")
    plt.hlines(0, 0, len(ave_grads) + 1, lw=2, color="k")
    plt.xticks(range(0, len(ave_grads), 1), layers, rotation="vertical")
    plt.xlim(left=0, right=len(ave_grads))
    plt.ylim(bottom=-0.001, top=1)  # Zoom in on lower gradient regions
    plt.xlabel("Layers")
    plt.ylabel("Average Gradient")
    plt.title("Gradient Flow")
    plt.grid(True)
    plt.legend([
        Line2D([0], [0], color="c", lw=4),
        Line2D([0], [0], color="b", lw=4),
        Line2D([0], [0], color="k", lw=4)
    ], ['max-gradient', 'mean-gradient', 'zero-gradient'])
    # Save the plot to the specified directory
    plt.savefig(os.path.join(grad_flow_directory, "weights_" + model + "_" + str(itera) + "_" + str(epoch) + ".png"), dpi=500, bbox_inches='tight')


def get_mol(smiles_or_mol):
    """
    Loads a SMILES string or molecule into an RDKit molecule object.
    
    Args:
        smiles_or_mol (str or RDKit Mol): SMILES string or RDKit molecule.
    
    Returns:
        RDKit Mol or None: Sanitized molecule object, or None if invalid.
    """
    if isinstance(smiles_or_mol, str):
        if len(smiles_or_mol) == 0:
            return None
        mol = Chem.MolFromSmiles(smiles_or_mol)
        if mol is None:
            return None
        try:
            Chem.SanitizeMol(mol)
        except ValueError:
            return None
        return mol
    return smiles_or_mol


def mapper(n_jobs):
    """
    Returns a mapping function for parallel or serial processing.
    
    If n_jobs == 1, returns the built-in map function.
    If n_jobs > 1, returns a function that uses a multiprocessing pool.
    
    Args:
        n_jobs (int or pool object): Number of jobs or a Pool instance.
    
    Returns:
        callable: A function that acts like map.
    """
    if n_jobs == 1:
        def _mapper(*args, **kwargs):
            return list(map(*args, **kwargs))
        return _mapper
    if isinstance(n_jobs, int):
        pool = Pool(n_jobs)
        def _mapper(*args, **kwargs):
            try:
                result = pool.map(*args, **kwargs)
            finally:
                pool.terminate()
            return result
        return _mapper
    return n_jobs.map


def remove_invalid(gen, canonize=True, n_jobs=1):
    """
    Removes invalid molecules from the provided dataset.
    
    Optionally canonizes the SMILES strings.
    
    Args:
        gen (list): List of SMILES strings.
        canonize (bool): Whether to convert to canonical SMILES.
        n_jobs (int): Number of parallel jobs.
    
    Returns:
        list: Filtered list of valid molecules.
    """
    if not canonize:
        mols = mapper(n_jobs)(get_mol, gen)
        return [gen_ for gen_, mol in zip(gen, mols) if mol is not None]
    return [x for x in mapper(n_jobs)(canonic_smiles, gen) if x is not None]


def fraction_valid(gen, n_jobs=1):
    """
    Computes the fraction of valid molecules in the dataset.
    
    Args:
        gen (list): List of SMILES strings.
        n_jobs (int): Number of parallel jobs.
    
    Returns:
        float: Fraction of molecules that are valid.
    """
    gen = mapper(n_jobs)(get_mol, gen)
    return 1 - gen.count(None) / len(gen)


def canonic_smiles(smiles_or_mol):
    """
    Converts a SMILES string or molecule to its canonical SMILES.
    
    Args:
        smiles_or_mol (str or RDKit Mol): Input molecule.
    
    Returns:
        str or None: Canonical SMILES string or None if invalid.
    """
    mol = get_mol(smiles_or_mol)
    if mol is None:
        return None
    return Chem.MolToSmiles(mol)


def fraction_unique(gen, k=None, n_jobs=1, check_validity=True):
    """
    Computes the fraction of unique molecules.
    
    Optionally computes unique@k, where only the first k molecules are considered.
    
    Args:
        gen (list): List of SMILES strings.
        k (int): Optional cutoff for unique@k computation.
        n_jobs (int): Number of parallel jobs.
        check_validity (bool): Whether to check for validity of molecules.
    
    Returns:
        float: Fraction of unique molecules.
    """
    if k is not None:
        if len(gen) < k:
            warnings.warn("Can't compute unique@{}.".format(k) +
                          " gen contains only {} molecules".format(len(gen)))
        gen = gen[:k]
    if check_validity:
        canonic = list(mapper(n_jobs)(canonic_smiles, gen))
        canonic = [i for i in canonic if i is not None]
    set_cannonic = set(canonic)
    return 0 if len(canonic) == 0 else len(set_cannonic) / len(canonic)


def novelty(gen, train, n_jobs=1):
    """
    Computes the novelty score of generated molecules.
    
    Novelty is defined as the fraction of generated molecules that do not appear in the training set.
    
    Args:
        gen (list): List of generated SMILES strings.
        train (list): List of training SMILES strings.
        n_jobs (int): Number of parallel jobs.
    
    Returns:
        float: Novelty score.
    """
    gen_smiles = mapper(n_jobs)(canonic_smiles, gen)
    gen_smiles_set = set(gen_smiles) - {None}
    train_set = set(train)
    return 0 if len(gen_smiles_set) == 0 else len(gen_smiles_set - train_set) / len(gen_smiles_set)


def internal_diversity(gen):
    """
    Computes the internal diversity of a set of molecules.
    
    Internal diversity is defined as one minus the average Tanimoto similarity between all pairs.
    
    Args:
        gen: Array-like representation of molecules.
    
    Returns:
        tuple: Mean and standard deviation of internal diversity.
    """
    diversity = [1 - x for x in average_agg_tanimoto(gen, gen, agg="mean", intdiv=True)]
    return np.mean(diversity), np.std(diversity)


def average_agg_tanimoto(stock_vecs, gen_vecs, batch_size=5000, agg='max', device='cpu', p=1, intdiv=False):
    """
    Computes the average aggregated Tanimoto similarity between two sets of molecular fingerprints.
    
    For each fingerprint in gen_vecs, finds the closest (max or mean) similarity with fingerprints in stock_vecs.
    
    Args:
        stock_vecs (numpy.ndarray): Array of fingerprint vectors from the reference set.
        gen_vecs (numpy.ndarray): Array of fingerprint vectors from the generated set.
        batch_size (int): Batch size for processing fingerprints.
        agg (str): Aggregation method, either 'max' or 'mean'.
        device (str): Device to perform computations on.
        p (int): Power for averaging.
        intdiv (bool): Whether to return individual similarities or the average.
    
    Returns:
        float or numpy.ndarray: Average aggregated Tanimoto similarity or array of individual scores.
    """
    assert agg in ['max', 'mean'], "Can aggregate only max or mean"
    agg_tanimoto = np.zeros(len(gen_vecs))
    total = np.zeros(len(gen_vecs))
    for j in range(0, stock_vecs.shape[0], batch_size):
        x_stock = torch.tensor(stock_vecs[j:j + batch_size]).to(device).float()
        for i in range(0, gen_vecs.shape[0], batch_size):
            y_gen = torch.tensor(gen_vecs[i:i + batch_size]).to(device).float()
            y_gen = y_gen.transpose(0, 1)
            tp = torch.mm(x_stock, y_gen)
            # Compute Jaccard/Tanimoto similarity
            jac = (tp / (x_stock.sum(1, keepdim=True) + y_gen.sum(0, keepdim=True) - tp)).cpu().numpy()
            jac[np.isnan(jac)] = 1
            if p != 1:
                jac = jac ** p
            if agg == 'max':
                agg_tanimoto[i:i + y_gen.shape[1]] = np.maximum(
                    agg_tanimoto[i:i + y_gen.shape[1]], jac.max(0))
            elif agg == 'mean':
                agg_tanimoto[i:i + y_gen.shape[1]] += jac.sum(0)
                total[i:i + y_gen.shape[1]] += jac.shape[0]
    if agg == 'mean':
        agg_tanimoto /= total
    if p != 1:
        agg_tanimoto = (agg_tanimoto) ** (1 / p)
    if intdiv:
        return agg_tanimoto
    else:
        return np.mean(agg_tanimoto)


def str2bool(v):
    """
    Converts a string to a boolean.
    
    Args:
        v (str): Input string.
    
    Returns:
        bool: True if the string is 'true' (case insensitive), else False.
    """
    return v.lower() in ('true')


def obey_lipinski(mol):
    """
    Checks if a molecule obeys Lipinski's Rule of Five.
    
    The function evaluates weight, hydrogen bond donors and acceptors, logP, and rotatable bonds.
    
    Args:
        mol (RDKit Mol): Molecule object.
    
    Returns:
        int: Number of Lipinski rules satisfied.
    """
    mol = deepcopy(mol)
    Chem.SanitizeMol(mol)
    rule_1 = Descriptors.ExactMolWt(mol) < 500
    rule_2 = Lipinski.NumHDonors(mol) <= 5
    rule_3 = Lipinski.NumHAcceptors(mol) <= 10
    rule_4 = (logp := Crippen.MolLogP(mol) >= -2) & (logp <= 5)
    rule_5 = Chem.rdMolDescriptors.CalcNumRotatableBonds(mol) <= 10
    return np.sum([int(a) for a in [rule_1, rule_2, rule_3, rule_4, rule_5]])


def obey_veber(mol):
    """
    Checks if a molecule obeys Veber's rules.
    
    Veber's rules focus on the number of rotatable bonds and topological polar surface area.
    
    Args:
        mol (RDKit Mol): Molecule object.
    
    Returns:
        int: Number of Veber's rules satisfied.
    """
    mol = deepcopy(mol)
    Chem.SanitizeMol(mol)
    rule_1 = rdMolDescriptors.CalcNumRotatableBonds(mol) <= 10
    rule_2 = rdMolDescriptors.CalcTPSA(mol) <= 140
    return np.sum([int(a) for a in [rule_1, rule_2]])


def load_pains_filters():
    """
    Loads the PAINS (Pan-Assay INterference compoundS) filters A, B, and C.
    
    Returns:
        FilterCatalog: An RDKit FilterCatalog object containing PAINS filters.
    """
    params = FilterCatalog.FilterCatalogParams()
    params.AddCatalog(FilterCatalog.FilterCatalogParams.FilterCatalogs.PAINS_A)
    params.AddCatalog(FilterCatalog.FilterCatalogParams.FilterCatalogs.PAINS_B)
    params.AddCatalog(FilterCatalog.FilterCatalogParams.FilterCatalogs.PAINS_C)
    catalog = FilterCatalog.FilterCatalog(params)
    return catalog


def is_pains(mol, catalog):
    """
    Checks if the given molecule is a PAINS compound.
    
    Args:
        mol (RDKit Mol): Molecule object.
        catalog (FilterCatalog): A catalog of PAINS filters.
    
    Returns:
        bool: True if the molecule matches a PAINS filter, else False.
    """
    entry = catalog.GetFirstMatch(mol)
    return entry is not None


def mapper(n_jobs):
    """
    Returns a mapping function for parallel or serial processing.
    
    If n_jobs == 1, returns the built-in map function.
    If n_jobs > 1, returns a function that uses a multiprocessing pool.
    
    Args:
        n_jobs (int or pool object): Number of jobs or a Pool instance.
    
    Returns:
        callable: A function that acts like map.
    """
    if n_jobs == 1:
        def _mapper(*args, **kwargs):
            return list(map(*args, **kwargs))
        return _mapper
    if isinstance(n_jobs, int):
        pool = Pool(n_jobs)
        def _mapper(*args, **kwargs):
            try:
                result = pool.map(*args, **kwargs)
            finally:
                pool.terminate()
            return result
        return _mapper
    return n_jobs.map


def fragmenter(mol):
    """
    Fragments a molecule using BRICS and returns a list of fragment SMILES.
    
    Args:
        mol (str or RDKit Mol): Input molecule.
    
    Returns:
        list: List of fragment SMILES strings.
    """
    fgs = AllChem.FragmentOnBRICSBonds(get_mol(mol))
    fgs_smi = Chem.MolToSmiles(fgs).split(".")
    return fgs_smi


def get_mol(smiles_or_mol):
    """
    Loads a SMILES string or molecule into an RDKit molecule object.
    
    Args:
        smiles_or_mol (str or RDKit Mol): SMILES string or molecule.
    
    Returns:
        RDKit Mol or None: Sanitized molecule object or None if invalid.
    """
    if isinstance(smiles_or_mol, str):
        if len(smiles_or_mol) == 0:
            return None
        mol = Chem.MolFromSmiles(smiles_or_mol)
        if mol is None:
            return None
        try:
            Chem.SanitizeMol(mol)
        except ValueError:
            return None
        return mol
    return smiles_or_mol


def compute_fragments(mol_list, n_jobs=1):
    """
    Fragments a list of molecules using BRICS and returns a counter of fragment occurrences.
    
    Args:
        mol_list (list): List of molecules (SMILES or RDKit Mol).
        n_jobs (int): Number of parallel jobs.
    
    Returns:
        Counter: A Counter dictionary mapping fragment SMILES to counts.
    """
    fragments = Counter()
    for mol_frag in mapper(n_jobs)(fragmenter, mol_list):
        fragments.update(mol_frag)
    return fragments


def compute_scaffolds(mol_list, n_jobs=1, min_rings=2):
    """
    Extracts scaffolds from a list of molecules as canonical SMILES.
    
    Only scaffolds with at least min_rings rings are considered.
    
    Args:
        mol_list (list): List of molecules.
        n_jobs (int): Number of parallel jobs.
        min_rings (int): Minimum number of rings required in a scaffold.
    
    Returns:
        Counter: A Counter mapping scaffold SMILES to counts.
    """
    scaffolds = Counter()
    map_ = mapper(n_jobs)
    scaffolds = Counter(map_(partial(compute_scaffold, min_rings=min_rings), mol_list))
    if None in scaffolds:
        scaffolds.pop(None)
    return scaffolds


def get_n_rings(mol):
    """
    Computes the number of rings in a molecule.
    
    Args:
        mol (RDKit Mol): Molecule object.
    
    Returns:
        int: Number of rings.
    """
    return mol.GetRingInfo().NumRings()


def compute_scaffold(mol, min_rings=2):
    """
    Computes the Murcko scaffold of a molecule and returns its canonical SMILES if it has enough rings.
    
    Args:
        mol (str or RDKit Mol): Input molecule.
        min_rings (int): Minimum number of rings required.
    
    Returns:
        str or None: Canonical SMILES of the scaffold if valid, else None.
    """
    mol = get_mol(mol)
    try:
        scaffold = MurckoScaffold.GetScaffoldForMol(mol)
    except (ValueError, RuntimeError):
        return None
    n_rings = get_n_rings(scaffold)
    scaffold_smiles = Chem.MolToSmiles(scaffold)
    if scaffold_smiles == '' or n_rings < min_rings:
        return None
    return scaffold_smiles


class Metric:
    """
    Abstract base class for chemical metrics.
    
    Derived classes should implement the precalc and metric methods.
    """
    def __init__(self, n_jobs=1, device='cpu', batch_size=512, **kwargs):
        self.n_jobs = n_jobs
        self.device = device
        self.batch_size = batch_size
        for k, v in kwargs.items():
            setattr(self, k, v)

    def __call__(self, ref=None, gen=None, pref=None, pgen=None):
        """
        Computes the metric between reference and generated molecules.
        
        Exactly one of ref or pref, and gen or pgen should be provided.
        
        Args:
            ref: Reference molecule list.
            gen: Generated molecule list.
            pref: Precalculated reference metric.
            pgen: Precalculated generated metric.
        
        Returns:
            Metric value computed by the metric method.
        """
        assert (ref is None) != (pref is None), "specify ref xor pref"
        assert (gen is None) != (pgen is None), "specify gen xor pgen"
        if pref is None:
            pref = self.precalc(ref)
        if pgen is None:
            pgen = self.precalc(gen)
        return self.metric(pref, pgen)

    def precalc(self, molecules):
        """
        Pre-calculates necessary representations from a list of molecules.
        Should be implemented by derived classes.
        """
        raise NotImplementedError

    def metric(self, pref, pgen):
        """
        Computes the metric given precalculated representations.
        Should be implemented by derived classes.
        """
        raise NotImplementedError


class FragMetric(Metric):
    """
    Metrics based on molecular fragments.
    """
    def precalc(self, mols):
        return {'frag': compute_fragments(mols, n_jobs=self.n_jobs)}

    def metric(self, pref, pgen):
        return cos_similarity(pref['frag'], pgen['frag'])


class ScafMetric(Metric):
    """
    Metrics based on molecular scaffolds.
    """
    def precalc(self, mols):
        return {'scaf': compute_scaffolds(mols, n_jobs=self.n_jobs)}

    def metric(self, pref, pgen):
        return cos_similarity(pref['scaf'], pgen['scaf'])


def cos_similarity(ref_counts, gen_counts):
    """
    Computes cosine similarity between two molecular vectors.
    
    Args:
        ref_counts (dict): Reference molecular vectors.
        gen_counts (dict): Generated molecular vectors.
    
    Returns:
        float: Cosine similarity between the two molecular vectors.
    """
    if len(ref_counts) == 0 or len(gen_counts) == 0:
        return np.nan
    keys = np.unique(list(ref_counts.keys()) + list(gen_counts.keys()))
    ref_vec = np.array([ref_counts.get(k, 0) for k in keys])
    gen_vec = np.array([gen_counts.get(k, 0) for k in keys])
    return 1 - cos_distance(ref_vec, gen_vec)