
from .._network import Network
from .._outputfiles import OutputFiles
from ._workspace import Workspace
from .._population import Population, Populations
from ._profiler import Profiler, NullProfiler
from ._iterate import iterate
from ..iterators._iterate_default import iterate_default
from ..extractors._extract_default import extract_default
from ._clear_all_infections import clear_all_infections
from ._extract_data import extract_data

__all__ = ["run_model"]


def run_model(network: Network,
              infections, play_infections,
              rngs, s: int,
              output_dir: OutputFiles,
              population: Population = Population(initial=57104043),
              nsteps: int = None,
              profile: bool = True,
              profiler: Profiler = None,
              nthreads: int = None,
              iterator=None,
              extractor=None):
    """Actually run the model... Real work happens here. The model
       will run until completion or until 'nsteps' have been
       completed (whichever happens first)

        Parameters
        ----------
        network: Network
            The network on which to run the model
        infections: list
            The space used to record the infections
        play_infections: list
            The space used to record the play infectionss
        rngs: list
            The list of random number generators to use, one per thread
        population: Population
            The initial population at the start of the model outbreak.
            This is also used to set the date and day of the start of
            the model outbreak
        seed: int
            The random number seed used for this model run. If this is
            None then a very random random number seed will be used
        output_dir: OutputFiles
            The directory to write all of the output into
        nsteps: int
            The maximum number of steps to run in the outbreak. If None
            then run until the outbreak has finished
        profile: bool
            Whether or not to profile the model run and print out the
            results
        profiler: Profiler
            The profiler to use to profile - a new one is created if
            one isn't passed
        s: int
            Index of the seeding parameter to use
        nthreads: int
            Number of threads over which to parallelise this model run
        iterator: function
            Function that will be used to dynamically get the functions
            that will be used at each iteration to advance the
            model. Any additional files or parameters needed by these
            functions should be included in the `network.params` object.
        extractor: function
            Function that will be used to dynamically get the functions
            that will be used at each iteration to extract data from
            the model run

        Returns
        -------
        trajectory: Populations
            The trajectory of the population for every day of the model run
    """
    if iterator is None:
        iterator = iterate_default
    elif isinstance(iterator, str):
        from ..iterators._iterate_custom import build_custom_iterator
        iterator = build_custom_iterator(iterator, __name__)

    if extractor is None:
        extractor = extract_default
    elif isinstance(iterator, str):
        from ..extractors._extract_custom import build_custom_extractor
        extractor = build_custom_extractor(extractor, __name__)

    if profile:
        if profiler:
            p = profiler
        else:
            p = Profiler()
    else:
        p = NullProfiler()

    p = p.start("run_model")

    params = network.params

    if params is None:
        return population

    from copy import deepcopy
    population = deepcopy(population)

    # create space to hold the population trajectory
    trajectory = Populations()

    p = p.start("clear_all_infections")
    clear_all_infections(infections=infections,
                         play_infections=play_infections,
                         nthreads=nthreads)
    p = p.stop()

    # get and call all of the functions that need to be called to set
    # up the model run
    p = p.start("setup_funcs")
    setup_funcs = iterator(nthreads=nthreads, setup=True)

    for setup_func in setup_funcs:
        setup_func(network=network, population=population,
                   infections=infections, play_infections=play_infections,
                   rngs=rngs, profiler=p, nthreads=nthreads)

    setup_funcs = extractor(nthreads=nthreads, setup=True)

    for setup_func in setup_funcs:
        setup_func(network=network, population=population,
                   output_dir=output_dir, profiler=p, nthreads=nthreads)
    p = p.stop()

    # create a workspace that is used as part of extract_data to
    # provide a scratch-pad while extracting data from the model
    workspace = Workspace(network=network)

    # Now get the population and network data for the first day of the
    # model ("day zero", unless a future day has been set by the user)
    p = p.start("extract_data")
    extract_data(network=network, population=population, workspace=workspace,
                 output_dir=output_dir, infections=infections,
                 play_infections=play_infections,
                 rngs=rngs, get_output_functions=extractor,
                 nthreads=nthreads, profiler=p)
    p = p.stop()

    infecteds = population.infecteds

    # save the initial population
    trajectory.append(population)

    p = p.start("run_model_loop")
    iteration_count = 0

    # keep looping until the outbreak is over or until we have completed
    # at least 5 loop iterations
    while (infecteds != 0) or (iteration_count < 5):
        if profile:
            p2 = Profiler()
        else:
            p2 = NullProfiler()

        p2 = p2.start(f"timing for day {population.day}")

        iterate(network=network, population=population,
                infections=infections, play_infections=play_infections,
                rngs=rngs, get_advance_functions=iterator,
                nthreads=nthreads, profiler=p2)

        print(f"\n {population.day} {infecteds}")

        p2 = p2.start("extract_data")
        extract_data(network=network, population=population,
                     workspace=workspace, output_dir=output_dir,
                     infections=infections, play_infections=play_infections,
                     rngs=rngs, get_output_functions=extractor,
                     nthreads=nthreads, profiler=p2)
        p2 = p2.stop()

        infecteds = population.infecteds

        iteration_count += 1
        population.day += 1

        if population.date:
            from datetime import timedelta
            population.date += timedelta(days=1)

        if nsteps is not None:
            if iteration_count >= nsteps:
                trajectory.append(population)
                print(f"Exiting model run early at nsteps = {nsteps}")
                break

        p2 = p2.stop()

        if not p2.is_null():
            print(f"\n{p2}\n")

        # save the population trajectory
        trajectory.append(population)

    # end of while loop
    p = p.stop()
    p.stop()

    if not p.is_null():
        print(f"\nOVERALL MODEL TIMING\n{p}")

    print(f"Infection died ... Ending on day {population.day}")

    return trajectory
