# Experiment combinator to run the same experiment several times.
#
# Copyright (C) 2016--2022 Simon Dobson
#
# This file is part of epyc, experiment management in Python.
#
# epyc is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# epyc is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with epyc. If not, see <http://www.gnu.org/licenses/gpl.html>.

from epyc import ExperimentCombinator, Experiment, ResultsDict
import sys
from typing import Any, Dict, List, Union
if sys.version_info >= (3, 8):
    from typing import Final
else:
    # backwards compatibility with Python 35, Python36, and Python37
    from typing_extensions import Final


class RepeatedExperiment(ExperimentCombinator):
    """A experiment combinator that takes a "base" experiment and runs it
    several times. This means you can define a single experiment separate
    from its repeating logic.

    When run, a repeated experiment runs a number of repetitions of the underlying
    experiment at the same point in the parameter space. The result of the
    repeated experiment is the list of results from the underlying experiment.
    If the underlying experiment itself returns a list of results, these are all
    flattened into a single list.

    :param ex: the underlying experiment
    :pamam N: the number of repetitions to perform"""

    # Extra metadata
    REPETITIONS: Final[str] = 'epyc.experiment.repetitions'   #: Metadata element for number of repetitions performed.
    I          : Final[str] = 'epyc.experiment.rep'           #: Metadata element for index in the repetitions.


    def __init__(self, ex: Experiment, N: int ):
        super().__init__(ex)
        self._N = N

    def repetitions(self) -> int:
        """Return the number of repetitions of the underlying experiment
        we expect to perform.

        :returns: the number of repetitions"""
        return self._N

    def do(self, params: Dict[str, Any]) -> List[ResultsDict]:
        """Perform the number of repetitions we want. The results returned
        will be a list of the results dicts generated by the repeated experiments.
        The metadata for each experiment will include an entry
        :attr:`RepeatedExperiment.REPETITIONS` for the number of
        repetitions that occurred (which will be the length of this list) and
        an entry :attr:`RepeatedExperiment.I` for the index of the result in that
        sequence.

        :param params: the parameters to the experiment
        :returns: a list of result dicts"""
        N = self.repetitions()
        e = self.experiment()
        results = []
        for i in range(N):
            rcs = e.run()

            # make sure we have a list to traverse
            if not isinstance(rcs, list):
                rcs = [rcs]

            # add repetition information to each results dict's metadata
            for j in range(len(rcs)):
                rc = rcs[j]
                rc[Experiment.METADATA][self.I] = i
                rc[Experiment.METADATA][self.REPETITIONS] = N
            results.extend(rcs)
        return results
