Python doesn't wait long enough for user input

I have a problem with a python program. When I prompt for user input the program waits only about 2-3 seconds, and if I input the answer in that timeframe it loads the rest, but if i don't it stops loading inputs and becomes unreponsive. I have tried adding a while loop but it doesnt change anything.

import pygame as pg

pg.init()
pg.display.set_caption("Struktura Atoma")
(sirina, visina) = (250, 300)
prozor = pg.display.set_mode((sirina, visina))

kraj = True

r = 20  # poluprecnik
x = sirina // 2  # koordinata x centra kruznice
y = visina // 2  # koordinara y centra kruznice

izbor = input("Unesite 1 za atomsku strukturu atoma ugljenika, unesite 2 za atomsku strukturu magnezijuma: ")

if izbor == "1":

    for i in range(3):
        roze = (243, 58, 106)
        siva = (220, 220, 220)
        bojaAtom = (0, 255, 255)
        font = pg.font.SysFont("Arial", 22)  # podesavanje fonta
        slC = font.render("C", True, pg.Color("black"))  # formiranje slova
        pg.draw.circle(prozor, bojaAtom, (x, y), r, 1)  # crtanje kruznice
        if r == 20:
            pg.draw.circle(prozor, bojaAtom, (x, y), r, 0)
            prozor.blit(slC, (x - 6, y - 12))

        if r == 50:
            pg.draw.circle(prozor, roze, (x - 45, y - 20), r - 40, 0)
            pg.draw.circle(prozor, roze, (x + 45, y + 20), r - 40, 0)

        if r == 80:
            pg.draw.circle(prozor, siva, (x - 80, y), r - 70, 0)
            pg.draw.circle(prozor, siva, (x + 80, y), r - 70, 0)
            pg.draw.circle(prozor, siva, (x, y - 80), r - 70, 0)
            pg.draw.circle(prozor, siva, (x, y + 80), r - 70, 0)
        pg.display.update()  # osvezavanje prozora

        r = r + 30  # promena poluprecnika sledece kruznice
        pg.time.wait(100)  # cekaj 100ms

        # provera da li je korisnik zatvorio prozor

if izbor == "2":
    for i in range(4):
        roze = (243, 58, 106)
        siva = (220, 220, 220)
        bojaAtom = (0, 255, 255)
        font = pg.font.SysFont("Arial", 22)  # podesavanje fonta
        slMg = font.render("Mg", True, pg.Color("black"))  # formiranje slova
        pg.draw.circle(prozor, bojaAtom, (x, y), r, 1)  # crtanje kruznice
        if r == 20:
            pg.draw.circle(prozor, bojaAtom, (x, y), r, 0)
            prozor.blit(slMg, (x - 12, y - 13))

        if r == 50:
            pg.draw.circle(prozor, roze, (x - 45, y - 20), r - 40, 0)
            pg.draw.circle(prozor, roze, (x + 45, y + 20), r - 40, 0)

        if r == 80:
            pg.draw.circle(prozor, siva, (x - 80, y), r - 70, 0)
            pg.draw.circle(prozor, siva, (x + 80, y), r - 70, 0)
            pg.draw.circle(prozor, siva, (x, y - 80), r - 70, 0)
            pg.draw.circle(prozor, siva, (x, y + 80), r - 70, 0)
            pg.draw.circle(prozor, siva, (x - 56, y - 56), r - 70, 0)
            pg.draw.circle(prozor, siva, (x + 56, y - 56), r - 70, 0)
            pg.draw.circle(prozor, siva, (x + 56, y + 56), r - 70, 0)
            pg.draw.circle(prozor, siva, (x - 56, y + 56), r - 70, 0)

        if r == 110:
            pg.draw.circle(prozor, siva, (x - 105, y), r - 100, 0)
            pg.draw.circle(prozor, siva, (x + 105, y), r - 100, 0)
        pg.display.update()  # osvezavanje prozora

        r = r + 30  # promena poluprecnika sledece kruznice
        pg.time.wait(100)  # cekaj 100ms


while kraj:
    for dogadjaj in pg.event.get():
        if dogadjaj.type == pg.QUIT:
            kraj = False



pg.quit()