#!/usr/bin/env python2.7

"""
 Simulation of outbreaks (SIMO)
"""

from SimPy.Simulation import *
import rngstreams
import simo_pb2
import csv

class State(object):

    def update_data(self, infection):
        infection.sim.n_infectives[infection.type] += self.delta_infectives
        infection.sim.n_susceptibles[infection.type] += self.delta_susceptibles
        state_name = self.__class__.__name__
        infection.entry_times[state_name] = infection.sim.now()


class Recovered(State):

    def __init__(self, infection):
        self.wait = None
        self.delta_susceptibles = 0
        self.delta_infectives = -1
        State.__init__(self)

    def update_data(self, infection):
        infection.sim.n_infectives[infection.type] += self.delta_infectives
        infection.sim.n_susceptibles[infection.type] += self.delta_susceptibles
        state_name = self.__class__.__name__
        infection.entry_times[state_name] = infection.sim.now()
        infection.entry_times["Isolated"] = infection.sim.now()

    next_state = None


class Isolated(State):

    def __init__(self, infection):
        self.wait = None
        self.delta_susceptibles = 0
        self.delta_infectives = -1
        State.__init__(self)

    def update_data(self, infection):
        infection.sim.n_infectives[infection.type] += self.delta_infectives
        infection.sim.n_susceptibles[infection.type] += self.delta_susceptibles
        state_name = self.__class__.__name__
        infection.entry_times[state_name] = infection.sim.now()
        infection.entry_times["Recovered"] = infection.sim.now() + infection.isolation_time

    next_state = None


class Infectious(State):

    def __init__(self, infection):
        alpha = infection.sim.alphas[infection.type]
        beta = infection.sim.betas[infection.type]
        maxinftime = infection.sim.maxinftime[infection.type]
        rgamma = infection.sim.rng.gammavariate(alpha, beta)
        if maxinftime >= 0 and maxinftime < rgamma:
            self.wait = maxinftime
            self.next_state = Isolated
            infection.isolation_time = rgamma - maxinftime
        else:
            self.wait = rgamma
            self.next_state = Recovered
        self.delta_susceptibles = 0
        self.delta_infectives = 1
        State.__init__(self)


class Latent(State):

    def __init__(self, infection):
        self.wait = infection.sim.latent_periods[infection.type]
        self.delta_susceptibles = -1
        self.delta_infectives = 0
        State.__init__(self)

    next_state = Infectious


class InfectionCourse(Process):
    def __init__(self, type, sim):
        Process.__init__(self, sim=sim)
        self.type = type
        self.entry_times = {}
  
    def progress(self):
        """Make host progress through series of infection states."""
        sim = self.sim
        self.state = Latent(self)
        while True:
            self.state.update_data(self)
            self.interrupt(sim.transmission)            
            if self.state.wait != None:
                wait = self.state.wait
                yield hold, self, wait
                self.state = self.state.next_state(self)
            else:
                sim.case_histories.obn.append(sim.obn)
                sim.case_histories.outbreak_id.append(sim.id)
                sim.case_histories.type.append(self.type)
                val = self.entry_times["Latent"]
                sim.case_histories.latent_start.append(val)
                val = self.entry_times["Infectious"]
                sim.case_histories.infective_start.append(val)
                val = self.entry_times["Recovered"]
                sim.case_histories.recovered_start.append(val)
                val = self.entry_times["Isolated"]
                sim.case_histories.isolated_start.append(val)
                break


class Transmission(Process):

    def __init__(self, sim):
        Process.__init__(self, sim=sim)
        self.rates = dict.fromkeys(sim.types, 0)

    def run(self):
        sim = self.sim
        while True:
            n_infectives = sum(sim.n_infectives.itervalues())
            for type in self.rates:
                self.rates[type] = (sim.proportionalities[type]
                                    * sim.n_susceptibles[type]
                                    * n_infectives)
            cum_sum = sum(self.rates.itervalues())
            if cum_sum > 0:
                wait = sim.rng.expovariate(cum_sum)
                X = sim.rng.random() * cum_sum
                running_sum = 0
                for type in self.rates:
                    running_sum += self.rates[type]
                    if running_sum >= X:
                        yield hold, self, wait
                        if not self.interrupted():
                            infection = InfectionCourse(type=type,
                                                        sim=sim)
                            sim.activate(infection,
                                         infection.progress())
                            break
                        else:
                            self.interruptReset()
                            break
            else: # Sit and wait to be interuppted
                yield hold, self, sim.max_time + 1
                self.interruptReset()


class PopulationModel(Simulation):

    def __init__(self, rng, id=0, types=['Patient', 'Staff'],
                 initial_cases=[1, 0], susceptibles=[100, 100],
                 proportionalities=[0.005, 0.0025], max_time=100,
                 maxinftime=[-1, -1],
                 obn=0L, alphas=[3, 3], betas=[0.67, 0.67],
                 latent_periods=[1, 1], output=simo_pb2.SimOutput()):
        Simulation.__init__(self)
        self.types = types
        self.n_susceptibles = dict(zip(types, susceptibles))
        self.n_infectives = dict.fromkeys(types, 0)
        self.initial_cases = dict(zip(types, initial_cases))
        self.proportionalities = dict(zip(types, proportionalities))
        self.alphas = dict(zip(types, alphas))
        self.betas = dict(zip(types, betas))
        self.latent_periods = dict(zip(types, latent_periods))
        self.case_histories = output
        self.max_time = max_time
        self.maxinftime = dict(zip(types, maxinftime))
        self.obn = obn
        self.id = id
        self.rng = rng

    def run(self):
        self.initialize()
        self.transmission = Transmission(sim=self)
        self.activate(self.transmission, self.transmission.run())
        for type in self.initial_cases:
            for count in range(self.initial_cases[type]):
                infection = InfectionCourse(type=type, sim=self)
                self.n_susceptibles[type] += 1
                self.activate(infection, infection.progress())
        self.simulate(until=self.max_time)

def get_input(path):
    f = open(path, "rb")
    ss = simo_pb2.SimSet()
    ss.ParseFromString(f.read())
    f.close()
    return ss

def put_output(path, output):
    f = open(path, "wb")
    f.write(output.SerializeToString())
    f.close()

def main():
    simset = get_input("simo-in.pb")
    output = simo_pb2.SimOutput()
    rng = rngstreams.sprng()
    rng.jumpasubstream(simset.jump)
    i = 1
    for si in simset.siminput:
        experi = PopulationModel(rng, id=i, types=si.type,
                                 initial_cases=si.initial_cases,
                                 susceptibles=si.susceptibles,
                                 proportionalities=si.proportionalities,
                                 latent_periods=si.latent_periods,
                                 max_time=si.max_time, obn=si.obn,
                                 alphas=si.alphas, maxinftime=si.maxinftime,
                                 betas=si.betas, output=output)
        experi.run()
        i += 1
    put_output("simo-out.pb", output)

if __name__ == '__main__':
  main()
