import random
import sys
import pygame

class MondrianRectangle(object):
    def __init__(self, rect, colour = (255, 255, 255)):
        self.rect = rect
        self.colour = colour

    def __getitem__(self, key):
        return self.rect[key]

    def __setitem__(self, key, value):
        self.rect[key] = value
    
    def draw(self, surface):
        surface.fill(self.colour, self.rect)
        

class MondrianLine(object):
    def __init__(self, direction = "horizontal", start_pos = (0, 0), end_pos = (0, 0), width = 0):
        self.direction = direction
        self.start_pos = start_pos
        self.end_pos = end_pos
        self.width = width

    def draw(self, surface):
        pygame.draw.line(surface, (0, 0, 0), self.start_pos, self.end_pos, self.width)


def generate_random_line():
    width = 5
    if random.random() > 0.5:
        line_pos = random.randint(0, 1024)
        return MondrianLine("vertical", (line_pos, 0), (line_pos, 768), width)
    else:                     
        line_pos = random.randint(0, 768)
        return MondrianLine("horizontal", (0, line_pos), (1024, line_pos), width)


def split_rectangles(rectangles, line):
    new_rectangles = []
    for j, rectangle in enumerate(rectangles):
        if (line.direction == "horizontal" and 
            rectangle[1] < line.start_pos[1] < rectangle[3] and 
            min(line.start_pos[0], line.end_pos[0]) <= rectangle[0] and 
            max(line.start_pos[0], line.end_pos[0]) >= rectangle[2]
        ):
            new_rectangles.extend([
                MondrianRectangle((rectangle[0], rectangle[1], rectangle[2], line.start_pos[1])),
                MondrianRectangle((rectangle[0], line.start_pos[1], rectangle[2], rectangle[3]))
            ])

        elif (line.direction == "vertical" and 
            rectangle[0] < line.start_pos[0] < rectangle[2] and
            min(line.start_pos[1], line.end_pos[1]) <= rectangle[1] and
            max(line.start_pos[1], line.end_pos[1]) >= rectangle[3]
        ):
            new_rectangles.extend([
                MondrianRectangle((rectangle[0], rectangle[1], line.start_pos[0], rectangle[3])),
                MondrianRectangle((line.start_pos[0], rectangle[1], rectangle[2], rectangle[3]))
            ])

        else:
            new_rectangles.append(rectangle)

    return new_rectangles


def main():
    screen = pygame.display.set_mode((1024, 768))

    # Step 1
    lines = []
    rectangles = [MondrianRectangle((0, 0, 1024, 768))]
    for i in range(random.randint(5, 10)):
        line = generate_random_line()
        lines.append(line)
        rectangles = split_rectangles(rectangles, line)

    #Step 2
    horizontal_lines = [line for line in lines if line.direction == "horizontal"]
    vertical_lines = [line for line in lines if line.direction == "vertical"]

    for i in range(random.randint(1, 10)):
        if random.random() > 0.5 and len(horizontal_lines) >= 2:
            lines_to_join = random.sample(horizontal_lines, 2)
            line_pos = random.randint(0, 1024)
            new_line = MondrianLine(
                "vertical", 
                start_pos = (line_pos, lines_to_join[0].start_pos[1]),
                end_pos = (line_pos, lines_to_join[1].start_pos[1]),
                width = 5
            )
            rectangles = split_rectangles(rectangles, new_line)
            lines.append(new_line)
        elif len(vertical_lines) >= 2:
            lines_to_join = random.sample(vertical_lines, 2)
            line_pos = random.randint(0, 768)
            new_line = MondrianLine(
                "horizontal", 
                start_pos = (lines_to_join[0].start_pos[0], line_pos),
                end_pos = (lines_to_join[1].start_pos[0], line_pos),
                width = 5
            )
            rectangles = split_rectangles(rectangles, new_line)
            lines.append(new_line)
     
    #Step 3
    for rectangle in random.sample(rectangles, random.randint(1, 10)):
        rectangle.colour = (
            random.randint(0, 255), 
            random.randint(0, 255),
            random.randint(0, 255)
        )
    
    while True:
        screen.fill((255, 255, 255))
        for event in pygame.event.get():
            if event.type == pygame.QUIT: sys.exit()
        [rectangle.draw(screen) for rectangle in rectangles]
        [line.draw(screen) for line in lines]
        pygame.display.flip()
 

if __name__ == "__main__":
    main()