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)