Зависание при отрисовке квадратов

# -*- coding: utf-8 -*-
 
from Tkinter import *
import time
from random import choice
 
 
def draw(i, y, count, color, number_in_string, number, refresh_rate):
        size = 800//number_in_string
        size1 = 600//number_in_string
        x = 0 + i * size
        canvas.create_rectangle(x,y, x+size, y+size1, fill=color, outline=color)
        if x+size == 800:
            i = -1
            y = y + size1
        #if y+size == 600:
             #удалить и сдвинуть
        if count == number:
            return
        count = count + 1
        canvas.after(refresh_rate, lambda: draw(i+1, y, count, choice(colors), \
        number_in_string, number, refresh_rate))
 
if __name__ == "__main__":
    print 'Enter the number of elements in string:'
    number_in_string=input()
    print 'Enter the number of elements:'
    number=input()
    print 'Enter refresh rate:'
    refresh_rate=input()
    print 'The number of elements in string:', number_in_string
    print 'The number of elements:', number
    print 'Refresh rate:', refresh_rate
 
    tk = Tk()
    tk.title("laboratory_number_5")
    tk.wm_geometry("800x600+20+40")
    canvas = Canvas(tk, width=800, height=600)
    colors = "Red Orange Yellow Green LightBlue Blue Violet".split()
    draw(0, 0, 0, choice(colors), number_in_string, number, refresh_rate)
    canvas.pack()
    tk.mainloop()

    if x+size == size * number_in_string:

# -*- coding: utf-8 -*-
 
from Tkinter import *
import time
from random import choice
 
 
def draw(row, col, color, count, last_row):
 
        x1 = col * sizeX
        y1 = row * sizeY
        x2 = x1 + sizeX
        y2 = y1 + sizeY
        rectId = canvas.create_rectangle(x1,y1,x2,y2, fill=color, outline=color)
 
        last_row.append(rectId)
        if count == number_total:
            print "Done!"
            return
 
        count += 1
        col += 1
        if (col > cols - 1):
            col = 0
            row += 1
            last_row = []
 
        
        canvas.after(refresh_rate, lambda: draw(row, col, choice(colors), count, last_row))
 
 
if __name__ == "__main__":
    width = 800
    height = 600
 
    print 'Enter the number of quadrats in row:'
    cols=input()
    print 'Enter the total number of quadrats:'
    number_total=input()
    print 'Enter refresh rate (in msec):'
    refresh_rate=input()
    print 'The number of elements in row:', cols
    print 'The number of elements:', number_total
    print 'Refresh rate:', refresh_rate
 
    rows=number_total//cols
    print "rows: ", rows
 
    sizeX = width//cols
    sizeY = height//cols
    print 'sizeX: %d; sizeY: %d' % (sizeX, sizeY)
 
    tk = Tk()
    tk.title("laboratory_number_5")
    tk.wm_geometry("%dx%d+20+40" % (width, height))
    canvas = Canvas(tk, width=width, height=height)
    colors = ["Red","Orange","Yellow","Green","LightBlue","Blue","Violet"]
    draw(0, 0, choice(colors), 1, [])
    canvas.pack()
    tk.mainloop()