# Monitor the progress of an epidemic
#
# Copyright (C) 2020--2021 Simon Dobson
#
# This file is part of epydemic, epidemic network simulations in Python.
#
# epydemic 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.
#
# epydemic 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 epydemic. If not, see <http://www.gnu.org/licenses/gpl.html>.

import sys
from typing import Dict, Any
if sys.version_info >= (3, 8):
    from typing import Final
else:
    # backport compatibility with older typing
    from typing_extensions import Final
from epydemic import Process


class Monitor(Process):
    '''Add progress monitoring to a process. This class captures the sizes
    of all loci in the process at regular intervals, returning them as
    time series. Each time series appears as an array in the results dict, keyed
    by :attr:`OBSERVATIONS` (for the simulation times at which the observations
    were taken) or by a name related to the locus name, generated by
    :meth:`timeSeriesForLocus`. This means that each locus has its own
    time series tracking its evolving size.

    '''

    # Experimental parameters
    DELTA: Final[str] = "epydemic.monitor.time_delta"             #: Parameter for the time interval for observations.

    # Results
    OBSERVATIONS: Final[str] = 'epydemic.monitor.observations'    #: Result holding the times of the observations.

    # Result name stem
    TIMESERIES_STEM: Final[str] = "epydemic.monitor.timeseries"   #: Stem for names of the timeseries for different loci.

    @staticmethod
    def timeSeriesForLocus(l: str) -> str:
        '''Generate the name of the time series corresponding to the
        given locus.

        :param l: the locus
        :returns: the result name'''
        return '{stem}-{locus}'.format(stem=Monitor.TIMESERIES_STEM, locus=l)

    def __init__(self):
        super().__init__()

    def reset(self):
        '''Reset the process.'''
        super().reset()
        self._timeSeries = None

    def build(self, params: Dict[str, Any]):
        '''Build the observation process.

        :param params: the experimental parameters

        '''
        super().build(params)

        # post a repeating event to observe the process
        delta = params[self.DELTA]
        self.postRepeatingEvent(0.0, delta, None, self.observe)

    def observe(self, t: float, e: Any):
        '''Observe the network, capturing the sizes of all loci which are then
        stored into individual time series. Another time series,
        tagged :attr:`OBSERVATIONS`, is created to store the sequence
        of times at which observations were made.

        :param t: the current simulation time
        :param e: the element (ignored)

        '''

        # if this is the first run, build the initial dicts: one
        # per locus and one for the observation times
        # (We can't do this as part of build() as it depends on which
        # other processes we're composed with, and in what order)
        if self._timeSeries is None:
            self._timeSeries = dict()
            self._timeSeries[self.OBSERVATIONS] = []
            for n in self.dynamics().loci().keys():   # loci for all processes in this experiment
                self._timeSeries[Monitor.timeSeriesForLocus(n)] = []

        # make the observation
        self._timeSeries[self.OBSERVATIONS].append(t)
        for (n, l) in self.dynamics().loci().items():
            self._timeSeries[Monitor.timeSeriesForLocus(n)].append(len(l))

    def results(self) -> Dict[str, Any]:
        '''Return the time series tagged as specified by
        :meth:`timeSeriesForLocus`. There is one time series *per*
        locus, plus one for the sequence of times at which the
        observations were made.

        :returns: the results

        '''
        res = super().results()

        # store the time series
        for n in self._timeSeries:
            res[n] = self._timeSeries[n]
        return res
