from random import randint, choice
from time import sleep
from kandinsky import *
from ion import *

class Maze:
  def __init__(self, W, H, SQUARE_SIZE, draw, t=2):
    self.W = W
    self.H = H
    self.SQUARE_SIZE = SQUARE_SIZE
    self.gen(draw, t)

  def gen(self, draw, t):
    self.maze = [[[1, 1] for x in range(self.W)] for y in range(self.H)] # [east, south]
    self.visited = [[0 for x in range(self.W)] for y in range(self.H)]
    self.pos = (0, 0)
    self.path = [self.pos]
    self.visited_n = 1
    self.visited[self.pos[1]][self.pos[0]] = 1

    if draw:
      fill_rect(0, 0, 320, 222, (0, 0, 0))

    while self.visited_n < self.W*self.H:
      while not self.next_square(draw, t):
        self.pos = self.path.pop()
      if draw:
        self.draw_square(self.pos[0], self.pos[1])
      sleep(0.01)

  def next_square(self, draw, t):
    possible_dirs = [1, 1, 1, 1]
    x, y = self.pos[0], self.pos[1]

    if x == 0 or self.visited[y][x-1]:
      possible_dirs[0] = 0
    if x == self.W-1 or self.visited[y][x+1]:
      possible_dirs[1] = 0
    if y == 0 or self.visited[y-1][x]:
      possible_dirs[2] = 0
    if y == self.H-1 or self.visited[y+1][x]:
      possible_dirs[3] = 0

    break_x = x
    break_y = y
    break_wall = 0

    if 1 in possible_dirs:
      next_dir = randint(0, 3)
      while possible_dirs[next_dir] == 0:
        next_dir = randint(0, 3)

      if next_dir == 0:
        break_x -= 1
        self.pos = (x-1, y)
      elif next_dir == 1:
        self.pos = (x+1, y)
      elif next_dir == 2:
        break_y -= 1
        break_wall = 1
        self.pos = (x, y-1)
      else:
        break_wall = 1
        self.pos = (x, y+1)
      
      self.maze[break_y][break_x][break_wall] = 0
      if draw:
        self.draw_square(break_x, break_y, t)

      self.path.append(self.pos)
      self.visited[self.pos[1]][self.pos[0]] = 1
      self.visited_n += 1
      return True

    return False

  def draw_square(self, x, y, t=2):
    if self.maze[y][x][0]:
      east = (0, 0, 0)
    else:
      east = (255, 255, 255)
    if self.maze[y][x][1]:
      south = (0, 0, 0)
    else:
      south = (255, 255, 255)
    fill_rect(x*self.SQUARE_SIZE, y*self.SQUARE_SIZE, self.SQUARE_SIZE - t, self.SQUARE_SIZE - t, (255, 255, 255))
    fill_rect((x+1)*self.SQUARE_SIZE, y*self.SQUARE_SIZE, -t, self.SQUARE_SIZE, east)
    fill_rect(x*self.SQUARE_SIZE, (y+1)*self.SQUARE_SIZE, self.SQUARE_SIZE, -t, south)

  def draw_maze(self, thick=2):
    fill_rect(0, 0, 320, 222, (255, 255, 255))
    for y, row in enumerate(self.maze):
      for x, walls in enumerate(row):
        self.draw_square(x, y, thick)

SQUARE_SIZE = int(input("Input tile size (px): "))
W= 320 // SQUARE_SIZE
H= 222 // SQUARE_SIZE

maze = Maze(W, H, SQUARE_SIZE, True)