# -*- coding: utf-8 -*-
"""Created on 23.10"""
__author__ = "Trygve B. Nomeland"
__email__ = "trygve.borte.nomeland@nmbu.no"

""""
SIR modell av en pandemi:
S (Susceptible) – folk som ikke er eller har vært syke og som kan smittes
I (Infectious) – folk som er syke og som kan smitte andre
R (Recoverd) – folk har vært syke og har blitt friske igjen. Modellen antar at disse ikke kan bli syke igjen.
_y betyr verdi fra forrige steg
"""
def updateSIR(s_y, i_y, r_y, b, k):
    return {'s': s_y - b*s_y*i_y, 'i': i_y + b*s_y*i_y-k*i_y, 'r': r_y + k*i_y}

# Generator som kjører simulasjonen frem til "stop" verdien
def genSIR(s_0, i_0, r_0, stop, b=1/3, k=1/10,):
    SIR = {'s': s_0, 'i': i_0, 'r': r_0}
    yield SIR
    n = 0
    while n < stop:
        n += 1
        SIR = updateSIR(SIR.get('s'), SIR.get('i'), SIR.get('r'), b, k)
        yield SIR

def main():
    f = open('pan.csv', 'w')
    f.write('Day,s@day,i@day,r@day\n')
    for n, S in enumerate(genSIR(1, 10/(5*10**6), 0, stop=120)):
        f.write(f'{n},{S.get('s'):.3e},{S.get('i'):.3e},{S.get('r'):.3e}\n')
    f.close()

if __name__ == '__main__':
    main()