from common import *

import numpy as np
from typing import *

# concepts:
# repeating patterns, colors as indicators, scaling

# description:
# In the input you will see a nxm sprite with black background. 
# Construct an output grid with n^2 x m^2 black pixels. Divide the output grid into subgrids, 
# and look at the corresponding pixel in the nxm input grid. If the corresponding pixel is not black, 
# then copy the nxm input grid into the subgrid. Else, the subgrid does not change. 

def transform(input_grid):
    # creates an empty 9x9 output grid 
    output_grid = np.zeros((input_grid.shape[0]**2,input_grid.shape[1]**2),dtype=int)

    input_sprite = input_grid

    # Go through the input grid. If an input grid pixel is not black, 
    # then copy the input grid to the corresponding location on the output grid
    for n in range(input_grid.shape[0]):
      for m in range(input_grid.shape[1]):
        if input_grid[n,m] != Color.BLACK:
            blit_sprite(output_grid, input_sprite, n*input_grid.shape[0], m*input_grid.shape[1])
    
    return output_grid

def generate_input():
  n,m = random.randint(3, 6), random.randint(3, 6)
  random_color = random.choice(list(Color.NOT_BLACK))
  return random_sprite(n, m, color_palette=[random_color])

from common import *

import numpy as np
from typing import *

# concepts:
# topology

# description:
# The input grid is a square grid with black and green pixels. The input grid should have regions that are enclosed by the green pixels. 
# To produce the output, you need to find the enclosed regions in the input grid, and then color them yellow. 
                
def transform(input_grid):
    # Create initial output grid template based on input grid.
    output_grid = input_grid.copy()

    # Find enclosed regions
    interior_mask = object_interior(input_grid)
    boundary_mask = object_boundary(input_grid)
    inside_but_not_on_edge = interior_mask & ~boundary_mask

    # Color enclosed regions
    for x, y in np.argwhere(inside_but_not_on_edge):
        if output_grid[x, y] == Color.BLACK:
            output_grid[x, y] = Color.YELLOW

    return output_grid


def generate_input():
    # Generate a square grid of arbitrary size with black background, size from 5x5 to 20x20
    n = random.randint(10, 20)
    grid = np.zeros((n, n), dtype=int)

    # Generate some random green sprites, and then hollow out the interior
    n_objects = random.randint(1, 3)
    for _ in range(n_objects):
        n, m = random.randint(4, 10), random.randint(4, 10)
        sprite = random_sprite(n, m, color_palette=[Color.GREEN], connectivity=8)
        interior_mask = object_interior(sprite)
        boundary_mask = object_boundary(sprite)
        interior_but_not_edges = interior_mask & ~boundary_mask
        sprite[interior_but_not_edges] = Color.BLACK

        try:
            x, y = random_free_location_for_sprite(grid, sprite, border_size=1, padding=1)
        except:
            continue

        blit_sprite(grid, sprite, x, y, background=Color.BLACK)
    
    return grid

from common import *

import numpy as np
from typing import *

# concepts:
# translational symmetry, symmetry detection

# description:
# In the input you will see a grid consisting of a blue sprite that is repeatedly translated vertically, forming a stack of the same sprite.
# To make the output, expand the input to have height 9, and continue to repeatedly translate the sprite vertically. Change color to red.
 
def transform(input_grid):
    # Plan:
    # 1. Find the repeated translation, which is a symmetry
    # 2. Extend the pattern by copying the sprite and its symmetric copies
    # 3. Change the color from blue to red
    
    symmetries = detect_translational_symmetry(input_grid, ignore_colors=[], background=Color.BLACK)
    assert len(symmetries) > 0, "No translational symmetry found"

    # make the output (the height is now 9)
    output_grid = np.full((input_grid.shape[0], 9), Color.BLACK)
    
    # Copy all of the input pixels to the output, INCLUDING their symmetric copies (i.e. their orbit)
    for x, y in np.argwhere(input_grid != Color.BLACK):
        # Compute the orbit into the output grid
        for x2, y2 in orbit(output_grid, x, y, symmetries):
            output_grid[x2, y2] = input_grid[x, y]
    
    # Color change: blue -> red
    output_grid[output_grid == Color.BLUE] = Color.RED

    return output_grid


def generate_input():
    # grid is always 3x6
    grid = np.zeros((3, 6),dtype = int)

    # The input is always blue
    color = Color.BLUE

    # Creates a random smaller sprite, where the height (period) is chosen randomly
    height = random.randint(2, 7)
    sprite = random_sprite(3, height, symmetry="not_symmetric", color_palette=[color], density=0.4, connectivity=8)

    # place the smaller pattern, tiling it so that it is repeated vertically
    # tile "infinitely" (x100)
    vertically_repeated = np.tile(sprite, (1, 100))
    # crop to the size of the grid
    vertically_repeated = vertically_repeated[:, :grid.shape[1]]
    # copy to the grid
    grid[:,:] = vertically_repeated

    return grid

from common import *

import numpy as np
from typing import *

# concepts:
# objects, pixel manipulation

# description:
# In the input you will see a set of objects, each consisting of a horizontal top/bottom and diagonal left/right edges (but that structure is not important)
# To make the output shift right each pixel in the object *except* when there are no other pixels down and to the right

def transform(input_grid: np.ndarray) -> np.ndarray:
    # find the connected components, which are monochromatic objects
    objects = find_connected_components(input_grid, background=Color.BLACK, connectivity=8, monochromatic=True)

    output_grid = np.zeros_like(input_grid)

    for obj in objects:
        transformed_object = np.zeros_like(obj)

        for x in range(obj.shape[0]):
            for y in range(obj.shape[1]):
                if obj[x, y] != Color.BLACK:
                    # check that there are other colored pixels down and to the right
                    down_and_to_the_right = obj[x+1:, y+1:]
                    if np.any(down_and_to_the_right != Color.BLACK):
                        transformed_object[x+1, y] = obj[x, y]
                    else:
                        transformed_object[x, y] = obj[x, y]

        blit_object(output_grid, transformed_object, background=Color.BLACK)

    return output_grid


def generate_input():
    n, m = np.random.randint(10, 30), np.random.randint(10, 30)
    grid = np.full((n, m), Color.BLACK)

    n_objects = np.random.randint(1, 3)

    for _ in range(n_objects):
        color = random.choice(Color.NOT_BLACK)

        bar_width = np.random.randint(3, n//2)
        side_height = np.random.randint(3, m - bar_width)

        width, height = bar_width + side_height, side_height
        obj = np.zeros((width, height), dtype=int)

        # make the horizontal top edge
        obj[:bar_width+1, 0] = color
        # make the horizontal bottom edge
        obj[-bar_width:, -1] = color
        # make the diagonal left edge
        for i in range(side_height):
            obj[i, i] = color
        # make the diagonal right edge
        for i in range(side_height-1):
            obj[bar_width+i+1, i] = color

        # place the object randomly on the grid, assuming we can find a spot
        try:
            x, y = random_free_location_for_sprite(grid, obj, background=Color.BLACK, padding=2, padding_connectivity=8, border_size=2)
        except:
            continue

        blit_sprite(grid, obj, x=x, y=y, background=Color.BLACK)

    # Make sure that we actually generated something
    if np.all(grid == Color.BLACK):
        return generate_input()
    
    return grid

from common import *

import numpy as np
from typing import *

# concepts:
# sprites, color change, collision detection, repetition, overlap

# description:
# In the input you will see a 3x3 object with a few other objects around it.
# For each of the other sprites around the central 3x3 object:
# 1. Slide the central sprite so it completely overlaps the other sprite (slide it as much as you can to do so)
# 2. Change the color of the central sprite to match the color of the other sprite
# 3. Repeat the slide (by the same displacement vector) indefinitely until it falls off the canvas

def transform(input_grid: np.ndarray) -> np.ndarray:
    # find the objects, which are monochromatic connected components
    objects = find_connected_components(input_grid, background=Color.BLACK, connectivity=8, monochromatic=True)

    # find the central object, which is the biggest
    central_object = max(objects, key=lambda obj: np.sum(obj != Color.BLACK))

    # find the other objects
    other_objects = [obj for obj in objects if not np.array_equal(obj, central_object)]

    output_grid = np.copy(input_grid)

    for other_object in other_objects:
        # find the biggest displacement vector that will make the central object completely overlap the other object
        biggest_displacement_vector = (0,0)
        displacement_vectors = [ (i, j) for i in range(-10, 10) for j in range(-10, 10) ]
        for displacement_vector in displacement_vectors:
            # translate the central object by the displacement vector
            translated_central_object = translate(central_object, displacement_vector[0], displacement_vector[1], background=Color.BLACK)

            # check if the translated object completely overlaps the other object
            translated_mask, other_mask = translated_central_object != Color.BLACK, other_object != Color.BLACK
            overlaps = np.all(translated_mask & other_mask == other_mask)

            if overlaps:
                # but is it the biggest?
                if biggest_displacement_vector[0] ** 2 + biggest_displacement_vector[1] ** 2 < displacement_vector[0] ** 2 + displacement_vector[1] ** 2:
                    biggest_displacement_vector = displacement_vector

        displacement_vector = biggest_displacement_vector

        # color change
        color_of_other_object = np.unique(other_object[other_object != Color.BLACK])[0]
        central_object[central_object != Color.BLACK] = color_of_other_object

        # repeat the displacement indefinitely until it falls off the canvas
        for i in range(1, 10):
            displaced_central_object = translate(central_object, displacement_vector[0] * i, displacement_vector[1] * i, background=Color.BLACK)
            blit_object(output_grid, displaced_central_object, background=Color.BLACK)
            
    return output_grid



def generate_input() -> np.ndarray:
    # make a black grid first as background
    n, m = 21, 21
    grid = np.zeros((n, m), dtype=int)

    # make a 3x3 object
    central_color = np.random.choice(Color.NOT_BLACK)
    central_sprite = random_sprite(3, 3, color_palette=[central_color])

    # place the central object near the center
    x, y = np.random.randint(int(0.3*n), int(0.7*n)), np.random.randint(int(0.3*m), int(0.7*m))
    blit_sprite(grid, central_sprite, x, y, background=Color.BLACK)

    # possible displacement vectors can range in any of the eight different directions (cardinal directions and in between them)
    # they should be close to just a little more than the length of the central object, however
    eight_cardinal_directions = [(0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1), (-1, 0), (-1, 1)]
    # pick a random subset of them, between 1-4
    displacement_vectors = random.sample(eight_cardinal_directions, k=np.random.randint(1, 5))

    for vector in displacement_vectors:
        vector_length = np.random.randint(4, 5)
        vector = (vector[0] * vector_length, vector[1] * vector_length)

        # make a random object by recoloring the central object, translating by the vector,
        # and then randomly removing parts of it by flipping random pixels to black
        other_sprite = np.copy(central_sprite)
        other_sprite[other_sprite != Color.BLACK] = np.random.choice(Color.NOT_BLACK)

        # flip some random pixels to black:
        # first find the foreground (nonblack) pixels,
        # then randomly sample a subset of them to color black
        nonblack_pixels = np.argwhere(other_sprite != Color.BLACK)
        num_nonblack = len(nonblack_pixels)
        num_to_flip = np.random.randint(1, num_nonblack-3)
        random_subset_of_nonblack_pixels = random.sample(list(nonblack_pixels), k=num_to_flip)

        # color black
        for pixel in random_subset_of_nonblack_pixels:
            other_sprite[pixel[0], pixel[1]] = Color.BLACK        

        # place the new object near the center, but offset by the vector
        blit_sprite(grid, other_sprite, x + vector[0], y + vector[1], background=Color.BLACK)

    return grid

import numpy as np
from typing import *
from common import *

# concepts:
# boolean logical operations, bitmasks with separator

# description:
# In the input you will see two blue bitmasks separated by a grey vertical bar
# To make the output, color teal the red that are set in both bitmasks (logical AND)

def transform(input_grid: np.ndarray) -> np.ndarray:
    # Find the grey vertical bar. Vertical means constant X
    for x_bar in range(input_grid.shape[0]):
        if np.all(input_grid[x_bar, :] == Color.GREY):
            break

    left_mask = input_grid[:x_bar, :]
    right_mask = input_grid[x_bar+1:, :]

    output_grid = np.zeros_like(left_mask)
    output_grid[(left_mask == Color.BLUE) & (right_mask == Color.BLUE)] = Color.RED
    
    return output_grid


def generate_input() -> np.ndarray:
    # create a pair of equally sized maroon bitmasks
    width, height = np.random.randint(2, 10), np.random.randint(2, 10)

    grid1 = np.zeros((width, height), dtype=int)
    grid2 = np.zeros((width, height), dtype=int)

    for x in range(width):
        for y in range(height):
            grid1[x, y] = np.random.choice([Color.BLUE, Color.BLACK])
            grid2[x, y] = np.random.choice([Color.BLUE, Color.BLACK])
    
    # create a blue vertical bar
    bar = np.zeros((1, height), dtype=int)
    bar[0, :] = Color.GREY

    grid = np.concatenate((grid1, bar, grid2), axis=0)

    return grid

from common import *

import numpy as np
from typing import *

# concepts:
# diagonal lines, repetition

# description:
# In the input you will see a 7x7 grid, with three diagonal lines that stretch from one end of the canvas to the other
# Each line is a different color, and the colors are not black
# The output should be the result of repeating every diagonal line on multiples of 3 offset from the original, which gives an interlacing pattern filling the output canvas


def transform(input_grid: np.ndarray) -> np.ndarray:
    output_grid = np.zeros((7, 7), dtype=int)

    # Loop over the input looking for any of the three diagonals
    # If we find one, we will fill the output with the same color in the same pattern
    for i in range(output_grid.shape[0]):
        for j in range(output_grid.shape[1]):
            c = input_grid[i][j]
            if c != Color.BLACK:
                # Fill the output with the same color in the same pattern
                # Loop by multiples of 3 to create the pattern
                # Loop way beyond the canvas (double the canvas size) to make sure we cover everything
                for distance in range(0, output_grid.shape[0]*2, 3):
                    draw_diagonal(output_grid, i-distance, j, c)
                    draw_diagonal(output_grid, i+distance, j, c)
    
    return output_grid

def draw_diagonal(grid, x, y, c):
    # create diagonal line that stretches from one side of the canvas to the other
    # to do this, draw infinite rays pointing in opposite directions
    draw_line(grid, x, y, length=None, color=c, direction=(1, -1))
    draw_line(grid, x, y, length=None, color=c, direction=(-1, 1))

def generate_input() -> np.ndarray:

    # create a 7x7 grid of black (0)
    grid = np.zeros((7, 7), dtype=int)
    # pick 3 random distinct colors
    c1, c2, c3 = np.random.choice(Color.NOT_BLACK, 3, replace=False)

    # put down the three diagonal lines
    draw_diagonal(grid, np.random.choice(range(7)), np.random.choice(range(7)), c1)
    draw_diagonal(grid, np.random.choice(range(7)), np.random.choice(range(7)), c2)
    draw_diagonal(grid, np.random.choice(range(7)), np.random.choice(range(7)), c3)

    # make sure that the output is completely filled up with the pattern meaning it doesn't contain any black pixels
    if np.any(main(grid) == Color.BLACK):
        return generate_input()
    
    return grid

from common import *

import numpy as np
from typing import *

# concepts:
# collision detection, sliding objects

# description:
# In the input you will see a teal 2x2 square and a red object (the red object might be irregular in its shape)
# Slide the red object in any of the four directions until it just touches the teal square

def transform(input_grid):

    # get just the teal object
    teal_object = np.zeros_like(input_grid)
    teal_object[input_grid == Color.TEAL] = Color.TEAL

    # get just the red object
    red_object = np.zeros_like(input_grid)
    red_object[input_grid == Color.RED] = Color.RED

    # the output grid starts with just the teal object, because we still need to figure out where the red object will be by sliding it
    output_grid = np.copy(teal_object)
    
    # consider sliding in the 4 cardinal directions, and consider sliding as far as possible
    possible_displacements = [ (slide_distance*dx, slide_distance*dy)
                               for slide_distance in range(max(input_grid.shape))
                               for dx, dy in [(0, 1), (0, -1), (1, 0), (-1, 0)] ]
    for x, y in possible_displacements:
        # check if the objects are touching after sliding
        translated_red_object = translate(red_object, x, y, background=Color.BLACK)
        if contact(object1=teal_object, object2=translated_red_object):
            # put the red object where it belongs
            blit_object(output_grid, translated_red_object, background=Color.BLACK)
            return output_grid
            
    assert 0, "No valid slide found"

def generate_input():
    # make a black grid first as background, roughly 5 x 5 to 10x10 works
    n, m = random.randint(5, 20), random.randint(5, 20)
    grid = np.full((n, m), Color.BLACK)

    # make a 2x2 teal square, put it somewhere random on the grid
    square_sprite = np.full((2, 2), Color.TEAL)
    x, y = random_free_location_for_sprite(grid, square_sprite, background=Color.BLACK, padding=1, border_size=1)
    blit_sprite(grid, square_sprite, x, y, background=Color.BLACK)

    # make a random sprite of [3,4] x [3,4] with a random symmetry type and the color red
    sprite = random_sprite([3,4], [3,4], symmetry="not_symmetric", color_palette=[Color.RED])

    # put the sprite somewhere random on the grid
    x, y = random_free_location_for_sprite(grid, sprite, background=Color.BLACK, padding=1, border_size=1)
    blit_sprite(grid, sprite, x, y, background=Color.BLACK)    

    # check that we could slide the object either vertically or horizontally in order to touch the red square
    # this will be true if there is a row or column that has both red and blue
    for x in range(n):
        if Color.TEAL in grid[x, :] and Color.RED in grid[x, :]:
            return grid
    for y in range(m):
        if Color.TEAL in grid[:, y] and Color.RED in grid[:, y]:
            return grid
    
    # if not, try again
    return generate_input()

from common import *

import numpy as np
from typing import *

# concepts:
# rectangular cells, flood fill, connecting same color

# description:
# In the input you will see horizontal and vertical bars that divide the grid into rectangular cells
# To make the output, find any pair of rectangular cells that are in the same row and column and have the same color, then color all the rectangular cells between them with that color

def transform(input_grid: np.ndarray) -> np.ndarray:

    # find the color of the horizontal and vertical bars that divide the rectangular cells
    # this is the color of any line that extends all the way horizontally or vertically
    jail_color = None
    for i in range(input_grid.shape[0]):
        for j in range(input_grid.shape[1]):
            color = input_grid[i][j]
            if np.all(input_grid[i, :] == color) or np.all(input_grid[:, j] == color):
                jail_color = color
                break
    
    assert jail_color is not None, "No jail color found"

    output_grid = input_grid.copy()

    # color all the cells between the same color pixels
    for x in range(input_grid.shape[0]):
        for y in range(input_grid.shape[1]):
            color = input_grid[x][y]
            if color == jail_color or color == Color.BLACK:
                continue

            # check if there is a cell with the same color in the same X value
            for y2 in range(y+1, input_grid.shape[1]):
                if input_grid[x][y2] == color:
                    for y3 in range(y+1, y2):
                        if input_grid[x][y3] == Color.BLACK:
                            output_grid[x][y3] = color
                    break

            # check if there is a cell with the same color in the same Y value
            for x2 in range(x+1, input_grid.shape[0]):
                if input_grid[x2][y] == color:
                    for x3 in range(x+1, x2):
                        if input_grid[x3][y] == Color.BLACK:
                            output_grid[x3][y] = color
                    break
                
    return output_grid

def generate_input() -> np.ndarray:

    grid_size = 32
    cell_size = 2

    # First create the array of rectangular cells, each of which is separated by horizontal and vertical bars dividing cells

    # pick a non-black color for the divider
    divider_color = random.choice(Color.NOT_BLACK)
    grid = np.zeros((grid_size, grid_size), dtype=int)
    r_offset_x, r_offset_y = np.random.randint(0, cell_size), np.random.randint(0, cell_size)

    # make horizontal bars with cell_size gaps, but +1 because we need to include the divider, which is one pixel wide
    for x in range(r_offset_x, grid_size, cell_size+1):
        grid[x, :] = divider_color
    # make vertical bars with cell_size gaps
    for y in range(r_offset_y, grid_size, cell_size+1):
        grid[:, y] = divider_color

    # Second we will color some of the cells with a random color

    # random number of cells to color
    number_to_color = np.random.randint(1, 4)
    for _ in range(number_to_color):
        # pick what we're going to color the inside of the cell, which needs to be a different color from the divider
        other_color = np.random.choice([c for c in Color.ALL_COLORS if c != divider_color and c != Color.BLACK])

        # get all coords of black cells
        black_coords = np.argwhere(grid == Color.BLACK)
        # pick a random black cell
        x, y = random.choice(black_coords)
        flood_fill(grid, x, y, other_color)

        # sometimes skip coloring the other side of the divider
        if random.random() <= 0.2:
            continue 

        # flip a coin to decide if horizontal or vertical
        h_or_v = random.random() < 0.5
        if h_or_v:
            # horizontal
            # get all the black cells in the same row
            black_coords = np.argwhere(grid[x, :] == Color.BLACK)
            # pick a random black cell
            other_y = random.choice(black_coords)
            flood_fill(grid, x, other_y, other_color)
        else:
            # vertical
            # get all the black cells in the same column
            black_coords = np.argwhere(grid[:, y] == Color.BLACK)
            # pick a random black cell
            other_x = random.choice(black_coords)
            flood_fill(grid, other_x, y, other_color)

    return grid

from common import *

import numpy as np
from typing import *

# concepts:
# sorting, color change, size

# description:
# In the input you will see a row of exactly 4 grey bars of different heights, each starting at the bottom of the canvas, and each separated by 1 pixel (so they are two pixels apart)
# Color the tallest one blue, the second tallest one red, the third tallest one green, and the shortest one yellow.

def transform(input_grid):

    # extract the bars, each of which is a connected component
    bars = find_connected_components(input_grid, background=Color.BLACK)

    # sort the bars by height
    bars = list(sorted(bars, key=lambda bar: np.sum(bar != Color.BLACK), reverse=True))

    # color the bars
    output_grid = input_grid.copy()

    biggest_bar = bars[0]
    biggest_bar_mask = biggest_bar != Color.BLACK
    output_grid[biggest_bar_mask] = Color.BLUE

    second_biggest_bar = bars[1]
    second_biggest_bar_mask = second_biggest_bar != Color.BLACK
    output_grid[second_biggest_bar_mask] = Color.RED

    third_biggest_bar = bars[2]
    third_biggest_bar_mask = third_biggest_bar != Color.BLACK
    output_grid[third_biggest_bar_mask] = Color.GREEN

    smallest_bar = bars[3]
    smallest_bar_mask = smallest_bar != Color.BLACK
    output_grid[smallest_bar_mask] = Color.YELLOW

    return output_grid



def generate_input():
    # make a black 9x9 grid
    n, m = 9, 9
    grid = np.zeros((9, 9), dtype=int)

    # pick 4 distinct heights (can't reuse the same height)
    heights = random.sample(range(1, 9), 4)

    # draw the bars
    # space them by 2 so that there is a black pixel in between each pair
    for i, height in enumerate(heights):
        grid[1+i*2, -height:] = Color.GREY
        
    return grid

from common import *

import numpy as np
from typing import *

# concepts:
# rectangular cells, color guide

# description:
# In the input you will see grey horizontal and vertical bars that divide the grid into nine 3x3 rectangular regions, each of which contains 4-5 colored pixels
# To make the output, find the region that has exactly 4 colored pixels, and use its colors as a guide to fill in all the other cells

def transform(input_grid: np.ndarray) -> np.ndarray:

    # First identify 

    # Trick for decomposing inputs divided into rectangular regions by horizontal/vertical bars:
    # Treat the bar color as the background, and break the input up into connected components with that background color

    # The divider color is the color of the horizontal and vertical bars
    divider_colors = [ input_grid[x,y] for x in range(input_grid.shape[0]) for y in range(input_grid.shape[1])
                     if np.all(input_grid[x,:] == input_grid[x,0]) or np.all(input_grid[:,y] == input_grid[0,y]) ]
    assert len(set(divider_colors)) == 1, "There should be exactly one divider color"
    divider_color = divider_colors[0] # background=divider_color

    # Find multicolored regions, which are divided by divider_color, so we treat that as background, because it separates objects
    # Within each region there can be multiple colors
    regions = find_connected_components(input_grid, background=divider_color, monochromatic=False)
    # Tag the regions with their location within the 2D grid of (divided) regions
    # First get the bounding-box locations...
    locations = []
    for region in regions:
        x, y, w, h = bounding_box(region, background=divider_color)
        locations.append((x, y, region))
    # ...then re-index them so that (x, y) is the coordinate within the grid of rectangular regions
    grid_of_regions = []
    for x, y, region in locations:
        num_left_of_region = len({other_x for other_x, other_y, other_region in locations if other_x < x})
        num_above_region = len({other_y for other_x, other_y, other_region in locations if other_y < y})
        grid_of_regions.append((num_left_of_region, num_above_region, region))

    # Find the region with exactly 4 colors
    special_region = None
    for region in regions:
        not_divider_and_not_black = (region != divider_color) & (region != Color.BLACK)
        if np.sum(not_divider_and_not_black) == 4:
            assert special_region is None, "More than one special region found"
            special_region = region
    
    # Convert to a sprite
    special_sprite = crop(special_region, background=divider_color)
    
    # Create the output grid
    output_grid = np.zeros_like(input_grid)

    # Put the dividers back in
    output_grid[input_grid == divider_color] = divider_color

    # Fill in the cells with the special colors
    for x, y, region in grid_of_regions:
        output_grid[region != divider_color] = special_sprite[x, y]

    return output_grid



def generate_input() -> np.ndarray:
    
    divider_color = Color.GRAY

    # make the dividers, which comprise horizontal/vertical bars creating 3x3 cells, with 3 cells in each direction
    cell_size = 3
    n_cells = 3
    divider_size = 1 # the divider is a single pixel
    n_dividers = n_cells - 1
    distance_between_cells = cell_size + divider_size
    m = cell_size*n_cells + divider_size*n_dividers
    grid = np.full((m, m), Color.BLACK)
    for i in range(n_dividers):
        # horizontal dividers
        grid[cell_size + i*(cell_size + divider_size), :] = divider_color
        # vertical dividers
        grid[:, cell_size + i*(cell_size + divider_size)] = divider_color
    
    # pick one of the cells to have exactly 4 colors (the others will have 5)
    special_cell_x, special_cell_y = np.random.randint(3), np.random.randint(3)

    for x in range(3):
        for y in range(3):
            if x == special_cell_x and y == special_cell_y:
                n_colors = 4
            else:
                n_colors = 5

            # extract view of the cell
            # each of the cells is 3x3, but there is a divider in between them, so they are actually 4x4 apart
            cell = grid[x*distance_between_cells : x*distance_between_cells + cell_size,
                        y*distance_between_cells : y*distance_between_cells + cell_size]

            # color the cell by picking random positions and random colors until we have enough colored pixels
            while np.sum(cell!=Color.BLACK) < n_colors:
                # pick a random spot to color
                cell_x, cell_y = np.random.randint(cell_size), np.random.randint(cell_size)
                cell[cell_x, cell_y] = random.choice([color for color in Color.ALL_COLORS if color != Color.BLACK and color != divider_color])

    return grid

from common import *

import numpy as np
from typing import *

# concepts:
# pixel manipulation, growing

# description:
# In the input you will see some number of colored crosses, each of which is 3 pixels tall, 3 pixels wide, and has a single pixel in the center of the cross that is a different color.
# Make the output by growing the cross by 1 pixel north/south/east/west, and growing the center pixel by 2 pixels along each of the 4 diagonals.

def transform(input_grid):

    # extract the 3x3 crosses
    crosses = find_connected_components(input_grid, background=Color.BLACK, monochromatic=False)

    output_grid = input_grid.copy()

    for cross in crosses:
        # find the center
        x, y, w, h = bounding_box(cross)
        center_x, center_y = x + w//2, y + h//2

        # extract the relevant colors
        center_color = cross[center_x, center_y]
        cross_color = cross[cross != Color.BLACK][0]

        # grow the cross
        for output_x in range(x-1, x+w+1):
            for output_y in range(y-1, y+h+1):
                # skip if out of bounds
                if output_x < 0 or output_y < 0 or output_x >= input_grid.shape[0] or output_y >= input_grid.shape[1]:
                    continue
                
                # grow the cross north/south/east/west
                if output_x == center_x or output_y == center_y:
                    output_grid[output_x, output_y] = cross_color
                
                # grow the center diagonally
                if (output_x - center_x) == (output_y - center_y) or (output_x - center_x) == (center_y - output_y):
                    output_grid[output_x, output_y] = center_color

    return output_grid



def generate_input():
    input_grid = np.zeros((20, 20), dtype=int)

    # create 2 crosses
    for cross_number in range(2):
        # create a separate sprite
        cross = np.zeros((3, 3), dtype=int)
        cross_color, center_color = random.sample(list(Color.NOT_BLACK), 2)
        cross[1, :] = cross_color
        cross[:, 1] = cross_color
        cross[1, 1] = center_color

        # find a free place on the grid
        x, y = random_free_location_for_sprite(input_grid, cross)

        # blit the cross to the canvas
        blit_sprite(input_grid, cross, x, y)

    return input_grid