'''
Execute by running: ``python vivarium/process/template_process.py``

TODO: Replace the template code to implement your own process.
'''

import os

from vivarium.core.process import Process
from vivarium.core.directories import PROCESS_OUT_DIR
from vivarium.core.engine import Engine
from vivarium.plots.simulation_output import plot_simulation_output


# a global NAME is used for the output directory and for the process name
NAME = 'template'


class Template(Process):
    '''
    This mock process provides a basic template that can be used for a new process
    '''

    # give the process a name, so that it can register in the process_repository
    name = NAME

    # declare default parameters as class variables
    defaults = {
        'uptake_rate': 0.1,
    }

    def __init__(self, parameters=None):
        # parameters passed into the constructor merge with the defaults
        # and can be access through the self.parameters class variable
        super(Template, self).__init__(parameters)

    def ports_schema(self):
        '''
        ports_schema returns a dictionary that declares how each state will behave.
        Each key can be assigned settings for the schema_keys declared in Store:

        * `_default`
        * `_updater`
        * `_divider`
        * `_value`
        * `_properties`
        * `_emit`
        * `_serializer`
        '''

        return {
            'internal': {
                'A': {
                    '_default': 1.0,
                    '_updater': 'accumulate',
                    '_emit': True,
                }
            },
            'external': {
                'A': {
                    '_default': 1.0,
                    '_updater': 'accumulate',
                    '_emit': True,
                }
            },
        }

    def next_update(self, timestep, states):

        # get the states
        internal_a = states['internal']['A']
        external_a = states['external']['A']

        # calculate timestep-dependent updates
        internal_update = self.parameters['uptake_rate'] * external_a * timestep
        external_update = -1 * internal_update

        # return an update that mirrors the ports structure
        return {
            'internal': {
                'A': internal_update},
            'external': {
                'A': external_update}
        }


# functions to configure and run the process
def run_template_process():
    '''Run a simulation of the process.

    Returns:
        The simulation output.
    '''

    # initialize the process by passing in parameters
    parameters = {}
    template_process = Template(parameters)

    # declare the initial state, mirroring the ports structure
    initial_state = {
        'internal': {
            'A': 0.0
        },
        'external': {
            'A': 1.0
        },
    }

    # run the simulation
    composite = template_process.generate()
    sim = Engine(composite=composite, initial_state=initial_state)
    sim.update(10)
    output = sim.emitter.get_timeseries()

    return output


def test_template_process():
    '''Test that the process runs correctly.

    This will be executed by pytest.
    '''
    output = run_template_process()
    # TODO: Add assert statements to ensure correct performance.


def main():
    '''Simulate the process and plot results.'''
    # make an output directory to save plots
    out_dir = os.path.join(PROCESS_OUT_DIR, NAME)
    os.makedirs(out_dir, exist_ok=True)

    output = run_template_process()

    # plot the simulation output
    plot_settings = {}
    plot_simulation_output(output, plot_settings, out_dir)


# run module with python vivarium/process/template_process.py
if __name__ == '__main__':
    main()
