<commit_before>"""The engine."""
import os
import os.path
import argparse
from time import time
from collections import Iterable
from datetime import datetime, timedelta
from pprint import pformat
import gc
from bdb import BdbQuit

import pandas as pd

from vivarium import config

from vivarium.framework.values import ValuesManager
from vivarium.framework.event import EventManager, Event, emits
from vivarium.framework.population import PopulationManager, creates_simulants
from vivarium.framework.lookup import InterpolatedDataManager
from vivarium.framework.components import load, read_component_configuration
from vivarium.framework.randomness import RandomnessStream
from vivarium.framework.util import collapse_nested_dict

import logging
_log = logging.getLogger(__name__)


class Builder:
    def __init__(self, context):
        self.lookup = context.tables.build_table
        self.value = context.values.get_value
        self.rate = context.values.get_rate
        self.modifies_value = context.values.mutator
        self.emitter = context.events.get_emitter
        self.population_view = context.population.get_view
        self.clock = lambda: lambda: context.current_time
        input_draw_number = config.run_configuration.draw_number
        model_draw_number = config.run_configuration.model_draw_number
        self.randomness = lambda key: RandomnessStream(key, self.clock(), (input_draw_number, model_draw_number))

    def __repr__(self):
        return ("Builder(\nlookup: {},\nvalue: {},\nrate: {},\n".format(self.lookup, self.value, self.rate)
                + "emitter: {},\npopulation_view: {},\n".format(self.emitter, self.population_view)
                + "clock: {},\nrandomness: {}\n)".format(self.clock, self.randomness))


class SimulationContext:
    """context"""
    def __init__(self, components):
        self.components = components
        self.values = ValuesManager()
        self.events = EventManager()
        self.population = PopulationManager()
        self.tables = InterpolatedDataManager()
        self.components.extend([self.tables, self.values, self.events, self.population])
        self.current_time = None

    def setup(self):
        builder = Builder(self)
        components = [self.values, self.events, self.population, self.tables] + list(self.components)
        done = set()

        i = 0
        while i < len(components):
            component = components[i]
            if component is None:
                raise ValueError('None in component list. This likely indicates a bug in a factory function')

            if isinstance(component, Iterable):
                # Unpack lists of components so their constituent components get initialized
                components.extend(component)
            if component not in done:
                if hasattr(component, 'configuration_defaults'):
                    # This reapplies configuration from some components but
                    # that shouldn't be a problem.
                    config.read_dict(component.configuration_defaults, layer='component_configs', source=component)
                if hasattr(component, 'setup'):
                    sub_components = component.setup(builder)
                    done.add(component)
                    if sub_components:
                        components.extend(sub_components)
            i += 1
        self.values.setup_components(components)
        self.events.setup_components(components)
        self.population.setup_components(components)

        self.events.get_emitter('post_setup')(None)

    def __repr__(self):
        return ("SimulationContext(\ncomponents: {},\nvalues: {},\n".format(self.components, self.values)
                + "events: {},\npopulation: {},\ntables: {},\n".format(self.events, self.population, self.tables)
                + "current_time: {})".format(self.current_time))


@emits('time_step')
@emits('time_step__prepare')
@emits('time_step__cleanup')
def _step(simulation, time_step, time_step_emitter, time_step__prepare_emitter, time_step__cleanup_emitter):
    _log.debug(simulation.current_time)
    time_step__prepare_emitter(Event(simulation.population.population.index))
    time_step_emitter(Event(simulation.population.population.index))
    time_step__cleanup_emitter(Event(simulation.population.population.index))
    simulation.current_time += time_step


def _get_time(suffix):
    params = config.simulation_parameters
    month, day = 'month' + suffix, 'day' + suffix
    if month in params or day in params:
        if not (month in params and day in params):
            raise ValueError("you must either specify both a month {0} and a day {0} or neither".format(suffix))
        return datetime(params['year_{}'.format(suffix)], params[month], params.day_start)
    else:
        return datetime(params['year_{}'.format(suffix)], 7, 2)


@creates_simulants
@emits('simulation_end')
def event_loop(simulation, simulant_creator, end_emitter):
    start = _get_time('start')
    stop = _get_time('end')
    time_step = config.simulation_parameters.time_step
    time_step = timedelta(days=time_step)

    simulation.current_time = start

    population_size = config.simulation_parameters.population_size

    if config.simulation_parameters.initial_age is not None and config.simulation_parameters.pop_age_start is None:
        simulant_creator(population_size, population_configuration={
            'initial_age': config.simulation_parameters.initial_age})
    else:
        simulant_creator(population_size)

    while simulation.current_time < stop:
        gc.collect()  # TODO: Actually figure out where the memory leak is.
        _step(simulation, time_step)

    end_emitter(Event(simulation.population.population.index))


def setup_simulation(components):
    if not components:
        components = []
    simulation = SimulationContext(load(components + [_step, event_loop]))

    simulation.setup()

    return simulation


def run_simulation(simulation):
    start = time()

    event_loop(simulation)

    metrics = simulation.values.get_value('metrics')
    metrics.source = lambda index: {}
    metrics = metrics(simulation.population.population.index)
    metrics['simulation_run_time'] = time() - start
    return metrics


def configure(input_draw_number=None, model_draw_number=None, verbose=False, simulation_config=None):
    if simulation_config:
        if isinstance(simulation_config, dict):
            config.read_dict(simulation_config)
        else:
            config.read(simulation_config)

    if input_draw_number is not None:
        config.run_configuration.set_with_metadata('draw_number', input_draw_number,
                                                   layer='override', source='command_line_argument')
    if model_draw_number is None:<commit_msg>Revert commits on the wrong branch<commit_after>"""The engine."""
import os
import os.path
import argparse
from time import time
from collections import Iterable
from datetime import datetime, timedelta
from pprint import pformat
import gc
from bdb import BdbQuit

import pandas as pd

from vivarium import config

from vivarium.framework.values import ValuesManager
from vivarium.framework.event import EventManager, Event, emits
from vivarium.framework.population import PopulationManager, creates_simulants
from vivarium.framework.lookup import InterpolatedDataManager
from vivarium.framework.components import load, read_component_configuration
from vivarium.framework.randomness import RandomnessStream
from vivarium.framework.util import collapse_nested_dict

import logging
_log = logging.getLogger(__name__)


class Builder:
    def __init__(self, context):
        self.lookup = context.tables.build_table
        self.value = context.values.get_value
        self.rate = context.values.get_rate
        self.modifies_value = context.values.mutator
        self.emitter = context.events.get_emitter
        self.population_view = context.population.get_view
        self.clock = lambda: lambda: context.current_time
        input_draw_number = config.run_configuration.draw_number
        model_draw_number = config.run_configuration.model_draw_number
        self.randomness = lambda key: RandomnessStream(key, self.clock(), (input_draw_number, model_draw_number))

    def __repr__(self):
        return ("Builder(\nlookup: {},\nvalue: {},\nrate: {},\n".format(self.lookup, self.value, self.rate)
                + "emitter: {},\npopulation_view: {},\n".format(self.emitter, self.population_view)
                + "clock: {},\nrandomness: {}\n)".format(self.clock, self.randomness))


class SimulationContext:
    """context"""
    def __init__(self, components):
        self.components = components
        self.values = ValuesManager()
        self.events = EventManager()
        self.population = PopulationManager()
        self.tables = InterpolatedDataManager()
        self.components.extend([self.tables, self.values, self.events, self.population])
        self.current_time = None

    def setup(self):
        builder = Builder(self)
        components = [self.values, self.events, self.population, self.tables] + list(self.components)
        done = set()

        i = 0
        while i < len(components):
            component = components[i]
            if component is None:
                raise ValueError('None in component list. This likely indicates a bug in a factory function')

            if isinstance(component, Iterable):
                # Unpack lists of components so their constituent components get initialized
                components.extend(component)
            if component not in done:
                if hasattr(component, 'configuration_defaults'):
                    # This reapplies configuration from some components but
                    # that shouldn't be a problem.
                    config.read_dict(component.configuration_defaults, layer='component_configs', source=component)
                if hasattr(component, 'setup'):
                    sub_components = component.setup(builder)
                    done.add(component)
                    if sub_components:
                        components.extend(sub_components)
            i += 1
        self.values.setup_components(components)
        self.events.setup_components(components)
        self.population.setup_components(components)

        self.events.get_emitter('post_setup')(None)

    def __repr__(self):
        return ("SimulationContext(\ncomponents: {},\nvalues: {},\n".format(self.components, self.values)
                + "events: {},\npopulation: {},\ntables: {},\n".format(self.events, self.population, self.tables)
                + "current_time: {})".format(self.current_time))


@emits('time_step')
@emits('time_step__prepare')
@emits('time_step__cleanup')
def _step(simulation, time_step, time_step_emitter, time_step__prepare_emitter, time_step__cleanup_emitter):
    _log.debug(simulation.current_time)
    time_step__prepare_emitter(Event(simulation.population.population.index))
    time_step_emitter(Event(simulation.population.population.index))
    time_step__cleanup_emitter(Event(simulation.population.population.index))
    simulation.current_time += time_step


def _get_time(suffix):
    params = config.simulation_parameters
    month, day = 'month' + suffix, 'day' + suffix
    if month in params or day in params:
        if not (month in params and day in params):
            raise ValueError("you must either specify both a month {0} and a day {0} or neither".format(suffix))
        return datetime(params['year_{}'.format(suffix)], params[month], params.day_start)
    else:
        return datetime(params['year_{}'.format(suffix)], 7, 2)


@creates_simulants
@emits('simulation_end')
def event_loop(simulation, simulant_creator, end_emitter):
    start = _get_time('start')
    stop = _get_time('end')
    time_step = config.simulation_parameters.time_step
    time_step = timedelta(days=time_step)

    simulation.current_time = start

    population_size = config.simulation_parameters.population_size

    if config.simulation_parameters.initial_age is not None and config.simulation_parameters.pop_age_start is None:
        simulant_creator(population_size, population_configuration={
            'initial_age': config.simulation_parameters.initial_age})
    else:
        simulant_creator(population_size)

    while simulation.current_time < stop:
        gc.collect()  # TODO: Actually figure out where the memory leak is.
        _step(simulation, time_step)

    end_emitter(Event(simulation.population.population.index))


def setup_simulation(components):
    if not components:
        components = []
    simulation = SimulationContext(load(components + [_step, event_loop]))

    simulation.setup()

    return simulation


def run_simulation(simulation):
    start = time()

    event_loop(simulation)

    metrics = simulation.values.get_value('metrics')
    metrics.source = lambda index: {}
    metrics = metrics(simulation.population.population.index)
    metrics['simulation_run_time'] = time() - start
    return metrics


def configure(input_draw_number=None, model_draw_number=None, verbose=False, simulation_config=None):
    if simulation_config:
        if isinstance(simulation_config, dict):
            config.read_dict(simulation_config)
        else:
            config.read(simulation_config)

    if input_draw_number is not None:
        config.run_configuration.set_with_metadata('draw_number', input_draw_number,
                                                   layer='override', source='command_line_argument')
    if model_draw_number None:
{'fraction_non_alphanumeric': 0.09352623118515271, 'fraction_numerical': 0.0008768084173608067, 'mean_word_length': 3.3135833595965964, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3802621', 'n_tokens_mistral': 3796, 'n_tokens_neox': 3610, 'n_words': 845}