# Path: guacamol/distribution_matching_generator.py
# class DistributionMatchingGenerator(metaclass=ABCMeta):
#     """
#     Interface for molecule generators.
#     """
# 
#     @abstractmethod
#     def generate(self, number_samples: int) -> List[str]:
#         """
#         Samples SMILES strings from a molecule generator.
# 
#         Args:
#             number_samples: number of molecules to generate
# 
#         Returns:
#             A list of SMILES strings.
#         """
# 
# Path: guacamol/utils/chemistry.py
# def is_valid(smiles: str):
#     """
#     Verifies whether a SMILES string corresponds to a valid molecule.
# 
#     Args:
#         smiles: SMILES string
# 
#     Returns:
#         True if the SMILES strings corresponds to a valid, non-empty molecule.
#     """
# 
#     mol = Chem.MolFromSmiles(smiles)
# 
#     return smiles != '' and mol is not None and mol.GetNumAtoms() > 0
# 
# def canonicalize(smiles: str, include_stereocenters=True) -> Optional[str]:
#     """
#     Canonicalize the SMILES strings with RDKit.
# 
#     The algorithm is detailed under https://pubs.acs.org/doi/full/10.1021/acs.jcim.5b00543
# 
#     Args:
#         smiles: SMILES string to canonicalize
#         include_stereocenters: whether to keep the stereochemical information in the canonical SMILES string
# 
#     Returns:
#         Canonicalized SMILES string, None if the molecule is invalid.
#     """
# 
#     mol = Chem.MolFromSmiles(smiles)
# 
#     if mol is not None:
#         return Chem.MolToSmiles(mol, isomericSmiles=include_stereocenters)
#     else:
#         return None


from typing import List, Set
from guacamol.distribution_matching_generator import DistributionMatchingGenerator
from guacamol.utils.chemistry import is_valid, canonicalize
def sample_valid_molecules(model: DistributionMatchingGenerator, number_molecules: int, max_tries=10) -> List[str]:
    """
    Sample from the given generator until the desired number of valid molecules
    has been sampled (i.e., ignore invalid molecules).

    Args:
        model: model to sample from
        number_molecules: number of valid molecules to generate
        max_tries: determines the maximum number N of samples to draw, N = number_molecules * max_tries

    Returns:
        A list of number_molecules valid molecules. If this was not possible with the given max_tries, the list may be shorter.
    """

    max_samples = max_tries * number_molecules
    number_already_sampled = 0

    valid_molecules: List[str] = []

    while len(valid_molecules) < number_molecules and number_already_sampled < max_samples:
        remaining_to_sample = number_molecules - len(valid_molecules)

        samples = model.generate(remaining_to_sample)
        number_already_sampled += remaining_to_sample

        valid_molecules += [m for m in samples if is_valid(m)]

    return valid_molecules



# Path: guacamol/distribution_matching_generator.py
# class DistributionMatchingGenerator(metaclass=ABCMeta):
#     """
#     Interface for molecule generators.
#     """
# 
#     @abstractmethod
#     def generate(self, number_samples: int) -> List[str]:
#         """
#         Samples SMILES strings from a molecule generator.
# 
#         Args:
#             number_samples: number of molecules to generate
# 
#         Returns:
#             A list of SMILES strings.
#         """
# 
# Path: guacamol/utils/chemistry.py
# def is_valid(smiles: str):
#     """
#     Verifies whether a SMILES string corresponds to a valid molecule.
# 
#     Args:
#         smiles: SMILES string
# 
#     Returns:
#         True if the SMILES strings corresponds to a valid, non-empty molecule.
#     """
# 
#     mol = Chem.MolFromSmiles(smiles)
# 
#     return smiles != '' and mol is not None and mol.GetNumAtoms() > 0
# 
# def canonicalize(smiles: str, include_stereocenters=True) -> Optional[str]:
#     """
#     Canonicalize the SMILES strings with RDKit.
# 
#     The algorithm is detailed under https://pubs.acs.org/doi/full/10.1021/acs.jcim.5b00543
# 
#     Args:
#         smiles: SMILES string to canonicalize
#         include_stereocenters: whether to keep the stereochemical information in the canonical SMILES string
# 
#     Returns:
#         Canonicalized SMILES string, None if the molecule is invalid.
#     """
# 
#     mol = Chem.MolFromSmiles(smiles)
# 
#     if mol is not None:
#         return Chem.MolToSmiles(mol, isomericSmiles=include_stereocenters)
#     else:
#         return None


# Path: guacamol/utils/sampling_helpers.py
from typing import List, Set
from guacamol.distribution_matching_generator import DistributionMatchingGenerator
from guacamol.utils.chemistry import is_valid, canonicalize

def sample_valid_molecules(model: DistributionMatchingGenerator, number_molecules: int, max_tries=10) -> List[str]:
    """
    Sample from the given generator until the desired number of valid molecules
    has been sampled (i.e., ignore invalid molecules).

    Args:
        model: model to sample from
        number_molecules: number of valid molecules to generate
        max_tries: determines the maximum number N of samples to draw, N = number_molecules * max_tries

    Returns:
        A list of number_molecules valid molecules. If this was not possible with the given max_tries, the list may be shorter.
    """

    max_samples = max_tries * number_molecules
    number_already_sampled = 0

    valid_molecules: List[str] = []

    while len(valid_molecules) < number_molecules and number_already_sampled < max_samples:
        remaining_to_sample = number_molecules - len(valid_molecules)

        samples = model.generate(remaining_to_sample)
        number_already_sampled += remaining_to_sample

        valid_molecules += [m for m in samples if is_valid(m)]

    return valid_molecules



def sample_unique_molecules(model: DistributionMatchingGenerator, number_molecules: int, max_tries=10) -> List[str]: