import csv
import math
import random
import matplotlib.pyplot as plt

class Mensch(object):

    def __init__(self, alter, geschlecht, bdl):
        self.Alter = alter
        self.Geschlecht = geschlecht #0 Mann, 1 Frau
        self.Bundesland = bdl
        self.Koordinaten = self.getCoordinates(bdl)
        self.Infiziert = False
        self.Stichtag = None


    def __del__(self):
        self.Alter = None
        self.Geschlecht = None
        self.Bundesland = None
        self.Koordinaten = None

    def getCoordinates(self, bdl):
        bdlDict = {'Vorarlberg':[0,0], 'Tirol':[1,0], 'Kärnten':[2,0], 'Salzburg':[0,1], 'Steiermark':[1,1], 'Burgenland':[2,1], 'Oberösterreich':[0,2] ,'Niederösterreich':[1,2], 'Wien':[2,2]}

        return [int(bdlDict[bdl][0]) + random.random(), int(bdlDict[bdl][1]) + random.random()]

    def die(self):
        randDie = random.random()
        infectFaktor = 1
        genderFaktor = 0.493
        if self.Geschlecht == 1:
            genderFaktor = 0.507
        if self.Infiziert:
            infectFaktor = 5
        if (randDie < (int(gestDict[self.Bundesland][int(self.Geschlecht)])*infectFaktor/(int(bevDict[self.Bundesland])*365*genderFaktor))):
            return True
        else:
            return False

    def infect(self, day):
        if (int(day) - int(self.Stichtag)) > 10:
            self.Infiziert = False
        else:
            for i in range(20):
                randIndex = random.randint(0, len(Population)-1)
                randInfect = random.random()
                if(self.Bundesland == Population[randIndex].Bundesland):
                    Distanz = math.sqrt((self.Koordinaten[0] - Population[randIndex].Koordinaten[0]) * (
                                self.Koordinaten[0] - Population[randIndex].Koordinaten[0]) + (
                                                    self.Koordinaten[1] - Population[randIndex].Koordinaten[1]) * (
                                                    self.Koordinaten[1] - Population[randIndex].Koordinaten[1]))
                    InfectProb = 0.3 * (1 - Distanz)
                else:
                    InfectProb = 0.0015
                if(randInfect< InfectProb):
                    Population[randIndex].Infiziert = True
                    Population[randIndex].Stichtag = day

class Simulation(object):

    def __init__(self, population):
        self.Population = population

    def simulate(self):

        print(len(Population))

        RandomStartInfekt = random.randint(0, len(Population)-1)
        Population[RandomStartInfekt].Infiziert = True
        Population[RandomStartInfekt].Stichtag = 0
        RandomStartInfekt = random.randint(0, len(Population)-1)
        Population[RandomStartInfekt].Infiziert = True
        Population[RandomStartInfekt].Stichtag = 0
        RandomStartInfekt = random.randint(0, len(Population)-1)
        Population[RandomStartInfekt].Infiziert = True
        Population[RandomStartInfekt].Stichtag = 0

        plt.ion()

        for i in range(365):
            print("Today is day number "+ str(i))
            for key in gebDict.keys():
                for j in range(int(abs(random.gauss((int(gebDict[key][0])*Populationsfaktor/(365)), 1)))):
                     self.Population.append(Mensch(0, 0, key))
                for j in range(int(abs(random.gauss((int(gebDict[key][1])*Populationsfaktor / (365)), 1)))):
                    self.Population.append(Mensch(0, 1, key))

            print(len(Population))

            deathlist = []
            for human in Population:
                if(human.die()):
                    deathlist.append(Population.index(human))
                if(human.Infiziert):
                    human.infect(i)

            deathlist.sort(reverse=True)
            for index in deathlist:
                Population.remove(Population[index])



            #Visualisierung
            #Scatter Plot
            x_inf = []
            y_inf = []
            x_fit = []
            y_fit = []

            for human in Population:
                if(human.Infiziert):
                    #plt.plot(human.Koordinaten[0], human.Koordinaten[1], "o", color="red",  markersize=2)
                    x_inf.append(human.Koordinaten[0])
                    y_inf.append(human.Koordinaten[1])
                else:
                    #plt.plot(human.Koordinaten[0], human.Koordinaten[1], "o", color="blue", markersize=1)
                    x_fit.append(human.Koordinaten[0])
                    y_fit.append(human.Koordinaten[1])

            plt.title("Population")
            plt.scatter(x_inf, y_inf, c="red", s=5)
            plt.scatter(x_fit, y_fit, c="blue", s=2)

            plt.draw()
            if(i % 5 == 0):
                plt.savefig('result/populationOnDay' + str(i) + '.png')
            plt.pause(0.0001)
            plt.clf()



        print(len(Population))



if __name__ == '__main__':


    print("Generate Population:")
    Populationsfaktor = 0.01
    Population = []
    bevAnzahl = 0
    bevDict = {}
    gebDict = {}
    gestDict = {}
    AnteilFrauen = 0.507


    #read data from Statistik Austria
    with open('./data/Bevölkerung.csv') as bevoelkerung:
        csv_reader = csv.reader(bevoelkerung, delimiter = ';')
        rowcnt = 0
        for row in csv_reader:
            rowcnt = rowcnt + 1
            if rowcnt == 2:
                bevAnzahl = int(int(row[-1])*Populationsfaktor)
            if rowcnt > 2:
                bevDict[row[0]] = row[-1]
    with open('./data/Geburten2021.csv') as geburten:
        csv_reader = csv.reader(geburten, delimiter=';')
        rowcnt = 0
        for row in csv_reader:
            rowcnt = rowcnt + 1
            if rowcnt > 2:
                gebDict[row[0]] = [row[-2], row[-1]]

    with open('./data/Gestorbene2021.csv') as gestorbene:
        csv_reader = csv.reader(gestorbene, delimiter=';')
        rowcnt = 0
        for row in csv_reader:
            rowcnt = rowcnt + 1
            if rowcnt > 2:
                gestDict[row[0]] = [row[-2], row[-1]]

    print(bevDict)
    print(gebDict)
    print(gestDict)
    print(bevAnzahl)

    bdls = []
    weights = []
    for key in bevDict.keys():
        bdls.append(str(key))
        weights.append(int(bevDict[key]) * Populationsfaktor / int(bevAnzahl))


    for i in range(int(bevAnzahl)):
        gender = 1
        if(random.random()>0.507):
            gender = 0
        age = random.randint(0, 101)

        bdl = random.choices(bdls, weights)[0]

        Population.append(Mensch(age, gender, bdl))

    Simulation = Simulation(Population)
    Simulation.simulate()