import sys
#import and init pygame
import pygame
from random import randint as nahodne 
pygame.init() 

sirka = 800
vyska = 800

#create the screen
window = pygame.display.set_mode((sirka+200, vyska)) 
pygame.font.init() 
myfont = pygame.font.SysFont('Comic Sans MS', 100)
myfont_maly = pygame.font.SysFont('Comic Sans MS', 40)


#draw it to the screen
pygame.display.flip() 

#################################### konštanty #########################################

zavaznost = 10 # sanca v promile, ze v dalsom cykle nakazeny vazne ochorie. default 10
smrtnost = 1 # sanca v promile, ze v dalsom cykle chory clovek zomrie. Mení sa počas hry podľa počtu nakazených.
liecitelnost = 2 # sanca v promile, ze v dalsom cykle chory, alebo nakazeny vyzdravie. default 2
inkubacna_doba = 200 #po kolkych cykloch prechadza na chorobu. Default 200
pocet_ludi = 500 #### default 500. Pri číslach nad pár tisíc začne lagovať. Teraz mám zložitosť n*log(n) zatiaľ bez multithreadingu
su_doma = 0.21 ## pomer ľudí, ktorí ostávajú doma a nehýbu sa (číslo medzi 0.0 a 1) default 0.2

########################################################################################
################################# Generovanie počiatku hry #############################

ludkovia = []
ludkovia.append([nahodne(0,sirka),nahodne(0,vyska),nahodne(0,6)-3,nahodne(0,6)-3,'chory',0])
for i in range(pocet_ludi-1):
    if nahodne(0,100)>su_doma*100:
        ludkovia.append([nahodne(0,sirka),nahodne(0,vyska),nahodne(0,4)-2,nahodne(0,4)-2,'zdravy',0])
    else:
        ludkovia.append([nahodne(0,sirka),nahodne(0,vyska),0,0,'zdravy',0])

##################### Kreslí súčastný stav, rôzne chorých rôznymi farbami ##############
def vykresli():
    for clovek in ludkovia:
        if clovek[4] == 'zdravy':
            farba = (200,200,200)
        elif clovek[4] == 'nakazeny':
            farba = (255,150,0)
        elif clovek[4] == 'vylieceny':
            farba = (50,255,50)
        elif clovek[4] == 'mrtvy':
            farba = (0,0,0)
        else:
            farba = (255,0,0)
        pygame.draw.circle(window, farba, (clovek[0],clovek[1]), 7, 7)

##################### Hýbe ľuďmi a rieši kolízie so stenami ############################
def pohyb():
    for clovek in ludkovia:
        if clovek[0]>=sirka-2:
            clovek[2]=-abs(clovek[2])
        if clovek[0]<=2:
            clovek[2]=abs(clovek[2])
        if clovek[1]>=vyska-2:
            clovek[3]=-abs(clovek[3])
        if clovek[1]<=2:
            clovek[3]=abs(clovek[3])
        clovek[0] = clovek[0]+clovek[2]
        clovek[1] = clovek[1]+clovek[3]


##################### Usporiada ľudí podľa x-ovej súradnice ############################
def usporiadaj():
    ludkovia.sort(key=lambda x: x[1])


########################### Rieši čo sa má stať, ak sa dvaja ľudia stretnú #############
########################### geometriu aj nákazu ########################################
# a áno, toto by bolo oveľa ľahšie a krajšie, keby to bolo objektovo spravené niekde o 
# level vyššie. :) som lenivá lama. Začal som to robiť od brucha a nemal som to veľmi
# premyslené. 

def kolizie():
    for i in range(0,len(ludkovia)-1):
        for j in range(i+1,len(ludkovia)):
            if abs(ludkovia[i][0]-ludkovia[j][0])<9:
                if abs(ludkovia[i][1]-ludkovia[j][1])<9:
                    if ludkovia[i][4]=='chory' and ludkovia[j][4]=='zdravy':
                        ludkovia[j][4]='nakazeny'
                    if ludkovia[j][4]=='chory' and ludkovia[i][4]=='zdravy':
                        ludkovia[i][4]='nakazeny'
                    if abs(ludkovia[i][2])+abs(ludkovia[i][3])!=0:                    
                        ludkovia[i][2]=round((ludkovia[i][0]-ludkovia[j][0])/2)
                        ludkovia[i][3]=round((ludkovia[i][1]-ludkovia[j][1])/2)
                        if ludkovia[i][2]==0:
                            ludkovia[i][2]=nahodne(0,6)-3
                        if ludkovia[i][3]==0:
                            ludkovia[i][3]=nahodne(0,6)-3
                    if abs(ludkovia[j][2])+abs(ludkovia[j][3])!=0:                    
                        ludkovia[j][2]=round((ludkovia[j][0]-ludkovia[i][0])/2)
                        ludkovia[j][3]=round((ludkovia[j][1]-ludkovia[i][1])/2)
                        if ludkovia[j][2]==0:
                            ludkovia[j][2]=nahodne(0,6)-3
                        if ludkovia[j][3]==0:
                            ludkovia[j][3]=nahodne(0,6)-3
                else: 
                    break

########################### Rieši pre každého nakazeného, či mu prepukne choroba ########
########################### a pre každého chorého, či sa vylieči, alebo zomrie ##########
def postup_choroby():
    for clovek in ludkovia:
        if clovek[4] == 'nakazeny':
            clovek[5] += 1
            if clovek[5] > inkubacna_doba:
                if nahodne(0,1000)< zavaznost:
                    clovek[4] = 'chory'
                    clovek[5] = 0
        elif clovek[4] == 'chory':
            clovek[5] += 1
            if clovek[5] > inkubacna_doba:
                if nahodne(0,1000)< smrtnost:
                    clovek[4] = 'mrtvy'
                    clovek[5] = 0 
                    clovek[2] = 0
                    clovek[3] = 0   
            if nahodne(0,1000)< liecitelnost:
                clovek[4] = 'vylieceny'  

########################### Počíta, koľko je chorých, nakazených a mŕtvych ľudí ##########
def spocitaj():
    chorych = 0
    zdravych = 0
    nakazenych = 0
    mrtvych = 0
    vyliecenych = 0
    for clovek in ludkovia:
        if clovek[4] == 'zdravy':
            zdravych += 1
        elif clovek[4] == 'chory':
            chorych += 1
        elif clovek[4] == 'nakazeny':
            nakazenych += 1
        elif clovek[4] == 'mrtvy':
            mrtvych += 1 
        if clovek[4] == 'vylieceny':
            vyliecenych += 1
    global smrtnost 
    smrtnost = 5*chorych/pocet_ludi
    text = myfont.render(str(zdravych), True, (200,200,200))
    window.blit(text,(sirka,0))
    text = myfont.render(str(nakazenych), True, (255,150,0))
    window.blit(text,(sirka,vyska/6))
    text = myfont.render(str(chorych), True, (255,0,0))
    window.blit(text,(sirka,vyska/6*2))
    text = myfont.render(str(vyliecenych), True, (50,255,50))
    window.blit(text,(sirka,vyska/6*3))
    text = myfont.render(str(mrtvych), True, (0,0,0))
    window.blit(text,(sirka,vyska/6*4))   
    text = myfont_maly.render(str(round(smrtnost*300))+"%", True, (100,100,100))
    window.blit(text,(sirka,vyska/6*5))   
    text = myfont_maly.render("Pretaženosť", True, (100,100,100))
    window.blit(text,(sirka,vyska/6*5+50)) 
    text = myfont_maly.render("nemocníc", True, (100,100,100))
    window.blit(text,(sirka,vyska/6*5+100))      

#input handling (somewhat boilerplate code):
while True: 
    pohyb()                     # pohne každým človekom
    postup_choroby()            # ľudia sa ozdravujú, umierajú, zhoršuje sa im stav ,...
    usporiadaj()               # zíde sa, ak by sme optimalizovali kolízie
    kolizie()                   # Tu sa srážajú a nakazujú ľudia
    window.fill((255,255,255))  # Zresetuje sa obrazovka
    spocitaj()                  # Spočítam, koľko je chorých, zdravých ,..
    vykresli()                  # vykreslia sa nové pozície
    pygame.time.delay(20)       # obmedzujem fps (treba zlepšiť, nech sa adaptuje na výkonové a výpočtové požiadavky)
    pygame.display.flip()       #aktualizuje sa obrazovka s novými pozíciami.
    for event in pygame.event.get(): 
        if event.type == pygame.QUIT: 
            sys.exit(0) 
    ##    else: 
    ##        print (event)