# Path: guacamol/distribution_learning_benchmark.py
# class DistributionLearningBenchmark:
#     """
#     Base class for assessing how well a model is able to generate molecules matching a molecule distribution.
# 
#     Derived class should implement the assess_molecules function
#     """
# 
#     def __init__(self, name: str, number_samples: int) -> None:
#         self.name = name
#         self.number_samples = number_samples
# 
#     @abstractmethod
#     def assess_model(self, model: DistributionMatchingGenerator) -> DistributionLearningBenchmarkResult:
#         """
#         Assess a distribution-matching generator model.
# 
#         Args:
#             model: model to assess
#         """
# 
# class DistributionLearningBenchmarkResult:
#     """
#     Contains the results of a distribution learning benchmark.
# 
#     NB: timing does not make sense since training happens outside of DistributionLearningBenchmark.
#     """
# 
#     def __init__(self, benchmark_name: str, score: float, sampling_time: float, metadata: Dict[str, Any]) -> None:
#         """
#         Args:
#             benchmark_name: name of the distribution-learning benchmark
#             score: benchmark score
#             sampling_time: time for sampling the molecules in seconds
#             metadata: benchmark-specific information
#         """
#         self.benchmark_name = benchmark_name
#         self.score = score
#         self.sampling_time = sampling_time
#         self.metadata = metadata
# 
# 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/data.py
# def get_random_subset(dataset: List[Any], subset_size: int, seed: Optional[int] = None) -> List[Any]:
#     """
#     Get a random subset of some dataset.
# 
#     For reproducibility, the random number generator seed can be specified.
#     Nevertheless, the state of the random number generator is restored to avoid side effects.
# 
#     Args:
#         dataset: full set to select a subset from
#         subset_size: target size of the subset
#         seed: random number generator seed. Defaults to not setting the seed.
# 
#     Returns:
#         subset of the original dataset as a list
#     """
#     if len(dataset) < subset_size:
#         raise Exception(f'The dataset to extract a subset from is too small: '
#                         f'{len(dataset)} < {subset_size}')
# 
#     # save random number generator state
#     rng_state = np.random.get_state()
# 
#     if seed is not None:
#         # extract a subset (for a given training set, the subset will always be identical).
#         np.random.seed(seed)
# 
#     subset = np.random.choice(dataset, subset_size, replace=False)
# 
#     if seed is not None:
#         # reset random number generator state, only if needed
#         np.random.set_state(rng_state)
# 
#     return list(subset)
# 
# Path: guacamol/utils/sampling_helpers.py
# 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


import logging
import os
import pkgutil
import tempfile
import time
import fcd
import numpy as np
from typing import List
from guacamol.distribution_learning_benchmark import DistributionLearningBenchmark, DistributionLearningBenchmarkResult
from guacamol.distribution_matching_generator import DistributionMatchingGenerator
from guacamol.utils.data import get_random_subset
from guacamol.utils.sampling_helpers import sample_valid_molecules
        """
        self.chemnet_model_filename = chemnet_model_filename
        self.sample_size = sample_size
        super().__init__(name='Frechet ChemNet Distance', number_samples=self.sample_size)

        self.reference_molecules = get_random_subset(training_set, self.sample_size, seed=42)

    def assess_model(self, model: DistributionMatchingGenerator) -> DistributionLearningBenchmarkResult:
        chemnet = self._load_chemnet()

        start_time = time.time()
        generated_molecules = sample_valid_molecules(model=model, number_molecules=self.number_samples)
        end_time = time.time()

        if len(generated_molecules) != self.number_samples:
            logger.warning('The model could not generate enough valid molecules.')

        mu_ref, cov_ref = self._calculate_distribution_statistics(chemnet, self.reference_molecules)
        mu, cov = self._calculate_distribution_statistics(chemnet, generated_molecules)

        FCD = fcd.calculate_frechet_distance(mu1=mu_ref, mu2=mu,
                                             sigma1=cov_ref, sigma2=cov)
        score = np.exp(-0.2 * FCD)

        metadata = {
            'number_reference_molecules': len(self.reference_molecules),
            'number_generated_molecules': len(generated_molecules),
            'FCD': FCD
        }


# Path: guacamol/distribution_learning_benchmark.py
# class DistributionLearningBenchmark:
#     """
#     Base class for assessing how well a model is able to generate molecules matching a molecule distribution.
# 
#     Derived class should implement the assess_molecules function
#     """
# 
#     def __init__(self, name: str, number_samples: int) -> None:
#         self.name = name
#         self.number_samples = number_samples
# 
#     @abstractmethod
#     def assess_model(self, model: DistributionMatchingGenerator) -> DistributionLearningBenchmarkResult:
#         """
#         Assess a distribution-matching generator model.
# 
#         Args:
#             model: model to assess
#         """
# 
# class DistributionLearningBenchmarkResult:
#     """
#     Contains the results of a distribution learning benchmark.
# 
#     NB: timing does not make sense since training happens outside of DistributionLearningBenchmark.
#     """
# 
#     def __init__(self, benchmark_name: str, score: float, sampling_time: float, metadata: Dict[str, Any]) -> None:
#         """
#         Args:
#             benchmark_name: name of the distribution-learning benchmark
#             score: benchmark score
#             sampling_time: time for sampling the molecules in seconds
#             metadata: benchmark-specific information
#         """
#         self.benchmark_name = benchmark_name
#         self.score = score
#         self.sampling_time = sampling_time
#         self.metadata = metadata
# 
# 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/data.py
# def get_random_subset(dataset: List[Any], subset_size: int, seed: Optional[int] = None) -> List[Any]:
#     """
#     Get a random subset of some dataset.
# 
#     For reproducibility, the random number generator seed can be specified.
#     Nevertheless, the state of the random number generator is restored to avoid side effects.
# 
#     Args:
#         dataset: full set to select a subset from
#         subset_size: target size of the subset
#         seed: random number generator seed. Defaults to not setting the seed.
# 
#     Returns:
#         subset of the original dataset as a list
#     """
#     if len(dataset) < subset_size:
#         raise Exception(f'The dataset to extract a subset from is too small: '
#                         f'{len(dataset)} < {subset_size}')
# 
#     # save random number generator state
#     rng_state = np.random.get_state()
# 
#     if seed is not None:
#         # extract a subset (for a given training set, the subset will always be identical).
#         np.random.seed(seed)
# 
#     subset = np.random.choice(dataset, subset_size, replace=False)
# 
#     if seed is not None:
#         # reset random number generator state, only if needed
#         np.random.set_state(rng_state)
# 
#     return list(subset)
# 
# Path: guacamol/utils/sampling_helpers.py
# 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/frechet_benchmark.py
import logging
import os
import pkgutil
import tempfile
import time
import fcd
import numpy as np
from typing import List
from guacamol.distribution_learning_benchmark import DistributionLearningBenchmark, DistributionLearningBenchmarkResult
from guacamol.distribution_matching_generator import DistributionMatchingGenerator
from guacamol.utils.data import get_random_subset
from guacamol.utils.sampling_helpers import sample_valid_molecules

        """
        self.chemnet_model_filename = chemnet_model_filename
        self.sample_size = sample_size
        super().__init__(name='Frechet ChemNet Distance', number_samples=self.sample_size)

        self.reference_molecules = get_random_subset(training_set, self.sample_size, seed=42)

    def assess_model(self, model: DistributionMatchingGenerator) -> DistributionLearningBenchmarkResult:
        chemnet = self._load_chemnet()

        start_time = time.time()
        generated_molecules = sample_valid_molecules(model=model, number_molecules=self.number_samples)
        end_time = time.time()

        if len(generated_molecules) != self.number_samples:
            logger.warning('The model could not generate enough valid molecules.')

        mu_ref, cov_ref = self._calculate_distribution_statistics(chemnet, self.reference_molecules)
        mu, cov = self._calculate_distribution_statistics(chemnet, generated_molecules)

        FCD = fcd.calculate_frechet_distance(mu1=mu_ref, mu2=mu,
                                             sigma1=cov_ref, sigma2=cov)
        score = np.exp(-0.2 * FCD)

        metadata = {
            'number_reference_molecules': len(self.reference_molecules),
            'number_generated_molecules': len(generated_molecules),
            'FCD': FCD
        }


        return DistributionLearningBenchmarkResult(benchmark_name=self.name,