0813.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# import "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% Setup\n",
    "\n",
    "import numpy as np # linear algebra\n",
    "import pandas as pd\n",
    "import json\n",
    "from pathlib import Path\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import operator\n",
    "from collections import Counter\n",
    "import copy\n",
    "from itertools import product, permutations, combinations, combinations_with_replacement\n",
    "from functools import partial\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import colors\n",
    "\n",
    "data_path = Path('/kaggle/input/abstraction-and-reasoning-challenge/')\n",
    "#data_path = Path('data')\n",
    "train_path = data_path / 'training'\n",
    "eval_path = data_path / 'evaluation'\n",
    "test_path = data_path / 'test'\n",
    "\n",
    "train_tasks = { task.stem: json.load(task.open()) for task in train_path.iterdir() } \n",
    "eval_tasks = { task.stem: json.load(task.open()) for task in eval_path.iterdir() }\n",
    "test_tasks = { task.stem: json.load(task.open()) for task in test_path.iterdir() }\n",
    "\n",
    "data = test_tasks\n",
    "\n",
    "cmap = colors.ListedColormap(\n",
    "        ['#000000', '#0074D9','#FF4136','#2ECC40','#FFDC00',\n",
    "         '#AAAAAA', '#F012BE', '#FF851B', '#7FDBFF', '#870C25'])\n",
    "norm = colors.Normalize(vmin=0, vmax=9)\n",
    "\n",
    "def plot_pictures(pictures, labels):\n",
    "    fig, axs = plt.subplots(1, len(pictures), figsize=(2*len(pictures),32))\n",
    "    for i, (pict, label) in enumerate(zip(pictures, labels)):\n",
    "        axs[i].imshow(np.array(pict), cmap=cmap, norm=norm)\n",
    "        axs[i].set_title(label)\n",
    "    plt.show()\n",
    "\n",
    "def plot_sample(sample, predict=None):\n",
    "    \"\"\"\n",
    "    This function plots a sample. sample is an object of the class Sample.\n",
    "    predict is any matrix (numpy ndarray).\n",
    "    \"\"\"\n",
    "    if predict is None:\n",
    "        plot_pictures([sample.inMatrix.m, sample.outMatrix.m], ['Input', 'Output'])\n",
    "    else:\n",
    "        plot_pictures([sample.inMatrix.m, sample.outMatrix.m, predict], ['Input', 'Output', 'Predict'])\n",
    "\n",
    "def plot_task(task):\n",
    "    \"\"\"\n",
    "    Given a task (in its original format), this function plots all of its\n",
    "    matrices.\n",
    "    \"\"\"\n",
    "    len_train = len(task['train'])\n",
    "    len_test  = len(task['test'])\n",
    "    len_max   = max(len_train, len_test)\n",
    "    length    = {'train': len_train, 'test': len_test}\n",
    "    fig, axs  = plt.subplots(len_max, 4, figsize=(15, 15*len_max//4))\n",
    "    for col, mode in enumerate(['train', 'test']):\n",
    "        for idx in range(length[mode]):\n",
    "            axs[idx][2*col+0].axis('off')\n",
    "            axs[idx][2*col+0].imshow(task[mode][idx]['input'], cmap=cmap, norm=norm)\n",
    "            axs[idx][2*col+0].set_title(f\"Input {mode}, {np.array(task[mode][idx]['input']).shape}\")\n",
    "            try:\n",
    "                axs[idx][2*col+1].axis('off')\n",
    "                axs[idx][2*col+1].imshow(task[mode][idx]['output'], cmap=cmap, norm=norm)\n",
    "                axs[idx][2*col+1].set_title(f\"Output {mode}, {np.array(task[mode][idx]['output']).shape}\")\n",
    "            except:\n",
    "                pass\n",
    "        for idx in range(length[mode], len_max):\n",
    "            axs[idx][2*col+0].axis('off')\n",
    "            axs[idx][2*col+1].axis('off')\n",
    "    plt.tight_layout()\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "def flattener(pred):\n",
    "    str_pred = str([row for row in pred])\n",
    "    str_pred = str_pred.replace(', ', '')\n",
    "    str_pred = str_pred.replace('[[', '|')\n",
    "    str_pred = str_pred.replace('][', '|')\n",
    "    str_pred = str_pred.replace(']]', '|')\n",
    "    return str_pred\n",
    "\n",
    "##############################################################################\n",
    "# %% CORE OBJECTS\n",
    "\n",
    "# %% Frontiers\n",
    "class Frontier:\n",
    "    \"\"\"\n",
    "    A Frontier is defined as a straight line with a single color that crosses\n",
    "    all of the matrix. For example, if the matrix has shape MxN, then a\n",
    "    Frontier will have shape Mx1 or 1xN. See the function \"detectFrontiers\"\n",
    "    for details in the implementation.\n",
    "    \n",
    "    ...\n",
    "    \n",
    "    Attributes\n",
    "    ----------\n",
    "    color: int\n",
    "        The color of the frontier\n",
    "    directrion: str\n",
    "        A character ('h' or 'v') determining whether the frontier is horizontal\n",
    "        or vertical\n",
    "    position: tuple\n",
    "        A 2-tuple of ints determining the position of the upper-left pixel of\n",
    "        the frontier\n",
    "    \"\"\"\n",
    "    def __init__(self, color, direction, position):\n",
    "        \"\"\"\n",
    "        direction can be 'h' or 'v' (horizontal, vertical)\n",
    "        color, position and are all integers\n",
    "        \"\"\"\n",
    "        self.color = color\n",
    "        self.direction = direction\n",
    "        self.position = position\n",
    "        \n",
    "    def __eq__(self, other):\n",
    "        if isinstance(other, self.__class__):\n",
    "            return self.__dict__ == other.__dict__\n",
    "        else:\n",
    "            return False\n",
    "        \n",
    "def detectFrontiers(m):\n",
    "    \"\"\"\n",
    "    Returns a list of the Frontiers detected in the matrix m (numpy.ndarray).\n",
    "    \"\"\"\n",
    "    frontiers = []\n",
    "    \n",
    "    # Horizontal lines\n",
    "    if m.shape[0]>1:\n",
    "        for i in range(m.shape[0]):\n",
    "            color = m[i, 0]\n",
    "            isFrontier = True\n",
    "            for j in range(m.shape[1]):\n",
    "                if color != m[i,j]:\n",
    "                    isFrontier = False\n",
    "                    break\n",
    "            if isFrontier:\n",
    "                frontiers.append(Frontier(color, 'h', i))\n",
    "            \n",
    "    # Vertical lines\n",
    "    if m.shape[1]>1:\n",
    "        for j in range(m.shape[1]):\n",
    "            color = m[0, j]\n",
    "            isFrontier = True\n",
    "            for i in range(m.shape[0]):\n",
    "                if color != m[i,j]:\n",
    "                    isFrontier = False\n",
    "                    break\n",
    "            if isFrontier:\n",
    "                frontiers.append(Frontier(color, 'v', j))\n",
    "            \n",
    "    return frontiers\n",
    "\n",
    "# %% Grids\n",
    "class Grid:\n",
    "    \"\"\"\n",
    "    An object of the class Grid is basically a collection of frontiers that\n",
    "    have all the same color.\n",
    "    It is useful to check, for example, whether the cells defined by the grid\n",
    "    always have the same size or not.\n",
    "    \n",
    "    ...\n",
    "    \n",
    "    Attributes\n",
    "    ----------\n",
    "    color: int\n",
    "        The color of the grid\n",
    "    m: numpy.ndarray\n",
    "        The whole matrix\n",
    "    frontiers: list\n",
    "        A list of all the frontiers the grid is composed of\n",
    "    cells: list of list of 2-tuples\n",
    "        cells can be viewed as a 2-dimensional matrix of 2-tuples (Matrix, \n",
    "        position). The first element is an object of the class Matrix, and the\n",
    "        second element is the position of the cell in m.\n",
    "        Each element represents a cell of the grid.\n",
    "    shape: tuple\n",
    "        A 2-tuple of ints representing the number of cells of the grid\n",
    "    nCells: int\n",
    "        Number of cells of the grid\n",
    "    cellList: list\n",
    "        A list of all the cells\n",
    "    allCellsSameShape: bool\n",
    "        Determines whether all the cells of the grid have the same shape (as\n",
    "        matrices).\n",
    "    cellShape: tuple\n",
    "        Only defined if allCellsSameShape is True. Shape of the cells.\n",
    "    allCellsHaveOneColor: bool\n",
    "        Determines whether the ALL of the cells of the grid are composed of\n",
    "        pixels of the same color\n",
    "    \"\"\"\n",
    "    def __init__(self, m, frontiers):\n",
    "        self.color = frontiers[0].color\n",
    "        self.m = m\n",
    "        self.frontiers = frontiers\n",
    "        hPositions = [f.position for f in frontiers if f.direction == 'h']\n",
    "        hPositions.append(-1)\n",
    "        hPositions.append(m.shape[0])\n",
    "        hPositions.sort()\n",
    "        vPositions = [f.position for f in frontiers if f.direction == 'v']\n",
    "        vPositions.append(-1)\n",
    "        vPositions.append(m.shape[1])\n",
    "        vPositions.sort()\n",
    "        # cells is a matrix (list of lists) of 2-tuples (Matrix, position)\n",
    "        self.cells = []\n",
    "        hShape = 0\n",
    "        vShape = 0\n",
    "        for h in range(len(hPositions)-1):\n",
    "            if hPositions[h]+1 == hPositions[h+1]:\n",
    "                continue\n",
    "            self.cells.append([])\n",
    "            for v in range(len(vPositions)-1):\n",
    "                if vPositions[v]+1 == vPositions[v+1]:\n",
    "                    continue\n",
    "                if hShape == 0:\n",
    "                    vShape += 1\n",
    "                self.cells[hShape].append((Matrix(m[hPositions[h]+1:hPositions[h+1], \\\n",
    "                                                   vPositions[v]+1:vPositions[v+1]], \\\n",
    "                                                 detectGrid=False), \\\n",
    "                                          (hPositions[h]+1, vPositions[v]+1)))\n",
    "            hShape += 1\n",
    "            \n",
    "        self.shape = (hShape, vShape) # N of h cells x N of v cells\n",
    "        self.cellList = []\n",
    "        for cellRow in range(len(self.cells)):\n",
    "            for cellCol in range(len(self.cells[0])):\n",
    "                self.cellList.append(self.cells[cellRow][cellCol])\n",
    "        self.allCellsSameShape = len(set([c[0].shape for c in self.cellList])) == 1\n",
    "        if self.allCellsSameShape:\n",
    "            self.cellShape = self.cells[0][0][0].shape\n",
    "            \n",
    "        self.nCells = len(self.cellList)\n",
    "            \n",
    "        # Check whether each cell has one and only one color\n",
    "        self.allCellsHaveOneColor = True\n",
    "        for c in self.cellList:\n",
    "            if c[0].nColors!=1:\n",
    "                self.allCellsHaveOneColor = False\n",
    "                break\n",
    "        \n",
    "        \n",
    "    def __eq__(self, other):\n",
    "        if isinstance(other, self.__class__):\n",
    "            return all([f in other.frontiers for f in self.frontiers])\n",
    "        else:\n",
    "            return False\n",
    "\n",
    "# %% Shapes and subclasses\n",
    "class Shape:\n",
    "    \"\"\"\n",
    "    An object of the class Shape is meant to represent a connected entity of a\n",
    "    Matrix. Its main attribute is \"m\", a 2-dimensional numpy-ndarray of type\n",
    "    np.uint8, in which all the entries have the color of the entity from the\n",
    "    Matrix to be represented, and the rest of the elements are equal to 255.\n",
    "    The attribute m has the smallest shape possible.\n",
    "    For example, in the matrix [[1, 1], [1, 3]] we might detect two different\n",
    "    Shapes: [[1, 1], [1, 255]] and [[3]].\n",
    "    ...\n",
    "    Main attributes\n",
    "    ---------------\n",
    "    m: numpy.ndarray\n",
    "        Described above.\n",
    "    nPixels: int\n",
    "        Number of pixels of the Shape. This is, number of elements that are\n",
    "        different from 255 in m.\n",
    "    position: tuple (int, int)\n",
    "        Position of the Shape in the original Matrix (upper-left corner of m).\n",
    "    background: int\n",
    "        Background color of the original Matrix.\n",
    "    isBorder: bool\n",
    "        Checks whether the Shape touches the border of the original Matrix.\n",
    "    colors: set\n",
    "        Set of colors present in the Shape.\n",
    "    nColors: int\n",
    "        Number of colors present in the Shape.\n",
    "    colorCount: collections.Counter\n",
    "        Dictionary whose keys are the colors that appear in the Shape, and\n",
    "        their values represent the amount of pixels with that color.\n",
    "    isSquare: bool\n",
    "        Determines whether the Shape is a full square or not.\n",
    "    isRectangle: bool\n",
    "        Determines whether the Shape is a full rectangle or not.\n",
    "        \n",
    "    Methods\n",
    "    -------\n",
    "    hasSameShape(other, sameColor=False, samePosition=False, rotation=False,\n",
    "                 mirror=False, scaling=False)\n",
    "        Checks whether the Shape \"other\" is the same Shape as the current one,\n",
    "        modulo the given parameters.\n",
    "    \n",
    "    \"\"\"\n",
    "    def __init__(self, m, xPos, yPos, background, isBorder):\n",
    "        # pixels is a 2xn numpy array, where n is the number of pixels\n",
    "        self.m = m\n",
    "        self.nPixels = m.size - np.count_nonzero(m==255)\n",
    "        self.background = background\n",
    "        self.shape = m.shape\n",
    "        self.position = (xPos, yPos)\n",
    "        self.pixels = set([(i,j) for i,j in np.ndindex(m.shape) if m[i,j]!=255])\n",
    "            \n",
    "        # Is the shape in the border?\n",
    "        self.isBorder = isBorder\n",
    "        \n",
    "        # Which colors does the shape have?\n",
    "        self.colors = set(np.unique(m)) - set([255])\n",
    "        self.nColors = len(self.colors)\n",
    "        if self.nColors==1:\n",
    "            self.color = next(iter(self.colors))\n",
    "\n",
    "        self.colorCount = Counter(self.m.flatten()) + Counter({0:0, 1:0, 2:0, 3:0, 4:0, 5:0, 6:0, 7:0, 8:0, 9:0})\n",
    "        del self.colorCount[255]\n",
    "\n",
    "        # Symmetries\n",
    "        self.lrSymmetric = np.array_equal(self.m, np.fliplr(self.m))\n",
    "        self.udSymmetric = np.array_equal(self.m, np.flipud(self.m))\n",
    "        if self.m.shape[0] == self.m.shape[1]:\n",
    "            self.d1Symmetric = np.array_equal(self.m, self.m.T)\n",
    "            self.d2Symmetric = np.array_equal(np.fliplr(self.m), (np.fliplr(self.m)).T)\n",
    "        else:\n",
    "            self.d1Symmetric = False\n",
    "            self.d2Symmetric = False\n",
    "\n",
    "        self.isRectangle = 255 not in np.unique(m)\n",
    "        self.isSquare = self.isRectangle and self.shape[0]==self.shape[1]\n",
    "        \n",
    "        if self.isRectangle and self.nColors > 1:\n",
    "            self.subshapes = detectShapes(self.m, background=self.colorCount.most_common(1)[0][0],\\\n",
    "                                              singleColor=True, diagonals=False)\n",
    "        self.nHoles = self.getNHoles()\n",
    "        \n",
    "        if self.nColors==1:\n",
    "            self.isFullFrame = self.isFullFrame()\n",
    "            self.isPartialFrame = self.isPartialFrame()\n",
    "        else:\n",
    "            self.isFullFrame = False\n",
    "            self.isPartialFrame = False\n",
    "        \n",
    "        if self.nColors==1:\n",
    "            self.boolFeatures = []\n",
    "            for c in range(10):\n",
    "                self.boolFeatures.append(self.color==c)\n",
    "            self.boolFeatures.append(self.isBorder)\n",
    "            self.boolFeatures.append(not self.isBorder)\n",
    "            self.boolFeatures.append(self.lrSymmetric)\n",
    "            self.boolFeatures.append(self.udSymmetric)\n",
    "            self.boolFeatures.append(self.d1Symmetric)\n",
    "            self.boolFeatures.append(self.d2Symmetric)\n",
    "            self.boolFeatures.append(self.isSquare)\n",
    "            self.boolFeatures.append(self.isRectangle)\n",
    "            for nPix in range(1,30):\n",
    "                self.boolFeatures.append(self.nPixels==nPix)\n",
    "            self.boolFeatures.append((self.nPixels%2)==0)\n",
    "            self.boolFeatures.append((self.nPixels%2)==1)\n",
    "    \n",
    "    def hasSameShape(self, other, sameColor=False, samePosition=False, rotation=False, \\\n",
    "                     mirror=False, scaling=False):\n",
    "        \"\"\"\n",
    "        Checks whether the Shape \"other\" is the same Shape as the current one,\n",
    "        modulo the given parameters.\n",
    "        ...\n",
    "        Parameters\n",
    "        ----------\n",
    "        other: Shape\n",
    "            Shape to be compared with the current one.\n",
    "        sameColor: bool\n",
    "            True if we require the colors to be the same. Default is False.\n",
    "        samePosition: bool\n",
    "            True if we require both Shapes to be in the same position of the\n",
    "            Matrix. Default is False.\n",
    "        rotation: bool\n",
    "            True if we allow one Shape to be the rotated version of the other\n",
    "            Shape. Rotations of 90, 180 and 270 degrees are considered.\n",
    "            Default is False.\n",
    "        mirror: bool\n",
    "            True if we allow one Shape to be the other one mirrored. Only\n",
    "            Left-Right and Up-Down mirrorings are always considered. Default\n",
    "            is False.\n",
    "        scaling: bool\n",
    "            True if we allow one shape to be equal to the other Shape modulo\n",
    "            scaling. Default is False.\n",
    "        \"\"\"\n",
    "        if samePosition:\n",
    "            if self.position != other.position:\n",
    "                return False\n",
    "        if sameColor:\n",
    "            m1 = self.m\n",
    "            m2 = other.m\n",
    "        else:\n",
    "            m1 = self.shapeDummyMatrix()\n",
    "            m2 = other.shapeDummyMatrix()\n",
    "        if scaling and m1.shape!=m2.shape:\n",
    "            def multiplyPixels(matrix, factor):\n",
    "                m = np.zeros(tuple(s * f for s, f in zip(matrix.shape, factor)), dtype=np.uint8)\n",
    "                for i,j in np.ndindex(matrix.shape):\n",
    "                    for k,l in np.ndindex(factor):\n",
    "                        m[i*factor[0]+k, j*factor[1]+l] = matrix[i,j]\n",
    "                return m\n",
    "            \n",
    "            if (m1.shape[0]%m2.shape[0])==0 and (m1.shape[1]%m2.shape[1])==0:\n",
    "                factor = (int(m1.shape[0]/m2.shape[0]), int(m1.shape[1]/m2.shape[1]))\n",
    "                m2 = multiplyPixels(m2, factor)\n",
    "            elif (m2.shape[0]%m1.shape[0])==0 and (m2.shape[1]%m1.shape[1])==0:\n",
    "                factor = (int(m2.shape[0]/m1.shape[0]), int(m2.shape[1]/m1.shape[1]))\n",
    "                m1 = multiplyPixels(m1, factor)\n",
    "            elif rotation and (m1.shape[0]%m2.shape[1])==0 and (m1.shape[1]%m2.shape[0])==0:\n",
    "                factor = (int(m1.shape[0]/m2.shape[1]), int(m1.shape[1]/m2.shape[0]))\n",
    "                m2 = multiplyPixels(m2, factor)\n",
    "            elif rotation and (m2.shape[0]%m1.shape[1])==0 and (m2.shape[1]%m1.shape[0])==0:\n",
    "                factor = (int(m2.shape[0]/m1.shape[1]), int(m2.shape[1]/m1.shape[0]))\n",
    "                m1 = multiplyPixels(m1, factor)\n",
    "            else:\n",
    "                return False\n",
    "        if rotation and not mirror:\n",
    "            if any([np.array_equal(m1, np.rot90(m2,x)) for x in range(1,4)]):\n",
    "                return True\n",
    "        if mirror and not rotation:\n",
    "            if np.array_equal(m1, np.fliplr(m2)) or np.array_equal(m1, np.flipud(m2)):\n",
    "                return True\n",
    "        if mirror and rotation:\n",
    "            for x in range(1, 4):\n",
    "                if any([np.array_equal(m1, np.rot90(m2,x))\\\n",
    "                        or np.array_equal(m1, np.fliplr(np.rot90(m2,x))) for x in range(0,4)]):\n",
    "                    return True               \n",
    "                \n",
    "        return np.array_equal(m1,m2)\n",
    "    \n",
    "    def __eq__(self, other):\n",
    "        \"\"\"\n",
    "        Two Shapes are considered equal if their matrices m are equal.\n",
    "        \"\"\"\n",
    "        if isinstance(other, self.__class__):\n",
    "            if self.shape != other.shape:\n",
    "                return False\n",
    "            return np.array_equal(self.m, other.m)\n",
    "        else:\n",
    "            return False\n",
    " \n",
    "    def isSubshape(self, other, sameColor=False, rotation=False, mirror=False):\n",
    "        \"\"\"\n",
    "        Checks whether the Shape \"other\" is a subShape of the current one,\n",
    "        modulo the given parameters.\n",
    "        ...\n",
    "        Parameters\n",
    "        ----------\n",
    "        other: Shape\n",
    "            Shape to be compared with the current one.\n",
    "        sameColor: bool\n",
    "            True if we require the colors to be the same. Default is False.\n",
    "        rotation: bool\n",
    "            True if we allow one Shape to be the rotated version of the other\n",
    "            Shape. Rotations of 90, 180 and 270 degrees are considered.\n",
    "            Default is False.\n",
    "        mirror: bool\n",
    "            True if we allow one Shape to be the other one mirrored. Only\n",
    "            Left-Right and Up-Down mirrorings are always considered. Default\n",
    "            is False.\n",
    "        \"\"\"\n",
    "        #return positions\n",
    "        if rotation:\n",
    "            m1 = self.m\n",
    "            for x in range(1,4):\n",
    "                if Shape(np.rot90(m1,x), 0, 0, 0, self.isBorder).isSubshape(other, sameColor, False, mirror):\n",
    "                    return True\n",
    "        if mirror == 'lr':\n",
    "            if Shape(self.m[::,::-1], 0, 0, 0, self.isBorder).isSubshape(other, sameColor, rotation, False):\n",
    "                return True\n",
    "        if mirror == 'ud':\n",
    "            if Shape(self.m[::-1,::], 0, 0, 0, self.isBorder).isSubshape(other, sameColor, rotation, False):\n",
    "                return True\n",
    "        if sameColor:\n",
    "            if hasattr(self,'color') and hasattr(other,'color') and self.color != other.color:\n",
    "                return False\n",
    "        if any(other.shape[i] < self.shape[i] for i in [0,1]):\n",
    "            return False\n",
    "        \n",
    "        for yIn in range(other.shape[1] - self.shape[1] + 1):\n",
    "            for xIn in range(other.shape[0] - self.shape[0] + 1):\n",
    "                if sameColor:\n",
    "                    if np.all(np.logical_or((self.m == other.m[xIn: xIn + self.shape[0], yIn: yIn + self.shape[1]]),\\\n",
    "                                            self.m==255)):\n",
    "                        return True\n",
    "                else:\n",
    "                    if set([tuple(np.add(ps,[xIn,yIn])) for ps in self.pixels]) <= other.pixels:\n",
    "                        return True\n",
    "        return False\n",
    "    \n",
    "    def shapeDummyMatrix(self):\n",
    "        \"\"\"\n",
    "        Returns the smallest possible matrix containing the shape. The values\n",
    "        of the matrix are ones and zeros, depending on whether the pixel is a\n",
    "        Shape pixel or not.\n",
    "        \"\"\"\n",
    "        return (self.m!=255).astype(np.uint8) \n",
    "    \n",
    "    def hasFeatures(self, features):\n",
    "        \"\"\"\n",
    "        Given a list of features, this function returns True if the current\n",
    "        Shape has the given features. Otherwise, it returns False.\n",
    "        The parameter features is a list of boolean values.\n",
    "        \"\"\"\n",
    "        for i in range(len(features)):\n",
    "            if features[i] and not self.boolFeatures[i]:\n",
    "                return False\n",
    "        return True\n",
    "\n",
    "    def getNHoles(self):\n",
    "        \"\"\"\n",
    "        Returns the number of holes of the Shape.\n",
    "        \"\"\"\n",
    "        nHoles = 0\n",
    "        m = self.m\n",
    "        seen = np.zeros((self.shape[0], self.shape[1]), dtype=np.bool)\n",
    "        def isInHole(i,j):\n",
    "            if i<0 or j<0 or i>self.shape[0]-1 or j>self.shape[1]-1:\n",
    "                return False\n",
    "            if seen[i,j] or m[i,j] != 255:\n",
    "                return True\n",
    "            seen[i,j] = True\n",
    "            ret = isInHole(i+1,j)*isInHole(i-1,j)*isInHole(i,j+1)*isInHole(i,j-1)\n",
    "            return ret\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            if m[i,j] == 255 and not seen[i,j]:\n",
    "                if isInHole(i,j):\n",
    "                    nHoles += 1\n",
    "        return nHoles        \n",
    "\n",
    "    def isPartialFrame(self):\n",
    "        \"\"\"\n",
    "        Checks whether the Shape is a partial Frame.\n",
    "        \"\"\"\n",
    "        if self.shape[0] < 4 or self.shape[1] < 4 or len(self.pixels) < 4:\n",
    "            return False\n",
    "        if len(np.unique(self.m[1:-1,1:-1])) > 1 or self.color in np.unique(self.m[1:-1,1:-1]):\n",
    "            return False\n",
    "        return True\n",
    "    \n",
    "    def isFullFrame(self):\n",
    "        \"\"\"\n",
    "        Checks whether the Shape is a full Frame.\n",
    "        \"\"\"\n",
    "        if self.shape[0]<3 or self.shape[1]<3:\n",
    "            return False\n",
    "        for i in range(self.shape[0]):\n",
    "            if self.m[i,0]==255 or self.m[i,self.shape[1]-1]==255:\n",
    "                return False\n",
    "        for j in range(self.shape[1]):\n",
    "            if self.m[0,j]==255 or self.m[self.shape[0]-1,j]==255:\n",
    "                return False\n",
    "            \n",
    "        # We require fullFrames to have less than 20% of the pixels inside the\n",
    "        # frame of the same color of the frame\n",
    "        \n",
    "        if self.nPixels - 2*(self.shape[0]+self.shape[1]-2) < 0.2*(self.shape[0]-2)*(self.shape[1]-2):\n",
    "            return True\n",
    "        \n",
    "        return False\n",
    "\n",
    "def detectShapesByColor(x, background):\n",
    "    shapes = []\n",
    "    for c in range(10):\n",
    "        if c == background or c not in x:\n",
    "            continue\n",
    "        mc = np.zeros(x.shape, dtype=int)\n",
    "        mc[x==c] = c\n",
    "        mc[x!=c] = 255\n",
    "        x1, x2, y1, y2 = 0, mc.shape[0]-1, 0, mc.shape[1]-1\n",
    "        while x1 <= x2 and np.all(mc[x1,:] == 255):\n",
    "            x1 += 1 \n",
    "        while x2 >= x1 and np.all(mc[x2,:] == 255):\n",
    "            x2 -= 1\n",
    "        while y1 <= y2 and np.all(mc[:,y1] == 255):\n",
    "            y1 += 1\n",
    "        while y2 >= y1 and np.all(mc[:,y2] == 255):\n",
    "            y2 -= 1\n",
    "        m = mc[x1:x2+1,y1:y2+1]\n",
    "        s = Shape(m.copy(), x1, y1, background, False)\n",
    "        shapes.append(s)\n",
    "    return shapes\n",
    "\n",
    "def detectShapes(x, background, singleColor=False, diagonals=False):\n",
    "    \"\"\"\n",
    "    Given a numpy array x (2D), returns a list of the Shapes present in x.\n",
    "    ...\n",
    "    Parameters\n",
    "    ----------\n",
    "    x: numpy.ndarray\n",
    "        The matrix we want to detect Shapes from.\n",
    "    background: int\n",
    "        The background color of the matrix\n",
    "    singleColor: bool\n",
    "        True if we want all the shapes to have only one color. Default is\n",
    "        False.\n",
    "    diagonals: bool\n",
    "        True if we allow pixels of the Shape to be connected diagonally.\n",
    "        Default is False.\n",
    "    \"\"\"\n",
    "    # Helper function to add pixels to a shape\n",
    "    def addPixelsAround(i,j):\n",
    "        def addPixel(i,j):\n",
    "            if i < 0 or j < 0 or i > iMax or j > jMax or seen[i,j] == True:\n",
    "                return\n",
    "            if singleColor:\n",
    "                if x[i,j] != color:\n",
    "                    return\n",
    "                newShape[i,j] = color\n",
    "            else:\n",
    "                if x[i,j] == background:\n",
    "                    return\n",
    "                newShape[i,j] = x[i,j]\n",
    "            seen[i,j] = True                \n",
    "            addPixelsAround(i,j)\n",
    "        \n",
    "        addPixel(i-1,j)\n",
    "        addPixel(i+1,j)\n",
    "        addPixel(i,j-1)\n",
    "        addPixel(i,j+1)\n",
    "        \n",
    "        if diagonals:\n",
    "            addPixel(i-1,j-1)\n",
    "            addPixel(i-1,j+1)\n",
    "            addPixel(i+1,j-1)\n",
    "            addPixel(i+1,j+1)\n",
    "            \n",
    "    def crop(matrix):\n",
    "        ret = matrix.copy()\n",
    "        for k in range(x.shape[0]):\n",
    "            if any(matrix[k,:] != 255): # -1==255 for dtype=np.uint8\n",
    "                x0 = k\n",
    "                break\n",
    "        for k in reversed(range(x.shape[0])):\n",
    "            if any(matrix[k,:] != 255): # -1==255 for dtype=np.uint8\n",
    "                x1 = k\n",
    "                break\n",
    "        for k in range(x.shape[1]):\n",
    "            if any(matrix[:,k] != 255): # -1==255 for dtype=np.uint8\n",
    "                y0 = k\n",
    "                break\n",
    "        for k in reversed(range(x.shape[1])):\n",
    "            if any(matrix[:,k] != 255): # -1==255 for dtype=np.uint8\n",
    "                y1 = k\n",
    "                break\n",
    "        return ret[x0:x1+1,y0:y1+1], x0, y0\n",
    "                \n",
    "    shapes = []\n",
    "    seen = np.zeros(x.shape, dtype=bool)\n",
    "    iMax = x.shape[0]-1\n",
    "    jMax = x.shape[1]-1\n",
    "    for i, j in np.ndindex(x.shape):\n",
    "        if seen[i,j] == False:\n",
    "            seen[i,j] = True\n",
    "            if not singleColor and x[i,j]==background:\n",
    "                continue\n",
    "            newShape = np.full((x.shape), -1, dtype=np.uint8)\n",
    "            newShape[i,j] = x[i,j]\n",
    "            if singleColor:\n",
    "                color = x[i][j]\n",
    "            addPixelsAround(i,j)\n",
    "            m, xPos, yPos = crop(newShape)\n",
    "            isBorder = xPos==0 or yPos==0 or (xPos+m.shape[0]==x.shape[0]) or (yPos+m.shape[1]==x.shape[1])\n",
    "            s = Shape(m.copy(), xPos, yPos, background, isBorder)\n",
    "            shapes.append(s)\n",
    "    return shapes\n",
    "\n",
    "def detectIsolatedPixels(matrix, dShapeList):\n",
    "    pixList = []\n",
    "    for sh in dShapeList:\n",
    "        if sh.nPixels > 1 or sh.color == matrix.backgroundColor:\n",
    "            continue\n",
    "        else:\n",
    "            cc = set()\n",
    "            for i,j in np.ndindex(3, 3):\n",
    "                if i - 1 + sh.position[0] < matrix.shape[0] and i - 1 + sh.position[0] >= 0 \\\n",
    "                        and j - 1 + sh.position[1] < matrix.shape[1] and j - 1 + sh.position[1] >= 0:\n",
    "                    cc  = cc.union(set([matrix.m[i - 1 + sh.position[0],j - 1 + sh.position[1]]]))\n",
    "            if len(cc) == 2:\n",
    "                pixList.append(sh)\n",
    "    return pixList\n",
    "\n",
    "# %% Class Matrix\n",
    "class Matrix():\n",
    "    \"\"\"\n",
    "    An object of the class Matrix stores all the relevant information about\n",
    "    any matrix, be it input matrix, output matrix, from the training samples\n",
    "    or from the test samples.\n",
    "    ...\n",
    "    Main Attributes\n",
    "    ----------\n",
    "    m: numpy.ndarray\n",
    "        The matrix\n",
    "    shape: tuple (int, int)\n",
    "        Shape of m\n",
    "    colors: set\n",
    "        Colors present in the matrix\n",
    "    colorCount: collections.Counter\n",
    "        Dictionary containing the number of appearances of each color.\n",
    "    backgroundColor: int\n",
    "        The background color of the matrix\n",
    "    shapes: list (of Shapes)\n",
    "        List of Shapes present in the matrix. The Shapes have only one color,\n",
    "        and cannot contain diagonal connections.\n",
    "    dShapes: list (of Shapes)\n",
    "        List of Shapes present in the matrix. The Shapes have only one color,\n",
    "        and allow diagonal connections.\n",
    "    multicolorShapes: list (of Shapes)\n",
    "        List of Shapes present in the matrix. The Shapes can have many colors,\n",
    "        and cannot contain diagonal connections. They just cannot contain the\n",
    "        background color.\n",
    "    multicolorDShapes: list (of Shapes)\n",
    "        List of Shapes present in the matrix. The Shapes can have many colors,\n",
    "        and allow diagonal connections. They just cannot contain the\n",
    "        background color.\n",
    "    frontiers: list (of Frontiers)\n",
    "        List of Frontiers of the matrix.\n",
    "    isGrid: bool\n",
    "        True if the matrix contains a symmetric grid. False otherwise.\n",
    "    grid: Grid\n",
    "        The Grid present in the matrix (there can only be one). This attribute\n",
    "        is only defined if isGrid.\n",
    "    isAsymmetricGrid: bool\n",
    "        True if the matrix contains an asymmetric grid. False otherwise.\n",
    "    asymmetricGrid: Grid\n",
    "        The asymmetric Grid present in the matrix (there can only be one).\n",
    "        This attribute is only defined if isAsymmetricGrid.\n",
    "    \"\"\"\n",
    "    def __init__(self, m, detectGrid=True, backgroundColor=None):\n",
    "        if type(m) == Matrix:\n",
    "            return m\n",
    "        \n",
    "        self.m = np.array(m)\n",
    "        \n",
    "        # interesting properties:\n",
    "        \n",
    "        # Dimensions\n",
    "        self.shape = self.m.shape\n",
    "        self.nElements = self.m.size\n",
    "        \n",
    "        # Counter of colors\n",
    "        self.colorCount = self.getColors()\n",
    "        self.colors = set(self.colorCount.keys())\n",
    "        self.nColors = len(self.colorCount)\n",
    "        \n",
    "        # Background color\n",
    "        if backgroundColor==None:\n",
    "            self.backgroundColor = max(self.colorCount, key=self.colorCount.get)\n",
    "        else:\n",
    "            self.backgroundColor = backgroundColor\n",
    "        \n",
    "        # Shapes\n",
    "        self.shapes = detectShapes(self.m, self.backgroundColor, singleColor=True)\n",
    "        self.nShapes = len(self.shapes)\n",
    "        self.dShapes = detectShapes(self.m, self.backgroundColor, singleColor=True, diagonals=True)\n",
    "        self.nDShapes = len(self.dShapes)\n",
    "        self.fullFrames = [shape for shape in self.shapes if shape.isFullFrame]\n",
    "        self.fullFrames = sorted(self.fullFrames, key=lambda x: x.shape[0]*x.shape[1], reverse=True)\n",
    "        self.shapesByColor = detectShapesByColor(self.m, self.backgroundColor)\n",
    "        self.partialFrames = [shape for shape in self.shapesByColor if shape.isPartialFrame]\n",
    "        self.isolatedPixels = detectIsolatedPixels(self, self.dShapes)\n",
    "        self.nIsolatedPixels = len(self.isolatedPixels)\n",
    "        \n",
    "        self.shapeColorCounter = Counter([s.color for s in self.shapes])\n",
    "        self.blanks = []\n",
    "        for s in self.shapes:\n",
    "            if s.isRectangle and self.shapeColorCounter[s.color]==1:\n",
    "                self.blanks.append(s)\n",
    "            \n",
    "        # Frontiers\n",
    "        self.frontiers = detectFrontiers(self.m)\n",
    "        self.frontierColors = [f.color for f in self.frontiers]\n",
    "        if len(self.frontiers) == 0:\n",
    "            self.allFrontiersEqualColor = False\n",
    "        else: self.allFrontiersEqualColor = (self.frontierColors.count(self.frontiers[0]) ==\\\n",
    "                                         len(self.frontiers))\n",
    "        # Check if it's a grid and the dimensions of the cells\n",
    "        self.isGrid = False\n",
    "        self.isAsymmetricGrid = False\n",
    "        if detectGrid:\n",
    "            for fc in set(self.frontierColors):\n",
    "                possibleGrid = [f for f in self.frontiers if f.color==fc]\n",
    "                possibleGrid = Grid(self.m, possibleGrid)\n",
    "                if possibleGrid.nCells>1:\n",
    "                    if possibleGrid.allCellsSameShape:\n",
    "                        self.grid = copy.deepcopy(possibleGrid)\n",
    "                        self.isGrid = True\n",
    "                        self.asymmetricGrid = copy.deepcopy(possibleGrid)\n",
    "                        self.isAsymmetricGrid = True\n",
    "                        break\n",
    "                    else:\n",
    "                        self.asymmetricGrid = copy.deepcopy(possibleGrid)\n",
    "                        self.isAsymmetricGrid=True\n",
    "                        \n",
    "        # Shape-based backgroundColor\n",
    "        if not self.isGrid:\n",
    "            for shape in self.shapes:\n",
    "                if shape.shape==self.shape:\n",
    "                    self.backgroundColor = shape.color\n",
    "                    break\n",
    "        # Define multicolor shapes based on the background color\n",
    "        self.multicolorShapes = detectShapes(self.m, self.backgroundColor)\n",
    "        self.multicolorDShapes = detectShapes(self.m, self.backgroundColor, diagonals=True)\n",
    "        self.dummyMatrix = (self.m!=self.backgroundColor).astype(np.uint8) \n",
    "        # Symmetries\n",
    "        self.lrSymmetric = np.array_equal(self.m, np.fliplr(self.m))\n",
    "        # Up-Down\n",
    "        self.udSymmetric = np.array_equal(self.m, np.flipud(self.m))\n",
    "        # Diagonals (only if square)\n",
    "        if self.m.shape[0] == self.m.shape[1]:\n",
    "            self.d1Symmetric = np.array_equal(self.m, self.m.T)\n",
    "            self.d2Symmetric = np.array_equal(np.fliplr(self.m), (np.fliplr(self.m)).T)\n",
    "        else:\n",
    "            self.d1Symmetric = False\n",
    "            self.d2Symmetric = False\n",
    "        self.totalSymmetric = self.lrSymmetric and self.udSymmetric and \\\n",
    "        self.d1Symmetric and self.d2Symmetric\n",
    "        \n",
    "        self.fullBorders = []\n",
    "        for f in self.frontiers:\n",
    "            if f.color != self.backgroundColor:\n",
    "                if f.position==0:\n",
    "                    self.fullBorders.append(f)\n",
    "                elif (f.direction=='h' and f.position==self.shape[0]-1) or\\\n",
    "                (f.direction=='v' and f.position==self.shape[1]-1):\n",
    "                    self.fullBorders.append(f)\n",
    "               \n",
    "        self.isVertical = False\n",
    "        self.isHorizontal = False\n",
    "        if len(self.frontiers)!=0:\n",
    "            self.isVertical = all([f.direction=='v' for f in self.frontiers])\n",
    "            self.isHorizontal = all([f.direction=='h' for f in self.frontiers])\n",
    "    \n",
    "    def getColors(self):\n",
    "        unique, counts = np.unique(self.m, return_counts=True)\n",
    "        return dict(zip(unique, counts))\n",
    "    \n",
    "    def getShapes(self, color=None, bigOrSmall=None, isBorder=None, diag=False):\n",
    "        \"\"\"\n",
    "        Return a list of the shapes meeting the required specifications.\n",
    "        \"\"\"\n",
    "        if diag:\n",
    "            candidates = self.dShapes\n",
    "        else:\n",
    "            candidates = self.shapes\n",
    "        if color != None:\n",
    "            candidates = [c for c in candidates if c.color == color]\n",
    "        if isBorder==True:\n",
    "            candidates = [c for c in candidates if c.isBorder]\n",
    "        if isBorder==False:\n",
    "            candidates = [c for c in candidates if not c.isBorder]\n",
    "        if len(candidates) ==  0:\n",
    "            return []\n",
    "        sizes = [c.nPixels for c in candidates]\n",
    "        if bigOrSmall == \"big\":\n",
    "            maxSize = max(sizes)\n",
    "            return [c for c in candidates if c.nPixels==maxSize]\n",
    "        elif bigOrSmall == \"small\":\n",
    "            minSize = min(sizes)\n",
    "            return [c for c in candidates if c.nPixels==minSize]\n",
    "        else:\n",
    "            return candidates\n",
    "        \n",
    "    def followsColPattern(self):\n",
    "        \"\"\"\n",
    "        This function checks whether the matrix follows a pattern of lines or\n",
    "        columns being always the same (task 771 for example).\n",
    "        Meant to be used for the output matrix mainly.\n",
    "        It returns a number (length of the pattern) and \"row\" or \"col\".\n",
    "        \"\"\"\n",
    "        m = self.m.copy()\n",
    "        col0 = m[:,0]\n",
    "        for i in range(1,int(m.shape[1]/2)+1):\n",
    "            if np.all(col0 == m[:,i]):\n",
    "                isPattern=True\n",
    "                for j in range(i):\n",
    "                    k=0\n",
    "                    while k*i+j < m.shape[1]:\n",
    "                        if np.any(m[:,j] != m[:,k*i+j]):\n",
    "                            isPattern=False\n",
    "                            break\n",
    "                        k+=1\n",
    "                    if not isPattern:\n",
    "                        break\n",
    "                if isPattern:\n",
    "                    return i\n",
    "        return False\n",
    "    \n",
    "    def followsRowPattern(self):\n",
    "        m = self.m.copy()\n",
    "        row0 = m[0,:]\n",
    "        for i in range(1,int(m.shape[0]/2)+1):\n",
    "            if np.all(row0 == m[i,:]):\n",
    "                isPattern=True\n",
    "                for j in range(i):\n",
    "                    k=0\n",
    "                    while k*i+j < m.shape[0]:\n",
    "                        if np.any(m[j,:] != m[k*i+j,:]):\n",
    "                            isPattern=False\n",
    "                            break\n",
    "                        k+=1\n",
    "                    if not isPattern:\n",
    "                        break\n",
    "                if isPattern:\n",
    "                    return i\n",
    "        return False\n",
    "    \n",
    "    def isUniqueShape(self, shape):\n",
    "        count = 0\n",
    "        for sh in self.shapes:\n",
    "            if sh.hasSameShape(shape):\n",
    "                count += 1\n",
    "        if count==1:\n",
    "            return True\n",
    "        return False\n",
    "    \n",
    "    def getShapeAttributes(self, backgroundColor=0, singleColor=True, diagonals=True):\n",
    "        '''\n",
    "        Returns list of shape attributes that matches list of shapes\n",
    "        Add:\n",
    "            - is border\n",
    "            - has neighbors\n",
    "            - is reference\n",
    "            - is referenced\n",
    "        '''\n",
    "        if singleColor: \n",
    "            if diagonals:   \n",
    "                shapeList = [sh for sh in self.dShapes]\n",
    "            else:   \n",
    "                shapeList = [sh for sh in self.shapes]\n",
    "            if len([sh for sh in shapeList if sh.color != backgroundColor]) == 0:\n",
    "                return [set() for sh in shapeList]\n",
    "        else:\n",
    "            if diagonals: \n",
    "                shapeList = [sh for sh in self.multicolorDShapes]\n",
    "            else:\n",
    "                shapeList = [sh for sh in self.multicolorShapes]\n",
    "            if len(shapeList) == 0:\n",
    "                return [set()]\n",
    "        attrList =[[] for i in range(len(shapeList))]\n",
    "        if singleColor:\n",
    "            cc = Counter([sh.color for sh in shapeList])\n",
    "        if singleColor:\n",
    "            sc = Counter([sh.nPixels for sh in shapeList if sh.color != backgroundColor])\n",
    "        else:\n",
    "            sc = Counter([sh.nPixels for sh in shapeList])\n",
    "        largest, smallest, mcopies, mcolors = -1, 1000, 0, 0\n",
    "        if singleColor:\n",
    "            maxH, minH = max([sh.nHoles for sh in shapeList if sh.color != backgroundColor]),\\\n",
    "                            min([sh.nHoles for sh in shapeList if sh.color != backgroundColor])\n",
    "        ila, ism = [], []\n",
    "        for i in range(len(shapeList)):\n",
    "            #color count\n",
    "            if singleColor:\n",
    "                if shapeList[i].color == backgroundColor:\n",
    "                    attrList[i].append(-1)\n",
    "                    continue\n",
    "                else:\n",
    "                    attrList[i].append(shapeList[i].color)\n",
    "            else:\n",
    "                attrList[i].append(shapeList[i].nColors)\n",
    "                if shapeList[i].nColors > mcolors:\n",
    "                    mcolors = shapeList[i].nColors\n",
    "            #copies\n",
    "            if singleColor:\n",
    "                attrList[i] = [np.count_nonzero([np.all(shapeList[i].pixels == osh.pixels) for osh in shapeList])] + attrList[i]\n",
    "                if attrList[i][0] > mcopies:\n",
    "                    mcopies = attrList[i][0]\n",
    "            else: \n",
    "                attrList[i] = [np.count_nonzero([shapeList[i] == osh for osh in shapeList])] + attrList[i]\n",
    "                if attrList[i][0] > mcopies:\n",
    "                    mcopies = attrList[i][0]\n",
    "            #unique color?\n",
    "            if singleColor:\n",
    "                if cc[shapeList[i].color] == 1:\n",
    "                    attrList[i].append('UnCo')\n",
    "            #more of x color?\n",
    "            if not singleColor:\n",
    "                for c in range(10):\n",
    "                    if shapeList[i].colorCount[c] > 0 and  shapeList[i].colorCount[c] == max([sh.colorCount[c] for sh in shapeList]):\n",
    "                        attrList[i].append('mo'+str(c))    \n",
    "            #largest?\n",
    "            if len(shapeList[i].pixels) >= largest:\n",
    "                ila += [i]\n",
    "                if len(shapeList[i].pixels) > largest:\n",
    "                    largest = len(shapeList[i].pixels)\n",
    "                    ila = [i]\n",
    "            #smallest?\n",
    "            if len(shapeList[i].pixels) <= smallest:\n",
    "                ism += [i]\n",
    "                if len(shapeList[i].pixels) < smallest:\n",
    "                    smallest = len(shapeList[i].pixels)\n",
    "                    ism = [i]\n",
    "            #unique size\n",
    "            if sc[shapeList[i].nPixels] == 1 and len(sc) == 2:\n",
    "                attrList[i].append('UnSi')\n",
    "            #symmetric?\n",
    "            if shapeList[i].lrSymmetric:\n",
    "                attrList[i].append('LrSy')\n",
    "            else:\n",
    "                attrList[i].append('NlrSy')\n",
    "            if shapeList[i].udSymmetric:\n",
    "                attrList[i].append('UdSy')\n",
    "            else:\n",
    "                attrList[i].append('NudSy')\n",
    "            if shapeList[i].d1Symmetric: \n",
    "                attrList[i].append('D1Sy')\n",
    "            else:\n",
    "                attrList[i].append('ND1Sy')\n",
    "            if shapeList[i].d2Symmetric:\n",
    "                attrList[i].append('D2Sy')\n",
    "            else:\n",
    "                attrList[i].append('ND2Sy')\n",
    "            attrList[i].append(shapeList[i].position)\n",
    "            #pixels\n",
    "            if len(shapeList[i].pixels) == 1:\n",
    "                attrList[i].append('PiXl')\n",
    "            #holes\n",
    "            if singleColor:\n",
    "                if maxH>minH:\n",
    "                    if shapeList[i].nHoles == maxH:\n",
    "                        attrList[i].append('MoHo')\n",
    "                    elif shapeList[i].nHoles == minH:\n",
    "                        attrList[i].append('LeHo')           \n",
    "            #is referenced by a full/partial frame?\n",
    "                if any((shapeList[i].position[0] >= fr.position[0] and shapeList[i].position[1] >= fr.position[1]\\\n",
    "                        and shapeList[i].position[0] + shapeList[i].shape[0] <= fr.position[0] + fr.shape[0] and\\\n",
    "                        shapeList[i].position[1] + shapeList[i].shape[1] <= fr.position[1] + fr.shape[1] and\\\n",
    "                        shapeList[i].color != fr.color) for fr in self.partialFrames):\n",
    "                    attrList[i].append('IsRef')\n",
    "                if any((shapeList[i].position[0] >= fr.position[0] and shapeList[i].position[1] >= fr.position[1]\\\n",
    "                        and shapeList[i].position[0] + shapeList[i].shape[0] <= fr.position[0] + fr.shape[0] and\\\n",
    "                        shapeList[i].position[1] + shapeList[i].shape[1] <= fr.position[1] + fr.shape[1] and\\\n",
    "                        shapeList[i].color != fr.color) for fr in self.fullFrames):\n",
    "                    attrList[i].append('IsFRef')\n",
    "    \n",
    "        if len(ism) == 1:\n",
    "            attrList[ism[0]].append('SmSh')\n",
    "        if len(ila) == 1:\n",
    "            attrList[ila[0]].append('LaSh')\n",
    "        for i in range(len(shapeList)):\n",
    "            if len(attrList[i]) > 0 and attrList[i][0] == mcopies:\n",
    "                attrList[i].append('MoCo')\n",
    "        if not singleColor:\n",
    "            for i in range(len(shapeList)):\n",
    "                if len(attrList[i]) > 0 and attrList[i][1] == mcolors:\n",
    "                    attrList[i].append('MoCl')\n",
    "        if [l[0] for l in attrList].count(1) == 1:\n",
    "            for i in range(len(shapeList)):\n",
    "                if len(attrList[i]) > 0 and attrList[i][0] == 1:\n",
    "                    attrList[i].append('UnSh')\n",
    "                    break\n",
    "        return [set(l[1:]) for l in attrList]\n",
    "            \n",
    "\n",
    "# %% Class Sample\n",
    "class Sample():\n",
    "    \"\"\"\n",
    "    An object of the class Sample stores all the information refering to a\n",
    "    training or test sample. Almost all of the attributes are only set if the\n",
    "    sample is a training sample or the init parameter \"submission\" is set to\n",
    "    False.\n",
    "    ...\n",
    "    Main Attributes\n",
    "    ---------------\n",
    "    inMatrix: Matrix\n",
    "        The input Matrix.\n",
    "    outMatrix: Matrix\n",
    "        The output Matrix.\n",
    "    sameShape: bool\n",
    "        True if inMatrix and outMatrix have the same shape. False otherwise.\n",
    "    colors: set\n",
    "        Colors present in either the input Matrix or the output Matrix.\n",
    "    commonColors: set\n",
    "        Colors present in both the inputMatrix and the output Matrix.\n",
    "    fixedColors: set\n",
    "        Colors that appear both in the input and in the output, and in the\n",
    "        exact same pixels. Only defined if sameShape.\n",
    "    changedInColors: set\n",
    "        Colors that are not fixed and are such that, if a pixel in the output\n",
    "        matrix has that color, then it has the same color in the input matrix.\n",
    "        Only defined if sameShape.\n",
    "    changedOutColors: set\n",
    "        Colors that are not fixed and are such that, if a pixel in the input\n",
    "        matrix has that color, then it has the same color in the output matrix.\n",
    "        Only defined if sameShape.\n",
    "    sameColorCount: bool\n",
    "        True if the input and the output matrices have the same color count.\n",
    "    commonShapes: list (of Shapes)\n",
    "        Shapes that appear both in the input matrix and the output matrix.\n",
    "    gridIsUnchanged: bool\n",
    "        True if there is the exact same Grid both in the input and in the\n",
    "        output matrices.\n",
    "    \"\"\"\n",
    "    def __init__(self, s, trainOrTest, submission=False, backgroundColor=None):\n",
    "        \n",
    "        self.inMatrix = Matrix(s['input'], backgroundColor=backgroundColor)\n",
    "        \n",
    "        if trainOrTest == \"train\" or submission==False:\n",
    "            self.outMatrix = Matrix(s['output'], backgroundColor=backgroundColor)\n",
    "                    \n",
    "            # We want to compare the input and the output\n",
    "            # Do they have the same dimensions?\n",
    "            self.sameHeight = self.inMatrix.shape[0] == self.outMatrix.shape[0]\n",
    "            self.sameWidth = self.inMatrix.shape[1] == self.outMatrix.shape[1]\n",
    "            self.sameShape = self.sameHeight and self.sameWidth\n",
    "            \n",
    "            # Is the input shape a factor of the output shape?\n",
    "            # Or the other way around?\n",
    "            if not self.sameShape:\n",
    "                if (self.inMatrix.shape[0] % self.outMatrix.shape[0]) == 0 and \\\n",
    "                (self.inMatrix.shape[1] % self.outMatrix.shape[1]) == 0 :\n",
    "                    self.outShapeFactor = (int(self.inMatrix.shape[0]/self.outMatrix.shape[0]),\\\n",
    "                                           int(self.inMatrix.shape[1]/self.outMatrix.shape[1]))\n",
    "                if (self.outMatrix.shape[0] % self.inMatrix.shape[0]) == 0 and \\\n",
    "                (self.outMatrix.shape[1] % self.inMatrix.shape[1]) == 0 :\n",
    "                    self.inShapeFactor = (int(self.outMatrix.shape[0]/self.inMatrix.shape[0]),\\\n",
    "                                          int(self.outMatrix.shape[1]/self.inMatrix.shape[1]))\n",
    "\n",
    "            # Is one a subset of the other? for now always includes diagonals\n",
    "            self.inSmallerThanOut = all(self.inMatrix.shape[i] <= self.outMatrix.shape[i] for i in [0,1]) and not self.sameShape\n",
    "            self.outSmallerThanIn = all(self.inMatrix.shape[i] >= self.outMatrix.shape[i] for i in [0,1]) and not self.sameShape\n",
    "    \n",
    "            #R: Is the output a shape (faster than checking if is a subset?\n",
    "    \n",
    "            if self.outSmallerThanIn:\n",
    "                #check if output is the size of a multicolored shape\n",
    "                self.outIsInMulticolorShapeSize = any((sh.shape == self.outMatrix.shape) for sh in self.inMatrix.multicolorShapes)\n",
    "                self.outIsInMulticolorDShapeSize = any((sh.shape == self.outMatrix.shape) for sh in self.inMatrix.multicolorDShapes)\n",
    "            self.commonShapes, self.commonDShapes, self.commonMulticolorShapes, self.commonMulticolorDShapes = [], [], [], []\n",
    "            if len(self.inMatrix.shapes) < 15 or len(self.outMatrix.shapes) < 10:\n",
    "                self.commonShapes = self.getCommonShapes(diagonal=False, sameColor=True,\\\n",
    "                                                     multicolor=False, rotation=True, scaling=True, mirror=True)\n",
    "            if len(self.inMatrix.dShapes) < 15 or len(self.outMatrix.dShapes) < 10:\n",
    "                self.commonDShapes = self.getCommonShapes(diagonal=True, sameColor=True,\\\n",
    "                                                      multicolor=False, rotation=True, scaling=True, mirror=True)\n",
    "            if len(self.inMatrix.multicolorShapes) < 15 or len(self.outMatrix.multicolorShapes) < 10:\n",
    "                self.commonMulticolorShapes = self.getCommonShapes(diagonal=False, sameColor=True,\\\n",
    "                                                               multicolor=True, rotation=True, scaling=True, mirror=True)\n",
    "            if len(self.inMatrix.multicolorDShapes) < 15 or len(self.outMatrix.multicolorDShapes) < 10:\n",
    "                self.commonMulticolorDShapes = self.getCommonShapes(diagonal=True, sameColor=True,\\\n",
    "                                                                multicolor=True, rotation=True, scaling=True, mirror=True)\n",
    "\n",
    "            # Which colors are there in the sample?\n",
    "            self.colors = set(self.inMatrix.colors | self.outMatrix.colors)\n",
    "            self.commonColors = set(self.inMatrix.colors & self.outMatrix.colors)\n",
    "            self.nColors = len(self.colors)\n",
    "            # Do they have the same colors?\n",
    "            self.sameColors = len(self.colors) == len(self.commonColors)\n",
    "            # Do they have the same number of colors?\n",
    "            self.sameNumColors = self.inMatrix.nColors == self.outMatrix.nColors\n",
    "            # Does output contain all input colors or viceversa?\n",
    "            self.inHasOutColors = self.outMatrix.colors <= self.inMatrix.colors  \n",
    "            self.outHasInColors = self.inMatrix.colors <= self.outMatrix.colors\n",
    "            if self.sameShape:\n",
    "                # Which pixels changes happened? How many times?\n",
    "                self.changedPixels = Counter()\n",
    "                self.sameColorCount = self.inMatrix.colorCount == self.outMatrix.colorCount\n",
    "                for i, j in np.ndindex(self.inMatrix.shape):\n",
    "                    if self.inMatrix.m[i,j] != self.outMatrix.m[i,j]:\n",
    "                        self.changedPixels[(self.inMatrix.m[i,j], self.outMatrix.m[i,j])] += 1\n",
    "                # Are any of these changes complete? (i.e. all pixels of one color are changed to another one)\n",
    "                self.completeColorChanges = set(change for change in self.changedPixels.keys() if\\\n",
    "                                             self.changedPixels[change]==self.inMatrix.colorCount[change[0]] and\\\n",
    "                                             change[0] not in self.outMatrix.colorCount.keys())\n",
    "                self.allColorChangesAreComplete = len(self.changedPixels) == len(self.completeColorChanges)\n",
    "                # Does any color never change?\n",
    "                self.changedInColors = set(change[0] for change in self.changedPixels.keys())\n",
    "                self.changedOutColors = set(change[1] for change in self.changedPixels.keys())\n",
    "                self.unchangedColors = set(x for x in self.colors if x not in set.union(self.changedInColors, self.changedOutColors))\n",
    "                # Colors that stay unchanged\n",
    "                self.fixedColors = set(x for x in self.colors if x not in set.union(self.changedInColors, self.changedOutColors))\n",
    "            \n",
    "            if self.sameShape and self.sameColorCount:\n",
    "                self.sameRowCount = True\n",
    "                for r in range(self.inMatrix.shape[0]):\n",
    "                    _,inCounts = np.unique(self.inMatrix.m[r,:], return_counts=True)\n",
    "                    _,outCounts = np.unique(self.outMatrix.m[r,:], return_counts=True)\n",
    "                    if not np.array_equal(inCounts, outCounts):\n",
    "                        self.sameRowCount = False\n",
    "                        break\n",
    "                self.sameColCount = True\n",
    "                for c in range(self.inMatrix.shape[1]):\n",
    "                    _,inCounts = np.unique(self.inMatrix.m[:,c], return_counts=True)\n",
    "                    _,outCounts = np.unique(self.outMatrix.m[:,c], return_counts=True)\n",
    "                    if not np.array_equal(inCounts, outCounts):\n",
    "                        self.sameColCount = False\n",
    "                        break\n",
    "                    \n",
    "            # Shapes in the input that are fixed\n",
    "            if self.sameShape:\n",
    "                self.fixedShapes = []\n",
    "                for sh in self.inMatrix.shapes:\n",
    "                    if sh.color in self.fixedColors:\n",
    "                        continue\n",
    "                    shapeIsFixed = True\n",
    "                    for i,j in np.ndindex(sh.shape):\n",
    "                        if sh.m[i,j] != 255:\n",
    "                            if self.outMatrix.m[sh.position[0]+i,sh.position[1]+j]!=sh.m[i,j]:\n",
    "                                shapeIsFixed=False\n",
    "                                break\n",
    "                    if shapeIsFixed:\n",
    "                        self.fixedShapes.append(sh)\n",
    "                    \n",
    "            # Frames\n",
    "            self.commonFullFrames = [f for f in self.inMatrix.fullFrames if f in self.outMatrix.fullFrames]\n",
    "            if len(self.inMatrix.fullFrames)==1:\n",
    "                frameM = self.inMatrix.fullFrames[0].m.copy()\n",
    "                frameM[frameM==255] = self.inMatrix.fullFrames[0].background\n",
    "                if frameM.shape==self.outMatrix.shape:\n",
    "                    self.frameIsOutShape = True\n",
    "                elif frameM.shape==(self.outMatrix.shape[0]+1, self.outMatrix.shape[1]+1):\n",
    "                    self.frameInsideIsOutShape = True\n",
    "            \n",
    "            # Grids\n",
    "            # Is the grid the same in the input and in the output?\n",
    "            self.gridIsUnchanged = self.inMatrix.isGrid and self.outMatrix.isGrid \\\n",
    "            and self.inMatrix.grid == self.outMatrix.grid\n",
    "            # Does the shape of the grid cells determine the output shape?\n",
    "            if hasattr(self.inMatrix, \"grid\") and self.inMatrix.grid.allCellsSameShape:\n",
    "                self.gridCellIsOutputShape = self.outMatrix.shape == self.inMatrix.grid.cellShape\n",
    "            # Does the shape of the input determine the shape of the grid cells of the output?\n",
    "            if hasattr(self.outMatrix, \"grid\") and self.outMatrix.grid.allCellsSameShape:\n",
    "                self.gridCellIsInputShape = self.inMatrix.shape == self.outMatrix.grid.cellShape\n",
    "            # Do all the grid cells have one color?\n",
    "            if self.gridIsUnchanged:\n",
    "                self.gridCellsHaveOneColor = self.inMatrix.grid.allCellsHaveOneColor and\\\n",
    "                                             self.outMatrix.grid.allCellsHaveOneColor\n",
    "            # Asymmetric grids\n",
    "            self.asymmetricGridIsUnchanged = self.inMatrix.isAsymmetricGrid and self.outMatrix.isAsymmetricGrid \\\n",
    "            and self.inMatrix.asymmetricGrid == self.outMatrix.asymmetricGrid\n",
    "            if self.asymmetricGridIsUnchanged:\n",
    "                self.asymmetricGridCellsHaveOneColor = self.inMatrix.asymmetricGrid.allCellsHaveOneColor and\\\n",
    "                self.outMatrix.asymmetricGrid.allCellsHaveOneColor\n",
    "            \n",
    "            # Is there a blank to fill?\n",
    "            self.inputHasBlank = len(self.inMatrix.blanks)>0\n",
    "            if self.inputHasBlank:\n",
    "                for s in self.inMatrix.blanks:\n",
    "                    if s.shape == self.outMatrix.shape:\n",
    "                        self.blankToFill = s\n",
    "             \n",
    "            # Does the output matrix follow a pattern?\n",
    "            self.followsRowPattern = self.outMatrix.followsRowPattern()\n",
    "            self.followsColPattern = self.outMatrix.followsColPattern()\n",
    "            \n",
    "            # Full borders and horizontal/vertical\n",
    "            if self.sameShape:\n",
    "                self.commonFullBorders = []\n",
    "                for inBorder in self.inMatrix.fullBorders:\n",
    "                    for outBorder in self.outMatrix.fullBorders:\n",
    "                        if inBorder==outBorder:\n",
    "                            self.commonFullBorders.append(inBorder)\n",
    "                \n",
    "                self.isHorizontal = self.inMatrix.isHorizontal and self.outMatrix.isHorizontal\n",
    "                self.isVertical = self.inMatrix.isVertical and self.outMatrix.isVertical\n",
    "\n",
    "    def getCommonShapes(self, diagonal=True, multicolor=False, sameColor=False, samePosition=False, rotation=False, \\\n",
    "                     mirror=False, scaling=False):\n",
    "        comSh = []\n",
    "        if diagonal:\n",
    "            if not multicolor:\n",
    "                ishs = self.inMatrix.dShapes\n",
    "                oshs = self.outMatrix.dShapes\n",
    "            else:\n",
    "                ishs = self.inMatrix.multicolorDShapes\n",
    "                oshs = self.outMatrix.multicolorDShapes\n",
    "        else:\n",
    "            if not multicolor:\n",
    "                ishs = self.inMatrix.shapes\n",
    "                oshs = self.outMatrix.shapes\n",
    "            else:\n",
    "                ishs = self.inMatrix.multicolorShapes\n",
    "                oshs = self.outMatrix.multicolorShapes\n",
    "        #Arbitrary: shapes have size < 100 and > 3\n",
    "        for ish in ishs:\n",
    "            outCount = 0\n",
    "            if len(ish.pixels) < 4 or len(ish.pixels) > 100:\n",
    "                continue\n",
    "            for osh in oshs:\n",
    "                if len(osh.pixels) < 4 or len(osh.pixels) > 100:\n",
    "                    continue\n",
    "                if ish.hasSameShape(osh, sameColor=sameColor, samePosition=samePosition,\\\n",
    "                                    rotation=rotation, mirror=mirror, scaling=scaling):\n",
    "                    outCount += 1\n",
    "            if outCount > 0:\n",
    "                comSh.append((ish, np.count_nonzero([ish.hasSameShape(ish2, sameColor=sameColor, samePosition=samePosition,\\\n",
    "                                    rotation=rotation, mirror=mirror, scaling=scaling) for ish2 in ishs]), outCount))\n",
    "        return comSh\n",
    "\n",
    "# %% Class Task\n",
    "class Task():\n",
    "    \"\"\"\n",
    "    An object of the class Task stores all the relevant information about an\n",
    "    ARC task.\n",
    "    ...\n",
    "    Main Attributes\n",
    "    ---------------\n",
    "    task: dict\n",
    "        The task given in the standard ARC format.\n",
    "    index: str\n",
    "        The name of the task, in hexadecimal characters.\n",
    "    submission: bool\n",
    "        True if we don't know the output matrices of the test sample.\n",
    "    trainSamples: list (of Samples)\n",
    "        List of the training Samples of the task.\n",
    "    testSamples: list (of Samples)\n",
    "        List of the test Samples of the task.\n",
    "    nTrain: int\n",
    "        Number of train Samples.\n",
    "    nTest: int\n",
    "        Number of test Samples.\n",
    "    sameInShape: bool\n",
    "        True if all the input matrices have the same shape. False otherwise.\n",
    "    sameOutShape: bool\n",
    "        True if all the output matrices have the same shape. False otherwise.\n",
    "    sameIOShapes: bool\n",
    "        True if the input matrices have the same shape as the output matrices\n",
    "        for every sample. False otherwise.\n",
    "    colors: set\n",
    "        The colors that appear in the task.\n",
    "    totalInColors: set\n",
    "        The colors that appear in at least one of the input matrices.\n",
    "    commonInColors: set\n",
    "        The colors that appear in all the input matrices.\n",
    "    totalOutColors: set\n",
    "        The colors that appear in at least one of the output matrices.\n",
    "    commonOutColors: set\n",
    "        The colors that appear in all the output matrices.\n",
    "    fixedColors: set\n",
    "        The fixedColors that are common to every training Sample.\n",
    "    commonChangedInColors: set\n",
    "        The changedInColors that are common to every training Sample.\n",
    "    commonChangedOutColors: set\n",
    "        The changedOutColors that are common to every training Sample.\n",
    "    backgroundColor: int\n",
    "        If all the input matrices have the same background color, this\n",
    "        attribute represents it. Otherwise, it is set to -1.\n",
    "    \"\"\"\n",
    "    def __init__(self, t, i, submission=False, backgrounds=None):\n",
    "        self.task = t\n",
    "        self.index = i\n",
    "        self.submission = submission\n",
    "        \n",
    "        if backgrounds==None:\n",
    "            self.trainSamples = [Sample(s, \"train\", submission) for s in t['train']]\n",
    "            self.testSamples = [Sample(s, \"test\", submission) for s in t['test']]\n",
    "        else:\n",
    "            self.trainSamples = [Sample(t[\"train\"][s], \"train\", submission, backgrounds[\"train\"][s]) for s in range(len(t['train']))]\n",
    "            self.testSamples = [Sample(t[\"test\"][s], \"test\", submission, backgrounds[\"test\"][s]) for s in range(len(t['test']))]\n",
    "        \n",
    "        self.nTrain = len(self.trainSamples)\n",
    "        self.nTest = len(self.testSamples)\n",
    "        \n",
    "        # Common properties I want to know:\n",
    "        \n",
    "        # Dimension:\n",
    "        # Do all input/output matrices have the same shape?\n",
    "        inShapes = [s.inMatrix.shape for s in self.trainSamples]\n",
    "        self.sameInShape = self.allEqual(inShapes)\n",
    "        if self.sameInShape:\n",
    "            self.inShape = self.trainSamples[0].inMatrix.shape\n",
    "        outShapes = [s.outMatrix.shape for s in self.trainSamples]\n",
    "        self.sameOutShape = self.allEqual(outShapes)\n",
    "        if self.sameOutShape:\n",
    "            self.outShape = self.trainSamples[0].outMatrix.shape\n",
    "            \n",
    "        # Do all output matrices have the same shape as the input matrix?\n",
    "        self.sameIOShapes = all([s.sameShape for s in self.trainSamples])\n",
    "        \n",
    "        # Are the input/output matrices always squared?\n",
    "        self.inMatricesSquared = all([s.inMatrix.shape[0] == s.inMatrix.shape[1] \\\n",
    "                                      for s in self.trainSamples+self.testSamples])\n",
    "        self.outMatricesSquared = all([s.outMatrix.shape[0] == s.outMatrix.shape[1] \\\n",
    "                                       for s in self.trainSamples])\n",
    "    \n",
    "        # Are shapes of in (out) matrices always a factor of the shape of the \n",
    "        # out (in) matrices?\n",
    "        if all([hasattr(s, 'inShapeFactor') for s in self.trainSamples]):\n",
    "            if self.allEqual([s.inShapeFactor for s in self.trainSamples]):\n",
    "                self.inShapeFactor = self.trainSamples[0].inShapeFactor\n",
    "            elif all([s.inMatrix.shape[0]**2 == s.outMatrix.shape[0] and \\\n",
    "                      s.inMatrix.shape[1]**2 == s.outMatrix.shape[1] \\\n",
    "                      for s in self.trainSamples]):\n",
    "                self.inShapeFactor = \"squared\"\n",
    "            elif all([s.inMatrix.shape[0]**2 == s.outMatrix.shape[0] and \\\n",
    "                      s.inMatrix.shape[1] == s.outMatrix.shape[1] \\\n",
    "                      for s in self.trainSamples]):\n",
    "                self.inShapeFactor = \"xSquared\"\n",
    "            elif all([s.inMatrix.shape[0] == s.outMatrix.shape[0] and \\\n",
    "                      s.inMatrix.shape[1]**2 == s.outMatrix.shape[1] \\\n",
    "                      for s in self.trainSamples]):\n",
    "                self.inShapeFactor = \"ySquared\"\n",
    "            elif all([s.inMatrix.shape[0]*s.inMatrix.nColors == s.outMatrix.shape[0] and \\\n",
    "                     s.inMatrix.shape[1]*s.inMatrix.nColors == s.outMatrix.shape[1] \\\n",
    "                     for s in self.trainSamples]):\n",
    "                self.inShapeFactor = \"nColors\"\n",
    "            elif all([s.inMatrix.shape[0]*(s.inMatrix.nColors-1) == s.outMatrix.shape[0] and \\\n",
    "                     s.inMatrix.shape[1]*(s.inMatrix.nColors-1) == s.outMatrix.shape[1] \\\n",
    "                     for s in self.trainSamples]):\n",
    "                self.inShapeFactor = \"nColors-1\"\n",
    "        if all([hasattr(s, 'outShapeFactor') for s in self.trainSamples]):\n",
    "            if self.allEqual([s.outShapeFactor for s in self.trainSamples]):\n",
    "                self.outShapeFactor = self.trainSamples[0].outShapeFactor\n",
    "                \n",
    "        # Is the output always smaller?\n",
    "        self.outSmallerThanIn = all(s.outSmallerThanIn for s in self.trainSamples)\n",
    "        self.inSmallerThanOut = all(s.inSmallerThanOut for s in self.trainSamples)\n",
    "        \n",
    "        # Symmetries:\n",
    "        # Are all outputs LR, UD, D1 or D2 symmetric?\n",
    "        self.lrSymmetric = all([s.outMatrix.lrSymmetric for s in self.trainSamples])\n",
    "        self.udSymmetric = all([s.outMatrix.udSymmetric for s in self.trainSamples])\n",
    "        self.d1Symmetric = all([s.outMatrix.d1Symmetric for s in self.trainSamples])\n",
    "        self.d2Symmetric = all([s.outMatrix.d2Symmetric for s in self.trainSamples])\n",
    "        \n",
    "        # Colors\n",
    "        # How many colors are there in the input? Is it always the same number?\n",
    "        # How many colors are there in the output? Is it always the same number?\n",
    "        self.sameNumColors = all([s.sameNumColors for s in self.trainSamples])\n",
    "        self.nInColors = [s.inMatrix.nColors for s in self.trainSamples] + \\\n",
    "        [s.inMatrix.nColors for s in self.testSamples]\n",
    "        self.sameNInColors = self.allEqual(self.nInColors)\n",
    "        self.nOutColors = [s.outMatrix.nColors for s in self.trainSamples]\n",
    "        self.sameNOutColors = self.allEqual(self.nOutColors)\n",
    "        # Which colors does the input have? Union and intersection.\n",
    "        self.inColors = [s.inMatrix.colors for s in self.trainSamples+self.testSamples]\n",
    "        self.commonInColors = set.intersection(*self.inColors)\n",
    "        self.totalInColors = set.union(*self.inColors)\n",
    "        # Which colors does the output have? Union and intersection.\n",
    "        self.outColors = [s.outMatrix.colors for s in self.trainSamples]\n",
    "        self.commonOutColors = set.intersection(*self.outColors)\n",
    "        self.totalOutColors = set.union(*self.outColors)\n",
    "        # Which colors appear in every sample?\n",
    "        self.sampleColors = [s.colors for s in self.trainSamples]\n",
    "        self.commonSampleColors = set.intersection(*self.sampleColors)\n",
    "        self.almostCommonColors = self.commonSampleColors.copy()\n",
    "        if self.nTrain > 3:\n",
    "            for color in range(10):\n",
    "                if color not in self.almostCommonColors:\n",
    "                    apps = 0\n",
    "                    for i in range(self.nTrain):\n",
    "                        if color in self.trainSamples[i].colors:\n",
    "                            apps += 1\n",
    "                    if apps == self.nTrain-1:\n",
    "                        self.almostCommonColors.add(color)\n",
    "        # Input colors of the test samples\n",
    "        self.testInColors = [s.inMatrix.colors for s in self.testSamples]\n",
    "        # Are there the same number of colors in every sample?\n",
    "        self.sameNSampleColors = self.allEqual([len(sc) for sc in self.sampleColors]) and\\\n",
    "        all([len(s.inMatrix.colors | self.commonOutColors) <= len(self.sampleColors[0]) for s in self.testSamples])\n",
    "        # How many colors are there in total? Which ones?\n",
    "        self.colors = self.totalInColors | self.totalOutColors\n",
    "        self.nColors = len(self.colors)\n",
    "        # Does the output always have the same colors as the input?\n",
    "        if self.sameNumColors:\n",
    "            self.sameIOColors = all([i==j for i,j in zip(self.inColors, self.outColors)])\n",
    "        if self.sameIOShapes:\n",
    "            # Do the matrices have the same color count?\n",
    "            self.sameColorCount = all([s.sameColorCount for s in self.trainSamples])\n",
    "            if self.sameColorCount:\n",
    "                self.sameRowCount = all([s.sameRowCount for s in self.trainSamples])\n",
    "                self.sameColCount = all([s.sameColCount for s in self.trainSamples])\n",
    "            # Which color changes happen? Union and intersection.\n",
    "            cc = [set(s.changedPixels.keys()) for s in self.trainSamples]\n",
    "            self.colorChanges = set.union(*cc)\n",
    "            self.commonColorChanges = set.intersection(*cc)\n",
    "            # Does any color always change? (to and from)\n",
    "            self.changedInColors = [s.changedInColors for s in self.trainSamples]\n",
    "            self.commonChangedInColors = set.intersection(*self.changedInColors)\n",
    "            self.changedOutColors = [s.changedOutColors for s in self.trainSamples]\n",
    "            self.commonChangedOutColors = set.intersection(*self.changedOutColors)\n",
    "            self.commonOnlyChangedInColors = self.commonChangedInColors - set.union(*self.changedOutColors)\n",
    "            # Complete color changes\n",
    "            self.completeColorChanges = [s.completeColorChanges for s in self.trainSamples]\n",
    "            self.commonCompleteColorChanges = set.intersection(*self.completeColorChanges)\n",
    "            self.allColorChangesAreComplete = all([s.allColorChangesAreComplete for s in self.trainSamples])\n",
    "            # Are there any fixed colors?\n",
    "            self.fixedColors = set.intersection(*[s.fixedColors for s in self.trainSamples])\n",
    "            self.fixedColors2 = set.union(*[s.fixedColors for s in self.trainSamples]) - \\\n",
    "            set.union(*[s.changedInColors for s in self.trainSamples]) -\\\n",
    "            set.union(*[s.changedOutColors for s in self.trainSamples])\n",
    "            # Does any color never change?\n",
    "            if self.commonChangedInColors == set(self.changedInColors[0]):\n",
    "                self.unchangedColors = set(range(10)) - self.commonChangedInColors\n",
    "            else:\n",
    "                self.unchangedColors = [s.unchangedColors for s in self.trainSamples]\n",
    "                self.unchangedColors = set.intersection(*self.unchangedColors)\n",
    "                \n",
    "        # Is the number of pixels changed always the same?\n",
    "        \"\"\"\n",
    "        if self.sameIOShapes:\n",
    "            self.sameChanges = self.allEqual([s.diffPixels for s in self.trainSamples])\n",
    "        \"\"\"\n",
    "            \n",
    "        #R: is output a shape in the input\n",
    "        self.outIsInMulticolorShapeSize = False\n",
    "        self.outIsInMulticolorDShapeSize = False\n",
    "\n",
    "        if all([(hasattr(s, \"outIsInMulticolorShapeSize\") and s.outIsInMulticolorShapeSize) for s in self.trainSamples]):\n",
    "             self.outIsInMulticolorShapeSize = True\n",
    "        if all([(hasattr(s, \"outIsInMulticolorDShapeSize\") and s.outIsInMulticolorDShapeSize) for s in self.trainSamples]):\n",
    "             self.outIsInMulticolorDShapeSize = True\n",
    "             \n",
    "        self.nCommonInOutShapes = min(len(s.commonShapes) for s in self.trainSamples)\n",
    "        self.nCommonInOutDShapes = min(len(s.commonDShapes) for s in self.trainSamples) \n",
    "        self.nCommonInOutMulticolorShapes = min(len(s.commonMulticolorShapes) for s in self.trainSamples)\n",
    "        self.nCommonInOutMulticolorDShapes = min(len(s.commonMulticolorDShapes) for s in self.trainSamples) \n",
    "        \n",
    "        if self.sameIOShapes:\n",
    "            self.fixedShapes = []\n",
    "            for s in self.trainSamples:\n",
    "                for shape in s.fixedShapes:\n",
    "                    self.fixedShapes.append(shape)\n",
    "            self.fixedShapeFeatures = []\n",
    "            nFeatures = len(self.trainSamples[0].inMatrix.shapes[0].boolFeatures)\n",
    "            for i in range(nFeatures):\n",
    "                self.fixedShapeFeatures.append(True)\n",
    "            for shape in self.fixedShapes:\n",
    "                self.fixedShapeFeatures = [shape.boolFeatures[i] and self.fixedShapeFeatures[i] \\\n",
    "                                             for i in range(nFeatures)]\n",
    "        \n",
    "        # Grids:\n",
    "        self.inputIsGrid = all([s.inMatrix.isGrid for s in self.trainSamples+self.testSamples])\n",
    "        self.outputIsGrid = all([s.outMatrix.isGrid for s in self.trainSamples])\n",
    "        self.hasUnchangedGrid = all([s.gridIsUnchanged for s in self.trainSamples])\n",
    "        if all([hasattr(s, \"gridCellIsOutputShape\") for s in self.trainSamples]):\n",
    "            self.gridCellIsOutputShape = all([s.gridCellIsOutputShape for s in self.trainSamples])\n",
    "        if all([hasattr(s, \"gridCellIsInputShape\") for s in self.trainSamples]):\n",
    "            self.gridCellIsInputShape = all([s.gridCellIsInputShape for s in self.trainSamples])\n",
    "        if self.hasUnchangedGrid:\n",
    "            self.gridCellsHaveOneColor = all([s.gridCellsHaveOneColor for s in self.trainSamples])\n",
    "            self.outGridCellsHaveOneColor = all([s.outMatrix.grid.allCellsHaveOneColor for s in self.trainSamples])\n",
    "        # Asymmetric grids\n",
    "        self.inputIsAsymmetricGrid = all([s.inMatrix.isAsymmetricGrid for s in self.trainSamples+self.testSamples])\n",
    "        self.hasUnchangedAsymmetricGrid = all([s.asymmetricGridIsUnchanged for s in self.trainSamples])\n",
    "        if self.hasUnchangedAsymmetricGrid:\n",
    "            self.assymmetricGridCellsHaveOneColor = all([s.asymmetricGridCellsHaveOneColor for s in self.trainSamples])\n",
    "            self.outAsymmetricGridCellsHaveOneColor = all([s.outMatrix.asymmetricGrid.allCellsHaveOneColor for s in self.trainSamples])\n",
    "        \n",
    "        # Background color\n",
    "        \n",
    "        # Is there always a background color? Which one?\n",
    "        if self.allEqual([s.inMatrix.backgroundColor for s in self.trainSamples]) and\\\n",
    "        self.trainSamples[0].inMatrix.backgroundColor == self.testSamples[0].inMatrix.backgroundColor:\n",
    "            self.backgroundColor = self.trainSamples[0].inMatrix.backgroundColor\n",
    "        elif self.hasUnchangedAsymmetricGrid and all([s.inMatrix.asymmetricGrid.nCells>6 for s in self.trainSamples]):\n",
    "            self.backgroundColor = self.trainSamples[0].inMatrix.asymmetricGrid.color\n",
    "            for sample in self.trainSamples:\n",
    "                sample.inMatrix.backgroundColor = self.backgroundColor\n",
    "                sample.outMatrix.backgroundColor = self.backgroundColor\n",
    "            for sample in self.testSamples:\n",
    "                sample.inMatrix.backgroundColor = self.backgroundColor\n",
    "        else:\n",
    "            self.backgroundColor = -1\n",
    "            \n",
    "        self.orderedColors = self.orderColors()\n",
    "        \n",
    "        # Shapes:\n",
    "        # Does the task ONLY involve changing colors of shapes?\n",
    "        if self.sameIOShapes:\n",
    "            self.onlyShapeColorChanges = True\n",
    "            for s in self.trainSamples:\n",
    "                nShapes = s.inMatrix.nShapes\n",
    "                if s.outMatrix.nShapes != nShapes:\n",
    "                    self.onlyShapeColorChanges = False\n",
    "                    break\n",
    "                for shapeI in range(nShapes):\n",
    "                    if not s.inMatrix.shapes[shapeI].hasSameShape(s.outMatrix.shapes[shapeI]):\n",
    "                        self.onlyShapeColorChanges = False\n",
    "                        break\n",
    "                if not self.onlyShapeColorChanges:\n",
    "                    break\n",
    "            \n",
    "            # Get a list with the number of pixels shapes have\n",
    "            if self.onlyShapeColorChanges:\n",
    "                nPixels = set()\n",
    "                for s in self.trainSamples:\n",
    "                    for shape in s.inMatrix.shapes:\n",
    "                        nPixels.add(shape.nPixels)\n",
    "                self.shapePixelNumbers =  list(nPixels)\n",
    "                \n",
    "        #R: Are there any common input shapes accross samples?\n",
    "        self.commonInShapes = []\n",
    "        for sh1 in self.trainSamples[0].inMatrix.shapes:\n",
    "            if sh1.color == self.trainSamples[0].inMatrix.backgroundColor:\n",
    "                continue\n",
    "            addShape = True\n",
    "            for s in range(1,self.nTrain):\n",
    "                if not any([sh1.pixels == sh2.pixels for sh2 in self.trainSamples[s].inMatrix.shapes]):\n",
    "                    addShape = False\n",
    "                    break\n",
    "            if addShape and sh1 not in self.commonInShapes:\n",
    "                self.commonInShapes.append(sh1)\n",
    "\n",
    "        self.commonInDShapes = []\n",
    "        for sh1 in self.trainSamples[0].inMatrix.dShapes:\n",
    "            if sh1.color == self.trainSamples[0].inMatrix.backgroundColor:\n",
    "                continue\n",
    "            addShape = True\n",
    "            for s in range(1,self.nTrain):\n",
    "                if not any([sh1.pixels == sh2.pixels for sh2 in self.trainSamples[s].inMatrix.dShapes]):\n",
    "                    addShape = False\n",
    "                    break\n",
    "            if addShape:\n",
    "                self.commonInDShapes.append(sh1)\n",
    "        #Does the task use the information of isolated pixels?\n",
    "        #if all(s.inMatrix.nIsolatedPixels)\n",
    "        #Does the input always consist in two shapes?\n",
    "        self.twoShapeTask = (False, False, False, False, 1)\n",
    "        if all(len(s.inMatrix.multicolorDShapes)==2 for s in self.trainSamples):\n",
    "            self.twoShapeTask = (True, True, True, False, 1)\n",
    "            if all(s.inMatrix.multicolorDShapes[0].shape == s.inMatrix.multicolorDShapes[1].shape for s in self.trainSamples):\n",
    "                self.twoShapeTask = (True, True, True, True, 1)\n",
    "                \n",
    "        elif all(len(s.inMatrix.multicolorShapes)==2 for s in self.trainSamples):\n",
    "            self.twoShapeTask = (True, True, False, False, 1)\n",
    "            if all(s.inMatrix.multicolorShapes[0].shape == s.inMatrix.multicolorShapes[1].shape for s in self.trainSamples):\n",
    "                self.twoShapeTask = (True, True, False, True, 1)\n",
    "        elif all(len(s.inMatrix.dShapes)==2 for s in self.trainSamples):\n",
    "            self.twoShapeTask = (True, False, True, False, 1)\n",
    "            if all(s.inMatrix.dShapes[0].shape == s.inMatrix.dShapes[1].shape for s in self.trainSamples):\n",
    "                self.twoShapeTask = (True, False, True, True, 1)\n",
    "        if self.inputIsGrid:\n",
    "            if all(s.inMatrix.grid.nCells == 2 for s in self.trainSamples):\n",
    "                self.twoShapeTask = (True, False, False, True, 2)\n",
    "        elif  all((s.inMatrix.shape[0]*2 == s.inMatrix.shape[1] or\\\n",
    "                  s.inMatrix.shape[0] == s.inMatrix.shape[1]*2) for s in self.trainSamples):\n",
    "                self.twoShapeTask = (True, False, False, True, 3)\n",
    "                \n",
    "        #Are all output matrices equal mod nonBackgroundColor?\n",
    "        if self.sameOutShape:\n",
    "            self.sameOutDummyMatrix = all(np.all(self.trainSamples[0].outMatrix.dummyMatrix==s.outMatrix.dummyMatrix) for s in self.trainSamples)\n",
    "        # Frames\n",
    "        self.hasFullFrame = all([len(s.inMatrix.fullFrames)>0 for s in self.trainSamples])\n",
    "        self.hasPartialFrame = all([len(s.inMatrix.partialFrames)>0 for s in self.trainSamples])\n",
    "        # Is the task about filling a blank?\n",
    "        self.fillTheBlank =  all([hasattr(s, 'blankToFill') for s in self.trainSamples])\n",
    "                \n",
    "        # Do all output matrices follow a pattern?\n",
    "        self.followsRowPattern = all([s.followsRowPattern != False for s in self.trainSamples])\n",
    "        self.followsColPattern = all([s.followsColPattern != False for s in self.trainSamples])\n",
    "        if self.followsRowPattern:\n",
    "            self.rowPatterns = [s.outMatrix.followsRowPattern() for s in self.trainSamples]\n",
    "        if self.followsColPattern:\n",
    "            self.colPatterns = [s.outMatrix.followsColPattern() for s in self.trainSamples]\n",
    "        \n",
    "        # Full Borders / Requires vertical-horizontal rotation\n",
    "        if self.sameIOShapes:\n",
    "            if self.submission:\n",
    "                self.hasOneFullBorder = all([len(s.commonFullBorders)==1 for s in self.trainSamples])\n",
    "            else:\n",
    "                self.hasOneFullBorder = all([hasattr(s, 'commonFullBorders') and len(s.commonFullBorders)==1 for s in self.trainSamples+self.testSamples])\n",
    "            self.requiresHVRotation = False\n",
    "            if not (self.allEqual([s.isHorizontal for s in self.trainSamples]) or \\\n",
    "                    self.allEqual([s.isVertical for s in self.trainSamples])):    \n",
    "                self.requiresHVRotation = all([s.isHorizontal or s.isVertical for s in self.trainSamples])\n",
    "        \n",
    "    def allEqual(self, x):\n",
    "        \"\"\"\n",
    "        x is a list.\n",
    "        Returns true if all elements of x are equal.\n",
    "        \"\"\"\n",
    "        if len(x) == 0:\n",
    "            return False\n",
    "        return x.count(x[0]) == len(x)\n",
    "    \n",
    "    def orderColors(self):\n",
    "        \"\"\"\n",
    "        The aim of this function is to give the colors a specific order, in\n",
    "        order to do the OHE in the right way for every sample.\n",
    "        \"\"\"\n",
    "        orderedColors = []\n",
    "        # 1: Colors that appear in every sample, input and output, and never\n",
    "        # change. Only valid if t.sameIOShapes\n",
    "        if self.sameIOShapes:\n",
    "            for c in self.fixedColors:\n",
    "                if all([c in sample.inMatrix.colors for sample in self.testSamples]):\n",
    "                    orderedColors.append(c)\n",
    "        # 2: Colors that appear in every sample and are always changed from,\n",
    "        # never changed to.\n",
    "            for c in self.commonChangedInColors:\n",
    "                if c not in self.commonChangedOutColors:\n",
    "                    if all([c in sample.inMatrix.colors for sample in self.testSamples]):\n",
    "                        if c not in orderedColors:\n",
    "                            orderedColors.append(c)\n",
    "        # 3: Colors that appear in every sample and are always changed to,\n",
    "        # never changed from.\n",
    "            for c in self.commonChangedOutColors:\n",
    "                if not all([c in sample.inMatrix.colors for sample in self.trainSamples]):\n",
    "                    if c not in orderedColors:\n",
    "                        orderedColors.append(c)\n",
    "        # 4: Add the background color.\n",
    "        if self.backgroundColor != -1:\n",
    "            if self.backgroundColor not in orderedColors:\n",
    "                orderedColors.append(self.backgroundColor)\n",
    "        # 5: Other colors that appear in every input.\n",
    "        for c in self.commonInColors:\n",
    "            if all([c in sample.inMatrix.colors for sample in self.testSamples]):\n",
    "                if c not in orderedColors:\n",
    "                    orderedColors.append(c)\n",
    "        # 6: Other colors that appear in every output.\n",
    "        for c in self.commonOutColors:\n",
    "            if not all([c in sample.inMatrix.colors for sample in self.trainSamples]):\n",
    "                if c not in orderedColors:\n",
    "                    orderedColors.append(c)\n",
    "                \n",
    "        # TODO Dealing with grids and frames\n",
    "        \n",
    "        return orderedColors   \n",
    "    \n",
    "#############################################################################\n",
    "# %% Models\n",
    "        \n",
    "class OneConvModel(nn.Module):\n",
    "    \"\"\"\n",
    "    Simple CNN model consisting of only one 2d convolution. Input and output\n",
    "    tensors have the same shape.\n",
    "    \"\"\"\n",
    "    def __init__(self, ch=10, kernel=3, padVal = -1):\n",
    "        super(OneConvModel, self).__init__()\n",
    "        self.conv = nn.Conv2d(ch, ch, kernel_size=kernel, bias=0)\n",
    "        self.pad = nn.ConstantPad2d(int((kernel-1)/2), padVal)\n",
    "        \n",
    "    def forward(self, x, steps=1):\n",
    "        for _ in range(steps):\n",
    "            x = self.conv(self.pad(x))\n",
    "        return x\n",
    "    \n",
    "class LinearModel(nn.Module):\n",
    "    \"\"\"\n",
    "    Model consisting only of a fully connected layer and a given number of\n",
    "    channels.\n",
    "    \"\"\"\n",
    "    def __init__(self, inSize, outSize, ch):\n",
    "        super(LinearModel, self).__init__()\n",
    "        self.inSize = inSize\n",
    "        self.outSize = outSize\n",
    "        self.ch = ch\n",
    "        self.fc = nn.Linear(inSize[0]*inSize[1]*ch, outSize[0]*outSize[1]*ch)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.view(1, self.inSize[0]*self.inSize[1]*self.ch)\n",
    "        x = self.fc(x)\n",
    "        x = x.view(1, self.ch, self.outSize[0]*self.outSize[1])\n",
    "        return x\n",
    "    \n",
    "class LinearModelDummy(nn.Module): #(dummy = 2 channels)\n",
    "    \"\"\"\n",
    "    Model consisting only of a fully connected layer and two channels.\n",
    "    \"\"\"\n",
    "    def __init__(self, inSize, outSize):\n",
    "        super(LinearModelDummy, self).__init__()\n",
    "        self.inSize = inSize\n",
    "        self.outSize = outSize\n",
    "        self.fc = nn.Linear(inSize[0]*inSize[1]*2, outSize[0]*outSize[1]*2, bias=0)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.view(1, self.inSize[0]*self.inSize[1]*2)\n",
    "        x = self.fc(x)\n",
    "        x = x.view(1, 2, self.outSize[0]*self.outSize[1])\n",
    "        return x\n",
    "    \n",
    "class SimpleLinearModel(nn.Module):\n",
    "    \"\"\"\n",
    "    Model consisting only of a fully connected layer.\n",
    "    \"\"\"\n",
    "    def __init__(self, inSize, outSize):\n",
    "        super(SimpleLinearModel, self).__init__()\n",
    "        self.fc = nn.Linear(inSize, outSize)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.fc(x)\n",
    "        return x\n",
    "    \n",
    "class LSTMTagger(nn.Module):\n",
    "    \"\"\"\n",
    "    Simple LSTM model.\n",
    "    \"\"\"\n",
    "    def __init__(self, embedding_dim, hidden_dim, vocab_size, tagset_size):\n",
    "        super(LSTMTagger, self).__init__()\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)\n",
    "\n",
    "        # The LSTM takes word embeddings as inputs, and outputs hidden states\n",
    "        # with dimensionality hidden_dim.\n",
    "        self.lstm = nn.LSTM(embedding_dim, hidden_dim)\n",
    "\n",
    "        # The linear layer that maps from hidden state space to tag space\n",
    "        self.hidden2tag = nn.Linear(hidden_dim, tagset_size)\n",
    "        \n",
    "    def forward(self, sentence):\n",
    "        embeds = self.word_embeddings(sentence)\n",
    "        lstm_out, _ = self.lstm(embeds.view(len(sentence), 1, -1))\n",
    "        tag_space = self.hidden2tag(lstm_out.view(len(sentence), -1))\n",
    "        tag_scores = F.log_softmax(tag_space, dim=1)\n",
    "        return tag_scores\n",
    "    \n",
    "def pixelCorrespondence(t):\n",
    "    \"\"\"\n",
    "    Returns a dictionary. Keys are positions of the output matrix. Values are\n",
    "    the pixel in the input matrix it corresponds to.\n",
    "    Function only valid if t.sameInSahpe and t.sameOutShape\n",
    "    \"\"\"\n",
    "    pixelsColoredAllSamples = []\n",
    "    # In which positions does each color appear?\n",
    "    for s in t.trainSamples:\n",
    "        pixelsColored = [[] for i in range(10)]\n",
    "        m = s.inMatrix.m\n",
    "        for i,j in np.ndindex(t.inShape):\n",
    "            pixelsColored[m[i,j]].append((i,j))\n",
    "        pixelsColoredAllSamples.append(pixelsColored)\n",
    "    # For each pixel in output matrix, find correspondent pixel in input matrix\n",
    "    pixelMap = {}\n",
    "    for i,j in np.ndindex(t.outShape):\n",
    "        candidates = set()\n",
    "        for s in range(t.nTrain):\n",
    "            m = t.trainSamples[s].outMatrix.m\n",
    "            if len(candidates) == 0:\n",
    "                candidates = set(pixelsColoredAllSamples[s][m[i,j]])\n",
    "            else:\n",
    "                candidates = set(pixelsColoredAllSamples[s][m[i,j]]) & candidates\n",
    "            if len(candidates) == 0:\n",
    "                return {}\n",
    "        pixelMap[(i,j)] = next(iter(candidates))\n",
    "    \n",
    "    return pixelMap\n",
    "\n",
    "###############################################################################\n",
    "# %% Utils\n",
    "\n",
    "def identityM(matrix):\n",
    "    \"\"\"\n",
    "    Function that, given Matrix, returns its corresponding numpy.ndarray m\n",
    "    \"\"\"\n",
    "    if isinstance(matrix, np.ndarray):\n",
    "        return matrix.copy()\n",
    "    else:\n",
    "        return matrix.m.copy()\n",
    "\n",
    "def correctFixedColors(inMatrix, x, fixedColors, onlyChangedInColors):\n",
    "    \"\"\"\n",
    "    Given an input matrix (inMatrix), an output matrix (x) and a set of colors\n",
    "    that should not change between the input and the output (fixedColors),\n",
    "    this function returns a copy of x, but correcting the pixels that \n",
    "    shouldn't have changed back into the original, unchanged color.\n",
    "    \n",
    "    inMatrix and x are required to have the same shape.\n",
    "    \"\"\"\n",
    "    m = x.copy()\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        if inMatrix[i,j] in fixedColors:\n",
    "            m[i,j] = inMatrix[i,j]\n",
    "        if m[i,j] in onlyChangedInColors:\n",
    "            m[i,j] = inMatrix[i,j]\n",
    "    return m\n",
    "                \n",
    "def incorrectPixels(m1, m2):\n",
    "    \"\"\"\n",
    "    Returns the number of incorrect pixels (0 is best).\n",
    "    \"\"\"\n",
    "    if m1.shape != m2.shape:\n",
    "        return 1000\n",
    "    return np.sum(m1!=m2)\n",
    "\n",
    "def deBackgroundizeMatrix(m, color):\n",
    "    \"\"\"\n",
    "    Given a matrix m and a color, this function returns a matrix whose elements\n",
    "    are 0 or 1, depending on whether the corresponding pixel is of the given\n",
    "    color or not.\n",
    "    \"\"\"\n",
    "    return np.uint8(m == color)\n",
    "\n",
    "def relDicts(colors):\n",
    "    \"\"\"\n",
    "    Given a list of colors (numbers from 0 to 9, no repetitions allowed), this\n",
    "    function returns two dictionaries giving the relationships between the\n",
    "    color and its index in the list.\n",
    "    It's just a way to map the colors to list(range(nColors)).\n",
    "    \"\"\"\n",
    "    rel = {}\n",
    "    for i in range(len(colors)):\n",
    "        rel[i] = colors[i]\n",
    "    invRel = {v: k for k,v in rel.items()}\n",
    "    for i in range(len(colors)):\n",
    "        rel[i] = [colors[i]]\n",
    "    return rel, invRel\n",
    "\n",
    "def dummify(x, nChannels, rel=None):\n",
    "    \"\"\"\n",
    "    Given a matrix and a relationship given by relDicts, this function returns\n",
    "    a nColors x shape(x) matrix consisting only of ones and zeros. For each\n",
    "    channel (corresponding to a color), each element will be 1 if in the \n",
    "    original matrix x that pixel is of the corresponding color.\n",
    "    If rel is not specified, it is expected that the values of x range from\n",
    "    0 to nChannels-1.\n",
    "    \"\"\"\n",
    "    img = np.full((nChannels, x.shape[0], x.shape[1]), 0, dtype=np.uint8)\n",
    "    if rel==None:\n",
    "        for i in range(nChannels):\n",
    "            img[i] = x==i\n",
    "    else:\n",
    "        for i in range(len(rel)):\n",
    "            img[i] = np.isin(x,rel[i])\n",
    "    return img\n",
    "\n",
    "def dummifyColor(x, color):\n",
    "    \"\"\"\n",
    "    Given a matrix x and a color, this function returns a 2-by-shape(x) matrix\n",
    "    of ones and zeros. In one channel, the elements will be 1 if the pixel is\n",
    "    of the given color. In the other channel, they will be 1 otherwise.\n",
    "    \"\"\"\n",
    "    img = np.full((2, x.shape[0], x.shape[1]), 0, dtype=np.uint8)\n",
    "    img[0] = x!=color\n",
    "    img[1] = x==color\n",
    "    return img\n",
    "\n",
    "def updateBestFunction(t, f, bestScore, bestFunction, checkPerfect=False, prevScore=None):\n",
    "    \"\"\"\n",
    "    Given a task t, a partial function f, a best score and a best function, \n",
    "    this function executes f to all the matrices in t.trainSamples. If the\n",
    "    resulting score is lower than bestScore, then it returns f and the new\n",
    "    best score. Otherwise, it returns bestFunction again.\n",
    "    If the parameter checkPerfect is set to True, it returns f if the\n",
    "    transformation is perfect, this is, if all the pixels that are changed\n",
    "    (cannot be zero) are changed correctly. \n",
    "    \"\"\"\n",
    "    fun = copy.deepcopy(f)\n",
    "    score = 0\n",
    "    if checkPerfect:\n",
    "        changedPixels = 0\n",
    "    for sample in t.trainSamples:\n",
    "        pred = fun(sample.inMatrix)\n",
    "        score += incorrectPixels(sample.outMatrix.m, pred)\n",
    "        if checkPerfect:\n",
    "            changedPixels += incorrectPixels(pred, sample.inMatrix.m)\n",
    "    if score < bestScore:\n",
    "        bestScore = score\n",
    "        bestFunction = fun\n",
    "    if checkPerfect:\n",
    "        if changedPixels != 0 and prevScore - changedPixels == score: \n",
    "            isPerfect = True\n",
    "        else:\n",
    "            isPerfect = False\n",
    "        return fun, score, isPerfect\n",
    "    return bestFunction, bestScore\n",
    "\n",
    "# %% Symmetrize\n",
    "\n",
    "# if t.lrSymmetric or t.udSymmetric or t.d1Symmetric:\n",
    "# if len(t.changingColors) == 1:\n",
    "def symmetrize(matrix, axis, color=None, outColor=None, refColor=None):\n",
    "    \"\"\"\n",
    "    Given a matrix and a color, this function tries to turn pixels of that\n",
    "    given color into some other one in order to make the matrix symmetric.\n",
    "    \"axis\" is a list or set specifying the symmetry axis (lr, ud, d1 or d2).\n",
    "    \"\"\"\n",
    "    # Left-Right\n",
    "    def LRSymmetrize(m):\n",
    "        width = m.shape[1] - 1\n",
    "        for i in range(m.shape[0]):\n",
    "            for j in range(int(m.shape[1] / 2)):\n",
    "                if m[i,j] != m[i,width-j]:\n",
    "                    if color==None:\n",
    "                        if m[i,j]==refColor and m[i,width-j]!=refColor:\n",
    "                            m[i,width-j] = outColor\n",
    "                        elif m[i,j]!=refColor and m[i,width-j]==refColor:\n",
    "                            m[i,j] = outColor\n",
    "                    else:\n",
    "                        if m[i,j] == color:\n",
    "                            m[i,j] = m[i,width-j]\n",
    "                        elif m[i,width-j]==color:\n",
    "                            m[i,width-j] = m[i,j]\n",
    "        return m\n",
    "    \n",
    "    # Up-Down\n",
    "    def UDSymmetrize(m):\n",
    "        height = m.shape[0] - 1\n",
    "        for i in range(int(m.shape[0] / 2)):\n",
    "            for j in range(m.shape[1]):\n",
    "                if m[i,j] != m[height-i,j]:\n",
    "                    if color==None:\n",
    "                        if m[i,j]==refColor and m[height-i,j]!=refColor:\n",
    "                            m[height-i,j] = outColor\n",
    "                        elif m[i,j]!=refColor and m[height-i,j]==refColor:\n",
    "                            m[i,j] = outColor\n",
    "                    else:\n",
    "                        if m[i,j] == color:\n",
    "                            m[i,j] = m[height-i,j]\n",
    "                        elif m[height-i,j]==color:\n",
    "                            m[height-i,j] = m[i,j]\n",
    "        return m\n",
    "\n",
    "    # Main diagonal\n",
    "    def D1Symmetrize(m):\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            if m[i,j] != m[j,i]:\n",
    "                if color==None:\n",
    "                    if m[i,j]==refColor and m[j,i]!=refColor:\n",
    "                        m[j,i] = outColor\n",
    "                    elif m[i,j]!=refColor and m[j,i]==refColor:\n",
    "                        m[i,j] = outColor\n",
    "                else:\n",
    "                    if m[i,j] == color:\n",
    "                        m[i,j] = m[j,i]\n",
    "                    elif m[j,i]==color:\n",
    "                        m[j,i] = m[i,j]\n",
    "        return m\n",
    "    \n",
    "    def D2Symmetrize(matrix):\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            if m[i,j] != m[m.shape[0]-j-1, m.shape[1]-i-1]:\n",
    "                if color==None:\n",
    "                    if m[i,j]==refColor and m[m.shape[0]-j-1, m.shape[1]-i-1]!=refColor:\n",
    "                        m[m.shape[0]-j-1, m.shape[1]-i-1] = outColor\n",
    "                    elif m[i,j]!=refColor and m[m.shape[0]-j-1, m.shape[1]-i-1]==refColor:\n",
    "                        m[i,j] = outColor\n",
    "                else:\n",
    "                    if m[i,j] == color:\n",
    "                        m[i,j] = m[m.shape[0]-j-1, m.shape[1]-i-1]\n",
    "                    elif m[m.shape[0]-j-1, m.shape[1]-i-1]==color:\n",
    "                        m[m.shape[0]-j-1, m.shape[1]-i-1] = m[i,j]\n",
    "        return m\n",
    "    \n",
    "    m = matrix.m.copy()\n",
    "    while True:\n",
    "        prevMatrix = m.copy()\n",
    "        if \"lr\" in axis:\n",
    "            m = LRSymmetrize(m)\n",
    "        if \"ud\" in axis:\n",
    "            m = UDSymmetrize(m)\n",
    "        if \"d1\" in axis:\n",
    "            m = D1Symmetrize(m)\n",
    "        if \"d2\" in axis:\n",
    "            m = D2Symmetrize(m)\n",
    "        if np.array_equal(prevMatrix, m):\n",
    "            break\n",
    "            \n",
    "    return m\n",
    "\n",
    "# %% Color symmetric pixels (task 653)\n",
    "\n",
    "def colorSymmetricPixels(matrix, inColor, outColor, axis, includeAxis=False):\n",
    "    \"\"\"\n",
    "    This function finds the pixels of color inColor that are symmetric\n",
    "    according to the given axis in the input matrix, and colors them with the\n",
    "    color given by outColor.\n",
    "    Axis can be \"lr\", \"ud\", \"d1\" or \"d2\".\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    if axis==\"lr\":\n",
    "        for i,j in np.ndindex((m.shape[0], int(m.shape[1]/2))):\n",
    "            if m[i,j]==inColor and m[i,m.shape[1]-1-j]==inColor:\n",
    "                m[i,j] = outColor\n",
    "                m[i,m.shape[1]-1-j] = outColor\n",
    "        if includeAxis and ((m.shape[1]%2)==1):\n",
    "            j = int(m.shape[1]/2)\n",
    "            for i in range(m.shape[0]):\n",
    "                if m[i,j]==inColor:\n",
    "                    m[i,j] = outColor\n",
    "    if axis==\"ud\":\n",
    "        for i,j in np.ndindex((int(m.shape[0]/2), m.shape[1])):\n",
    "            if m[i,j]==inColor and m[m.shape[0]-1-i,j]==inColor:\n",
    "                m[i,j] = outColor\n",
    "                m[m.shape[0]-1-i,j] = outColor\n",
    "        if includeAxis and ((m.shape[0]%2)==1):\n",
    "            i = int(m.shape[0]/2)\n",
    "            for j in range(m.shape[1]):\n",
    "                if m[i,j]==inColor:\n",
    "                    m[i,j] = outColor\n",
    "    if axis==\"d1\":\n",
    "        for i in range(m.shape[0]):\n",
    "            for j in range(i):\n",
    "                if m[i,j]==inColor and m[j,i]==inColor:\n",
    "                    m[i,j] = outColor\n",
    "                    m[j,i] = outColor\n",
    "        if includeAxis:\n",
    "            for i in range(m.shape[0]):\n",
    "                if m[i,i]==inColor:\n",
    "                    m[i,i] = outColor\n",
    "    if axis==\"d2\":\n",
    "        for i in range(m.shape[0]):\n",
    "            for j in range(m.shape[0]-i-1):\n",
    "                if m[i,j]==inColor and m[m.shape[1]-j-1,m.shape[0]-i-1]==inColor:\n",
    "                    m[i,j] = outColor\n",
    "                    m[m.shape[1]-j-1,m.shape[0]-i-1] = outColor\n",
    "        if includeAxis:\n",
    "            for i in range(m.shape[0]):\n",
    "                if m[i, m.shape[0]-i-1]==inColor:\n",
    "                   m[i, m.shape[0]-i-1] = outColor\n",
    "                \n",
    "    return m\n",
    "\n",
    "def getBestColorSymmetricPixels(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"colorSymmetricPixels\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    \n",
    "    for cic in t.commonChangedInColors:\n",
    "        for coc in t.commonChangedOutColors:\n",
    "            f = partial(colorSymmetricPixels, inColor=cic, outColor=coc, \\\n",
    "                        axis=\"lr\", includeAxis=True)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            f = partial(colorSymmetricPixels, inColor=cic, outColor=coc, \\\n",
    "                        axis=\"lr\", includeAxis=False)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            f = partial(colorSymmetricPixels, inColor=cic, outColor=coc, \\\n",
    "                        axis=\"ud\", includeAxis=True)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            f = partial(colorSymmetricPixels, inColor=cic, outColor=coc, \\\n",
    "                        axis=\"ud\", includeAxis=False)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            if all([s.inMatrix.shape[0]==s.inMatrix.shape[1] for s in t.trainSamples+t.testSamples]):\n",
    "                f = partial(colorSymmetricPixels, inColor=cic, outColor=coc, \\\n",
    "                            axis=\"d1\", includeAxis=True)\n",
    "                bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                if bestScore==0:\n",
    "                    return bestFunction\n",
    "                f = partial(colorSymmetricPixels, inColor=cic, outColor=coc, \\\n",
    "                            axis=\"d1\", includeAxis=False)\n",
    "                bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                if bestScore==0:\n",
    "                    return bestFunction\n",
    "                f = partial(colorSymmetricPixels, inColor=cic, outColor=coc, \\\n",
    "                            axis=\"d2\", includeAxis=True)\n",
    "                bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                if bestScore==0:\n",
    "                    return bestFunction\n",
    "                f = partial(colorSymmetricPixels, inColor=cic, outColor=coc, \\\n",
    "                            axis=\"d2\", includeAxis=False)\n",
    "                bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                if bestScore==0:\n",
    "                    return bestFunction\n",
    "                \n",
    "    return bestFunction\n",
    "\n",
    "# %% Train and predict models  \n",
    "\n",
    "def trainCNNDummyColor(t, k, pad):\n",
    "    \"\"\"\n",
    "    This function trains a CNN with only one convolution of filter k and with\n",
    "    padding values equal to pad.\n",
    "    The training samples will have two channels: the background color and any\n",
    "    other color. The training loop loops through all the non-background colors\n",
    "    of each sample, treating them independently.\n",
    "    This is useful for tasks like number 3.\n",
    "    \"\"\"\n",
    "    model = OneConvModel(2, k, pad)\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    for e in range(50): # numEpochs            \n",
    "        optimizer.zero_grad()\n",
    "        loss = 0.0\n",
    "        for s in t.trainSamples:\n",
    "            for c in s.colors:\n",
    "                if c != t.backgroundColor:\n",
    "                    x = dummifyColor(s.inMatrix.m, c)\n",
    "                    x = torch.tensor(x).unsqueeze(0).float()\n",
    "                    y = deBackgroundizeMatrix(s.outMatrix.m, c)\n",
    "                    y = torch.tensor(y).unsqueeze(0).long()\n",
    "                    y_pred = model(x)\n",
    "                    loss += criterion(y_pred, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    return model\n",
    "\n",
    "@torch.no_grad()\n",
    "def predictCNNDummyColor(matrix, model):\n",
    "    \"\"\"\n",
    "    Predict function for a model trained using trainCNNDummyColor.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    pred = np.ones(m.shape, dtype=np.uint8) * matrix.backgroundColor\n",
    "    for c in matrix.colors:\n",
    "        if c != matrix.backgroundColor:\n",
    "            x = dummifyColor(m, c)\n",
    "            x = torch.tensor(x).unsqueeze(0).float()\n",
    "            x = model(x).argmax(1).squeeze(0).numpy()\n",
    "            for i,j in np.ndindex(m.shape):\n",
    "                if x[i,j] != 0:\n",
    "                    pred[i,j] = c\n",
    "    return pred\n",
    "\n",
    "def trainCNN(t, commonColors, nChannels, k=5, pad=0):\n",
    "    \"\"\"\n",
    "    This function trains a CNN model with kernel k and padding value pad.\n",
    "    It is required that all the training samples have the same number of colors\n",
    "    (adding the colors in the input and in the output).\n",
    "    It is also required that the output matrix has always the same shape as the\n",
    "    input matrix.\n",
    "    The colors are tried to be order in a specific way: first the colors that\n",
    "    are common to every sample (commonColors), and then the others.\n",
    "    \"\"\"\n",
    "    model = OneConvModel(nChannels, k, pad)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)\n",
    "    #losses = np.zeros(100)\n",
    "    for e in range(100):\n",
    "        optimizer.zero_grad()\n",
    "        loss = 0.0\n",
    "        for s in t.trainSamples:\n",
    "            sColors = commonColors.copy()\n",
    "            for c in s.colors:\n",
    "                if c not in sColors:\n",
    "                    sColors.append(c)\n",
    "            rel, invRel = relDicts(sColors)\n",
    "            x = dummify(s.inMatrix.m, nChannels, rel)\n",
    "            x = torch.tensor(x).unsqueeze(0).float()\n",
    "            y = s.outMatrix.m.copy()\n",
    "            for i,j in np.ndindex(y.shape):\n",
    "                y[i,j] = invRel[y[i,j]]\n",
    "            y = torch.tensor(y).unsqueeze(0).long()\n",
    "            y_pred = model(x)\n",
    "            loss += criterion(y_pred, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        #losses[e] = loss\n",
    "    return model#, losses\n",
    "\n",
    "@torch.no_grad()\n",
    "def predictCNN(matrix, model, commonColors, nChannels):\n",
    "    \"\"\"\n",
    "    Predict function for a model trained using trainCNN.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    pred = np.zeros(m.shape, dtype=np.uint8)\n",
    "    sColors = commonColors.copy()\n",
    "    for c in matrix.colors:\n",
    "        if c not in sColors:\n",
    "            sColors.append(c)\n",
    "    rel, invRel = relDicts(sColors)\n",
    "    if len(sColors) > nChannels:\n",
    "        return m\n",
    "    x = dummify(m, nChannels, rel)\n",
    "    x = torch.tensor(x).unsqueeze(0).float()\n",
    "    x = model(x).argmax(1).squeeze(0).numpy()\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        if x[i,j] not in rel.keys():\n",
    "            pred[i,j] = x[i,j]\n",
    "        else:\n",
    "            pred[i,j] = rel[x[i,j]][0]\n",
    "    return pred\n",
    "\n",
    "def getBestCNN(t):\n",
    "    \"\"\"\n",
    "    This function returns the best CNN with only one convolution, after trying\n",
    "    different kernel sizes and padding values.\n",
    "    There are as many channels as total colors or the minimum number of\n",
    "    channels that is necessary.\n",
    "    \"\"\"\n",
    "    kernel = [3,5,7]\n",
    "    pad = [0,-1]    \n",
    "    bestScore = 100000\n",
    "    for k, p in product(kernel, pad):\n",
    "        cc = list(range(10))\n",
    "        model = trainCNN(t, commonColors=cc, nChannels=10, k=k, pad=p)\n",
    "        score = sum([incorrectPixels(predictCNN(t.trainSamples[s].inMatrix, model, cc, 10), \\\n",
    "                                     t.trainSamples[s].outMatrix.m) for s in range(t.nTrain)])\n",
    "        if score < bestScore:\n",
    "            bestScore=score\n",
    "            ret = partial(predictCNN, model=model, commonColors=cc, nChannels=10)\n",
    "            if score==0:\n",
    "                return ret\n",
    "    return ret\n",
    "\n",
    "def getBestSameNSampleColorsCNN(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"predictCNN\" and works best for the training samples, after\n",
    "    training a CNN using the training Samples for training.\n",
    "    \"\"\"\n",
    "    kernel = [3,5,7]\n",
    "    pad = [0,-1]    \n",
    "    bestScore = 100000\n",
    "    for k, p in product(kernel, pad):\n",
    "        cc = list(t.commonSampleColors)\n",
    "        nc = t.trainSamples[0].nColors\n",
    "        model = trainCNN(t, commonColors=cc, nChannels=nc, k=k, pad=p)\n",
    "        score = sum([incorrectPixels(predictCNN(t.trainSamples[s].inMatrix, model, cc, nc), \\\n",
    "                                     t.trainSamples[s].outMatrix.m) for s in range(t.nTrain)])\n",
    "        if score < bestScore:\n",
    "            bestScore=score\n",
    "            ret = partial(predictCNN, model=model, commonColors=cc, nChannels=nc)\n",
    "            if score==0:\n",
    "                return ret\n",
    "            \n",
    "    return ret\n",
    "\n",
    "# %% CNN learning the output\n",
    "    \n",
    "def getNeighbourColors(m, i, j, border=0):\n",
    "    \"\"\"\n",
    "    Given a matrix m and a position i,j, this function returns a list of the\n",
    "    values of the neighbours of (i,j).\n",
    "    \"\"\"\n",
    "    x = []\n",
    "    y = m[i-1,j] if i>0 else border\n",
    "    x.append(y)\n",
    "    y = m[i,j-1] if j>0 else border\n",
    "    x.append(y)\n",
    "    y = m[i+1,j] if i<m.shape[0]-1 else border\n",
    "    x.append(y)\n",
    "    y = m[i,j+1] if j<m.shape[1]-1 else border\n",
    "    x.append(y)\n",
    "    return x\n",
    "\n",
    "def getDNeighbourColors(m, i, j, kernel=3, border=0):\n",
    "    \"\"\"\n",
    "    Given a matrix m and a position i,j, this function returns a list of the\n",
    "    values of the diagonal neighbours of (i,j).\n",
    "    \"\"\"\n",
    "    x = []\n",
    "    y = m[i-1,j-1] if (i>0 and j>0) else border\n",
    "    x.append(y)\n",
    "    y = m[i+1,j-1] if (i<m.shape[0]-1 and j>0) else border\n",
    "    x.append(y)\n",
    "    y = m[i-1,j+1] if (i>0 and j<m.shape[1]-1) else border\n",
    "    x.append(y)\n",
    "    y = m[i+1,j+1] if (i<m.shape[0]-1 and j<m.shape[1]-1) else border\n",
    "    x.append(y)\n",
    "    \n",
    "    if kernel==5:\n",
    "        y = m[i-2,j-2] if (i>1 and j>1) else border\n",
    "        x.append(y)\n",
    "        y = m[i-1,j-2] if (i>0 and j>1) else border\n",
    "        x.append(y)\n",
    "        y = m[i,j-2] if j>1 else border\n",
    "        x.append(y)\n",
    "        y = m[i+1,j-2] if (i<m.shape[0]-1 and j>1) else border\n",
    "        x.append(y)\n",
    "        y = m[i+2,j-2] if (i<m.shape[0]-2 and j>1) else border\n",
    "        x.append(y)\n",
    "        y = m[i+2,j-1] if (i<m.shape[0]-2 and j>0) else border\n",
    "        x.append(y)\n",
    "        y = m[i+2,j] if i<m.shape[0]-2 else border\n",
    "        x.append(y)\n",
    "        y = m[i+2,j+1] if (i<m.shape[0]-2 and j<m.shape[1]-1) else border\n",
    "        x.append(y)\n",
    "        y = m[i+2,j+2] if (i<m.shape[0]-2 and j<m.shape[1]-2) else border\n",
    "        x.append(y)\n",
    "        y = m[i+1,j+2] if (i<m.shape[0]-1 and j<m.shape[1]-2) else border\n",
    "        x.append(y)\n",
    "        y = m[i,j+2] if j<m.shape[1]-2 else border\n",
    "        x.append(y)\n",
    "        y = m[i-1,j+2] if (i>0 and j<m.shape[1]-2) else border\n",
    "        x.append(y)\n",
    "        y = m[i-2,j+2] if (i>1 and j<m.shape[1]-2) else border\n",
    "        x.append(y)\n",
    "        y = m[i-2,j+1] if (i>1 and j<m.shape[1]-1) else border\n",
    "        x.append(y)\n",
    "        y = m[i-2,j] if i>1 else border\n",
    "        x.append(y)\n",
    "        y = m[i-2,j-1] if (i>1 and j>0) else border\n",
    "        x.append(y)\n",
    "    return x\n",
    "\n",
    "def getAllNeighbourColors(m, i, j, kernel=3, border=0):\n",
    "    \"\"\"\n",
    "    This function returns a list the neighbour colors of the pixel i,j in the\n",
    "    matrix m.\n",
    "    \"\"\"\n",
    "    return getNeighbourColors(m,i,j,border) + getDNeighbourColors(m,i,j,kernel,border)\n",
    "\n",
    "def colorNeighbours(mIn, mOut ,i, j):\n",
    "    \"\"\"\n",
    "    Given matrices mIn and mOut, and coordinates i and j, this function colors\n",
    "    the neighbours of i and j in mIn with the colors of the neighbours of i and\n",
    "    j in mOut.\n",
    "    \"\"\"\n",
    "    if i>0:\n",
    "        mIn[i-1,j] = mOut[i-1,j]\n",
    "    if j>0:\n",
    "        mIn[i,j-1] = mOut[i,j-1]\n",
    "    if i<mIn.shape[0]-1:\n",
    "        mIn[i+1,j] = mOut[i+1,j]\n",
    "    if j<mIn.shape[1]-1:\n",
    "        mIn[i,j+1] = mOut[i,j+1]\n",
    "        \n",
    "def colorDNeighbours(mIn, mOut, i, j):\n",
    "    \"\"\"\n",
    "    Given matrices mIn and mOut, and coordinates i and j, this function colors\n",
    "    the neighbours of i and j in mIn with the colors of the neighbours of i and\n",
    "    j in mOut. It includes diagonal neighbours.\n",
    "    \"\"\"\n",
    "    colorNeighbours(mIn, mOut ,i, j)\n",
    "    if i>0 and j>0:\n",
    "        mIn[i-1,j-1] = mOut[i-1,j-1]\n",
    "    if i<mIn.shape[0]-1 and j>0:\n",
    "        mIn[i+1,j-1] = mOut[i+1,j-1]\n",
    "    if i>0 and j<mIn.shape[1]-1:\n",
    "        mIn[i-1,j+1] = mOut[i-1,j+1]\n",
    "    if i<mIn.shape[0]-1 and j<mIn.shape[1]-1:\n",
    "        mIn[i+1,j+1] = mOut[i+1,j+1]\n",
    "     \n",
    "def evolve(t, kernel=3, border=0, includeRotations=False):\n",
    "    \"\"\"\n",
    "    Given a task t, this function returns a list of dictionaries. Each\n",
    "    dictionary contains as keys lists of numbers representing colors of \n",
    "    neighbouring pixels. The value is the color of the pixel that is surrounded\n",
    "    by these neighbours. A key only exist if the value can be unique.\n",
    "    There are 4 dictionaries returned: For the first one, only the direct\n",
    "    neighbours are considered. For the second one, also the diagonal neighbours\n",
    "    are considered. For the third one, a kernel of 5 is considered (this is,\n",
    "    24 neighbours), and for the fourth one is of kernel 3 and also considers\n",
    "    the background color.\n",
    "    \"\"\"\n",
    "    def evolveInputMatrices(mIn, mOut, changeCIC=False):\n",
    "        reference = [m.copy() for m in mIn]\n",
    "        for m in range(len(mIn)):\n",
    "            if changeCIC:\n",
    "                for i,j in np.ndindex(mIn[m].shape):\n",
    "                    if mIn[m][i,j] not in set.union(fixedColors, changedOutColors): \n",
    "                        colorDNeighbours(mIn[m], mOut[m], i, j)\n",
    "                        break\n",
    "            else:\n",
    "                for i,j in np.ndindex(mIn[m].shape):\n",
    "                    if referenceIsFixed and reference[m][i,j] in fixedColors:\n",
    "                        colorDNeighbours(mIn[m], mOut[m], i, j)\n",
    "                    elif reference[m][i,j] in changedOutColors:\n",
    "                        colorDNeighbours(mIn[m], mOut[m], i, j)\n",
    "                    \n",
    "    nColors = t.trainSamples[0].nColors\n",
    "    \n",
    "    if not t.allEqual(t.sampleColors):\n",
    "        sampleRel = []\n",
    "        sampleInvRel = []\n",
    "        commonColors = t.orderedColors\n",
    "        for s in t.trainSamples:\n",
    "            colors = commonColors.copy()\n",
    "            for i,j in np.ndindex(s.inMatrix.shape):\n",
    "                if s.inMatrix.m[i,j] not in colors:\n",
    "                    colors.append(s.inMatrix.m[i,j])\n",
    "                    if len(colors) == nColors:\n",
    "                        break\n",
    "            if len(colors) != nColors:\n",
    "                for i,j in np.ndindex(s.outMatrix.shape):\n",
    "                    if s.outMatrix.m[i,j] not in colors:\n",
    "                        colors.append(s.outMatrix.m[i,j])\n",
    "                        if len(colors) == nColors:\n",
    "                            break\n",
    "            rel, invRel = relDicts(colors)\n",
    "            sampleRel.append(rel)\n",
    "            sampleInvRel.append(invRel)\n",
    "            \n",
    "        fixedColors = set()\n",
    "        for c in t.fixedColors:\n",
    "            fixedColors.add(sampleInvRel[0][c])\n",
    "        changedOutColors = set()\n",
    "        for c in t.commonChangedOutColors:\n",
    "            changedOutColors.add(sampleInvRel[0][c])\n",
    "        for c in range(len(commonColors), nColors):\n",
    "            changedOutColors.add(c)\n",
    "    else:\n",
    "        fixedColors = t.fixedColors\n",
    "        changedOutColors = t.commonChangedOutColors\n",
    "        \n",
    "    referenceIsFixed = t.trainSamples[0].inMatrix.nColors == len(fixedColors)+1\n",
    "    \n",
    "    outMatrices = [s.outMatrix.m.copy() for s in t.trainSamples]\n",
    "    referenceOutput = [s.inMatrix.m.copy() for s in t.trainSamples]\n",
    "    \n",
    "    if includeRotations:\n",
    "        for i in range(1,4):\n",
    "            for m in range(t.nTrain):\n",
    "                outMatrices.append(np.rot90(outMatrices[m].copy(), i))\n",
    "                referenceOutput.append(np.rot90(referenceOutput[m].copy(), i))\n",
    "                    \n",
    "    \n",
    "    if not t.allEqual(t.sampleColors):\n",
    "        for m in range(len(outMatrices)):\n",
    "            for i,j in np.ndindex(outMatrices[m].shape):\n",
    "                outMatrices[m][i,j] = sampleInvRel[m%t.nTrain][outMatrices[m][i,j]]\n",
    "                referenceOutput[m][i,j] = sampleInvRel[m%t.nTrain][referenceOutput[m][i,j]]\n",
    "    \n",
    "    colorFromNeighboursK2 = {}\n",
    "    colorFromNeighboursK3 = {}\n",
    "    colorFromNeighboursK5 = {}\n",
    "    for i in range(10):\n",
    "        referenceInput = [m.copy() for m in referenceOutput]\n",
    "        evolveInputMatrices(referenceOutput, outMatrices)\n",
    "        if np.all([np.array_equal(referenceInput[m], referenceOutput[m]) for m in range(len(referenceInput))]):\n",
    "            evolveInputMatrices(referenceOutput, outMatrices, True)\n",
    "        for m in range(len(outMatrices)):\n",
    "            for i,j in np.ndindex(referenceInput[m].shape):\n",
    "                if referenceInput[m][i,j] != referenceOutput[m][i,j]:\n",
    "                    neighbourColors = tuple(getNeighbourColors(referenceInput[m],i,j,border))\n",
    "                    colorFromNeighboursK2[neighbourColors] = referenceOutput[m][i,j]\n",
    "                    neighbourColors = tuple(getAllNeighbourColors(referenceInput[m],i,j,3,border))\n",
    "                    colorFromNeighboursK3[neighbourColors] = referenceOutput[m][i,j]\n",
    "                    neighbourColors = tuple(getAllNeighbourColors(referenceInput[m],i,j,5,border))\n",
    "                    colorFromNeighboursK5[neighbourColors] = referenceOutput[m][i,j]\n",
    "       \n",
    "    colorFromNeighboursK2 = {k:v for k,v in colorFromNeighboursK2.items() if \\\n",
    "                             not all([x not in set.union(changedOutColors, fixedColors) for x in k])}\n",
    "    colorFromNeighboursK3 = {k:v for k,v in colorFromNeighboursK3.items() if \\\n",
    "                             not all([x not in set.union(changedOutColors, fixedColors) for x in k])}\n",
    "    colorFromNeighboursK5 = {k:v for k,v in colorFromNeighboursK5.items() if \\\n",
    "                             not all([x not in set.union(changedOutColors, fixedColors) for x in k])}\n",
    "                \n",
    "    colorfromNeighboursK3Background = {}\n",
    "    for m in outMatrices:\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            if m[i,j] in t.commonChangedInColors:\n",
    "                neighbourColors = tuple(getAllNeighbourColors(m,i,j,3,border))\n",
    "                colorfromNeighboursK3Background[neighbourColors] = m[i,j]\n",
    "        \n",
    "    \n",
    "    return [colorFromNeighboursK2, colorFromNeighboursK3,\\\n",
    "            colorFromNeighboursK5, colorfromNeighboursK3Background]\n",
    "\n",
    "def applyEvolve(matrix, cfn, nColors, changedOutColors=set(), fixedColors=set(),\\\n",
    "                changedInColors=set(), referenceIsFixed=False, commonColors=set(),\\\n",
    "                kernel=None, border=0, nIterations=1000):\n",
    "    \"\"\"\n",
    "    Given a matrix and the \"colors from neighbours\" (cfn) list of dictionaries\n",
    "    returned by the function evolve, this function colors the pixels of the\n",
    "    matrix according to the rules given by cfn.\n",
    "    \"\"\"    \n",
    "    \n",
    "    def colorPixel(m,newM,i,j):\n",
    "        if newM[i,j] not in cic and colorAroundCIC==False:\n",
    "            return\n",
    "        if kernel==None:\n",
    "            tup3 = tuple(getAllNeighbourColors(m,i,j,3,border))\n",
    "            tup5 = tuple(getAllNeighbourColors(m,i,j,5,border))\n",
    "            #tup = getMostSimilarTuple(tup)\n",
    "            if tup3 in cfn[1].keys():\n",
    "                if tup3 in cfn[3].keys():\n",
    "                    if tup5 in cfn[2].keys():\n",
    "                        newM[i,j] = cfn[2][tup5]\n",
    "                else:\n",
    "                    newM[i,j] = cfn[1][tup3] \n",
    "            elif tup5 in cfn[2].keys():\n",
    "                newM[i,j] = cfn[2][tup5]\n",
    "        elif kernel==2:\n",
    "            tup2 = tuple(getNeighbourColors(m,i,j,border))\n",
    "            if tup2 in cfn[0].keys():\n",
    "                newM[i,j] = cfn[0][tup2]\n",
    "        elif kernel==3:\n",
    "            tup3 = tuple(getAllNeighbourColors(m,i,j,3,border))\n",
    "            if tup3 in cfn[1].keys():\n",
    "                newM[i,j] = cfn[1][tup3]\n",
    "        elif kernel==5:\n",
    "            tup5 = tuple(getAllNeighbourColors(m,i,j,5,border))\n",
    "            if tup5 in cfn[2].keys():\n",
    "                newM[i,j] = cfn[2][tup5]\n",
    "        \n",
    "    def colorPixelsAround(m,newM,i,j):\n",
    "        if i>0:\n",
    "            colorPixel(m,newM,i-1,j)\n",
    "        if j>0:\n",
    "            colorPixel(m,newM,i,j-1)\n",
    "        if i<m.shape[0]-1:\n",
    "            colorPixel(m,newM,i+1,j)\n",
    "        if j<m.shape[1]-1:\n",
    "            colorPixel(m,newM,i,j+1)\n",
    "        if i>0 and j>0:\n",
    "            colorPixel(m,newM,i-1,j-1)\n",
    "        if i<m.shape[0]-1 and j>0:\n",
    "            colorPixel(m,newM,i+1,j-1)\n",
    "        if i>0 and j<m.shape[1]-1:\n",
    "            colorPixel(m,newM,i-1,j+1)\n",
    "        if i<m.shape[0]-1 and j<m.shape[1]-1:\n",
    "            colorPixel(m,newM,i+1,j+1)\n",
    "    \n",
    "    m = matrix.m.copy()\n",
    "    \n",
    "    if len(commonColors) > 0:\n",
    "        colors = list(commonColors.copy())\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            if m[i,j] not in colors:\n",
    "                colors.append(m[i,j])\n",
    "        rel, invRel = relDicts(colors)\n",
    "        \n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            m[i,j] = invRel[m[i,j]]\n",
    "            \n",
    "        fc = set()\n",
    "        for c in fixedColors:\n",
    "            fc.add(invRel[c])        \n",
    "        coc = set()\n",
    "        for c in changedOutColors:\n",
    "            coc.add(invRel[c])\n",
    "        for c in range(len(commonColors), nColors):\n",
    "            coc.add(c)\n",
    "        cic = set()\n",
    "        for c in changedInColors:\n",
    "            cic.add(invRel[c])\n",
    "    else:\n",
    "        fc = fixedColors\n",
    "        coc = changedOutColors\n",
    "        cic = changedInColors\n",
    "    \n",
    "    it = 0\n",
    "    colorAroundCIC=False\n",
    "    while it<nIterations:\n",
    "        it += 1\n",
    "        newM = m.copy()\n",
    "        #seen = np.zeros(m.shape, dtype=np.bool)\n",
    "        if colorAroundCIC:\n",
    "            for i,j in np.ndindex(m.shape):\n",
    "                colorPixelsAround(m,newM,i,j)\n",
    "            colorAroundCIC=False\n",
    "            m=newM.copy()\n",
    "            continue\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            if referenceIsFixed and m[i,j] in fixedColors:\n",
    "                colorPixelsAround(m,newM,i,j)\n",
    "            elif m[i,j] in coc:\n",
    "                colorPixelsAround(m,newM,i,j)\n",
    "        if np.array_equal(newM,m):\n",
    "            if it==1:\n",
    "                colorAroundCIC=True\n",
    "            else:\n",
    "                break\n",
    "        m = newM.copy()\n",
    "        \n",
    "    if len(commonColors) > 0:\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            if m[i,j] in rel.keys():\n",
    "                m[i,j] = rel[m[i,j]][0]\n",
    "            else:\n",
    "                m[i,j] = rel[0][0] # Patch for bug in task 22\n",
    "        \n",
    "    return m\n",
    "    \n",
    "def getBestEvolve(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"evolve\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    nColors = t.trainSamples[0].nColors\n",
    "    fc = t.fixedColors\n",
    "    cic = t.commonChangedInColors\n",
    "    coc = t.commonChangedOutColors\n",
    "    refIsFixed = t.trainSamples[0].inMatrix.nColors == len(fc)+1\n",
    "    \n",
    "    bestScore = 1000\n",
    "    bestFunction = None\n",
    "    \n",
    "    cfn = evolve(t)\n",
    "    if t.allEqual(t.sampleColors):\n",
    "        f = partial(applyEvolve, cfn=cfn, nColors=nColors, changedOutColors=coc,\\\n",
    "                    fixedColors=fc, changedInColors=cic, referenceIsFixed=refIsFixed,\\\n",
    "                    kernel=None, border=0)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        \n",
    "        f =  partial(applyEvolve, cfn=cfn, nColors=nColors, changedOutColors=coc,\\\n",
    "                     fixedColors=fc, changedInColors=cic, referenceIsFixed=refIsFixed,\\\n",
    "                     kernel=5, border=0)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "\n",
    "    else:\n",
    "        f = partial(applyEvolve, cfn=cfn, nColors=nColors, changedOutColors=coc,\\\n",
    "                    fixedColors=fc, changedInColors=cic, referenceIsFixed=refIsFixed,\\\n",
    "                    kernel=None, border=0, commonColors=t.orderedColors)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction    \n",
    "        \n",
    "        f =  partial(applyEvolve, cfn=cfn, nColors=nColors, changedOutColors=coc,\\\n",
    "                     fixedColors=fc, changedInColors=cic, referenceIsFixed=refIsFixed,\\\n",
    "                     kernel=5, border=0, commonColors=t.orderedColors)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        \n",
    "    cfn = evolve(t, includeRotations=True)\n",
    "    if t.allEqual(t.sampleColors):\n",
    "        f = partial(applyEvolve, cfn=cfn, nColors=nColors, changedOutColors=coc,\\\n",
    "                    fixedColors=fc, changedInColors=cic, referenceIsFixed=refIsFixed,\\\n",
    "                    kernel=None, border=0)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction    \n",
    "        \n",
    "        f =  partial(applyEvolve, cfn=cfn, nColors=nColors, changedOutColors=coc,\\\n",
    "                     fixedColors=fc, changedInColors=cic, referenceIsFixed=refIsFixed,\\\n",
    "                     kernel=5, border=0)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "\n",
    "    else:\n",
    "        f = partial(applyEvolve, cfn=cfn, nColors=nColors, changedOutColors=coc,\\\n",
    "                    fixedColors=fc, changedInColors=cic, referenceIsFixed=refIsFixed,\\\n",
    "                    kernel=None, border=0, commonColors=t.orderedColors)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction    \n",
    "        \n",
    "        f =  partial(applyEvolve, cfn=cfn, nColors=nColors, changedOutColors=coc,\\\n",
    "                     fixedColors=fc, changedInColors=cic, referenceIsFixed=refIsFixed,\\\n",
    "                     kernel=5, border=0, commonColors=t.orderedColors)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        \n",
    "    return bestFunction\n",
    "\n",
    "\n",
    "# Good examples: 790,749,748,703,679,629,605,585,575,573,457,344,322,\n",
    "#                283,236,231,201,198,59,23\n",
    "\n",
    "class EvolvingLine():\n",
    "    \"\"\"\n",
    "    The object of this class include information on how to draw an evolving \n",
    "    line.\n",
    "    ...\n",
    "    Attributes\n",
    "    ----------\n",
    "    color: int\n",
    "        The color of the line\n",
    "    direction: str\n",
    "        The direction of the line. It can be one of 4: 'l', 'r', 'u', 'd'.\n",
    "    position: tupe (int, int)\n",
    "        The current position of the line. It marks where to draw the next\n",
    "        pixel.\n",
    "    cic: set\n",
    "        Short for \"changedInColors\". Colors of the input matrix that can be\n",
    "        modified.\n",
    "    colorRules: dict\n",
    "        Keys are colors. Values are strings encoding what the line should do\n",
    "        when it encounters this color. The values can be 'r' (turn right),\n",
    "        'l' (turn left), 'stop', 'skip' (skip over the pixel and continue), \n",
    "        'split' (generate two evolvingLines, one turning right and the other\n",
    "        one left) or 'convert' (change the color of the line)\n",
    "    dealWith: dict\n",
    "        Same as colorRules, but including borders. Borders are considered as\n",
    "        colors number 10 (left), 11 (right), 12 (top) and 13 (bottom).\n",
    "    fixedDirection: bool\n",
    "        True if the direction cannot be modified. False if it is modified when\n",
    "        turning.\n",
    "    stepSize: int\n",
    "        Number of pixels to color at every step. If not specified, it is 1.\n",
    "    turning: bool\n",
    "        True if the line is currently turning. False otherwise.\n",
    "    \n",
    "    Methods\n",
    "    -------\n",
    "    draw(self, m, direction=None):\n",
    "        Draw the evolving line in the matrix m.\n",
    "    \"\"\"\n",
    "    def __init__(self, color, direction, position, cic, source=None, \\\n",
    "                 colorRules=None, stepSize=None, fixedDirection=True, turning=False,\\\n",
    "                 turnAfterNSteps=[None, None], stepIncrease=None, \\\n",
    "                 alternateStepIncrease=False):\n",
    "        \"\"\"\n",
    "        cic = changedInColors\n",
    "        \"\"\"\n",
    "        self.source = source # Shape\n",
    "        self.color = color\n",
    "        self.direction = direction\n",
    "        self.position = position\n",
    "        self.cic = cic\n",
    "        self.colorRules = colorRules\n",
    "        self.fixedDirection = fixedDirection\n",
    "        self.dealWith = {}\n",
    "        self.stepSize = stepSize\n",
    "        self.turning = turning\n",
    "        # left=10, right=11, top=12, bot=13\n",
    "        for color in range(14): # 10 colors + 4 borders\n",
    "            self.dealWith[color] = 'stop'\n",
    "        if type(colorRules)==str:\n",
    "            for i in range(10):\n",
    "                if i not in cic:\n",
    "                    self.dealWith[i] = colorRules\n",
    "        elif colorRules!=None:\n",
    "            for cr in colorRules:\n",
    "                self.dealWith[cr[0]] = cr[1]    \n",
    "        self.dealWith[self.color] = \"skip\"\n",
    "        self.maxSteps = turnAfterNSteps.copy() # [number of maximum steps, direction]\n",
    "        self.stepIncrease = stepIncrease\n",
    "        self.step = -1\n",
    "        self.alternateStepIncrease = alternateStepIncrease\n",
    "        self.increaseStep = False\n",
    "        \n",
    "    def draw(self, m, direction=None):\n",
    "        # If we reached the maximum number of steps, turn\n",
    "        if self.maxSteps[0]!=None:\n",
    "            if self.maxSteps[0]==self.step:\n",
    "                if self.maxSteps[1]==\"stop\":\n",
    "                    return\n",
    "                \n",
    "                #self.turning=True\n",
    "                \n",
    "                if self.direction=='u' and self.maxSteps[1]=='r' or\\\n",
    "                self.direction=='d' and self.maxSteps[1]=='l':\n",
    "                    direction = 'r'\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'r'\n",
    "                        \n",
    "                elif self.direction=='u' and self.maxSteps[1]=='l' or\\\n",
    "                self.direction=='d' and self.maxSteps[1]=='r':\n",
    "                    direction = 'l'\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'l'   \n",
    "                        \n",
    "                elif self.direction=='r' and self.maxSteps[1]=='l' or\\\n",
    "                self.direction=='l' and self.maxSteps[1]=='r':\n",
    "                    direction = 'u'\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'u'\n",
    "                        \n",
    "                elif self.direction=='l' and self.maxSteps[1]=='l' or\\\n",
    "                self.direction=='r' and self.maxSteps[1]=='r':\n",
    "                    direction = 'd'\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'd' \n",
    "                        \n",
    "                if self.stepIncrease!=None:\n",
    "                    if self.increaseStep:\n",
    "                        self.maxSteps[0]+=self.stepIncrease\n",
    "                        if self.alternateStepIncrease:\n",
    "                            self.increaseStep=False\n",
    "                    else:\n",
    "                        self.increaseStep=True\n",
    "                        \n",
    "                self.step = -1\n",
    "                    \n",
    "        self.step += 1\n",
    "        \n",
    "        if direction==None:\n",
    "            direction=self.direction\n",
    "                    \n",
    "        # Left\n",
    "        if direction=='l':\n",
    "            if self.position[1]==0:\n",
    "                if not self.turning:\n",
    "                    self.turning=True\n",
    "                    self.dealWithColor(10, m)\n",
    "                return\n",
    "            newColor = m[self.position[0], self.position[1]-1]\n",
    "            if newColor in self.cic:\n",
    "                if self.turning:\n",
    "                    self.turning=False\n",
    "                m[self.position[0], self.position[1]-1] = self.color\n",
    "                self.position[1] -= 1\n",
    "                self.draw(m)\n",
    "            else:\n",
    "                if not self.turning:\n",
    "                    self.turning=True\n",
    "                    self.dealWithColor(newColor, m)\n",
    "    \n",
    "        # Right\n",
    "        if direction=='r':\n",
    "            if self.position[1]==m.shape[1]-1:\n",
    "                if not self.turning:\n",
    "                    self.turning=True\n",
    "                    self.dealWithColor(11, m)\n",
    "                return\n",
    "            newColor = m[self.position[0], self.position[1]+1]\n",
    "            if newColor in self.cic:\n",
    "                if self.turning:\n",
    "                    self.turning=False\n",
    "                m[self.position[0], self.position[1]+1] = self.color\n",
    "                self.position[1] += 1\n",
    "                self.draw(m)\n",
    "            else:\n",
    "                if not self.turning:\n",
    "                    self.turning=True\n",
    "                    self.dealWithColor(newColor, m)\n",
    "                \n",
    "        # Up\n",
    "        if direction=='u':\n",
    "            if self.position[0]==0:\n",
    "                if not self.turning:\n",
    "                    self.turning=True\n",
    "                    self.dealWithColor(12, m)\n",
    "                return\n",
    "            newColor = m[self.position[0]-1, self.position[1]]\n",
    "            if newColor in self.cic:\n",
    "                if self.turning:\n",
    "                    self.turning=False\n",
    "                m[self.position[0]-1, self.position[1]] = self.color\n",
    "                self.position[0] -= 1\n",
    "                self.draw(m)\n",
    "            else:\n",
    "                if not self.turning:\n",
    "                    self.turning=True\n",
    "                    self.dealWithColor(newColor, m)\n",
    "        \n",
    "        # Down\n",
    "        if direction=='d':\n",
    "            if self.position[0]==m.shape[0]-1:\n",
    "                if not self.turning:\n",
    "                    self.turning=True\n",
    "                    self.dealWithColor(13, m)\n",
    "                return\n",
    "            newColor = m[self.position[0]+1, self.position[1]]\n",
    "            if newColor in self.cic:\n",
    "                if self.turning:\n",
    "                    self.turning=False\n",
    "                m[self.position[0]+1, self.position[1]] = self.color\n",
    "                self.position[0] += 1\n",
    "                self.draw(m)\n",
    "            else:\n",
    "                if not self.turning:\n",
    "                    self.turning=True\n",
    "                    self.dealWithColor(newColor, m)\n",
    "        \n",
    "    def dealWithColor(self, color, m):\n",
    "        if self.dealWith[color] == \"stop\":\n",
    "            return\n",
    "        \n",
    "        if self.dealWith[color] == \"convert\":\n",
    "            if self.direction=='l':\n",
    "                if self.position[1]!=0:\n",
    "                    self.color = color\n",
    "                    self.position[1]-=1\n",
    "                    self.draw(m)\n",
    "                else:\n",
    "                    return\n",
    "            if self.direction=='r':\n",
    "                if self.position[1]!=m.shape[1]-1:\n",
    "                    self.color = color\n",
    "                    self.position[1]+=1\n",
    "                    self.draw(m)\n",
    "                else:\n",
    "                    return\n",
    "            if self.direction=='u':\n",
    "                if self.position[0]!=0:\n",
    "                    self.color = color\n",
    "                    self.position[0]-=1\n",
    "                    self.draw(m)\n",
    "                else:\n",
    "                    return\n",
    "            if self.direction=='d':\n",
    "                if self.position[0]!=m.shape[0]-1:\n",
    "                    self.color = color\n",
    "                    self.position[0]+=1\n",
    "                    self.draw(m)\n",
    "                else:\n",
    "                    return\n",
    "            \n",
    "        if self.dealWith[color] == \"split\":\n",
    "            if self.direction=='l' or self.direction=='r':\n",
    "                if self.position[0]!=0:\n",
    "                    l1 = EvolvingLine(self.color, self.direction, self.position.copy(), self.cic,\\\n",
    "                                      colorRules=self.colorRules, fixedDirection=self.fixedDirection, \\\n",
    "                                      turning=True)\n",
    "                    if self.fixedDirection==False:\n",
    "                        l1.direction='u'\n",
    "                    l1.draw(m, direction='u')\n",
    "                if self.position[0]!=m.shape[0]-1:\n",
    "                    l2 = EvolvingLine(self.color, self.direction, self.position.copy(), self.cic,\\\n",
    "                                      colorRules=self.colorRules, fixedDirection=self.fixedDirection, \\\n",
    "                                      turning=True)\n",
    "                    if self.fixedDirection==False:\n",
    "                        l2.direction='d'\n",
    "                    l2.draw(m, direction='d')\n",
    "            if self.direction=='u' or self.direction=='d':\n",
    "                if self.position[1]!=0:\n",
    "                    l1 = EvolvingLine(self.color, self.direction, self.position.copy(), self.cic,\\\n",
    "                                      colorRules=self.colorRules, fixedDirection=self.fixedDirection, \\\n",
    "                                      turning=True)\n",
    "                    if self.fixedDirection==False:\n",
    "                        l1.direction='l'\n",
    "                    l1.draw(m, direction='l')\n",
    "                if self.position[1]!=m.shape[1]-1:\n",
    "                    l2 = EvolvingLine(self.color, self.direction, self.position.copy(), self.cic,\\\n",
    "                                      colorRules=self.colorRules, fixedDirection=self.fixedDirection, \\\n",
    "                                      turning=True)\n",
    "                    if self.fixedDirection==False:\n",
    "                        l2.direction='r'\n",
    "                    l2.draw(m, direction='r')\n",
    "                    \n",
    "        if self.dealWith[color] == \"skip\":\n",
    "            if self.direction=='l':\n",
    "                if self.position[1]!=0:\n",
    "                    self.position[1]-=1\n",
    "                    self.draw(m)\n",
    "                else:\n",
    "                    return\n",
    "            if self.direction=='r':\n",
    "                if self.position[1]!=m.shape[1]-1:\n",
    "                    self.position[1]+=1\n",
    "                    self.draw(m)\n",
    "                else:\n",
    "                    return\n",
    "            if self.direction=='u':\n",
    "                if self.position[0]!=0:\n",
    "                    self.position[0]-=1\n",
    "                    self.draw(m)\n",
    "                else:\n",
    "                    return\n",
    "            if self.direction=='d':\n",
    "                if self.position[0]!=m.shape[0]-1:\n",
    "                    self.position[0]+=1\n",
    "                    self.draw(m)\n",
    "                else:\n",
    "                    return\n",
    "                    \n",
    "        # Left\n",
    "        if self.dealWith[color] == 'l':\n",
    "            if self.direction=='u':\n",
    "                if self.position[1]!=0:\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'l'\n",
    "                    self.draw(m, direction='l')\n",
    "                return\n",
    "            if self.direction=='d':\n",
    "                if self.position[1]!=m.shape[1]-1:\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'r'\n",
    "                    self.draw(m, direction='r')\n",
    "                return\n",
    "            if self.direction=='l':\n",
    "                if self.position[0]!=m.shape[0]-1:\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'd'\n",
    "                    self.draw(m, direction='d')\n",
    "                return\n",
    "            if self.direction=='r':\n",
    "                if self.position[0]!=0:\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'u'\n",
    "                    self.draw(m, direction='u')\n",
    "                return\n",
    "            \n",
    "        # Right\n",
    "        if self.dealWith[color] == 'r':\n",
    "            if self.direction=='u':\n",
    "                if self.position[1]!=m.shape[1]-1:\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'r'\n",
    "                    self.draw(m, direction='r')\n",
    "                return\n",
    "            if self.direction=='d':\n",
    "                if self.position[1]!=0:\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'l'\n",
    "                    self.draw(m, direction='l')\n",
    "                return\n",
    "            if self.direction=='l':\n",
    "                if self.position[0]!=0:\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'u'\n",
    "                    self.draw(m, direction='u')\n",
    "                return\n",
    "            if self.direction=='r':\n",
    "                if self.position[0]!=m.shape[0]-1:\n",
    "                    if not self.fixedDirection:\n",
    "                        self.direction = 'd'\n",
    "                    self.draw(m, direction='d')\n",
    "                return            \n",
    "        \n",
    "def detectEvolvingLineSources(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function detects the sources in which EvolvingLines\n",
    "    start. Currently it only supports pixels. It returns a set of pairs \n",
    "    (color (int), direction (str)).\n",
    "    \"\"\"\n",
    "    sources = set()\n",
    "    if len(t.commonChangedOutColors)==1:\n",
    "        coc = next(iter(t.commonChangedOutColors))\n",
    "    else:\n",
    "        coc = None\n",
    "    possibleSourceColors = set.intersection(t.commonChangedOutColors, t.commonInColors)\n",
    "    if len(possibleSourceColors) == 0:\n",
    "        possibleSourceColors = set.intersection(t.almostCommonColors, t.unchangedColors)\n",
    "    if len(possibleSourceColors) != 0:\n",
    "        firstIt = True\n",
    "        for sample in t.trainSamples:\n",
    "            sampleSources = set()\n",
    "            for color in possibleSourceColors:\n",
    "                if coc==None:\n",
    "                    targetColor=color\n",
    "                else:\n",
    "                    targetColor=coc\n",
    "                for shape in sample.inMatrix.shapes:\n",
    "                    if shape.color==color and shape.nPixels==1:                        \n",
    "                        # First special case: Corners\n",
    "                        if shape.position==(0,0):\n",
    "                            if sample.outMatrix.m[1][0]==targetColor and sample.outMatrix.m[0][1]==targetColor:\n",
    "                                sampleSources.add((color, \"away\"))\n",
    "                                sampleSources.add((color,'u'))\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                                sampleSources.add((color, 'r'))\n",
    "                            elif sample.outMatrix.m[1][0]==targetColor:\n",
    "                                sampleSources.add((color,'u'))\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                            elif sample.outMatrix.m[0][1]==targetColor:\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                                sampleSources.add((color, 'r'))\n",
    "                        elif shape.position==(0,sample.inMatrix.shape[1]-1):\n",
    "                            if sample.outMatrix.m[1][sample.outMatrix.shape[1]-1]==targetColor and sample.outMatrix.m[0][sample.outMatrix.shape[1]-2]==targetColor:\n",
    "                                sampleSources.add((color, \"away\"))\n",
    "                                sampleSources.add((color,'u'))\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                                sampleSources.add((color, 'r'))\n",
    "                            elif sample.outMatrix.m[1][sample.outMatrix.shape[1]-1]==targetColor:\n",
    "                                sampleSources.add((color,'u'))\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                            elif sample.outMatrix.m[0][sample.outMatrix.shape[1]-2]==targetColor:\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                                sampleSources.add((color, 'r'))\n",
    "                        elif shape.position==(sample.inMatrix.shape[0]-1,0):\n",
    "                            if sample.outMatrix.m[sample.outMatrix.shape[0]-2][0]==targetColor and sample.outMatrix.m[sample.outMatrix.shape[0]-1][1]==targetColor:\n",
    "                                sampleSources.add((color, \"away\"))\n",
    "                                sampleSources.add((color,'u'))\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                                sampleSources.add((color, 'r'))\n",
    "                            elif sample.outMatrix.m[sample.outMatrix.shape[0]-2][0]==targetColor:\n",
    "                                sampleSources.add((color,'u'))\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                            elif sample.outMatrix.m[sample.outMatrix.shape[0]-1][1]==targetColor:\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                                sampleSources.add((color, 'r'))\n",
    "                        elif shape.position==(sample.inMatrix.shape[0]-1,sample.inMatrix.shape[1]-1):\n",
    "                            if sample.outMatrix.m[sample.outMatrix.shape[0]-2][sample.outMatrix.shape[1]-1]==targetColor and sample.outMatrix.m[sample.outMatrix.shape[0]-1][sample.outMatrix.shape[1]-2]==targetColor:\n",
    "                                sampleSources.add((color, \"away\"))\n",
    "                                sampleSources.add((color,'u'))\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                                sampleSources.add((color, 'r'))\n",
    "                            elif sample.outMatrix.m[sample.outMatrix.shape[0]-2][sample.outMatrix.shape[1]-1]==targetColor:\n",
    "                                sampleSources.add((color,'u'))\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                            elif sample.outMatrix.m[sample.outMatrix.shape[0]-1][sample.outMatrix.shape[1]-2]==targetColor:\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                                sampleSources.add((color, 'r'))\n",
    "                        \n",
    "                        # Second special case: Border but not corner\n",
    "                        elif shape.position[0]== 0:\n",
    "                            if sample.outMatrix.m[1,shape.position[1]]==targetColor:\n",
    "                                sampleSources.add((color,\"away\"))\n",
    "                                sampleSources.add((color,'u'))\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                            if sample.outMatrix.m[0,shape.position[1]-1]==targetColor:\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                            if sample.outMatrix.m[0,shape.position[1]+1]==targetColor:\n",
    "                                sampleSources.add((color, 'r'))\n",
    "                        elif shape.position[0]== sample.inMatrix.shape[0]-1:\n",
    "                            if sample.outMatrix.m[sample.inMatrix.shape[0]-2,shape.position[1]]==targetColor:\n",
    "                                sampleSources.add((color,\"away\"))\n",
    "                                sampleSources.add((color,'u'))\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                            if sample.outMatrix.m[sample.inMatrix.shape[0]-1,shape.position[1]-1]==targetColor:\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                            if sample.outMatrix.m[sample.inMatrix.shape[0]-1,shape.position[1]+1]==targetColor:\n",
    "                                sampleSources.add((color, 'r'))\n",
    "                        elif shape.position[1]== 0:\n",
    "                            if sample.outMatrix.m[shape.position[0],1]==targetColor:\n",
    "                                sampleSources.add((color,\"away\"))\n",
    "                                sampleSources.add((color,'r'))\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                            if sample.outMatrix.m[shape.position[0]-1,0]==targetColor:\n",
    "                                sampleSources.add((color, 'u'))\n",
    "                            if sample.outMatrix.m[shape.position[0]+1,0]==targetColor:\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                        elif shape.position[1]== sample.inMatrix.shape[1]-1:\n",
    "                            if sample.outMatrix.m[shape.position[0],sample.inMatrix.shape[1]-2]==targetColor:\n",
    "                                sampleSources.add((color,\"away\"))\n",
    "                                sampleSources.add((color,'r'))\n",
    "                                sampleSources.add((color, 'l'))\n",
    "                            if sample.outMatrix.m[shape.position[0]-1,sample.inMatrix.shape[1]-1]==targetColor:\n",
    "                                sampleSources.add((color, 'u'))\n",
    "                            if sample.outMatrix.m[shape.position[0]+1,sample.inMatrix.shape[1]-1]==targetColor:\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                                \n",
    "                        # Third case: Not border\n",
    "                        else:\n",
    "                            if sample.outMatrix.m[shape.position[0]+1, shape.position[1]]==targetColor:\n",
    "                                sampleSources.add((color, 'd'))\n",
    "                            if sample.outMatrix.m[shape.position[0]-1, shape.position[1]]==targetColor:\n",
    "                                sampleSources.add((color, 'u'))\n",
    "                            if sample.outMatrix.m[shape.position[0], shape.position[1]+1]==targetColor:\n",
    "                                sampleSources.add((color, 'r'))\n",
    "                            if sample.outMatrix.m[shape.position[0], shape.position[1]-1]==targetColor:\n",
    "                                sampleSources.add((color, 'l'))\n",
    "            if firstIt:\n",
    "                sources = sampleSources\n",
    "                firstIt = False\n",
    "            else:\n",
    "                sources = set.intersection(sources, sampleSources) \n",
    "                \n",
    "    return sources\n",
    "\n",
    "def getBestEvolvingLines(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"drawEvolvingLines\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    if any([s.inMatrix.shape[0]==1 or s.inMatrix.shape[1]==1 for s in t.trainSamples]):\n",
    "        return partial(identityM)\n",
    "    \n",
    "    sources = detectEvolvingLineSources(t)\n",
    "    \n",
    "    fixedColorsList = list(t.fixedColors2)\n",
    "    cic=t.commonChangedInColors\n",
    "    #cic = [color for color in list(range(10)) if color not in fixedColorsList]\n",
    "    if len(t.commonChangedOutColors)==1:\n",
    "        coc = next(iter(t.commonChangedOutColors))\n",
    "    else:\n",
    "        coc = None\n",
    "    \n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    \n",
    "    mergeColors = t.commonOutColors-t.totalInColors\n",
    "    if len(mergeColors) == 1:\n",
    "        mergeColor = next(iter(mergeColors))\n",
    "    else:\n",
    "        mergeColor = None\n",
    "    \n",
    "    for actions in combinations_with_replacement([\"stop\", 'l', 'r', \"split\", \"skip\"],\\\n",
    "                                                 len(t.fixedColors2)):\n",
    "        rules = []\n",
    "        for c in range(len(fixedColorsList)):\n",
    "            rules.append([fixedColorsList[c], actions[c]])\n",
    "            \n",
    "        f = partial(drawEvolvingLines, sources=sources, rules=rules, cic=cic, \\\n",
    "                    fixedDirection=True, coc=coc, mergeColor=mergeColor)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        f = partial(drawEvolvingLines, sources=sources, rules=rules, cic=cic, \\\n",
    "                    fixedDirection=False, coc=coc, mergeColor=mergeColor)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if coc!=None:\n",
    "            f = partial(drawEvolvingLines, sources=sources, rules=rules, cic=cic, \\\n",
    "                        fixedDirection=True, coc=coc, mergeColor=mergeColor)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            f = partial(drawEvolvingLines, sources=sources, rules=rules, cic=cic, \\\n",
    "                        fixedDirection=False, coc=coc, mergeColor=mergeColor)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        \n",
    "    f = partial(drawEvolvingLines, sources=sources, rules=\"convert\", cic=cic, \\\n",
    "                    fixedDirection=False, coc=coc, mergeColor=mergeColor)  \n",
    "    bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        \n",
    "    for asi in [True, False]:\n",
    "        for stepIncrease in [None, 1, 2]:\n",
    "            for steps in [1, 2, 3, 4]:\n",
    "                for direction in ['r', 'l']:\n",
    "                    f = partial(drawEvolvingLines, sources=sources, rules=\"stop\", cic=cic, \\\n",
    "                    fixedDirection=False, alternateStepIncrease=asi, \\\n",
    "                    stepIncrease=stepIncrease, turnAfterNSteps=[steps, direction], mergeColor=mergeColor) \n",
    "                    bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            \n",
    "    return bestFunction\n",
    "\n",
    "def mergeMatrices(matrices, backgroundColor, mergeColor=None):\n",
    "    \"\"\"\n",
    "    All matrices are required to have the same shape.\n",
    "    \"\"\"\n",
    "    result = np.zeros(matrices[0].shape, dtype=np.uint8)\n",
    "    if mergeColor==None:\n",
    "        for i,j in np.ndindex(matrices[0].shape):\n",
    "            done=False\n",
    "            for m in matrices:\n",
    "                if m[i,j]!=backgroundColor:\n",
    "                    result[i,j] = m[i,j]\n",
    "                    done=True\n",
    "                    break\n",
    "            if not done:\n",
    "                result[i,j] = backgroundColor\n",
    "    else:\n",
    "        for i,j in np.ndindex(matrices[0].shape):\n",
    "            colors = set()\n",
    "            for m in matrices:\n",
    "                if m[i,j]!=backgroundColor:\n",
    "                    colors.add(m[i,j])\n",
    "            if len(colors)==0:\n",
    "                result[i,j] = backgroundColor\n",
    "            elif len(colors)==1:\n",
    "                result[i,j] = next(iter(colors))\n",
    "            else:\n",
    "                result[i,j] = mergeColor\n",
    "    return result\n",
    "        \n",
    "def drawEvolvingLines(matrix, sources, rules, cic, fixedDirection, coc=None, \\\n",
    "                      stepIncrease=None, alternateStepIncrease=False, \\\n",
    "                      turnAfterNSteps=[None, None], mergeColor=None):\n",
    "    \"\"\"\n",
    "    Given a set of sources, this function draws the evolving lines starting\n",
    "    at these sources following the given rules in the given matrix.\n",
    "    \"\"\"\n",
    "    if len(sources)==0:\n",
    "        return matrix.m.copy()\n",
    "    fd = fixedDirection\n",
    "    matrices = []\n",
    "    for source in sources:\n",
    "        newM = matrix.m.copy()\n",
    "        for i,j in np.ndindex(matrix.shape):\n",
    "            if matrix.m[i,j]==source[0]:\n",
    "                if source[1]==\"away\":\n",
    "                    if i==0:\n",
    "                        if coc==None:\n",
    "                            line = EvolvingLine(source[0], 'd', [i,j], cic, colorRules=rules, fixedDirection=fd,\\\n",
    "                                                stepIncrease=stepIncrease, alternateStepIncrease=alternateStepIncrease,\\\n",
    "                                                turnAfterNSteps=turnAfterNSteps)\n",
    "                        else:\n",
    "                            line = EvolvingLine(coc, 'd', [i,j], cic, colorRules=rules, fixedDirection=fd,\\\n",
    "                                                stepIncrease=stepIncrease, alternateStepIncrease=alternateStepIncrease,\\\n",
    "                                                turnAfterNSteps=turnAfterNSteps)\n",
    "                    elif i==matrix.m.shape[0]-1:\n",
    "                        if coc==None:\n",
    "                            line = EvolvingLine(source[0], 'u', [i,j], cic, colorRules=rules, fixedDirection=fd,\\\n",
    "                                                stepIncrease=stepIncrease, alternateStepIncrease=alternateStepIncrease,\\\n",
    "                                                turnAfterNSteps=turnAfterNSteps)\n",
    "                        else:\n",
    "                            line = EvolvingLine(coc, 'u', [i,j], cic, colorRules=rules, fixedDirection=fd,\\\n",
    "                                                stepIncrease=stepIncrease, alternateStepIncrease=alternateStepIncrease,\\\n",
    "                                                turnAfterNSteps=turnAfterNSteps)\n",
    "                    elif j==0:\n",
    "                        if coc==None:\n",
    "                            line = EvolvingLine(source[0], 'r', [i,j], cic, colorRules=rules, fixedDirection=fd,\\\n",
    "                                                stepIncrease=stepIncrease, alternateStepIncrease=alternateStepIncrease,\\\n",
    "                                                turnAfterNSteps=turnAfterNSteps)\n",
    "                        else:\n",
    "                            line = EvolvingLine(coc, 'r', [i,j], cic, colorRules=rules, fixedDirection=fd,\\\n",
    "                                                stepIncrease=stepIncrease, alternateStepIncrease=alternateStepIncrease,\\\n",
    "                                                turnAfterNSteps=turnAfterNSteps)\n",
    "                    elif j==matrix.m.shape[1]-1:\n",
    "                        if coc==None:\n",
    "                            line = EvolvingLine(source[0], 'l', [i,j], cic, colorRules=rules, fixedDirection=fd,\\\n",
    "                                                stepIncrease=stepIncrease, alternateStepIncrease=alternateStepIncrease,\\\n",
    "                                                turnAfterNSteps=turnAfterNSteps)\n",
    "                        else:\n",
    "                            line = EvolvingLine(coc, 'l', [i,j], cic, colorRules=rules, fixedDirection=fd,\\\n",
    "                                                stepIncrease=stepIncrease, alternateStepIncrease=alternateStepIncrease,\\\n",
    "                                                turnAfterNSteps=turnAfterNSteps)\n",
    "                    else:\n",
    "                        return matrix.m.copy()\n",
    "                else:\n",
    "                    if coc==None:\n",
    "                        line = EvolvingLine(source[0], source[1], [i,j], cic, colorRules=rules, fixedDirection=fd,\\\n",
    "                                            stepIncrease=stepIncrease, alternateStepIncrease=alternateStepIncrease,\\\n",
    "                                            turnAfterNSteps=turnAfterNSteps)\n",
    "                    else:\n",
    "                        line = EvolvingLine(coc, source[1], [i,j], cic, colorRules=rules, fixedDirection=fd,\\\n",
    "                                            stepIncrease=stepIncrease, alternateStepIncrease=alternateStepIncrease,\\\n",
    "                                            turnAfterNSteps=turnAfterNSteps)\n",
    "                line.draw(newM)\n",
    "        matrices.append(newM)\n",
    "    m = mergeMatrices(matrices, next(iter(cic)), mergeColor)\n",
    "    return m\n",
    "\n",
    "# %% Crossed coordinates\n",
    "    \n",
    "def paintCrossedCoordinates(matrix, refColor, outColor, fixedColors=set()):\n",
    "    \"\"\"\n",
    "    Given a Matrix, this function returns a matrix (numpy.ndarray) by coloring\n",
    "    the crossed coordinates that have a given refColor in the input matrix\n",
    "    with the color outColor.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    xCoord = set()\n",
    "    yCoord = set()\n",
    "    \n",
    "    for i,j in np.ndindex(matrix.shape):\n",
    "        if m[i,j]==refColor:\n",
    "            xCoord.add(i)\n",
    "            yCoord.add(j)\n",
    "    \n",
    "    for i in xCoord:\n",
    "        for j in yCoord:\n",
    "            if matrix.m[i,j] not in fixedColors:\n",
    "                m[i,j] = outColor\n",
    "    \n",
    "    return m\n",
    "\n",
    "# %% Linear Models\n",
    "\n",
    "# If input always has the same shape and output always has the same shape\n",
    "# And there is always the same number of colors in each sample    \n",
    "def trainLinearModel(t, commonColors, nChannels):\n",
    "    \"\"\"\n",
    "    This function trains a linear model.\n",
    "    It is required that all the training samples have the same number of colors\n",
    "    (adding the colors in the input and in the output).\n",
    "    It is also required that all the input matrices have the same shape, and\n",
    "    all the output matrices have the same shape.\n",
    "    The colors are tried to be order in a specific way: first the colors that\n",
    "    are common to every sample (commonColors), and then the others.\n",
    "    \"\"\"\n",
    "    model = LinearModel(t.inShape, t.outShape, nChannels)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)\n",
    "    for e in range(100):\n",
    "        optimizer.zero_grad()\n",
    "        loss = 0.0\n",
    "        for s in t.trainSamples:\n",
    "            sColors = commonColors.copy()\n",
    "            for c in s.colors:\n",
    "                if c not in sColors:\n",
    "                    sColors.append(c)\n",
    "            rel, invRel = relDicts(sColors)\n",
    "            x = dummify(s.inMatrix.m, nChannels, rel)\n",
    "            x = torch.tensor(x).unsqueeze(0).float()\n",
    "            y = s.outMatrix.m.copy()\n",
    "            for i,j in np.ndindex(y.shape):\n",
    "                y[i,j] = invRel[y[i,j]]\n",
    "            y = torch.tensor(y).unsqueeze(0).view(1,-1).long()\n",
    "            y_pred = model(x)\n",
    "            loss += criterion(y_pred, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    return model\n",
    "\n",
    "@torch.no_grad()\n",
    "def predictLinearModel(matrix, model, commonColors, nChannels, outShape):\n",
    "    \"\"\"\n",
    "    Predict function for a model trained using trainLinearModel.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    pred = np.zeros(outShape, dtype=np.uint8)\n",
    "    sColors = commonColors.copy()\n",
    "    for c in matrix.colors:\n",
    "        if c not in sColors:\n",
    "            sColors.append(c)\n",
    "    rel, invRel = relDicts(sColors)\n",
    "    if len(sColors) > nChannels:\n",
    "        return\n",
    "    x = dummify(m, nChannels, rel)\n",
    "    x = torch.tensor(x).unsqueeze(0).float()\n",
    "    x = model(x).argmax(1).squeeze(0).view(outShape).numpy()\n",
    "    for i,j in np.ndindex(outShape):\n",
    "        if x[i,j] not in rel.keys():\n",
    "            pred[i,j] = x[i,j]\n",
    "        else:\n",
    "            pred[i,j] = rel[x[i,j]][0]\n",
    "    return pred\n",
    "\n",
    "def trainLinearDummyModel(t):\n",
    "    \"\"\"\n",
    "    This function trains a linear model.\n",
    "    The training samples will have two channels: the background color and any\n",
    "    other color. The training loop loops through all the non-background colors\n",
    "    of each sample, treating them independently.\n",
    "    \"\"\"\n",
    "    model = LinearModelDummy(t.inShape, t.outShape)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.05)\n",
    "    for e in range(100):\n",
    "        optimizer.zero_grad()\n",
    "        loss = 0.0\n",
    "        for s in t.trainSamples:\n",
    "            for c in s.colors:\n",
    "                if c != t.backgroundColor:\n",
    "                    x = dummifyColor(s.inMatrix.m, c)\n",
    "                    x = torch.tensor(x).unsqueeze(0).float()\n",
    "                    y = deBackgroundizeMatrix(s.outMatrix.m, c)\n",
    "                    y = torch.tensor(y).unsqueeze(0).long()\n",
    "                    y = y.view(1, -1)\n",
    "                    y_pred = model(x)\n",
    "                    loss += criterion(y_pred, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    return model\n",
    "    \n",
    "@torch.no_grad()\n",
    "def predictLinearDummyModel(matrix, model, outShape, backgroundColor):\n",
    "    \"\"\"\n",
    "    Predict function for a model trained using trainLinearDummyModel.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    pred = np.zeros(outShape, dtype=np.uint8)\n",
    "    for c in matrix.colors:\n",
    "        if c != backgroundColor:\n",
    "            x = dummifyColor(m, c)\n",
    "            x = torch.tensor(x).unsqueeze(0).float()\n",
    "            x = model(x).argmax(1).squeeze().view(outShape).numpy()\n",
    "            for i,j in np.ndindex(outShape):\n",
    "                if x[i,j] != 0:\n",
    "                    pred[i,j] = c\n",
    "    return pred\n",
    "\n",
    "def trainLinearModelShapeColor(t):\n",
    "    \"\"\"\n",
    "    For trainLinearModelShapeColor we need to have the same shapes in the input\n",
    "    and in the output, and in the exact same positions. The training loop loops\n",
    "    through all the shapes of the task, and its aim is to predict the final\n",
    "    color of each shape.\n",
    "    The features of the linear model are:\n",
    "        - One feature per color in the task. Value of 1 if the shape has that\n",
    "        color, 0 otherwise.\n",
    "        - Several features representing the number of pixels of the shape.\n",
    "        Only one of these features can be equal to 1, the rest will be equal\n",
    "        to 0.\n",
    "        - 5 features to encode the number of holes of the shape (0,1,2,3 or 4)\n",
    "        - Feature encoding whether the shape is a square or not.\n",
    "        - Feature encoding whether the shape is a rectangle or not.\n",
    "        - Feature encoding whether the shape touches the border or not.\n",
    "    \"\"\"\n",
    "    inColors = set.union(*t.changedInColors+t.changedOutColors) - t.unchangedColors\n",
    "    colors = list(inColors) + list(set.union(*t.changedInColors+t.changedOutColors) - inColors)\n",
    "    rel, invRel = relDicts(list(colors))\n",
    "    shapePixelNumbers = t.shapePixelNumbers\n",
    "    _,nPixelsRel = relDicts(shapePixelNumbers)\n",
    "    # inFeatures: [colors that change], [number of pixels]+1, [number of holes] (0-4),\n",
    "    # isSquare, isRectangle, isBorder\n",
    "    nInFeatures = len(inColors) + len(shapePixelNumbers) + 1 + 5 + 3\n",
    "    model = SimpleLinearModel(nInFeatures, len(colors))\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)\n",
    "    num_epochs = 80\n",
    "    trainShapes = []\n",
    "    for s in t.trainSamples:\n",
    "        for shapeI in range(s.inMatrix.nShapes):\n",
    "            trainShapes.append((s.inMatrix.shapes[shapeI],\\\n",
    "                                s.outMatrix.shapes[shapeI].color))\n",
    "        for e in range(num_epochs):\n",
    "            optimizer.zero_grad()\n",
    "            loss = 0.0\n",
    "            for s,label in trainShapes:\n",
    "                inFeatures = torch.zeros(nInFeatures)\n",
    "                if s.color in inColors:\n",
    "                    inFeatures[invRel[s.color]] = 1\n",
    "                    inFeatures[len(inColors)+nPixelsRel[s.nPixels]] = 1\n",
    "                    inFeatures[len(inColors)+len(shapePixelNumbers)+1+min(s.nHoles, 4)] = 1\n",
    "                    inFeatures[nInFeatures-1] = int(s.isSquare)\n",
    "                    inFeatures[nInFeatures-2] = int(s.isRectangle)\n",
    "                    inFeatures[nInFeatures-1] = s.isBorder\n",
    "                    #inFeatures[nInFeatures-4] = s.nHoles\n",
    "                    #inFeatures[t.nColors+5] = s.position[0].item()\n",
    "                    #inFeatures[t.nColors+6] = s.position[1].item()\n",
    "                    y = torch.tensor(invRel[label]).unsqueeze(0).long()\n",
    "                    x = inFeatures.unsqueeze(0).float()\n",
    "                    y_pred = model(x)\n",
    "                    loss += criterion(y_pred, y)\n",
    "            if loss == 0:\n",
    "                continue\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            for p in model.parameters():\n",
    "                p.data.clamp_(min=0.05, max=1)\n",
    "    return model\n",
    "\n",
    "@torch.no_grad()\n",
    "def predictLinearModelShapeColor(matrix, model, colors, unchangedColors, shapePixelNumbers):\n",
    "    \"\"\"\n",
    "    Predict function for a model trained using trainLinearModelShapeColor.\n",
    "    \"\"\"\n",
    "    inColors = colors - unchangedColors\n",
    "    colors = list(inColors) + list(colors - inColors)\n",
    "    rel, invRel = relDicts(list(colors))\n",
    "    _,nPixelsRel = relDicts(shapePixelNumbers)\n",
    "    nInFeatures = len(inColors) + len(shapePixelNumbers) + 1 + 5 + 3\n",
    "    pred = matrix.m.copy()\n",
    "    for shape in matrix.shapes:\n",
    "        if shape.color in inColors:\n",
    "            inFeatures = torch.zeros(nInFeatures)\n",
    "            inFeatures[invRel[shape.color]] = 1\n",
    "            if shape.nPixels not in nPixelsRel.keys():\n",
    "                inFeatures[len(inColors)+len(shapePixelNumbers)] = 1\n",
    "            else:\n",
    "                inFeatures[len(inColors)+nPixelsRel[shape.nPixels]] = 1\n",
    "            inFeatures[len(inColors)+len(shapePixelNumbers)+1+min(shape.nHoles, 4)] = 1\n",
    "            inFeatures[nInFeatures-1] = int(shape.isSquare)\n",
    "            inFeatures[nInFeatures-2] = int(shape.isRectangle)\n",
    "            inFeatures[nInFeatures-3] = shape.isBorder\n",
    "            #inFeatures[nInFeatures-4] = shape.nHoles\n",
    "            #inFeatures[nColors+5] = shape.position[0].item()\n",
    "            #inFeatures[nColors+6] = shape.position[1].item()\n",
    "            x = inFeatures.unsqueeze(0).float()\n",
    "            y = model(x).squeeze().argmax().item()\n",
    "            pred = changeColorShapes(pred, [shape], rel[y][0])\n",
    "    return pred\n",
    "\n",
    "# %% LSTM\n",
    "def prepare_sequence(seq, to_ix):\n",
    "    \"\"\"\n",
    "    Utility function for LSTM.\n",
    "    \"\"\"\n",
    "    idxs = [to_ix[w] for w in seq]\n",
    "    return torch.tensor(idxs, dtype=torch.long)\n",
    "\n",
    "def trainLSTM(t, inColors, colors, inRel, outRel, reverse, order):\n",
    "    \"\"\"\n",
    "    This function tries to train a model that colors shapes according to a\n",
    "    sequence.\n",
    "    \"\"\"\n",
    "    EMBEDDING_DIM = 10\n",
    "    HIDDEN_DIM = 10\n",
    "    model = LSTMTagger(EMBEDDING_DIM, HIDDEN_DIM, len(inColors), len(colors))\n",
    "    loss_function = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "    num_epochs = 150\n",
    "    for epoch in range(num_epochs):\n",
    "        optimizer.zero_grad()\n",
    "        loss = 0.0\n",
    "        for s in t.trainSamples:\n",
    "            inShapes = [shape for shape in s.inMatrix.shapes if shape.color in inColors]\n",
    "            inSeq = sorted(inShapes, key=lambda x: (x.position[order[0]], x.position[order[1]]), reverse=reverse)\n",
    "            inSeq = [shape.color for shape in inSeq]\n",
    "            outShapes = [shape for shape in s.outMatrix.shapes if shape.color in colors]\n",
    "            targetSeq = sorted(outShapes, key=lambda x: (x.position[order[0]], x.position[order[1]]), reverse=reverse)\n",
    "            targetSeq = [shape.color for shape in targetSeq]\n",
    "            inSeq = prepare_sequence(inSeq, inRel)\n",
    "            targetSeq = prepare_sequence(targetSeq, outRel)\n",
    "            tag_scores = model(inSeq)\n",
    "            loss += loss_function(tag_scores, targetSeq)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    return model\n",
    "    \n",
    "@torch.no_grad()\n",
    "def predictLSTM(matrix, model, inColors, colors, inRel, rel, reverse, order):\n",
    "    \"\"\"\n",
    "    Predict function for a model trained using trainLSTM.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    inShapes = [shape for shape in matrix.shapes if shape.color in inColors]\n",
    "    if len(inShapes)==0:\n",
    "        return m\n",
    "    sortedShapes = sorted(inShapes, key=lambda x: (x.position[order[0]], x.position[order[1]]), reverse=reverse)\n",
    "    inSeq = [shape.color for shape in sortedShapes]\n",
    "    inSeq = prepare_sequence(inSeq, inRel)\n",
    "    pred = model(inSeq).argmax(1).numpy()\n",
    "    for shapeI in range(len(sortedShapes)):\n",
    "        m = changeColorShapes(m, [sortedShapes[shapeI]], rel[pred[shapeI]][0])\n",
    "    return m\n",
    "             \n",
    "def getBestLSTM(t):\n",
    "    \"\"\"\n",
    "    This function tries to find out which one is the best-fitting LSTM model\n",
    "    for the task t. The LSTM models try to change the color of shapes to fit\n",
    "    sequences. Examples are tasks 175, 331, 459 or 594.\n",
    "    4 LSTM models are trained, considering models that order shapes by X\n",
    "    coordinage, models that order them by Y coordinate, and considering both\n",
    "    directions of the sequence (normal and reverse).\n",
    "    \"\"\"\n",
    "    colors = set.union(*t.changedInColors+t.changedOutColors)\n",
    "    inColors = colors - t.unchangedColors\n",
    "    if len(inColors) == 0:\n",
    "        return partial(identityM)\n",
    "    _,inRel = relDicts(list(inColors))\n",
    "    colors = list(inColors) + list(colors - inColors)\n",
    "    rel, outRel = relDicts(colors)\n",
    "    \n",
    "    for s in t.trainSamples:\n",
    "        inShapes = [shape for shape in s.inMatrix.shapes if shape.color in inColors]\n",
    "        outShapes = [shape for shape in s.outMatrix.shapes if shape.color in colors]\n",
    "        if len(inShapes) != len(outShapes) or len(inShapes) == 0:\n",
    "            return partial(identityM)\n",
    "            \n",
    "    reverse = [True, False]\n",
    "    order = [(0,1), (1,0)]    \n",
    "    bestScore = 1000\n",
    "    for r, o in product(reverse, order):        \n",
    "        model = trainLSTM(t, inColors=inColors, colors=colors, inRel=inRel,\\\n",
    "                          outRel=outRel, reverse=r, order=o)\n",
    "        \n",
    "        score = 0\n",
    "        for s in t.trainSamples:\n",
    "            m = predictLSTM(s.inMatrix, model, inColors, colors, inRel, rel, r, o)\n",
    "            score += incorrectPixels(m, s.outMatrix.m)\n",
    "        if score < bestScore:\n",
    "            bestScore=score\n",
    "            ret = partial(predictLSTM, model=model, inColors=inColors,\\\n",
    "                          colors=colors, inRel=inRel, rel=rel, reverse=r, order=o) \n",
    "            if bestScore==0:\n",
    "                return ret\n",
    "    return ret\n",
    "\n",
    "# %% Other utility functions\n",
    "\n",
    "def insertShape(matrix, shape):\n",
    "    \"\"\"\n",
    "    Given a matrix (numpy.ndarray) and a Shape, this function returns the\n",
    "    same matrix but with the shape inserted.\n",
    "    \"\"\"\n",
    "    m = matrix.copy()\n",
    "    shapeM = shape.m.copy()\n",
    "    for i,j in np.ndindex(shape.shape):\n",
    "        if shapeM[i,j] != 255:\n",
    "            if shape.position[0]+i<matrix.shape[0] and shape.position[1]+j<matrix.shape[1]\\\n",
    "                    and shape.position[0]+i >= 0 and shape.position[1]+j >= 0:\n",
    "                m[tuple(map(operator.add, (i,j), shape.position))] = shapeM[i,j]\n",
    "    return m\n",
    "\n",
    "def deleteShape(matrix, shape, backgroundColor):\n",
    "    \"\"\"\n",
    "    Given a matrix (numpy.ndarray) and a Shape, this function substitutes\n",
    "    the shape by the background color of the matrix.\n",
    "    \"\"\"\n",
    "    m = matrix.copy()\n",
    "    for c in shape.pixels:\n",
    "        m[tuple(map(operator.add, c, shape.position))] = backgroundColor\n",
    "    return m\n",
    "\n",
    "def symmetrizeSubmatrix(matrix, ud=False, lr=False, rotation=False, newColor=None, subShape=None):\n",
    "    \"\"\"\n",
    "    Given a Matrix, make the non-background part symmetric\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    bC = matrix.backgroundColor\n",
    "    if np.all(m == bC):\n",
    "        return m\n",
    "    x1, x2, y1, y2 = 0, m.shape[0]-1, 0, m.shape[1]-1\n",
    "    while x1 <= x2 and np.all(m[x1,:] == bC):\n",
    "        x1 += 1\n",
    "    while x2 >= x1 and np.all(m[x2,:] == bC):\n",
    "        x2 -= 1\n",
    "    while y1 <= y2 and np.all(m[:,y1] == bC):\n",
    "        y1 += 1\n",
    "    while y2 >= y1 and np.all(m[:,y2] == bC):\n",
    "        y2 -= 1\n",
    "    subMat = m[x1:x2+1,y1:y2+1].copy()\n",
    "    \n",
    "    if subShape == None:\n",
    "        symList = []\n",
    "        if ud:\n",
    "            symList.append(np.flipud(subMat.copy()))\n",
    "            if lr:\n",
    "                symList.append(np.fliplr(subMat.copy()))\n",
    "                symList.append(np.fliplr(np.flipud(subMat.copy())))\n",
    "            elif lr:\n",
    "                symList.append(np.fliplr(subMat.copy()))\n",
    "        elif rotation:\n",
    "            for x in range(1,4):\n",
    "                symList.append(np.rot90(subMat.copy(),x))\n",
    "        for newSym in symList:\n",
    "            score, bestScore = 0, 0\n",
    "            bestX, bestY = 0, 0\n",
    "            for i, j in np.ndindex((m.shape[0]-newSym.shape[0]+1,m.shape[1]-newSym.shape[1]+1)):\n",
    "                score = np.count_nonzero(np.logical_or(m[i:i+newSym.shape[0],j:j+newSym.shape[1]] == newSym, newSym == bC))\n",
    "                if score > bestScore:\n",
    "                    bestX, bestY = i, j\n",
    "                    bestScore = score\n",
    "            for i, j in np.ndindex(newSym.shape):\n",
    "                if newSym[i,j] != bC:\n",
    "                    if newColor == None:\n",
    "                        m[bestX+i, bestY+j] = newSym[i,j]\n",
    "                    else:\n",
    "                        if m[bestX+i,bestY+j] == bC:\n",
    "                           m[bestX+i, bestY+j] = newColor \n",
    "    else:\n",
    "        found = False\n",
    "        for x in range(subMat.shape[0]-subShape.shape[0]+1):\n",
    "            for y in range(subMat.shape[1]-subShape.shape[1]+1):\n",
    "                #if np.all(m[x1+x:x1+x+subShape.shape[0],y1+y:y1+y+subShape.shape[1]]==subShape.m):\n",
    "                if np.all(np.equal(m[x1+x:x1+x+subShape.shape[0],y1+y:y1+y+subShape.shape[1]] == bC,subShape.m == 255)):\n",
    "                    found = True   \n",
    "                    break\n",
    "            if found:\n",
    "                break\n",
    "        if not found:\n",
    "            return m\n",
    "        if ud and lr:\n",
    "            if 2*x+x1+subShape.shape[0] > m.shape[0] or 2*y+y1+subShape.shape[0]> m.shape[1]:\n",
    "                return m\n",
    "            for i in range(subMat.shape[0]):\n",
    "                for j in range(subMat.shape[1]):\n",
    "                    if subMat[i][j] != bC:\n",
    "                        m[2*x+x1+subShape.shape[0]-i-1,y1+j] = subMat[i,j]\n",
    "                        m[x1+i,2*y+y1+subShape.shape[0]-j-1] = subMat[i,j]\n",
    "                        m[2*x+x1+subShape.shape[0]-i-1,2*y+y1+subShape.shape[0]-j-1] = subMat[i,j]\n",
    "        elif rotation:\n",
    "            if x1+y+x+subShape.shape[0] > m.shape[0] or y1+x+y+subShape.shape[1] > m.shape[1]\\\n",
    "                or y1+y-x+subMat.shape[0] >= m.shape[0] or x1+x-y+subMat.shape[1] >= m.shape[1]\\\n",
    "                or x1+2*x+subShape.shape[0] > m.shape[0] or y1+2*y+subShape.shape[0] > m.shape[1]:\n",
    "                return m\n",
    "            for i in range(subMat.shape[0]):\n",
    "                for j in range(subMat.shape[1]):\n",
    "                    if subMat[i,j] != bC:\n",
    "                        m[x1+x+subShape.shape[0]+y-j-1,y1+y-x+i] = subMat[i,j]\n",
    "                        m[x1+x-y+j,y1+y+subShape.shape[0]+x-i-1] = subMat[i,j]\n",
    "                        m[x1+2*x+subShape.shape[0]-i-1,y1+2*y+subShape.shape[0]-j-1] = subMat[i,j]        \n",
    "    return m\n",
    "\n",
    "def getBestSymmetrizeSubmatrix(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"symmetrizeSubmatrix\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    rotation, lr, ud = False, False, False\n",
    "    croppedSamples = [cropAllBackground(s.outMatrix) for s in t.trainSamples]\n",
    "    if all(np.all(np.flipud(m)==m) for m in croppedSamples):\n",
    "        lr = True\n",
    "    if all(np.all(np.fliplr(m)==m) for m in croppedSamples):    \n",
    "        ud = True    \n",
    "    if all(m.shape[0]==m.shape[1] and np.all(np.rot90(m)==m) for m in croppedSamples):\n",
    "        rotation = True\n",
    "    for sh in t.commonInDShapes:\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(symmetrizeSubmatrix,\\\n",
    "                                                        lr=lr,ud=ud,rotation=rotation,subShape=sh), bestScore, bestFunction)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(symmetrizeSubmatrix,lr=lr,\\\n",
    "                                                            ud=ud,rotation=rotation), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def colorMap(matrix, cMap):\n",
    "    \"\"\"\n",
    "    cMap is a dict of color changes. Each input color can map to one and only\n",
    "    one output color. Only valid if t.sameIOShapes.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        if m[i,j] in cMap.keys(): # Otherwise, it means m[i,j] unchanged\n",
    "            m[i,j] = cMap[matrix.m[i,j]]\n",
    "    return m\n",
    "\n",
    "def revertColorOrder(matrix):\n",
    "    m = matrix.m.copy()\n",
    "    colors = [color for color,count in sorted(matrix.colorCount.items(), key=lambda item:item[1])]\n",
    "    colorDict = {}\n",
    "    for i in range(len(colors)):\n",
    "        colorDict[colors[i]] = colors[len(colors)-i-1]\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        m[i,j] = colorDict[m[i,j]]\n",
    "    return m\n",
    "\n",
    "def changeColorShapes(matrix, shapes, color):\n",
    "    \"\"\"\n",
    "    Given a matrix (numpy.ndarray), a list of Shapes (they are expected to\n",
    "    be present in the matrix) and a color, this function returns the same\n",
    "    matrix, but with the shapes of the list having the given color.\n",
    "    \"\"\"\n",
    "    if len(shapes) == 0:\n",
    "        return matrix\n",
    "    m = matrix.copy()\n",
    "    if color not in list(range(10)):\n",
    "        return m\n",
    "    for s in shapes:\n",
    "        for c in s.pixels:\n",
    "            m[tuple(map(operator.add, c, s.position))] = color\n",
    "    return m\n",
    "\n",
    "def changeShapes(m, inColor, outColor, bigOrSmall=None, isBorder=None):\n",
    "    \"\"\"\n",
    "    Given a Matrix, this function changes the Shapes of the matrix\n",
    "    that have color inColor to having the color outColor, if they satisfy the\n",
    "    given conditions bigOrSmall (is the shape the smallest/biggest one?) and\n",
    "    isBorder.\n",
    "    \"\"\"\n",
    "    return changeColorShapes(m.m.copy(), m.getShapes(inColor, bigOrSmall, isBorder), outColor)\n",
    "\n",
    "def paintShapesInHalf(matrix, shapeColor, color, half, diagonal=False, middle=None):\n",
    "    \"\"\"\n",
    "    Half can be 'u', 'd', 'l' or 'r'.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    if diagonal:\n",
    "        shapesToPaint = [shape for shape in matrix.dShapes if shape.color==shapeColor]\n",
    "    else:\n",
    "        shapesToPaint = [shape for shape in matrix.shapes if shape.color==shapeColor]\n",
    "    \n",
    "    for shape in shapesToPaint:\n",
    "        if (shape.shape[0]%2)==0 or middle!=None:\n",
    "            if middle==True:\n",
    "                iLimit=int((shape.shape[0]+1)/2)\n",
    "            else:\n",
    "                iLimit=int(shape.shape[0]/2)\n",
    "            if half=='u':\n",
    "                for i,j in np.ndindex((iLimit, shape.shape[1])):\n",
    "                    if shape.m[i,j]==shape.color:\n",
    "                        m[shape.position[0]+i, shape.position[1]+j] = color\n",
    "            if half=='d':\n",
    "                for i,j in np.ndindex((iLimit, shape.shape[1])):\n",
    "                    if shape.m[shape.shape[0]-1-i,j]==shape.color:\n",
    "                        m[shape.position[0]+shape.shape[0]-1-i, shape.position[1]+j] = color\n",
    "        if (shape.shape[1]%2)==0 or middle!=None:\n",
    "            if middle==True:\n",
    "                jLimit=int((shape.shape[1]+1)/2)\n",
    "            else:\n",
    "                jLimit=int(shape.shape[1]/2)\n",
    "            if half=='l':\n",
    "                for i,j in np.ndindex((shape.shape[0], jLimit)):\n",
    "                    if shape.m[i,j]==shape.color:\n",
    "                        m[shape.position[0]+i, shape.position[1]+j] = color\n",
    "            if half=='r':\n",
    "                for i,j in np.ndindex((shape.shape[0], jLimit)):\n",
    "                    if shape.m[i,shape.shape[1]-1-j]==shape.color:\n",
    "                        m[shape.position[0]+i, shape.position[1]+shape.shape[1]-1-j] = color\n",
    "            \n",
    "    return m\n",
    "\n",
    "def getBestPaintShapesInHalf(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"paintShapesInHalf\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    for half in ['u', 'd', 'l', 'r']:\n",
    "        for cic in t.commonChangedInColors:\n",
    "            for coc in t.commonChangedOutColors:\n",
    "                for middle, diagonal in product([None, True, False], [True, False]):\n",
    "                    f = partial(paintShapesInHalf, shapeColor=cic, color=coc,\\\n",
    "                                half=half, diagonal=diagonal, middle=middle)\n",
    "                    bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                    if bestScore==0:\n",
    "                        return bestFunction\n",
    "    return bestFunction\n",
    "\n",
    "# %% Paint shapes from border color\n",
    "\n",
    "def paintShapeFromBorderColor(matrix, shapeColors, fixedColors, diagonals=False):\n",
    "    m = matrix.m.copy()\n",
    "    shapesToChange = [shape for shape in matrix.shapes if shape.color in shapeColors]\n",
    "    \n",
    "    for shape in shapesToChange:\n",
    "        neighbourColors = Counter()\n",
    "        for i,j in np.ndindex(shape.shape):\n",
    "            if shape.m[i,j]!=255:\n",
    "                if diagonals:\n",
    "                    neighbourColors += Counter(getAllNeighbourColors(matrix.m,\\\n",
    "                                               shape.position[0]+i, shape.position[1]+j,\\\n",
    "                                               kernel=3, border=-1))\n",
    "                else:\n",
    "                    neighbourColors += Counter(getNeighbourColors(matrix.m,\\\n",
    "                                               shape.position[0]+i, shape.position[1]+j,\\\n",
    "                                               border=-1))\n",
    "        for color in fixedColors|shapeColors|set([-1]):\n",
    "            neighbourColors.pop(color, None)\n",
    "        if len(neighbourColors)>0:\n",
    "            newColor = max(neighbourColors.items(), key=operator.itemgetter(1))[0]\n",
    "            m = changeColorShapes(m, [shape], newColor)\n",
    "    return m\n",
    "\n",
    "def getBestPaintShapeFromBorderColor(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"paintShapeFromBorderColor\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    shapeColors = t.commonChangedInColors\n",
    "    fixedColors = t.fixedColors\n",
    "    for diagonals in [True, False]:\n",
    "        f = partial(paintShapeFromBorderColor, shapeColors=shapeColors,\\\n",
    "                    fixedColors=fixedColors, diagonals=diagonals)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "    return bestFunction\n",
    "\n",
    "# %% Things with features\n",
    "def isFixedShape(shape, fixedShapeFeatures):\n",
    "    return shape.hasFeatures(fixedShapeFeatures)\n",
    "\n",
    "def hasFeatures(candidate, reference):\n",
    "    if all([i==False for i in reference]):\n",
    "        return False\n",
    "    for i in range(len(reference)):\n",
    "        if reference[i] and not candidate[i]:\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "# %% Change Colors with features\n",
    "\n",
    "def getClosestFixedShapeColor(shape, fixedShapes):\n",
    "    def getDistance(x1, x2):\n",
    "        x, y = sorted((x1, x2))\n",
    "        if x[0] <= x[1] < y[0] and all( y[0] <= y[1] for y in (x1,x2)):\n",
    "            return y[0] - x[1]\n",
    "        return 0\n",
    "    \n",
    "    color = 0\n",
    "    minDistance = 1000\n",
    "    for fs in fixedShapes:\n",
    "        xDist = getDistance([fs.position[0],fs.position[0]+fs.shape[0]-1], \\\n",
    "                            [shape.position[0],shape.position[0]+shape.shape[0]-1])\n",
    "        yDist = getDistance([fs.position[1],fs.position[1]+fs.shape[1]-1], \\\n",
    "                            [shape.position[1],shape.position[1]+shape.shape[1]-1])\n",
    "        \n",
    "        if xDist+yDist < minDistance:\n",
    "            minDistance = xDist+yDist\n",
    "            color = fs.color\n",
    "    return color\n",
    "\n",
    "def getShapeFeaturesForColorChange(t, fixedShapeFeatures=None, fixedColors=None,\\\n",
    "                                   predict=False):  \n",
    "    shapeFeatures = []\n",
    "    \n",
    "    if predict:\n",
    "        matrices = [t]\n",
    "    else:\n",
    "        matrices = [s.inMatrix for s in t.trainSamples]\n",
    "                  \n",
    "    for m in range(len(matrices)):\n",
    "        # Smallest and biggest shapes:\n",
    "        biggestShape = 0\n",
    "        smallestShape = 1000\n",
    "        for shape in matrices[m].shapes:\n",
    "            if shape.color not in fixedColors:\n",
    "                if shape.nPixels>biggestShape:\n",
    "                    biggestShape=shape.nPixels\n",
    "                if shape.nPixels<smallestShape:\n",
    "                    smallestShape=shape.nPixels\n",
    "                    \n",
    "        # Closest fixed shape\n",
    "        if predict:\n",
    "            fixedShapes = []\n",
    "            for shape in matrices[0].shapes:\n",
    "                if shape.hasFeatures(fixedShapeFeatures):\n",
    "                    fixedShapes.append(shape)\n",
    "        else:\n",
    "            fixedShapes = t.trainSamples[m].fixedShapes\n",
    "        \n",
    "        for shape in matrices[m].shapes:\n",
    "            shFeatures = []\n",
    "            for c in range(10):\n",
    "                shFeatures.append(shape.color==c)\n",
    "            shFeatures.append(shape.isBorder)\n",
    "            shFeatures.append(not shape.isBorder)\n",
    "            shFeatures.append(shape.lrSymmetric)\n",
    "            shFeatures.append(shape.udSymmetric)\n",
    "            shFeatures.append(shape.d1Symmetric)\n",
    "            shFeatures.append(shape.d2Symmetric)\n",
    "            shFeatures.append(shape.isSquare)\n",
    "            shFeatures.append(shape.isRectangle)\n",
    "            for nPix in range(1,30):\n",
    "                shFeatures.append(shape.nPixels==nPix)\n",
    "            for nPix in range(1,6):\n",
    "                shFeatures.append(shape.nPixels>nPix)\n",
    "            for nPix in range(2,7):\n",
    "                shFeatures.append(shape.nPixels<nPix)\n",
    "            shFeatures.append((shape.nPixels%2)==0)\n",
    "            shFeatures.append((shape.nPixels%2)==1)\n",
    "            for h in range(5):\n",
    "                shFeatures.append(shape.nHoles==h)\n",
    "            shFeatures.append(shape.nPixels==biggestShape)\n",
    "            shFeatures.append(shape.nPixels==smallestShape)\n",
    "            #shFeatures.append(self.isUniqueShape(sh))\n",
    "            #shFeatures.append(not self.isUniqueShape(sh))\n",
    "            closestFixedShapeColor = getClosestFixedShapeColor(shape, fixedShapes)\n",
    "            for c in range(10):\n",
    "                shFeatures.append(closestFixedShapeColor==c)\n",
    "                \n",
    "            shapeFeatures.append(shFeatures)\n",
    "            \n",
    "    return shapeFeatures\n",
    "\n",
    "def getColorChangesWithFeatures(t):\n",
    "    shapeFeatures = getShapeFeaturesForColorChange(t, fixedColors=t.fixedColors)\n",
    "    colorChangesWithFeatures = {}\n",
    "    nFeatures = len(shapeFeatures[0])\n",
    "    trueList = []\n",
    "    falseList = []\n",
    "    for i in range(nFeatures):\n",
    "        trueList.append(True)\n",
    "        falseList.append(False)\n",
    "    for c in t.totalOutColors:\n",
    "        colorChangesWithFeatures[c] = trueList\n",
    "    changeCounter = Counter() # How many shapes change to color c?\n",
    "    # First, initialise the colorChangesWithFeatures. For every outColor, we\n",
    "    # detect which features are True for all the shapes that change to that color.\n",
    "    shapeCounter = 0            \n",
    "    for sample in t.trainSamples:\n",
    "        for shape in sample.inMatrix.shapes:\n",
    "            shapeChanges = False\n",
    "            for i,j in np.ndindex(shape.shape):\n",
    "                if shape.m[i,j]!=255:\n",
    "                    if sample.outMatrix.m[shape.position[0]+i,shape.position[1]+j]!=shape.m[i,j]:\n",
    "                        color = sample.outMatrix.m[shape.position[0]+i,shape.position[1]+j]\n",
    "                        shapeChanges=True\n",
    "                        break\n",
    "            if shapeChanges:\n",
    "                changeCounter[color] += 1\n",
    "                colorChangesWithFeatures[color] = \\\n",
    "                [colorChangesWithFeatures[color][x] and shapeFeatures[shapeCounter][x]\\\n",
    "                 for x in range(nFeatures)]\n",
    "            shapeCounter += 1\n",
    "    # Now, there might be more True values than necessary in a certain entry\n",
    "    # of colorChangesWithFeatures. Therefore, we try to determine the minimum\n",
    "    # number of necessary True features.\n",
    "    for c in t.totalOutColors:\n",
    "        if colorChangesWithFeatures[c] == trueList:\n",
    "            continue\n",
    "        trueIndices = [i for i, x in enumerate(colorChangesWithFeatures[c]) if x]\n",
    "        # First, check if only one feature is enough\n",
    "        goodIndices = []\n",
    "        for index in trueIndices:\n",
    "            trueCount = 0\n",
    "            featureList = falseList.copy()\n",
    "            featureList[index] = True\n",
    "            for sf in shapeFeatures:\n",
    "                if hasFeatures(sf, featureList):\n",
    "                    trueCount += 1\n",
    "            # If the true count matches the number of changed shapes, we're done!\n",
    "            if trueCount == changeCounter[c]:\n",
    "                goodIndices.append(index)\n",
    "        if len(goodIndices) > 0:\n",
    "            featureList = falseList.copy()\n",
    "            for index in goodIndices:\n",
    "                featureList[index] = True\n",
    "            colorChangesWithFeatures[c] = featureList\n",
    "        # If we're not done, then check with combinations of 2 features\n",
    "        else:\n",
    "            for i,j in combinations(trueIndices, 2):\n",
    "                trueCount = 0\n",
    "                featureList = falseList.copy()\n",
    "                featureList[i] = True\n",
    "                featureList[j] = True\n",
    "                for sf in shapeFeatures:\n",
    "                    if hasFeatures(sf, featureList):\n",
    "                        trueCount += 1\n",
    "                # If the true count matches the number of changed shapes, we're done!\n",
    "                if trueCount == changeCounter[c]:\n",
    "                    colorChangesWithFeatures[c] = featureList\n",
    "                    break   \n",
    "                    \n",
    "    return colorChangesWithFeatures\n",
    "\n",
    "def changeShapesWithFeatures(matrix, ccwf, fixedColors, fixedShapeFeatures):\n",
    "    \"\"\"\n",
    "    ccwp stands for 'color change with properties'. It's a dictionary. Its keys\n",
    "    are integers encoding the color of the output shape, and its values are the\n",
    "    properties that the input shape has to satisfy in order to execute the\n",
    "    color change.\n",
    "    \"\"\"\n",
    "    featureList = getShapeFeaturesForColorChange(matrix, fixedColors=fixedColors,\\\n",
    "                                                 fixedShapeFeatures=fixedShapeFeatures,\\\n",
    "                                                 predict=True)\n",
    "    m = matrix.m.copy()\n",
    "    sortedCcwf = {k: v for k, v in sorted(ccwf.items(), key=lambda item: sum(item[1]))}\n",
    "    for color in sortedCcwf.keys():\n",
    "        for sh in range(len(matrix.shapes)):\n",
    "            if (matrix.shapes[sh].color in fixedColors):# or \\\n",
    "            #(matrix.shapes[sh].hasFeatures(fixedShapeFeatures)):\n",
    "                continue\n",
    "            if hasFeatures(featureList[sh], ccwf[color]):\n",
    "                m = changeColorShapes(m, [matrix.shapes[sh]], color)\n",
    "                #break\n",
    "    return m\n",
    "\n",
    "\n",
    "# %% Change pixels with features\n",
    "\n",
    "def pixelRecolor(t):\n",
    "    \"\"\"\n",
    "    if t.sameIOShapes\n",
    "    \"\"\"\n",
    "    Input = [s.inMatrix.m for s in t.trainSamples]\n",
    "    Output = [s.outMatrix.m for s in t.trainSamples]\n",
    "        \n",
    "    Best_Dict = -1\n",
    "    Best_Q1 = -1\n",
    "    Best_Q2 = -1\n",
    "    Best_v = -1\n",
    "    \n",
    "    # v ranges from 0 to 3. This gives an extra flexibility of measuring distance from any of the 4 corners\n",
    "    Pairs = []\n",
    "    for t in range(15):\n",
    "        for Q1 in range(1,8):\n",
    "            for Q2 in range(1,8):\n",
    "                if Q1+Q2 == t:\n",
    "                    Pairs.append((Q1,Q2))\n",
    "                    \n",
    "    for Q1, Q2 in Pairs:\n",
    "        for v in range(4):\n",
    "            if Best_Dict != -1:\n",
    "                continue\n",
    "            possible = True\n",
    "            Dict = {}\n",
    "            \n",
    "            for x, y in zip(Input, Output):\n",
    "                n = len(x)\n",
    "                k = len(x[0])\n",
    "                for i in range(n):\n",
    "                    for j in range(k):\n",
    "                        if v == 0 or v ==2:\n",
    "                            p1 = i%Q1\n",
    "                        else:\n",
    "                            p1 = (n-1-i)%Q1\n",
    "                        if v == 0 or v ==3:\n",
    "                            p2 = j%Q2\n",
    "                        else :\n",
    "                            p2 = (k-1-j)%Q2\n",
    "                        color1 = x[i][j]\n",
    "                        color2 = y[i][j]\n",
    "                        if color1 != color2:\n",
    "                            rule = (p1, p2, color1)\n",
    "                            if rule not in Dict:\n",
    "                                Dict[rule] = color2\n",
    "                            elif Dict[rule] != color2:\n",
    "                                possible = False\n",
    "            if possible:\n",
    "                \n",
    "                # Let's see if we actually solve the problem\n",
    "                for x, y in zip(Input, Output):\n",
    "                    n = len(x)\n",
    "                    k = len(x[0])\n",
    "                    for i in range(n):\n",
    "                        for j in range(k):\n",
    "                            if v == 0 or v ==2:\n",
    "                                p1 = i%Q1\n",
    "                            else:\n",
    "                                p1 = (n-1-i)%Q1\n",
    "                            if v == 0 or v ==3:\n",
    "                                p2 = j%Q2\n",
    "                            else :\n",
    "                                p2 = (k-1-j)%Q2\n",
    "                           \n",
    "                            color1 = x[i][j]\n",
    "                            rule = (p1,p2,color1)\n",
    "                            \n",
    "                            if rule in Dict:\n",
    "                                color2 = 0 + Dict[rule]\n",
    "                            else:\n",
    "                                color2 = 0 + y[i][j]\n",
    "                            if color2 != y[i][j]:\n",
    "                                possible = False \n",
    "                if possible:\n",
    "                    Best_Dict = Dict\n",
    "                    Best_Q1 = Q1\n",
    "                    Best_Q2 = Q2\n",
    "                    Best_v = v\n",
    "        \n",
    "    if Best_Dict == -1:\n",
    "        return [-1]#meaning that we didn't find a rule that works for the traning cases\n",
    "    else:\n",
    "        return [Best_Dict, Best_v, Best_Q1, Best_Q2]\n",
    "    \n",
    "def executePixelRecolor(matrix, Best_Dict, Best_v, Best_Q1, Best_Q2):\n",
    "    m = np.zeros(matrix.shape, dtype = np.uint8)\n",
    "    for i,j in np.ndindex(matrix.shape):\n",
    "        if Best_v == 0 or Best_v ==2:\n",
    "            p1 = i%Best_Q1\n",
    "        else:\n",
    "            p1 = (matrix.shape[0]-1-i)%Best_Q1\n",
    "        if Best_v == 0 or Best_v ==3:\n",
    "            p2 = j%Best_Q2\n",
    "        else :\n",
    "            p2 = (matrix.shape[1]-1-j)%Best_Q2\n",
    "       \n",
    "        color1 = matrix.m[i,j]\n",
    "        rule = (p1, p2, color1)\n",
    "        if (p1, p2, color1) in Best_Dict:\n",
    "            m[i][j] = 0 + Best_Dict[rule]\n",
    "        else:\n",
    "            m[i][j] = 0 + color1\n",
    " \n",
    "    return m\n",
    "\n",
    "def doRulesWithReference(m, reference, rules):\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        y = (m[i,j], reference[i,j])\n",
    "        if y in rules.keys():\n",
    "            m[i,j] = rules[y]\n",
    "    return m    \n",
    "\n",
    "def doPixelMod2Row(matrix, rules):\n",
    "    m = matrix.m.copy()\n",
    "    reference = np.zeros(m.shape, dtype=np.uint8)\n",
    "    onesRow = np.ones(m.shape[1], dtype=np.uint8)\n",
    "    for i in range(m.shape[0]):\n",
    "        if i%2 == 0:\n",
    "            reference[i,:] = onesRow.copy()\n",
    "    m = doRulesWithReference(m, reference, rules)\n",
    "    return m\n",
    "\n",
    "def doPixelMod3Row(matrix, rules):\n",
    "    m = matrix.m.copy()\n",
    "    reference = np.zeros(m.shape, dtype=np.uint8)\n",
    "    onesRow = np.ones(m.shape[1], dtype=np.uint8)\n",
    "    twosRow = np.full(m.shape[1], 2, dtype=np.uint8)\n",
    "    for i in range(m.shape[0]):\n",
    "        if i%3 == 0:\n",
    "            reference[i,:] = onesRow.copy()\n",
    "        elif i%3 == 1:\n",
    "            reference[i,:] = twosRow.copy()\n",
    "    m = doRulesWithReference(m, reference, rules)\n",
    "    return m\n",
    "\n",
    "def doPixelMod2RowReverse(matrix, rules):\n",
    "    m = matrix.m.copy()\n",
    "    reference = np.zeros(m.shape, dtype=np.uint8)\n",
    "    onesRow = np.ones(m.shape[1], dtype=np.uint8)\n",
    "    for i in range(m.shape[0]):\n",
    "        if i%2 == 0:\n",
    "            reference[m.shape[0]-i-1,:] = onesRow.copy()\n",
    "    m = doRulesWithReference(m, reference, rules)\n",
    "    return m\n",
    "\n",
    "def doPixelMod3RowReverse(matrix, rules):\n",
    "    m = matrix.m.copy()\n",
    "    reference = np.zeros(m.shape, dtype=np.uint8)\n",
    "    onesRow = np.ones(m.shape[1], dtype=np.uint8)\n",
    "    twosRow = np.full(m.shape[1], 2, dtype=np.uint8)\n",
    "    for i in range(m.shape[0]):\n",
    "        if i%3 == 0:\n",
    "            reference[m.shape[0]-i-1,:] = onesRow.copy()\n",
    "        elif i%3 == 1:\n",
    "            reference[m.shape[0]-i-1,:] = twosRow.copy()\n",
    "    m = doRulesWithReference(m, reference, rules)\n",
    "    return m\n",
    "\n",
    "def doPixelMod2Col(matrix, rules):\n",
    "    m = matrix.m.copy()\n",
    "    reference = np.zeros(m.shape, dtype=np.uint8)\n",
    "    onesCol = np.ones(m.shape[0], dtype=np.uint8)\n",
    "    for j in range(m.shape[1]):\n",
    "        if j%2 == 0:\n",
    "            reference[:,j] = onesCol.copy()\n",
    "    m = doRulesWithReference(m, reference, rules)\n",
    "    return m\n",
    "\n",
    "def doPixelMod3Col(matrix, rules):\n",
    "    m = matrix.m.copy()\n",
    "    reference = np.zeros(m.shape, dtype=np.uint8)\n",
    "    onesCol = np.ones(m.shape[0], dtype=np.uint8)\n",
    "    twosCol = np.full(m.shape[0], 2, dtype=np.uint8)\n",
    "    for j in range(m.shape[1]):\n",
    "        if j%3 == 0:\n",
    "            reference[:,j] = onesCol.copy()\n",
    "        elif j%3 == 1:\n",
    "            reference[:,j] = twosCol.copy()\n",
    "    m = doRulesWithReference(m, reference, rules)\n",
    "    return m\n",
    "\n",
    "def doPixelMod2ColReverse(matrix, rules):\n",
    "    m = matrix.m.copy()\n",
    "    reference = np.zeros(m.shape, dtype=np.uint8)\n",
    "    onesCol = np.ones(m.shape[0], dtype=np.uint8)\n",
    "    for j in range(m.shape[1]):\n",
    "        if j%2 == 0:\n",
    "            reference[:,m.shape[1]-j-1] = onesCol.copy()\n",
    "    m = doRulesWithReference(m, reference, rules)\n",
    "    return m\n",
    "\n",
    "def doPixelMod3ColReverse(matrix, rules):\n",
    "    m = matrix.m.copy()\n",
    "    reference = np.zeros(m.shape, dtype=np.uint8)\n",
    "    onesCol = np.ones(m.shape[0], dtype=np.uint8)\n",
    "    twosCol = np.full(m.shape[0], 2, dtype=np.uint8)\n",
    "    for j in range(m.shape[1]):\n",
    "        if j%3 == 0:\n",
    "            reference[:,m.shape[1]-j-1] = onesCol.copy()\n",
    "        elif j%3 == 1:\n",
    "            reference[:,m.shape[1]-j-1] = twosCol.copy()\n",
    "    m = doRulesWithReference(m, reference, rules)\n",
    "    return m\n",
    "\n",
    "def doPixelMod2Alternate(matrix, rules):\n",
    "    m = matrix.m.copy()\n",
    "    reference = np.zeros(m.shape, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        reference[i,j] = (i+j)%2\n",
    "    m = doRulesWithReference(m, reference, rules)\n",
    "    return m\n",
    "\n",
    "def doPixelMod3Alternate(matrix, rules):\n",
    "    m = matrix.m.copy()\n",
    "    reference = np.zeros(m.shape, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        reference[i,j] = (i+j)%3\n",
    "    m = doRulesWithReference(m, reference, rules)\n",
    "    return m\n",
    "\n",
    "def getPixelChangeCriteria(t):\n",
    "    # Row\n",
    "    # Mod 2\n",
    "    x = {}\n",
    "    for sample in t.trainSamples:\n",
    "        reference = np.zeros(sample.inMatrix.shape, dtype=np.uint8)\n",
    "        onesRow = np.ones(sample.inMatrix.shape[1], dtype=np.uint8)\n",
    "        for i in range(sample.inMatrix.shape[0]):\n",
    "            if i%2 == 0:\n",
    "                reference[i,:] = onesRow.copy()\n",
    "        for i,j in np.ndindex(reference.shape):\n",
    "            y = (sample.inMatrix.m[i,j], reference[i,j])\n",
    "            if y in x.keys():\n",
    "                x[y].add(sample.outMatrix.m[i,j])\n",
    "            else:\n",
    "                x[y] = set([sample.outMatrix.m[i,j]])\n",
    "    x = {k:next(iter(v)) for k,v in x.items() if len(v)==1}\n",
    "    x = {k:v for k,v in x.items() if v not in t.fixedColors}\n",
    "    if len(x)>0:\n",
    "        return partial(doPixelMod2Row, rules=x)\n",
    "    # Mod 3\n",
    "    x = {}\n",
    "    for sample in t.trainSamples:\n",
    "        reference = np.zeros(sample.inMatrix.shape, dtype=np.uint8)\n",
    "        onesRow = np.ones(sample.inMatrix.shape[1], dtype=np.uint8)\n",
    "        twosRow = np.full(sample.inMatrix.shape[1], 2, dtype=np.uint8)\n",
    "        for i in range(sample.inMatrix.shape[0]):\n",
    "            if i%3 == 0:\n",
    "                reference[i,:] = onesRow.copy()\n",
    "            elif i%3 == 1:\n",
    "                reference[i,:] = twosRow.copy()\n",
    "        for i,j in np.ndindex(reference.shape):\n",
    "            y = (sample.inMatrix.m[i,j], reference[i,j])\n",
    "            if y in x.keys():\n",
    "                x[y].add(sample.outMatrix.m[i,j])\n",
    "            else:\n",
    "                x[y] = set([sample.outMatrix.m[i,j]])\n",
    "    x = {k:next(iter(v)) for k,v in x.items() if len(v)==1}\n",
    "    x = {k:v for k,v in x.items() if v not in t.fixedColors}\n",
    "    if len(x)>0:\n",
    "        return partial(doPixelMod3Row, rules=x)\n",
    "    \n",
    "    # Row Reverse\n",
    "    # Mod 2\n",
    "    x = {}\n",
    "    for sample in t.trainSamples:\n",
    "        reference = np.zeros(sample.inMatrix.shape, dtype=np.uint8)\n",
    "        onesRow = np.ones(sample.inMatrix.shape[1], dtype=np.uint8)\n",
    "        for i in range(sample.inMatrix.shape[0]):\n",
    "            if i%2 == 0:\n",
    "                reference[sample.inMatrix.shape[0]-i-1,:] = onesRow.copy()\n",
    "        for i,j in np.ndindex(reference.shape):\n",
    "            y = (sample.inMatrix.m[i,j], reference[i,j])\n",
    "            if y in x.keys():\n",
    "                x[y].add(sample.outMatrix.m[i,j])\n",
    "            else:\n",
    "                x[y] = set([sample.outMatrix.m[i,j]])\n",
    "    x = {k:next(iter(v)) for k,v in x.items() if len(v)==1}\n",
    "    x = {k:v for k,v in x.items() if v not in t.fixedColors}\n",
    "    if len(x)>0:\n",
    "        return partial(doPixelMod2RowReverse, rules=x)\n",
    "    # Mod 3\n",
    "    x = {}\n",
    "    for sample in t.trainSamples:\n",
    "        reference = np.zeros(sample.inMatrix.shape, dtype=np.uint8)\n",
    "        onesRow = np.ones(sample.inMatrix.shape[1], dtype=np.uint8)\n",
    "        twosRow = np.full(sample.inMatrix.shape[1], 2, dtype=np.uint8)\n",
    "        for i in range(sample.inMatrix.shape[0]):\n",
    "            if i%3 == 0:\n",
    "                reference[sample.inMatrix.shape[0]-i-1,:] = onesRow.copy()\n",
    "            elif i%3 == 1:\n",
    "                reference[sample.inMatrix.shape[0]-i-1,:] = twosRow.copy()\n",
    "        for i,j in np.ndindex(reference.shape):\n",
    "            y = (sample.inMatrix.m[i,j], reference[i,j])\n",
    "            if y in x.keys():\n",
    "                x[y].add(sample.outMatrix.m[i,j])\n",
    "            else:\n",
    "                x[y] = set([sample.outMatrix.m[i,j]])\n",
    "    x = {k:next(iter(v)) for k,v in x.items() if len(v)==1}\n",
    "    x = {k:v for k,v in x.items() if v not in t.fixedColors}\n",
    "    if len(x)>0:\n",
    "        return partial(doPixelMod3RowReverse, rules=x)\n",
    "\n",
    "    # Col\n",
    "    # Mod 2\n",
    "    x = {}\n",
    "    for sample in t.trainSamples:\n",
    "        reference = np.zeros(sample.inMatrix.shape, dtype=np.uint8)\n",
    "        onesCol = np.ones(sample.inMatrix.shape[0], dtype=np.uint8)\n",
    "        for j in range(sample.inMatrix.shape[1]):\n",
    "            if j%2 == 0:\n",
    "                reference[:,j] = onesCol.copy()\n",
    "        for i,j in np.ndindex(reference.shape):\n",
    "            y = (sample.inMatrix.m[i,j], reference[i,j])\n",
    "            if y in x.keys():\n",
    "                x[y].add(sample.outMatrix.m[i,j])\n",
    "            else:\n",
    "                x[y] = set([sample.outMatrix.m[i,j]])\n",
    "    x = {k:next(iter(v)) for k,v in x.items() if len(v)==1}\n",
    "    x = {k:v for k,v in x.items() if v not in t.fixedColors}\n",
    "    if len(x)>0:\n",
    "        return partial(doPixelMod2Col, rules=x)\n",
    "    # Mod 3\n",
    "    x = {}\n",
    "    for sample in t.trainSamples:\n",
    "        reference = np.zeros(sample.inMatrix.shape, dtype=np.uint8)\n",
    "        onesCol = np.ones(sample.inMatrix.shape[0], dtype=np.uint8)\n",
    "        twosCol = np.full(sample.inMatrix.shape[0], 2, dtype=np.uint8)\n",
    "        for j in range(sample.inMatrix.shape[1]):\n",
    "            if j%3 == 0:\n",
    "                reference[:,j] = onesCol.copy()\n",
    "            elif j%3 == 1:\n",
    "                reference[:,j] = twosCol.copy()\n",
    "        for i,j in np.ndindex(reference.shape):\n",
    "            y = (sample.inMatrix.m[i,j], reference[i,j])\n",
    "            if y in x.keys():\n",
    "                x[y].add(sample.outMatrix.m[i,j])\n",
    "            else:\n",
    "                x[y] = set([sample.outMatrix.m[i,j]])\n",
    "    x = {k:next(iter(v)) for k,v in x.items() if len(v)==1}\n",
    "    x = {k:v for k,v in x.items() if v not in t.fixedColors}\n",
    "    if len(x)>0:\n",
    "        return partial(doPixelMod3Col, rules=x)\n",
    "    \n",
    "    # Col Reverse\n",
    "    # Mod 2\n",
    "    x = {}\n",
    "    for sample in t.trainSamples:\n",
    "        reference = np.zeros(sample.inMatrix.shape, dtype=np.uint8)\n",
    "        onesCol = np.ones(sample.inMatrix.shape[0], dtype=np.uint8)\n",
    "        for j in range(sample.inMatrix.shape[1]):\n",
    "            if j%2 == 0:\n",
    "                reference[:,sample.inMatrix.shape[1]-j-1] = onesCol.copy()\n",
    "        for i,j in np.ndindex(reference.shape):\n",
    "            y = (sample.inMatrix.m[i,j], reference[i,j])\n",
    "            if y in x.keys():\n",
    "                x[y].add(sample.outMatrix.m[i,j])\n",
    "            else:\n",
    "                x[y] = set([sample.outMatrix.m[i,j]])\n",
    "    x = {k:next(iter(v)) for k,v in x.items() if len(v)==1}\n",
    "    x = {k:v for k,v in x.items() if v not in t.fixedColors}\n",
    "    if len(x)>0:\n",
    "        return partial(doPixelMod2ColReverse, rules=x)\n",
    "    # Mod 3\n",
    "    x = {}\n",
    "    for sample in t.trainSamples:\n",
    "        reference = np.zeros(sample.inMatrix.shape, dtype=np.uint8)\n",
    "        onesCol = np.ones(sample.inMatrix.shape[0], dtype=np.uint8)\n",
    "        twosCol = np.full(sample.inMatrix.shape[0], 2, dtype=np.uint8)\n",
    "        for j in range(sample.inMatrix.shape[1]):\n",
    "            if j%3 == 0:\n",
    "                reference[:,sample.inMatrix.shape[1]-j-1] = onesCol.copy()\n",
    "            elif j%3 == 1:\n",
    "                reference[:,sample.inMatrix.shape[1]-j-1] = twosCol.copy()\n",
    "        for i,j in np.ndindex(reference.shape):\n",
    "            y = (sample.inMatrix.m[i,j], reference[i,j])\n",
    "            if y in x.keys():\n",
    "                x[y].add(sample.outMatrix.m[i,j])\n",
    "            else:\n",
    "                x[y] = set([sample.outMatrix.m[i,j]])\n",
    "    x = {k:next(iter(v)) for k,v in x.items() if len(v)==1}\n",
    "    x = {k:v for k,v in x.items() if v not in t.fixedColors}\n",
    "    if len(x)>0:\n",
    "        return partial(doPixelMod3ColReverse, rules=x)\n",
    "    \n",
    "    # Alternate\n",
    "    # Mod2\n",
    "    x = {}\n",
    "    for sample in t.trainSamples:\n",
    "        reference = np.zeros(sample.inMatrix.shape, dtype=np.uint8)\n",
    "        for i,j in np.ndindex(sample.inMatrix.shape):\n",
    "            reference[i,j] = (i+j)%2\n",
    "        for i,j in np.ndindex(reference.shape):\n",
    "            y = (sample.inMatrix.m[i,j], reference[i,j])\n",
    "            if y in x.keys():\n",
    "                x[y].add(sample.outMatrix.m[i,j])\n",
    "            else:\n",
    "                x[y] = set([sample.outMatrix.m[i,j]])\n",
    "    x = {k:next(iter(v)) for k,v in x.items() if len(v)==1}\n",
    "    x = {k:v for k,v in x.items() if v not in t.fixedColors}\n",
    "    if len(x)>0:\n",
    "        return partial(doPixelMod2Alternate, rules=x)\n",
    "    # Mod3\n",
    "    x = {}\n",
    "    for sample in t.trainSamples:\n",
    "        reference = np.zeros(sample.inMatrix.shape, dtype=np.uint8)\n",
    "        for i,j in np.ndindex(sample.inMatrix.shape):\n",
    "            reference[i,j] = (i+j)%3\n",
    "        for i,j in np.ndindex(reference.shape):\n",
    "            y = (sample.inMatrix.m[i,j], reference[i,j])\n",
    "            if y in x.keys():\n",
    "                x[y].add(sample.outMatrix.m[i,j])\n",
    "            else:\n",
    "                x[y] = set([sample.outMatrix.m[i,j]])\n",
    "    x = {k:next(iter(v)) for k,v in x.items() if len(v)==1}\n",
    "    x = {k:v for k,v in x.items() if v not in t.fixedColors}\n",
    "    if len(x)>0:\n",
    "        return partial(doPixelMod3Alternate, rules=x)\n",
    "    \n",
    "    return 0    \n",
    "\n",
    "# %% Surround Shape\n",
    "\n",
    "def surroundShape(matrix, shape, color, fixedColors, nSteps = None, forceFull=False, \\\n",
    "                  stepIsShape=False, stepIsNHoles=False):\n",
    "\n",
    "    \"\"\"\n",
    "    Given a matrix (numpy.ndarray), a Shape and a color, this function\n",
    "    surrounds the given Shape with the given color.\n",
    "    \"\"\"\n",
    "    \n",
    "    m = matrix.copy()\n",
    "    shapeMatrix = shape.m.copy()\n",
    "    \n",
    "    if nSteps==None:\n",
    "        if stepIsShape:\n",
    "            nSteps = int(shape.shape[0]/2)\n",
    "        elif stepIsNHoles:\n",
    "            nSteps = shape.nHoles\n",
    "        else:\n",
    "            nSteps = 15\n",
    "    \n",
    "    step = 0\n",
    "    \n",
    "    while step<nSteps:\n",
    "        step += 1\n",
    "        if forceFull:\n",
    "            if shape.position[0]-step<0 or shape.position[0]+shape.shape[0]+step>matrix.shape[0] or\\\n",
    "            shape.position[1]-step<0 or shape.position[1]+shape.shape[1]+step>matrix.shape[1]:\n",
    "                step -= 1\n",
    "                break\n",
    "            \n",
    "            done = False\n",
    "            for i in range(shape.position[0]-step, shape.position[0]+shape.shape[0]+step):\n",
    "                if matrix[i, shape.position[1]-step] in fixedColors:\n",
    "                    step -= 1\n",
    "                    done = True\n",
    "                    break\n",
    "                if matrix[i, shape.position[1]+shape.shape[1]+step-1] in fixedColors:\n",
    "                    step -= 1\n",
    "                    done = True\n",
    "                    break\n",
    "            if done:\n",
    "                break\n",
    "            for j in range(shape.position[1]-step, shape.position[1]+shape.shape[1]+step):\n",
    "                if matrix[shape.position[0]-step, j] in fixedColors:\n",
    "                    step -= 1\n",
    "                    done = True\n",
    "                    break\n",
    "                if matrix[shape.position[0]+shape.shape[0]+step-1, j] in fixedColors:\n",
    "                    step -= 1\n",
    "                    done = True\n",
    "                    break\n",
    "            if done:\n",
    "                break\n",
    "        \n",
    "        row = np.full(shapeMatrix.shape[1], -1, dtype=np.uint8)\n",
    "        col = np.full(shapeMatrix.shape[0]+2, -1, dtype=np.uint8)\n",
    "        newM = shapeMatrix.copy() \n",
    "        newM = np.vstack([row,newM,row])\n",
    "        newM = np.column_stack([col,newM,col])\n",
    "        \n",
    "        for i in range(newM.shape[0]):\n",
    "            for j in range(newM.shape[1]):\n",
    "                if newM[i,j] != 255:\n",
    "                    newM[i, j-1] = color\n",
    "                    break\n",
    "            for j in reversed(range(newM.shape[1])):\n",
    "                if newM[i,j] != 255:\n",
    "                    newM[i, j+1] = color\n",
    "                    break\n",
    "                    \n",
    "        for j in range(newM.shape[1]):\n",
    "            for i in range(newM.shape[0]):\n",
    "                if newM[i,j] != 255:\n",
    "                    newM[i-1, j] = color\n",
    "                    break\n",
    "            for i in reversed(range(newM.shape[0])):\n",
    "                if newM[i,j] != 255:\n",
    "                    newM[i+1, j] = color\n",
    "                    break\n",
    "                    \n",
    "        shapeMatrix = newM.copy()\n",
    "                    \n",
    "    for i,j in np.ndindex(shapeMatrix.shape):\n",
    "        if shape.position[0]-step+i<0 or shape.position[0]-step+i>=matrix.shape[0] or \\\n",
    "        shape.position[1]-step+j<0 or shape.position[1]-step+j>=matrix.shape[1]:\n",
    "            continue\n",
    "        if shapeMatrix[i,j] != 255:\n",
    "            m[shape.position[0]-step+i, shape.position[1]-step+j] = shapeMatrix[i,j]\n",
    "        \n",
    "    return m\n",
    "    \n",
    "def surroundAllShapes(matrix, shapeColor, surroundColor, fixedColors, nSteps=None,\\\n",
    "                      forceFull=False, stepIsShape=False, stepIsNHoles=False):\n",
    "    \"\"\"\n",
    "    Given a Matrix, a shapeColor and a surroundColor, this function surrounds\n",
    "    all the Shapes of color shapeColor with the given surroundColor.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    shapesToSurround = [s for s in matrix.shapes if s.color == shapeColor]\n",
    "    if stepIsShape:\n",
    "        shapesToSurround = [s for s in shapesToSurround if s.isSquare]\n",
    "    for s in shapesToSurround:\n",
    "        m = surroundShape(m, s, surroundColor, fixedColors, nSteps=nSteps,\\\n",
    "                          forceFull=forceFull, stepIsShape=stepIsShape,\\\n",
    "                          stepIsNHoles=stepIsNHoles)\n",
    "    return m\n",
    "\n",
    "def getBestSurroundShapes(t):    \n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"surroundShapes\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    \n",
    "    for fc in t.fixedColors:\n",
    "        for coc in t.commonChangedOutColors:\n",
    "            f = partial(surroundAllShapes, shapeColor=fc, surroundColor=coc, \\\n",
    "                        fixedColors=t.fixedColors, forceFull=True)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            \n",
    "            f = partial(surroundAllShapes, shapeColor=fc, surroundColor=coc, \\\n",
    "                            fixedColors=t.fixedColors, stepIsShape=True)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            \n",
    "            f = partial(surroundAllShapes, shapeColor=fc, surroundColor=coc, \\\n",
    "                            fixedColors=t.fixedColors, forceFull=True, stepIsShape=True)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            \n",
    "            f = partial(surroundAllShapes, shapeColor=fc, surroundColor=coc, \\\n",
    "                            fixedColors=t.fixedColors, forceFull=False, stepIsNHoles=True)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            \n",
    "            f = partial(surroundAllShapes, shapeColor=fc, surroundColor=coc, \\\n",
    "                            fixedColors=t.fixedColors, forceFull=True, stepIsNHoles=True)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            \n",
    "            for nSteps in range(1,4):\n",
    "                f = partial(surroundAllShapes, shapeColor=fc, surroundColor=coc, \\\n",
    "                            fixedColors=t.fixedColors, nSteps=nSteps)\n",
    "                bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                if bestScore==0:\n",
    "                    return bestFunction\n",
    "                \n",
    "                f = partial(surroundAllShapes, shapeColor=fc, surroundColor=coc, \\\n",
    "                            fixedColors=t.fixedColors, nSteps=nSteps, forceFull=True)\n",
    "                bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                if bestScore==0:\n",
    "                    return bestFunction\n",
    "            \n",
    "    return bestFunction\n",
    "\n",
    "# %% Extend Color\n",
    "\n",
    "def extendColor(matrix, direction, cic, fixedColors, color=None, sourceColor=None,\\\n",
    "                deleteExtensionColors=set(), deleteIfBorder=False, \\\n",
    "                breakAtFixedColor=False, mergeColor=None):\n",
    "    \"\"\"\n",
    "    Given a Matrix matrix and a direction, this function extends all the pixels\n",
    "    of the given color in the given direction direction can be any of:\n",
    "        'hv', 'h', 'v', 'u', 'd', 'l', 'r', 'all', 'diag', 'd1', 'd2'\n",
    "    If no color is specified, then all the colors but the changedInColors (cic)\n",
    "    and the fixedColors are extended.\n",
    "    \"\"\"\n",
    "    \n",
    "    if sourceColor==None:\n",
    "        sourceColor=color\n",
    "        \n",
    "    matrices = []\n",
    "    \n",
    "    # Vertical\n",
    "    if direction=='all' or direction=='hv' or direction=='v' or direction=='u':\n",
    "        for j in range(matrix.shape[1]):\n",
    "            colorCells=False\n",
    "            start = matrix.shape[0]-1\n",
    "            for i in reversed(range(matrix.shape[0])):\n",
    "                if color==None and sourceColor==None:\n",
    "                    if matrix.m[i,j] not in (fixedColors|cic):\n",
    "                        sourceColor = matrix.m[i,j]\n",
    "                if matrix.m[i,j]==sourceColor:\n",
    "                    m = matrix.m.copy()\n",
    "                    colorCells=True\n",
    "                    start = i\n",
    "                if colorCells:\n",
    "                    if matrix.m[i,j] in cic:\n",
    "                        if color==None:\n",
    "                            m[i,j] = sourceColor\n",
    "                        else:\n",
    "                            m[i,j] = color\n",
    "                    elif m[i,j]!=sourceColor and breakAtFixedColor==\"any\":\n",
    "                        sourceColor = m[i,j]\n",
    "                    elif matrix.m[i,j] in fixedColors and breakAtFixedColor:\n",
    "                        break\n",
    "                if colorCells and ((matrix.m[i,j] in deleteExtensionColors) or \\\n",
    "                                   i==0 and deleteIfBorder):\n",
    "                    currentI = i\n",
    "                    for i in range(currentI, start):\n",
    "                        m[i,j] = matrix.m[i,j]\n",
    "                    break\n",
    "            if color==None:\n",
    "                sourceColor=None\n",
    "            if colorCells:\n",
    "                matrices.append(m)\n",
    "    if direction=='all' or direction=='hv' or direction=='v' or direction=='d':\n",
    "        for j in range(matrix.shape[1]):\n",
    "            colorCells=False\n",
    "            start = 0\n",
    "            for i in range(matrix.shape[0]):\n",
    "                if color==None and sourceColor==None:\n",
    "                    if matrix.m[i,j] not in (fixedColors|cic):\n",
    "                        sourceColor = matrix.m[i,j]\n",
    "                if matrix.m[i,j]==sourceColor:\n",
    "                    m = matrix.m.copy()\n",
    "                    colorCells=True\n",
    "                    start = i\n",
    "                if colorCells and matrix.m[i,j] in fixedColors and breakAtFixedColor:\n",
    "                    break\n",
    "                if colorCells:\n",
    "                    if matrix.m[i,j] in cic:\n",
    "                        if color==None:\n",
    "                            m[i,j] = sourceColor\n",
    "                        else:\n",
    "                            m[i,j] = color\n",
    "                    elif m[i,j]!=sourceColor and breakAtFixedColor==\"any\":\n",
    "                        sourceColor = m[i,j]\n",
    "                    elif matrix.m[i,j] in fixedColors and breakAtFixedColor:\n",
    "                        break\n",
    "                        \n",
    "                if colorCells and ((matrix.m[i,j] in deleteExtensionColors) or \\\n",
    "                                   i==m.shape[0]-1 and deleteIfBorder):\n",
    "                    currentI = i+1\n",
    "                    for i in reversed(range(start, currentI)):\n",
    "                        m[i,j] = matrix.m[i,j]\n",
    "                    break\n",
    "            if color==None:\n",
    "                sourceColor=None\n",
    "            if colorCells:\n",
    "                matrices.append(m)\n",
    "             \n",
    "    # Horizontal\n",
    "    if direction=='all' or direction=='hv' or direction=='h' or direction=='l':\n",
    "        for i in range(matrix.shape[0]):\n",
    "            colorCells=False\n",
    "            start = matrix.shape[1]-1\n",
    "            for j in reversed(range(matrix.shape[1])):  \n",
    "                if color==None and sourceColor==None:\n",
    "                    if matrix.m[i,j] not in (fixedColors|cic):\n",
    "                        sourceColor = matrix.m[i,j]\n",
    "                if matrix.m[i,j]==sourceColor:\n",
    "                    m = matrix.m.copy()\n",
    "                    colorCells=True\n",
    "                    start = j\n",
    "                if colorCells:\n",
    "                    if matrix.m[i,j] in cic:\n",
    "                        if color==None:\n",
    "                            m[i,j] = sourceColor\n",
    "                        else:\n",
    "                            m[i,j] = color\n",
    "                    elif m[i,j]!=sourceColor and breakAtFixedColor==\"any\":\n",
    "                        sourceColor = m[i,j]\n",
    "                    elif matrix.m[i,j] in fixedColors and breakAtFixedColor:\n",
    "                        break\n",
    "                if colorCells and ((matrix.m[i,j] in deleteExtensionColors) or \\\n",
    "                                   j==0 and deleteIfBorder):\n",
    "                    currentJ = j\n",
    "                    for j in range(currentJ, start):\n",
    "                        m[i,j] = matrix.m[i,j]\n",
    "                    break\n",
    "            if color==None:\n",
    "                sourceColor=None\n",
    "            if colorCells:\n",
    "                matrices.append(m)\n",
    "    if direction=='all' or direction=='hv' or direction=='h' or direction=='r':\n",
    "        for i in range(matrix.shape[0]):\n",
    "            colorCells=False\n",
    "            start = 0\n",
    "            for j in range(matrix.shape[1]):\n",
    "                if color==None and sourceColor==None:\n",
    "                    if matrix.m[i,j] not in (fixedColors|cic):\n",
    "                        sourceColor = matrix.m[i,j]\n",
    "                if matrix.m[i,j]==sourceColor:\n",
    "                    m = matrix.m.copy()\n",
    "                    colorCells=True\n",
    "                    start = j\n",
    "                if colorCells: \n",
    "                    if matrix.m[i,j] in cic:\n",
    "                        if color==None:\n",
    "                            m[i,j] = sourceColor\n",
    "                        else:\n",
    "                            m[i,j] = color\n",
    "                    elif m[i,j]!=sourceColor and breakAtFixedColor==\"any\":\n",
    "                        sourceColor = m[i,j]\n",
    "                    elif matrix.m[i,j] in fixedColors and breakAtFixedColor:\n",
    "                        break\n",
    "                if colorCells and ((matrix.m[i,j] in deleteExtensionColors) or \\\n",
    "                                   j==m.shape[1]-1 and deleteIfBorder):\n",
    "                    currentJ = j+1\n",
    "                    for j in reversed(range(start, currentJ)):\n",
    "                        m[i,j] = matrix.m[i,j]\n",
    "                    break\n",
    "            if color==None:\n",
    "                sourceColor=None\n",
    "            if colorCells:\n",
    "                matrices.append(m)\n",
    "              \n",
    "    if len(matrices)>0:\n",
    "        m = mergeMatrices(matrices, next(iter(cic)), mergeColor=mergeColor)\n",
    "    else:\n",
    "        m = matrix.m.copy()\n",
    "                \n",
    "    # Diagonal\n",
    "    if direction=='all' or  direction=='diag' or direction=='d1' or direction=='d2':\n",
    "        if direction=='diag' or direction=='all':\n",
    "            directions = ['d1', 'd2']\n",
    "        else:\n",
    "            directions = [direction]\n",
    "        for direction in directions:\n",
    "            for transpose in [True, False]:\n",
    "                if transpose and direction=='d1':\n",
    "                    matrix.m = np.rot90(matrix.m, 2).T\n",
    "                    m = np.rot90(m, 2).T\n",
    "                if transpose and direction=='d2':\n",
    "                    matrix.m = matrix.m.T\n",
    "                    m = m.T\n",
    "                if direction=='d2':\n",
    "                    matrix.m = np.fliplr(matrix.m)\n",
    "                    m = np.fliplr(m)\n",
    "                for i in range(-matrix.shape[0]+1, matrix.shape[1]):\n",
    "                    diag = np.diagonal(matrix.m, i)\n",
    "                    colorCells=False\n",
    "                    for j in range(len(diag)):\n",
    "                        if color==None and sourceColor==None:\n",
    "                            if i<=0:\n",
    "                                if matrix.m[-i+j,j] not in (fixedColors|cic):\n",
    "                                    sourceColor = matrix.m[-i+j,j]\n",
    "                            else:\n",
    "                                if matrix.m[j,i+j] not in (fixedColors|cic):\n",
    "                                    sourceColor = matrix.m[j,i+j]\n",
    "                        if i<=0:\n",
    "                            if matrix.m[-i+j,j]==sourceColor:\n",
    "                                colorCells=True\n",
    "                            if colorCells:\n",
    "                                if matrix.m[-i+j,j] in cic:\n",
    "                                    if color==None:\n",
    "                                        m[-i+j,j] = sourceColor\n",
    "                                    else:\n",
    "                                        m[-i+j,j] = color\n",
    "                                elif matrix.m[-i+j,j]!=sourceColor and breakAtFixedColor==\"any\":\n",
    "                                    sourceColor = m[-i+j,j]\n",
    "                                elif matrix.m[-i+j,j] in fixedColors and breakAtFixedColor:\n",
    "                                    break\n",
    "                            if colorCells and ((matrix.m[-i+j,j] in deleteExtensionColors) or \\\n",
    "                                               j==len(diag)-1 and deleteIfBorder):\n",
    "                                for j in range(len(diag)):\n",
    "                                    m[-i+j,j] = matrix.m[-i+j,j]\n",
    "                                break\n",
    "                        else:\n",
    "                            if matrix.m[j,i+j]==sourceColor:\n",
    "                                colorCells=True\n",
    "                            if colorCells:\n",
    "                                if matrix.m[j,i+j] in cic:\n",
    "                                    if color==None:\n",
    "                                        m[j,i+j] = sourceColor\n",
    "                                    else:\n",
    "                                        m[j,i+j] = color\n",
    "                                elif matrix.m[j,i+j]!=sourceColor and breakAtFixedColor==\"any\":\n",
    "                                    sourceColor = m[j,i+j]\n",
    "                                elif matrix.m[j,i+j] in fixedColors and breakAtFixedColor:\n",
    "                                    break\n",
    "                            if colorCells and ((matrix.m[j,i+j] in deleteExtensionColors) or \\\n",
    "                                               j==len(diag)-1 and deleteIfBorder):\n",
    "                                for j in range(len(diag)):\n",
    "                                    m[j,i+j] = matrix.m[j,i+j]\n",
    "                                break\n",
    "                    if color==None:\n",
    "                        sourceColor=None\n",
    "                if direction=='d2':\n",
    "                    matrix.m = np.fliplr(matrix.m)\n",
    "                    m = np.fliplr(m)\n",
    "                if transpose and direction=='d2':\n",
    "                    matrix.m = matrix.m.T\n",
    "                    m = m.T\n",
    "                if transpose and direction=='d1':\n",
    "                    matrix.m = np.rot90(matrix.m, 2).T\n",
    "                    m = np.rot90(m, 2).T\n",
    "\n",
    "    return m    \n",
    "\n",
    "def getBestExtendColor(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"extendColor\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    \n",
    "    mergeColors = t.commonOutColors-t.totalInColors\n",
    "    if len(mergeColors) == 1:\n",
    "        mergeColor = next(iter(mergeColors))\n",
    "    else:\n",
    "        mergeColor = None\n",
    "    \n",
    "    cic = t.commonChangedInColors\n",
    "    if len(cic)==0:\n",
    "        return bestFunction\n",
    "    fixedColors = t.fixedColors\n",
    "    for d in ['r', 'l', 'h', 'u', 'd', 'v', 'hv', 'd1', 'd2', 'diag', 'all']:\n",
    "        for dib,bafc in product([True, False], [True, False, \"any\"]):\n",
    "            f = partial(extendColor, direction=d, cic=cic, fixedColors=fixedColors,\\\n",
    "                        deleteIfBorder=dib, breakAtFixedColor=bafc, mergeColor=mergeColor)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            f = partial(extendColor, direction=d, cic=cic, fixedColors=fixedColors,\\\n",
    "                        deleteExtensionColors=fixedColors,\\\n",
    "                        deleteIfBorder=dib, breakAtFixedColor=bafc, mergeColor=mergeColor)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            for coc in t.commonChangedOutColors:    \n",
    "                f = partial(extendColor, color=coc, direction=d, cic=cic, fixedColors=fixedColors,\\\n",
    "                            deleteIfBorder=dib, breakAtFixedColor=bafc, mergeColor=mergeColor)\n",
    "                bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                if bestScore==0:\n",
    "                    return bestFunction\n",
    "                f = partial(extendColor, color=coc, direction=d, cic=cic, fixedColors=fixedColors,\\\n",
    "                            deleteExtensionColors=fixedColors,\\\n",
    "                            deleteIfBorder=dib, breakAtFixedColor=bafc, mergeColor=mergeColor)\n",
    "                bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                if bestScore==0:\n",
    "                    return bestFunction\n",
    "                for fc in t.fixedColors:\n",
    "                    f = partial(extendColor, color=coc, direction=d, cic=cic, sourceColor=fc, \\\n",
    "                                fixedColors=fixedColors,\\\n",
    "                                deleteIfBorder=dib, breakAtFixedColor=bafc, mergeColor=mergeColor)\n",
    "                    bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                    if bestScore==0:\n",
    "                        return bestFunction\n",
    "                    f = partial(extendColor, color=coc, direction=d, cic=cic, sourceColor=fc, \\\n",
    "                                fixedColors=fixedColors, deleteExtensionColors=fixedColors,\\\n",
    "                                deleteIfBorder=dib, breakAtFixedColor=bafc, mergeColor=mergeColor)\n",
    "                    bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                    if bestScore==0:\n",
    "                        return bestFunction\n",
    "                \n",
    "    return bestFunction\n",
    "\n",
    "# %% Fill rectangleInside\n",
    "def fillRectangleInside(matrix, rectangleColor, fillColor):\n",
    "    m = matrix.m.copy()\n",
    "    for shape in matrix.shapes:\n",
    "        if shape.isRectangle and shape.color==rectangleColor:\n",
    "            if shape.shape[0] > 2 and shape.shape[1] > 2:\n",
    "                rect = np.full((shape.shape[0]-2, shape.shape[1]-2), fillColor, dtype=np.uint8)\n",
    "                m[shape.position[0]+1:shape.position[0]+shape.shape[0]-1,\\\n",
    "                  shape.position[1]+1:shape.position[1]+shape.shape[1]-1] = rect\n",
    "    return m\n",
    "\n",
    "# %% Color longest line\n",
    "def colorLongestLines(matrix, cic, coc, direction):\n",
    "    \"\"\"\n",
    "    cic stands for \"changedInColor\"\n",
    "    coc stands for \"changedOutColor\"\n",
    "    direction can be one of 4 strings: 'v', 'h', 'hv', 'd' (vertical,\n",
    "    horizontal, diagonal)\n",
    "    It is assumed t.sameIOShapes\n",
    "    \"\"\"    \n",
    "    m = matrix.m.copy()\n",
    "    \n",
    "    longest=0\n",
    "    positions = set()\n",
    "    if direction=='h':\n",
    "        for i in range(m.shape[0]):\n",
    "            count = 0\n",
    "            for j in range(m.shape[1]):\n",
    "                if m[i,j]==cic:\n",
    "                    if count!=0:\n",
    "                        count += 1\n",
    "                    else:\n",
    "                        count = 1\n",
    "                else:\n",
    "                    if count >= longest:\n",
    "                        if count > longest:\n",
    "                            positions = set()\n",
    "                        longest = count\n",
    "                        positions.add((i,j))\n",
    "                    count = 0\n",
    "            if count >= longest:\n",
    "                if count > longest:\n",
    "                    positions = set()\n",
    "                longest = count\n",
    "                positions.add((i,m.shape[1]-1))\n",
    "        for pos in positions:\n",
    "            for j in range(pos[1]-longest, pos[1]):\n",
    "                m[pos[0],j] = coc\n",
    "        return m                \n",
    "        \n",
    "    elif direction=='v':\n",
    "        for j in range(m.shape[1]):\n",
    "            count = 0\n",
    "            for i in range(m.shape[0]):\n",
    "                if m[i,j]==cic:\n",
    "                    if count!=0:\n",
    "                        count += 1\n",
    "                    else:\n",
    "                        count = 1\n",
    "                else:\n",
    "                    if count >= longest:\n",
    "                        if count > longest:\n",
    "                            positions = set()\n",
    "                        longest = count\n",
    "                        positions.add((i,j))\n",
    "                    count = 0\n",
    "            if count >= longest:\n",
    "                if count > longest:\n",
    "                    positions = set()\n",
    "                longest = count\n",
    "                positions.add((m.shape[0]-1,j))\n",
    "        for pos in positions:\n",
    "            for i in range(pos[0]-longest, pos[0]):\n",
    "                m[i,pos[1]] = coc\n",
    "        return m \n",
    "                        \n",
    "    elif direction=='hv':\n",
    "        longestH = 0\n",
    "        longestV = 0\n",
    "        positionsH = set()\n",
    "        positionsV = set()\n",
    "        for i in range(m.shape[0]):\n",
    "            count = 0\n",
    "            for j in range(m.shape[1]):\n",
    "                if m[i,j]==cic:\n",
    "                    if count!=0:\n",
    "                        count += 1\n",
    "                    else:\n",
    "                        count = 1\n",
    "                else:\n",
    "                    if count >= longestH:\n",
    "                        if count > longestH:\n",
    "                            positionsH = set()\n",
    "                        longestH = count\n",
    "                        positionsH.add((i,j))\n",
    "                    count = 0\n",
    "            if count >= longestH:\n",
    "                if count > longestH:\n",
    "                    positionsH = set()\n",
    "                longestH = count\n",
    "                positionsH.add((i,m.shape[1]-1))\n",
    "        for j in range(m.shape[1]):\n",
    "            count = 0\n",
    "            for i in range(m.shape[0]):\n",
    "                if m[i,j]==cic:\n",
    "                    if count!=0:\n",
    "                        count += 1\n",
    "                    else:\n",
    "                        count = 1\n",
    "                else:\n",
    "                    if count >= longestV:\n",
    "                        if count > longestV:\n",
    "                            positionsV = set()\n",
    "                        longestV = count\n",
    "                        positionsV.add((i,j))\n",
    "                    count = 0\n",
    "            if count >= longestV:\n",
    "                if count > longestV:\n",
    "                    positionsV = set()\n",
    "                longestV = count\n",
    "                positionsV.add((m.shape[0]-1,j))\n",
    "        for pos in positionsH:\n",
    "            for j in range(pos[1]-longestH, pos[1]):\n",
    "                m[pos[0],j] = coc\n",
    "        for pos in positionsV:\n",
    "            for i in range(pos[0]-longestV, pos[0]):\n",
    "                m[i,pos[1]] = coc\n",
    "        return m\n",
    "    \n",
    "    elif direction=='d':\n",
    "        # Direction of main diagonal\n",
    "        for i in reversed(range(m.shape[0])):\n",
    "            count = 0\n",
    "            jLimit = min(m.shape[1], m.shape[0]-i)\n",
    "            for j in range(jLimit):\n",
    "                if m[i+j,j]==cic:\n",
    "                    if count!=0:\n",
    "                        count += 1\n",
    "                    else:\n",
    "                        count = 1\n",
    "                else:\n",
    "                    if count >= longest:\n",
    "                        if count > longest:\n",
    "                            positions = set()\n",
    "                        longest = count\n",
    "                        positions.add(((i+j-1,j-1), 'main'))\n",
    "                    count = 0\n",
    "            if count >= longest:\n",
    "                if count > longest:\n",
    "                    positions = set()\n",
    "                longest = count\n",
    "                positions.add(((i+jLimit-1,jLimit-1), 'main'))\n",
    "        for j in range(1, m.shape[1]):\n",
    "            count = 0\n",
    "            iLimit = min(m.shape[0], m.shape[1]-j)\n",
    "            for i in range(iLimit):\n",
    "                if m[i,j+i]==cic:\n",
    "                    if count!=0:\n",
    "                        count += 1\n",
    "                    else:\n",
    "                        count = 1\n",
    "                else:\n",
    "                    if count >= longest:\n",
    "                        if count > longest:\n",
    "                            positions = set()\n",
    "                        longest = count\n",
    "                        positions.add(((i-1,j+i-1), 'main'))\n",
    "                    count = 0\n",
    "            if count >= longest:\n",
    "                if count > longest:\n",
    "                    positions = set()\n",
    "                longest = count\n",
    "                positions.add(((iLimit-1,j+iLimit-1), 'main'))\n",
    "                \n",
    "        # Direction of counterdiagonal\n",
    "        for i in range(m.shape[0]):\n",
    "            count = 0\n",
    "            jLimit = min(m.shape[1], i+1)\n",
    "            for j in range(jLimit):\n",
    "                if m[i-j,j]==cic:\n",
    "                    if count!=0:\n",
    "                        count += 1\n",
    "                    else:\n",
    "                        count = 1\n",
    "                else:\n",
    "                    if count >= longest:\n",
    "                        if count > longest:\n",
    "                            positions = set()\n",
    "                        longest = count\n",
    "                        positions.add(((i-j+1, j-1), 'counter'))\n",
    "                    count = 0\n",
    "            if count >= longest:\n",
    "                if count > longest:\n",
    "                    positions = set()\n",
    "                longest = count\n",
    "                positions.add(((i-jLimit+1,jLimit-1), 'counter'))\n",
    "        for j in range(m.shape[1]):\n",
    "            count = 0\n",
    "            iLimit = min(m.shape[0], m.shape[1]-j)\n",
    "            for i in range(iLimit):\n",
    "                if m[m.shape[0]-i-1,j+i]==cic:\n",
    "                    if count!=0:\n",
    "                        count += 1\n",
    "                    else:\n",
    "                        count = 1\n",
    "                else:\n",
    "                    if count >= longest:\n",
    "                        if count > longest:\n",
    "                            positions = set()\n",
    "                        longest = count\n",
    "                        positions.add(((m.shape[0]-i,j+i-1), 'counter'))\n",
    "                    count = 0\n",
    "            if count >= longest:\n",
    "                if count > longest:\n",
    "                    positions = set()\n",
    "                longest = count\n",
    "                positions.add(((m.shape[0]-iLimit,j+iLimit-1), 'counter'))\n",
    "        \n",
    "        # Draw the lines\n",
    "        for pos in positions:\n",
    "            if pos[1]=='main':\n",
    "                for x in range(longest):\n",
    "                    m[pos[0][0]-x, pos[0][1]-x] = coc\n",
    "            else:\n",
    "                for x in range(longest):\n",
    "                    m[pos[0][0]+x, pos[0][1]-x] = coc\n",
    "        return m\n",
    "    return m\n",
    "    \n",
    "# %% Move shapes    \n",
    "\n",
    "def moveShape(matrix, shape, background, direction, until = -1, nSteps = 100, \\\n",
    "              keepOriginal=False):\n",
    "    \"\"\"\n",
    "    'direction' can be l, r, u, d, ul, ur, dl, dr\n",
    "    (left, right, up, down, horizontal, vertical, diagonal1, diagonal2)\n",
    "    'until' can be a color or -1, which will be interpreted as border\n",
    "    If 'until'==-2, then move until the shape encounters anything\n",
    "    \"\"\"\n",
    "    m = matrix.copy()\n",
    "    if not keepOriginal:\n",
    "        m = changeColorShapes(m, [shape], background)\n",
    "    s = copy.deepcopy(shape)\n",
    "    if nSteps==\"shapeX\":\n",
    "        nSteps = shape.shape[0]\n",
    "    if nSteps==\"shapeY\":\n",
    "        nSteps = shape.shape[1]\n",
    "    step = 0\n",
    "    while True and step != nSteps:\n",
    "        step += 1\n",
    "        for c in s.pixels:\n",
    "            pos = (s.position[0]+c[0], s.position[1]+c[1])\n",
    "            if direction == \"l\":\n",
    "                newPos = (pos[0], pos[1]-1)\n",
    "            if direction == \"r\":\n",
    "                newPos = (pos[0], pos[1]+1)\n",
    "            if direction == \"u\":\n",
    "                newPos = (pos[0]-1, pos[1])\n",
    "            if direction == \"d\":\n",
    "                newPos = (pos[0]+1, pos[1])\n",
    "            if direction == \"ul\":\n",
    "                newPos = (pos[0]-1, pos[1]-1)\n",
    "            if direction == \"ur\":\n",
    "                newPos = (pos[0]-1, pos[1]+1)\n",
    "            if direction == \"dl\":\n",
    "                newPos = (pos[0]+1, pos[1]-1)\n",
    "            if direction == \"dr\":\n",
    "                newPos = (pos[0]+1, pos[1]+1)\n",
    "                \n",
    "            if newPos[0] not in range(m.shape[0]) or \\\n",
    "            newPos[1] not in range(m.shape[1]):\n",
    "                if until != -1 and until != -2:\n",
    "                    return matrix.copy()\n",
    "                else:\n",
    "                    return insertShape(m, s)\n",
    "            if until == -2 and m[newPos] != background:\n",
    "                return insertShape(m, s)\n",
    "            if m[newPos] == until:\n",
    "                return insertShape(m, s)\n",
    "            \n",
    "        if direction == \"l\":\n",
    "            s.position = (s.position[0], s.position[1]-1)\n",
    "        if direction == \"r\":\n",
    "            s.position = (s.position[0], s.position[1]+1)\n",
    "        if direction == \"u\":\n",
    "            s.position = (s.position[0]-1, s.position[1])\n",
    "        if direction == \"d\":\n",
    "            s.position = (s.position[0]+1, s.position[1])\n",
    "        if direction == \"ul\":\n",
    "            s.position = (s.position[0]-1, s.position[1]-1)\n",
    "        if direction == \"ur\":\n",
    "            s.position = (s.position[0]-1, s.position[1]+1)\n",
    "        if direction == \"dl\":\n",
    "            s.position = (s.position[0]+1, s.position[1]-1)\n",
    "        if direction == \"dr\":\n",
    "            s.position = (s.position[0]+1, s.position[1]+1)\n",
    "      \n",
    "    return insertShape(m, s) \n",
    "    \n",
    "def moveAllShapes(matrix, background, direction, until, nSteps=100, color=None, \\\n",
    "                  keepOriginal=False):\n",
    "    \"\"\"\n",
    "    direction can be l, r, u, d, ul, ur, dl, dr, h, v, d1, d2, all, any\n",
    "    \"\"\"\n",
    "    if color==None or color==\"multiColor\":\n",
    "        shapesToMove = matrix.multicolorShapes\n",
    "    elif color==\"diagonalMultiColor\":\n",
    "        shapesToMove=matrix.multicolorDShapes\n",
    "    elif color==\"singleColor\":\n",
    "        shapesToMove = [s for s in matrix.shapes if s.color!=background]\n",
    "    elif color==\"diagonalSingleColor\":\n",
    "        shapesToMove = [s for s in matrix.dShapes if s.color!=background]\n",
    "    else:\n",
    "        shapesToMove = [s for s in matrix.shapes if s.color in color]\n",
    "    if direction == 'l':\n",
    "        shapesToMove.sort(key=lambda x: x.position[1])\n",
    "    if direction == 'r':\n",
    "        shapesToMove.sort(key=lambda x: x.position[1]+x.shape[1], reverse=True)\n",
    "    if direction == 'u':\n",
    "        shapesToMove.sort(key=lambda x: x.position[0])  \n",
    "    if direction == 'd':\n",
    "        shapesToMove.sort(key=lambda x: x.position[0]+x.shape[0], reverse=True)\n",
    "    m = matrix.m.copy()\n",
    "    if len(shapesToMove) > 15:\n",
    "        return m\n",
    "    for s in shapesToMove:\n",
    "        newMatrix = m.copy()\n",
    "        if direction == \"any\":\n",
    "            for d in ['l', 'r', 'u', 'd', 'ul', 'ur', 'dl', 'dr']:\n",
    "                newMatrix = moveShape(m, s, background, d, until, keepOriginal=keepOriginal)\n",
    "                if not np.all(newMatrix == m):\n",
    "                    return newMatrix\n",
    "                    break\n",
    "        else:\n",
    "            m = moveShape(m, s, background, direction, until, nSteps, keepOriginal=keepOriginal)\n",
    "    return m\n",
    "    \n",
    "def moveShapeToClosest(matrix, shape, background, until=None, diagonals=False, restore=True):\n",
    "    \"\"\"\n",
    "    Given a matrix (numpy.ndarray) and a Shape, this function moves the\n",
    "    given shape until the closest shape with the color given by \"until\".\n",
    "    \"\"\"\n",
    "    m = matrix.copy()\n",
    "    s = copy.deepcopy(shape)\n",
    "    m = deleteShape(m, shape, background)\n",
    "    if until==None:\n",
    "        if hasattr(shape, \"color\"):\n",
    "            until=shape.color\n",
    "        else:\n",
    "            return matrix\n",
    "    if until not in m:\n",
    "        return matrix\n",
    "    nSteps = 0\n",
    "    while True:\n",
    "        for c in s.pixels:\n",
    "            pixelPos = tuple(map(operator.add, c, s.position))\n",
    "            if nSteps <= pixelPos[0] and m[pixelPos[0]-nSteps, pixelPos[1]] == until:\n",
    "                while nSteps>=0 and m[pixelPos[0]-nSteps, pixelPos[1]]!=background:\n",
    "                    nSteps-=1\n",
    "                s.position = (s.position[0]-nSteps, s.position[1])\n",
    "                return insertShape(m, s)\n",
    "            if pixelPos[0]+nSteps < m.shape[0] and m[pixelPos[0]+nSteps, pixelPos[1]] == until:\n",
    "                while nSteps>=0 and m[pixelPos[0]+nSteps, pixelPos[1]]!=background:\n",
    "                    nSteps-=1\n",
    "                s.position = (s.position[0]+nSteps, s.position[1])\n",
    "                return insertShape(m, s)\n",
    "            if nSteps <= pixelPos[1] and m[pixelPos[0], pixelPos[1]-nSteps] == until:\n",
    "                while nSteps>=0 and m[pixelPos[0], pixelPos[1]-nSteps]!=background:\n",
    "                    nSteps-=1\n",
    "                s.position = (s.position[0], s.position[1]-nSteps)\n",
    "                return insertShape(m, s)\n",
    "            if pixelPos[1]+nSteps < m.shape[1] and m[pixelPos[0], pixelPos[1]+nSteps] == until:\n",
    "                while nSteps>=0 and m[pixelPos[0], pixelPos[1]+nSteps]!=background:\n",
    "                    nSteps-=1\n",
    "                s.position = (s.position[0], s.position[1]+nSteps)\n",
    "                return insertShape(m, s)\n",
    "            if diagonals:\n",
    "                if nSteps <= pixelPos[0] and nSteps <= pixelPos[1] and \\\n",
    "                m[pixelPos[0]-nSteps, pixelPos[1]-nSteps] == until:\n",
    "                    s.position = (s.position[0]-nSteps+1, s.position[1]-nSteps+1)\n",
    "                    return insertShape(m, s)\n",
    "                if nSteps <= pixelPos[0] and pixelPos[1]+nSteps < m.shape[1] and \\\n",
    "                m[pixelPos[0]-nSteps, pixelPos[1]+nSteps] == until:\n",
    "                    s.position = (s.position[0]-nSteps+1, s.position[1]+nSteps-1)\n",
    "                    return insertShape(m, s)\n",
    "                if pixelPos[0]+nSteps < m.shape[0] and nSteps <= pixelPos[1] and \\\n",
    "                m[pixelPos[0]+nSteps, pixelPos[1]-nSteps] == until:\n",
    "                    s.position = (s.position[0]+nSteps-1, s.position[1]-nSteps+1)\n",
    "                    return insertShape(m, s)\n",
    "                if pixelPos[0]+nSteps < m.shape[0] and pixelPos[1]+nSteps < m.shape[1] and \\\n",
    "                m[pixelPos[0]+nSteps, pixelPos[1]+nSteps] == until:\n",
    "                    s.position = (s.position[0]+nSteps-1, s.position[1]+nSteps-1)\n",
    "                    return insertShape(m, s)\n",
    "        nSteps += 1\n",
    "        if nSteps > m.shape[0] and nSteps > m.shape[1]:\n",
    "            if restore:\n",
    "                return matrix\n",
    "            else:\n",
    "                return m\n",
    "        \n",
    "def moveAllShapesToClosest(matrix, background, colorsToMove=None, until=None, \\\n",
    "                           diagonals=False, restore=True, fixedShapeFeatures=None):\n",
    "    \"\"\"\n",
    "    This function moves all the shapes with color \"colorsToMove\" until the\n",
    "    closest shape with color \"until\".\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    if len(matrix.shapes) > 25:\n",
    "        return m\n",
    "    fixedShapes = []\n",
    "    if until == None:\n",
    "        colorsToMove = []\n",
    "        for shape in matrix.shapes:\n",
    "            if hasFeatures(shape.boolFeatures, fixedShapeFeatures):\n",
    "                fixedShapes.append(shape)\n",
    "                colorsToMove.append(shape.color)\n",
    "    elif colorsToMove==None:\n",
    "        colorsToMove = matrix.colors - set([background, until])\n",
    "    else:\n",
    "        colorsToMove = [colorsToMove]\n",
    "    for ctm in colorsToMove:\n",
    "        for shape in matrix.shapes:\n",
    "            if shape not in fixedShapes:\n",
    "                if shape.color == ctm:\n",
    "                    m = moveShapeToClosest(m, shape, background, until, diagonals, restore)\n",
    "    return m\n",
    "\n",
    "def getBestMoveShapes(t, candidate):\n",
    "    \"\"\"\n",
    "    This functions tries to find, for a given task t, the best way to move\n",
    "    shapes.\n",
    "    \"\"\"\n",
    "    directions = ['l', 'r', 'u', 'd', 'ul', 'ur', 'dl', 'dr', 'any']\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    \n",
    "    doSingleColor = all([len(s.inMatrix.shapes) < 50 for s in t.trainSamples])\n",
    "    doDSingleColor = all([len(s.inMatrix.dShapes) < 15 for s in t.trainSamples])\n",
    "    doMulticolor = all([len(s.inMatrix.multicolorShapes) < 15 for s in t.trainSamples])\n",
    "        \n",
    "    # Move all shapes in a specific direction, until a non-background thing is touched\n",
    "    for d in directions:\n",
    "        if doSingleColor:\n",
    "            f = partial(moveAllShapes, background=t.backgroundColor, until=-2,\\\n",
    "                        direction=d, color=\"singleColor\")\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "        if doDSingleColor:\n",
    "            f = partial(moveAllShapes, background=t.backgroundColor, until=-2,\\\n",
    "                        direction=d, color=\"diagonalSingleColor\")\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "        if doMulticolor:\n",
    "            f = partial(moveAllShapes, background=t.backgroundColor, until=-2,\\\n",
    "                        direction=d, color=\"multiColor\")\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "        f = partial(moveAllShapes, background=t.backgroundColor, until=-2,\\\n",
    "                    direction=d, color=\"diagonalMultiColor\")\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        f = partial(moveAllShapes, background=t.backgroundColor, until=-2,\\\n",
    "                    direction=d, color=\"diagonalMultiColor\", nSteps=\"shapeX\")\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        f = partial(moveAllShapes, background=t.backgroundColor, until=-2,\\\n",
    "                    direction=d, color=\"diagonalMultiColor\", nSteps=\"shapeY\")\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        if doDSingleColor:\n",
    "            f = partial(moveAllShapes, background=t.backgroundColor, until=-2,\\\n",
    "                        direction=d, color=\"diagonalSingleColor\", nSteps=\"shapeX\")\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "            f = partial(moveAllShapes, background=t.backgroundColor, until=-2,\\\n",
    "                        direction=d, color=\"diagonalSingleColor\", nSteps=\"shapeY\")\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "        \n",
    "    if doSingleColor:\n",
    "        colorsToChange = list(t.colors - t.fixedColors - set({t.backgroundColor}))\n",
    "        ctc = [[c] for c in colorsToChange] + [colorsToChange] # Also all colors\n",
    "        for c in ctc:\n",
    "            for d in directions:\n",
    "                moveUntil = colorsToChange + [-1] + [-2] #Border, any\n",
    "                for u in moveUntil:\n",
    "                    f = partial(moveAllShapes, color=c, background=t.backgroundColor,\\\n",
    "                                    direction=d, until=u)\n",
    "                    bestFunction, bestScore, isPerfect = updateBestFunction(t, f, bestScore, bestFunction, \\\n",
    "                                                                            checkPerfect=True, prevScore=candidate.score)\n",
    "                    if isPerfect:\n",
    "                        return bestFunction\n",
    "                    if bestScore==0:\n",
    "                        return bestFunction\n",
    "                for nSteps in range(1, 5):\n",
    "                    f = partial(moveAllShapes, color=c, background=t.backgroundColor,\\\n",
    "                                direction=d, until=-2, nSteps=nSteps)\n",
    "                    bestFunction, bestScore, isPerfect = updateBestFunction(t, f, bestScore, bestFunction,\\\n",
    "                                                                            checkPerfect=True, prevScore=candidate.score)\n",
    "                    if isPerfect:\n",
    "                        return bestFunction\n",
    "                    if bestScore==0:\n",
    "                        return bestFunction\n",
    "                    f = partial(moveAllShapes, color=c, background=t.backgroundColor,\\\n",
    "                                direction=d, until=-2, nSteps=nSteps, keepOriginal=True)\n",
    "                    bestFunction, bestScore, isPerfect = updateBestFunction(t, f, bestScore, bestFunction,\\\n",
    "                                                                            checkPerfect=True, prevScore=candidate.score)\n",
    "                    if isPerfect:\n",
    "                        return bestFunction\n",
    "                    if bestScore==0:\n",
    "                        return bestFunction\n",
    "                \n",
    "    \n",
    "        if t.backgroundColor != -1 and hasattr(t, 'fixedColors'):\n",
    "            colorsToMove = t.almostCommonColors - set([t.backgroundColor]) - t.fixedColors\n",
    "            for ctm in colorsToMove:\n",
    "                for uc in t.unchangedColors:\n",
    "                    f = partial(moveAllShapesToClosest, colorsToMove=ctm,\\\n",
    "                                     background=t.backgroundColor, until=uc)\n",
    "                    bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                    if bestScore==0:\n",
    "                        return bestFunction\n",
    "                    \n",
    "                    f = partial(moveAllShapesToClosest, colorsToMove=ctm,\\\n",
    "                                background=t.backgroundColor, until=uc, diagonals=True)\n",
    "                    bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "                    if bestScore==0:\n",
    "                        return bestFunction\n",
    "                \n",
    "        if all([len(sample.fixedShapes)>0 for sample in t.trainSamples]):\n",
    "            f = partial(moveAllShapesToClosest, background=t.backgroundColor,\\\n",
    "                        fixedShapeFeatures = t.fixedShapeFeatures)\n",
    "            bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "            if bestScore==0:\n",
    "                return bestFunction\n",
    "        \n",
    "    return bestFunction\n",
    "\n",
    "# %% Complete rectangles\n",
    "def completeRectangles(matrix, sourceColor, newColor):\n",
    "    \"\"\"\n",
    "    It is assumed that the background is clear.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    for s in matrix.multicolorDShapes:\n",
    "        if hasattr(s, 'color') and s.color==sourceColor:\n",
    "            newShape = copy.deepcopy(s)\n",
    "            newShape.m[newShape.m==255] = newColor\n",
    "            m = insertShape(m, newShape)\n",
    "    return m\n",
    "\n",
    "# %% Delete shapes\n",
    "# Like that this only solves task 96. It's clearly not general enough.\n",
    "def deletePixels(matrix, diagonals=False):\n",
    "    \"\"\"\n",
    "    Given a matrix, this functions deletes all the pixels. This means that all\n",
    "    the dShapes consisting of only 1 pixel are converted to the color that\n",
    "    surrounds most of that pixel.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    if m.shape[0]==1 and m.shape[1]==1:\n",
    "        return m\n",
    "    \n",
    "    if diagonals:\n",
    "        shapes = matrix.dShapes\n",
    "    else:\n",
    "        shapes = matrix.shapes\n",
    "    for s in shapes:\n",
    "        if s.nPixels==1:\n",
    "            surrColors = Counter()\n",
    "            if s.position[0]>0:\n",
    "                surrColors[m[s.position[0]-1, s.position[1]]] += 1\n",
    "            if s.position[1]<m.shape[1]-1:\n",
    "                surrColors[m[s.position[0], s.position[1]+1]] += 1\n",
    "            if s.position[0]<m.shape[0]-1:\n",
    "                surrColors[m[s.position[0]+1, s.position[1]]] += 1\n",
    "            if s.position[1]>0:\n",
    "                surrColors[m[s.position[0], s.position[1]-1]] += 1\n",
    "            if len(set(surrColors.values()))==1:\n",
    "                if s.position[0]>0 and s.position[1]>0:\n",
    "                    surrColors[m[s.position[0]-1, s.position[1]-1]] += 1\n",
    "                if s.position[0]>0 and s.position[1]<m.shape[1]-1:\n",
    "                    surrColors[m[s.position[0]-1, s.position[1]+1]] += 1\n",
    "                if s.position[0]<m.shape[0]-1 and s.position[1]<m.shape[1]-1:\n",
    "                    surrColors[m[s.position[0]+1, s.position[1]+1]] += 1\n",
    "                if s.position[0]<m.shape[0]-1 and s.position[1]>0:\n",
    "                    surrColors[m[s.position[0]+1, s.position[1]-1]] += 1\n",
    "            \n",
    "            m[s.position[0],s.position[1]] = max(surrColors.items(), key=operator.itemgetter(1))[0]\n",
    "    return m\n",
    "\n",
    "# %% Connect Pixels\n",
    "\n",
    "def connectPixels(matrix, pixelColor=None, connColor=None, fixedColors=set(),\\\n",
    "                  allowedChanges={}, lineExclusive=False, diagonal=False):\n",
    "    \"\"\"\n",
    "    Given a matrix, this function connects all the pixels that have the same\n",
    "    color. This means that, for example, all the pixels between two red pixels\n",
    "    will also be red.\n",
    "    If \"pixelColor\" is specified, then only the pixels with the specified color\n",
    "    will be connected.\n",
    "    Ìf \"connColor\" is specified, the color used to connect the pixels will be\n",
    "    the one given by this parameter.\n",
    "    If there are any colors in the \"fixedColors\" set, then it is made sure that\n",
    "    they remain unchanged.\n",
    "    \"allowedChanges\" is a dictionary determining which color changes are\n",
    "    allowed. It's exclusive with the options unchangedColors and connColor.\n",
    "    \"\"\"\n",
    "    m = matrix.copy()\n",
    "    # Row\n",
    "    for i in range(m.shape[0]):\n",
    "        lowLimit = 0\n",
    "        while lowLimit < m.shape[1] and matrix[i, lowLimit] != pixelColor:\n",
    "            lowLimit += 1\n",
    "        lowLimit += 1\n",
    "        upLimit = m.shape[1]-1\n",
    "        while upLimit > lowLimit and matrix[i, upLimit] != pixelColor:\n",
    "            upLimit -= 1\n",
    "        if upLimit > lowLimit:\n",
    "            if lineExclusive:\n",
    "                for j in range(lowLimit, upLimit):\n",
    "                    if matrix[i,j] == pixelColor:\n",
    "                        lowLimit = upLimit\n",
    "                        break\n",
    "            for j in range(lowLimit, upLimit):\n",
    "                if connColor != None:\n",
    "                    if matrix[i,j] != pixelColor and matrix[i,j] not in fixedColors:\n",
    "                        m[i,j] = connColor\n",
    "                else:\n",
    "                    if matrix[i,j] in allowedChanges.keys():\n",
    "                        m[i,j] = allowedChanges[matrix[i,j]]\n",
    "       \n",
    "    # Column             \n",
    "    for j in range(m.shape[1]):\n",
    "        lowLimit = 0\n",
    "        while lowLimit < m.shape[0] and matrix[lowLimit, j] != pixelColor:\n",
    "            lowLimit += 1\n",
    "        lowLimit += 1\n",
    "        upLimit = m.shape[0]-1\n",
    "        while upLimit > lowLimit and matrix[upLimit, j] != pixelColor:\n",
    "            upLimit -= 1\n",
    "        if upLimit > lowLimit:\n",
    "            if lineExclusive:\n",
    "                for i in range(lowLimit, upLimit):\n",
    "                    if matrix[i,j] == pixelColor:\n",
    "                        lowLimit = upLimit\n",
    "                        break\n",
    "            for i in range(lowLimit, upLimit):\n",
    "                if connColor != None:\n",
    "                    if matrix[i,j] != pixelColor and matrix[i,j] not in fixedColors:\n",
    "                        m[i,j] = connColor\n",
    "                else:\n",
    "                    if matrix[i,j] in allowedChanges.keys():\n",
    "                        m[i,j] = allowedChanges[matrix[i,j]]\n",
    "    \n",
    "    # Diagonal         \n",
    "    if diagonal:\n",
    "        for d in [1, 2]:\n",
    "            if d==2:\n",
    "                matrix = np.fliplr(matrix)\n",
    "                m = np.fliplr(m)\n",
    "            for i in range(-matrix.shape[0]-1, matrix.shape[1]):\n",
    "                diag = np.diagonal(matrix, i)\n",
    "                lowLimit = 0\n",
    "                while lowLimit < len(diag) and diag[lowLimit] != pixelColor:\n",
    "                    lowLimit += 1\n",
    "                lowLimit += 1\n",
    "                upLimit = len(diag)-1\n",
    "                while upLimit > lowLimit and diag[upLimit] != pixelColor:\n",
    "                    upLimit -= 1\n",
    "                if upLimit > lowLimit:\n",
    "                    if lineExclusive:\n",
    "                        for j in range(lowLimit, upLimit):\n",
    "                            if i<=0:\n",
    "                                if matrix[-i+j,j] == pixelColor:\n",
    "                                    lowLimit = upLimit\n",
    "                                    break\n",
    "                            else:\n",
    "                                if matrix[j,i+j] == pixelColor:\n",
    "                                    lowLimit = upLimit\n",
    "                                    break\n",
    "                    for j in range(lowLimit, upLimit):\n",
    "                        if i<=0:\n",
    "                            if connColor != None:\n",
    "                                if matrix[-i+j,j] != pixelColor and matrix[-i+j,j] not in fixedColors:\n",
    "                                    m[-i+j,j] = connColor\n",
    "                            else:\n",
    "                                if matrix[-i+j,j] in allowedChanges.keys():\n",
    "                                    m[-i+j,j] = allowedChanges[matrix[-i+j,j]]\n",
    "                        else:\n",
    "                            if connColor != None:\n",
    "                                if matrix[j,i+j] != pixelColor and matrix[j,i+j] not in fixedColors:\n",
    "                                    m[j,i+j] = connColor\n",
    "                            else:\n",
    "                                if matrix[j,i+j] in allowedChanges.keys():\n",
    "                                    m[j,i+j] = allowedChanges[matrix[j,i+j]]\n",
    "            if d==2:\n",
    "                matrix = np.fliplr(matrix)\n",
    "                m = np.fliplr(m)\n",
    "\n",
    "    return m\n",
    "\n",
    "def connectAnyPixels(matrix, pixelColor=None, connColor=None, fixedColors=set(),\\\n",
    "                     allowedChanges={}, lineExclusive=False, diagonal=False):\n",
    "    \"\"\"\n",
    "    Given a Matrix, this function draws a line connecting two pixels of the\n",
    "    same color (different from the background color). The color of the line\n",
    "    is the same as the color of the pixels, unless specified by \"connColor\".\n",
    "    \"\"\"\n",
    "    \n",
    "    m = matrix.m.copy()\n",
    "    if pixelColor==None:\n",
    "        if connColor==None:\n",
    "            for c in matrix.colors - set([matrix.backgroundColor]):\n",
    "                m = connectPixels(m, c, c, lineExclusive=lineExclusive, diagonal=diagonal)\n",
    "            return m\n",
    "        else:\n",
    "            for c in matrix.colors - set([matrix.backgroundColor]):\n",
    "                m = connectPixels(m, c, connColor, lineExclusive=lineExclusive, diagonal=diagonal)\n",
    "            return m\n",
    "    else:\n",
    "        if len(allowedChanges)>0:\n",
    "            m = connectPixels(m, pixelColor, allowedChanges=allowedChanges,\\\n",
    "                              lineExclusive=lineExclusive, diagonal=diagonal)\n",
    "        else:\n",
    "            m = connectPixels(m, pixelColor, connColor, fixedColors, lineExclusive=lineExclusive,\\\n",
    "                              diagonal=diagonal)\n",
    "    return m\n",
    "\n",
    "def rotate(matrix, angle):\n",
    "    \"\"\"\n",
    "    Angle can be 90, 180, 270\n",
    "    \"\"\"\n",
    "    assert angle in [90, 180, 270], \"Invalid rotation angle\"\n",
    "    if isinstance(matrix, np.ndarray):\n",
    "        m = matrix.copy()\n",
    "    else:\n",
    "        m = matrix.m.copy()\n",
    "    return np.rot90(m, int(angle/90))    \n",
    "    \n",
    "def mirror(matrix, axis):\n",
    "    \"\"\"\n",
    "    Axis can be lr, up, d1, d2\n",
    "    \"\"\"\n",
    "    if isinstance(matrix, np.ndarray):\n",
    "        m = matrix.copy()\n",
    "    else:\n",
    "        m = matrix.m.copy()\n",
    "    assert axis in [\"lr\", \"ud\", \"d1\", \"d2\"], \"Invalid mirror axis\"\n",
    "    if axis == \"lr\":\n",
    "        return np.fliplr(m)\n",
    "    if axis == \"ud\":\n",
    "        return np.flipud(m)\n",
    "    if axis == \"d1\":\n",
    "        return m.T\n",
    "    if axis == \"d2\":\n",
    "        return m[::-1,::-1].T\n",
    "\n",
    "def flipShape(matrix, shape, axis, background):\n",
    "    # Axis can be lr, ud\n",
    "    m = matrix.copy()\n",
    "    smallM = np.ones((shape.shape[0], shape.shape[1]), dtype=np.uint8) * background\n",
    "    for c in shape.pixels:\n",
    "        smallM[c] = shape.color\n",
    "    if axis == \"lr\":\n",
    "        smallM = np.fliplr(smallM)\n",
    "    if axis == \"ud\":\n",
    "        smallM = np.flipud(smallM)\n",
    "    for i,j in np.ndindex(smallM.shape):\n",
    "        m[shape.position[0]+i, shape.position[1]+j] = smallM[i,j]\n",
    "    return m\n",
    "\n",
    "def flipAllShapes(matrix, axis, color, background, byColor=False, diagonal=False):#, multicolor=False):\n",
    "    m = matrix.m.copy()\n",
    "    if byColor:\n",
    "        shapesToMirror = [s for s in matrix.shapesByColor if s.color in color]\n",
    "    #elif multicolor:\n",
    "    #    if diagonal:\n",
    "    #        shapesToMirror = [s for s in matrix.multicolorDShapes]\n",
    "    #    else:\n",
    "    #        shapesToMirror = [s for s in matrix.multicolorShapes]\n",
    "    else:\n",
    "        if diagonal:\n",
    "            shapesToMirror = [s for s in matrix.dShapes if s.color in color]\n",
    "        else:\n",
    "            shapesToMirror = [s for s in matrix.shapes if s.color in color]\n",
    "    for s in shapesToMirror:\n",
    "        m = flipShape(m, s, axis, background)\n",
    "    return m\n",
    "\n",
    "def getBestFlipAllShapes(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"flipAllShapes\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestFunction = partial(identityM)\n",
    "    if t.backgroundColor == -1:\n",
    "        return bestFunction\n",
    "    bestScore = 1000\n",
    "    colors = set([0,1,2,3,4,5,6,7,8,9]) - set([t.backgroundColor])\n",
    "    for d in [\"lr\", \"ud\"]:\n",
    "        #for multicolor in [True, False]:\n",
    "        for diagonal in [True, False]:\n",
    "            bestFunction, bestScore = updateBestFunction(t, partial(flipAllShapes, axis=d, color=colors,\\\n",
    "                                    background=t.backgroundColor, diagonal=diagonal), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(flipAllShapes, axis=d, color=colors,\\\n",
    "                                background=t.backgroundColor, byColor=True), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def mapPixels(matrix, pixelMap, outShape):\n",
    "    \"\"\"\n",
    "    Given a Matrix as input, this function maps each pixel of that matrix\n",
    "    to an outputMatrix, given by outShape.\n",
    "    The dictionary pixelMap determines which pixel in the input matrix maps\n",
    "    to each pixel in the output matrix.\n",
    "    \"\"\"\n",
    "    inMatrix = matrix.m.copy()\n",
    "    m = np.zeros(outShape, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(outShape):\n",
    "        m[i,j] = inMatrix[pixelMap[i,j]]\n",
    "    return m\n",
    "\n",
    "def switchColors(matrix, color1=None, color2=None):\n",
    "    \"\"\"\n",
    "    This function switches the color1 and the color2 in the matrix.\n",
    "    If color1 and color2 are not specified, then the matrix is expected to only\n",
    "    have 2 colors, and they will be switched.\n",
    "    \"\"\"\n",
    "    if type(matrix) == np.ndarray:\n",
    "        m = matrix.copy()\n",
    "    else:\n",
    "        m = matrix.m.copy()\n",
    "    if color1==None or color2==None:\n",
    "        color1 = m[0,0]\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            if m[i,j]!=color1:\n",
    "                color2 = m[i,j]\n",
    "                break\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        if m[i,j]==color1:\n",
    "            m[i,j] = color2\n",
    "        else:\n",
    "            m[i,j] = color1        \n",
    "    return m\n",
    "\n",
    "# %% Follow row/col patterns\n",
    "def identifyColor(m, pixelPos, c2c, rowStep=None, colStep=None):\n",
    "    \"\"\"\n",
    "    Utility function for followPattern.\n",
    "    \"\"\"\n",
    "    if colStep!=None and rowStep!=None:\n",
    "        i = 0\n",
    "        while i+pixelPos[0] < m.shape[0]:\n",
    "            j = 0\n",
    "            while j+pixelPos[1] < m.shape[1]:\n",
    "                if m[pixelPos[0]+i, pixelPos[1]+j] != c2c:\n",
    "                    return m[pixelPos[0]+i, pixelPos[1]+j]\n",
    "                j += colStep\n",
    "            i += rowStep\n",
    "        return c2c\n",
    "    \n",
    "def identifyColStep(m, c2c):\n",
    "    \"\"\"\n",
    "    Utility function for followPattern.\n",
    "    \"\"\"\n",
    "    colStep = 1\n",
    "    while colStep < int(m.shape[1]/2)+1:\n",
    "        isGood = True\n",
    "        for j in range(colStep):\n",
    "            for i in range(m.shape[0]):\n",
    "                block = 0\n",
    "                colors = set()\n",
    "                while j+block < m.shape[1]:\n",
    "                    colors.add(m[i,j+block])\n",
    "                    block += colStep\n",
    "                if c2c in colors:\n",
    "                    if len(colors) > 2:\n",
    "                        isGood = False\n",
    "                        break\n",
    "                else:\n",
    "                    if len(colors) > 1:\n",
    "                        isGood = False\n",
    "                        break\n",
    "            if not isGood:\n",
    "                break  \n",
    "        if isGood:\n",
    "            return colStep \n",
    "        colStep+=1 \n",
    "    return m.shape[1]\n",
    "\n",
    "def identifyRowStep(m, c2c):\n",
    "    \"\"\"\n",
    "    Utility function for followPattern.\n",
    "    \"\"\"\n",
    "    rowStep = 1\n",
    "    while rowStep < int(m.shape[0]/2)+1:\n",
    "        isGood = True\n",
    "        for i in range(rowStep):\n",
    "            for j in range(m.shape[1]):\n",
    "                block = 0\n",
    "                colors = set()\n",
    "                while i+block < m.shape[0]:\n",
    "                    colors.add(m[i+block,j])\n",
    "                    block += rowStep\n",
    "                if c2c in colors:\n",
    "                    if len(colors) > 2:\n",
    "                        isGood = False\n",
    "                        break\n",
    "                else:\n",
    "                    if len(colors) > 1:\n",
    "                        isGood = False\n",
    "                        break\n",
    "            if not isGood:\n",
    "                break  \n",
    "        if isGood:\n",
    "            return rowStep \n",
    "        rowStep+=1  \n",
    "    return m.shape[0]            \n",
    "\n",
    "def followPattern(matrix, rc, colorToChange=None, rowStep=None, colStep=None):\n",
    "    \"\"\"\n",
    "    Given a Matrix, this function turns it into a matrix that follows a\n",
    "    pattern. This will be made row-wise, column-wise or both, depending on the\n",
    "    parameter \"rc\". \"rc\" can be \"row\", \"column\" or \"both\".\n",
    "    'colorToChange' is the number corresponding to the only color that changes,\n",
    "    if any.\n",
    "    'rowStep' and 'colStep' are only to be given if the rowStep/colStep is the\n",
    "    same for every train sample.\n",
    "    \"\"\"  \n",
    "    m = matrix.m.copy()\n",
    "            \n",
    "    if colorToChange!=None:\n",
    "        if rc==\"col\":\n",
    "            rowStep=m.shape[0]\n",
    "            if colStep==None:\n",
    "                colStep=identifyColStep(m, colorToChange)\n",
    "        if rc==\"row\":\n",
    "            colStep=m.shape[1]\n",
    "            if rowStep==None:\n",
    "                rowStep=identifyRowStep(m, colorToChange)\n",
    "        if rc==\"both\":\n",
    "            if colStep==None and rowStep==None:\n",
    "                colStep=identifyColStep(m, colorToChange)\n",
    "                rowStep=identifyRowStep(m, colorToChange) \n",
    "            elif rowStep==None:\n",
    "                rowStep=m.shape[0]\n",
    "            elif colStep==None:\n",
    "                colStep=m.shape[1]                       \n",
    "        for i,j in np.ndindex((rowStep, colStep)):\n",
    "            color = identifyColor(m, (i,j), colorToChange, rowStep, colStep)\n",
    "            k = 0\n",
    "            while i+k < m.shape[0]:\n",
    "                l = 0\n",
    "                while j+l < m.shape[1]:\n",
    "                    m[i+k, j+l] = color\n",
    "                    l += colStep\n",
    "                k += rowStep\n",
    "            \n",
    "    return m\n",
    "\n",
    "# %% Fill the blank\n",
    "def fillTheBlankParameters(t):\n",
    "    matrices = []\n",
    "    for s in t.trainSamples:\n",
    "        m = s.inMatrix.m.copy()\n",
    "        blank = s.blankToFill\n",
    "        m[blank.position[0]:blank.position[0]+blank.shape[0],\\\n",
    "          blank.position[1]:blank.position[1]+blank.shape[1]] = s.outMatrix.m.copy()\n",
    "        matrices.append(Matrix(m))\n",
    "        \n",
    "    x = []\n",
    "    x.append(all([m.lrSymmetric for m in matrices]))\n",
    "    x.append(all([m.udSymmetric for m in matrices]))\n",
    "    x.append(all([m.d1Symmetric for m in matrices]))\n",
    "    x.append(all([m.d2Symmetric for m in matrices]))\n",
    "    return x\n",
    "\n",
    "def fillTheBlank(matrix, params):\n",
    "    m = matrix.m.copy()\n",
    "    if len(matrix.blanks) == 0:\n",
    "        return m\n",
    "    blank = matrix.blanks[0]\n",
    "    color = blank.color\n",
    "    pred = np.zeros(blank.shape, dtype=np.uint8)\n",
    "    \n",
    "    # lr\n",
    "    if params[0]:\n",
    "        for i,j in np.ndindex(blank.shape):\n",
    "            if m[blank.position[0]+i, m.shape[1]-1-(blank.position[1]+j)] != color:\n",
    "                pred[i,j] = m[blank.position[0]+i, m.shape[1]-1-(blank.position[1]+j)]\n",
    "    # ud\n",
    "    if params[1]:\n",
    "        for i,j in np.ndindex(blank.shape):\n",
    "            if m[m.shape[0]-1-(blank.position[0]+i), blank.position[1]+j] != color:\n",
    "                pred[i,j] = m[m.shape[0]-1-(blank.position[0]+i), blank.position[1]+j]\n",
    "    # d1\n",
    "    if params[2] and m.shape[0]==m.shape[1]:\n",
    "        for i,j in np.ndindex(blank.shape):\n",
    "            if m[blank.position[1]+j, blank.position[0]+i] != color:\n",
    "                pred[i,j] = m[blank.position[1]+j, blank.position[0]+i]\n",
    "    # d2 (persymmetric matrix)\n",
    "    if params[3] and m.shape[0]==m.shape[1]:\n",
    "        for i,j in np.ndindex(blank.shape):\n",
    "            if m[m.shape[1]-1-(blank.position[1]+j), m.shape[0]-1-(blank.position[0]+i)] != color:\n",
    "                pred[i,j] = m[m.shape[1]-1-(blank.position[1]+j), m.shape[0]-1-(blank.position[0]+i)]\n",
    "    \n",
    "    return pred\n",
    "    \n",
    "# %% Operations with more than one matrix\n",
    "\n",
    "# All the matrices need to have the same shape\n",
    "\n",
    "def pixelwiseAnd(matrices, falseColor, targetColor=None, trueColor=None):\n",
    "    \"\"\"\n",
    "    This function returns the result of executing the pixelwise \"and\" operation\n",
    "    in a list of matrices.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    matrices: list\n",
    "        A list of numpy.ndarrays of the same shape\n",
    "    falseColor: int\n",
    "        The color of the pixel in the output matrix if the \"and\" operation is\n",
    "        false.\n",
    "    targetColor: int\n",
    "        The color to be targeted by the \"and\" operation. For example, if\n",
    "        targetColor is red, then the \"and\" operation will be true for a pixel\n",
    "        if all that pixel is red in all of the input matrices.\n",
    "        If targetColor is None, then the \"and\" operation will return true if\n",
    "        the pixel has the same color in all the matrices, and false otherwise.\n",
    "    trueColor: int\n",
    "        The color of the pixel in the output matrix if the \"and\" operation is\n",
    "        true.\n",
    "        If trueColor is none, the output color if the \"and\" operation is true\n",
    "        will be the color of the evaluated pixel.\n",
    "    \"\"\"\n",
    "    m = np.zeros(matrices[0].shape, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        if targetColor == None:\n",
    "            if all([x[i,j] == matrices[0][i,j] for x in matrices]):\n",
    "                if trueColor == None:\n",
    "                    m[i,j] = matrices[0][i,j]\n",
    "                else:\n",
    "                    m[i,j] = trueColor\n",
    "            else:\n",
    "                m[i,j] = falseColor\n",
    "        else:\n",
    "            if all([x[i,j] == targetColor for x in matrices]):\n",
    "                if trueColor == None:\n",
    "                    m[i,j] = matrices[0][i,j]\n",
    "                else:\n",
    "                    m[i,j] = trueColor\n",
    "            else:\n",
    "                m[i,j] = falseColor\n",
    "    return m\n",
    "\n",
    "\"\"\"\n",
    "def pixelwiseOr(matrices, falseColor, targetColor=None, trueColor=None, \\\n",
    "                trueValues=None):\n",
    "    See pixelwiseAnd.\n",
    "    trueValues is a list with as many elements as matrices.\n",
    "    m = np.zeros(matrices[0].shape, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        if targetColor == None:\n",
    "            isFalse = True\n",
    "            for x in matrices:\n",
    "                if x[i,j] != falseColor:\n",
    "                    isFalse = False\n",
    "                    if trueColor == None:\n",
    "                        m[i,j] = x[i,j]\n",
    "                    else:\n",
    "                        m[i,j] = trueColor\n",
    "                    break\n",
    "            if isFalse:\n",
    "                m[i,j] = falseColor\n",
    "        else:\n",
    "            if any([x[i,j] == targetColor for x in matrices]):\n",
    "                if trueColor == None:\n",
    "                    m[i,j] = targetColor\n",
    "                else:\n",
    "                    m[i,j] = trueColor\n",
    "            else:\n",
    "                m[i,j] = falseColor\n",
    "    return m\n",
    "\"\"\"\n",
    "\n",
    "def pixelwiseOr(matrices, falseColor, targetColor=None, trueColor=None, \\\n",
    "                trueValues=None):\n",
    "    \"\"\"\n",
    "    See pixelwiseAnd.\n",
    "    trueValues is a list with as many elements as matrices.\n",
    "    \"\"\"\n",
    "    m = np.zeros(matrices[0].shape, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        if targetColor == None:\n",
    "            trueCount = 0\n",
    "            index = 0\n",
    "            for x in matrices:\n",
    "                if x[i,j] != falseColor:\n",
    "                    trueCount += 1\n",
    "                    trueIndex = index\n",
    "                index += 1\n",
    "            if trueCount==0:\n",
    "                m[i,j] = falseColor\n",
    "            else:\n",
    "                if trueColor!=None:\n",
    "                    m[i,j] = trueColor\n",
    "                elif trueValues!=None:\n",
    "                    if trueCount==1:\n",
    "                        m[i,j] = trueValues[trueIndex]\n",
    "                    else:\n",
    "                        m[i,j] = matrices[trueIndex][i,j]\n",
    "                else:\n",
    "                    m[i,j] = matrices[trueIndex][i,j]\n",
    "        else:\n",
    "            if any([x[i,j] == targetColor for x in matrices]):\n",
    "                if trueColor == None:\n",
    "                    m[i,j] = targetColor\n",
    "                else:\n",
    "                    m[i,j] = trueColor\n",
    "            else:\n",
    "                m[i,j] = falseColor\n",
    "    return m\n",
    "\n",
    "def pixelwiseXor(m1, m2, falseColor, targetColor=None, trueColor=None, \\\n",
    "                 firstTrue=None, secondTrue=None):\n",
    "    \"\"\"\n",
    "    See pixelwiseAnd. The difference is that the Xor operation only makes sense\n",
    "    with two input matrices.\n",
    "    \"\"\"\n",
    "    m = np.zeros(m1.shape, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        if targetColor == None:\n",
    "            if (m1[i,j] == falseColor) != (m2[i,j] == falseColor):\n",
    "                if trueColor == None:\n",
    "                    if firstTrue == None:\n",
    "                        if m1[i,j] != falseColor:\n",
    "                            m[i,j] = m1[i,j]\n",
    "                        else:\n",
    "                            m[i,j] = m2[i,j]\n",
    "                    else:\n",
    "                        if m1[i,j] != falseColor:\n",
    "                            m[i,j] = firstTrue\n",
    "                        else:\n",
    "                            m[i,j] = secondTrue\n",
    "                else:\n",
    "                    m[i,j] = trueColor     \n",
    "            else:\n",
    "                m[i,j] = falseColor\n",
    "        else:\n",
    "            if (m1[i,j] == targetColor) != (m2[i,j] == targetColor):\n",
    "                if trueColor == None:\n",
    "                    if firstTrue == None:\n",
    "                        if m1[i,j] != falseColor:\n",
    "                            m[i,j] = m1[i,j]\n",
    "                        else:\n",
    "                            m[i,j] = m2[i,j]\n",
    "                    else:\n",
    "                        if m1[i,j] != falseColor:\n",
    "                            m[i,j] = firstTrue\n",
    "                        else:\n",
    "                            m[i,j] = secondTrue\n",
    "                else:\n",
    "                    m[i,j] = trueColor     \n",
    "            else:\n",
    "                m[i,j] = falseColor\n",
    "    return m\n",
    "\n",
    "# %% Downsize and Minimize\n",
    "    \n",
    "def getDownsizeFactors(matrix):\n",
    "    \"\"\"\n",
    "    Still unused\n",
    "    \"\"\"\n",
    "    xDivisors = set()\n",
    "    for x in range(1, matrix.shape[0]):\n",
    "        if (matrix.shape[0]%x)==0:\n",
    "            xDivisors.add(x)\n",
    "    yDivisors = set()\n",
    "    for y in range(1, matrix.shape[1]):\n",
    "        if (matrix.shape[1]%y)==0:\n",
    "            yDivisors.add(y)\n",
    "    \n",
    "    downsizeFactors = set()\n",
    "    for x,y in product(xDivisors, yDivisors):\n",
    "        downsizeFactors.add((x,y))\n",
    " \n",
    "    return downsizeFactors\n",
    "\n",
    "def downsize(matrix, newShape, falseColor=None):\n",
    "    \"\"\"\n",
    "    Given a matrix and a shape, this function returns a new matrix with the\n",
    "    given shape. The elements of the return matrix are given by the colors of \n",
    "    each of the submatrices. Each submatrix is only allowed to have the\n",
    "    background color and at most another one (that will define the output\n",
    "    color of the corresponding pixel).\n",
    "    \"\"\"\n",
    "    if falseColor==None:\n",
    "        falseColor = matrix.backgroundColor\n",
    "    if (matrix.shape[0]%newShape[0])!=0 or (matrix.shape[1]%newShape[1])!=0:\n",
    "        return matrix.m.copy()\n",
    "    xBlock = int(matrix.shape[0]/newShape[0])\n",
    "    yBlock = int(matrix.shape[1]/newShape[1])\n",
    "    m = np.full(newShape, matrix.backgroundColor, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(newShape[0], newShape[1]):\n",
    "        color = -1\n",
    "        for x,y in np.ndindex(xBlock, yBlock):\n",
    "            if matrix.m[i*xBlock+x, j*yBlock+y] not in [matrix.backgroundColor, color]:\n",
    "                if color==-1:\n",
    "                    color = matrix.m[i*xBlock+x, j*yBlock+y]\n",
    "                else:\n",
    "                    return matrix.m.copy()\n",
    "        if color==-1:\n",
    "            m[i,j] = falseColor\n",
    "        else:\n",
    "            m[i,j] = color\n",
    "    return m\n",
    "\n",
    "def minimize(matrix):\n",
    "    \"\"\"\n",
    "    Given a matrix, this function returns the matrix resulting from the\n",
    "    following operations:\n",
    "        If two consecutive rows are equal, delete one of them\n",
    "        If two consecutive columns are equal, delete one of them\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    x = 1\n",
    "    for i in range(1, matrix.shape[0]):\n",
    "        if np.array_equal(m[x,:],m[x-1,:]):\n",
    "            m = np.delete(m, (x), axis=0)\n",
    "        else:\n",
    "            x+=1\n",
    "    x = 1\n",
    "    for i in range(1, matrix.shape[1]):\n",
    "        if np.array_equal(m[:,x],m[:,x-1]):\n",
    "            m = np.delete(m, (x), axis=1)\n",
    "        else:\n",
    "            x+=1\n",
    "    return m\n",
    "            \n",
    "        \n",
    "\n",
    "# %% Operations to extend matrices\n",
    "    \n",
    "def extendMatrix(matrix, color, position=\"tl\", xShape=None, yShape=None, isSquare=False, goodDimension=None):\n",
    "    \"\"\"\n",
    "    Given a matrix, xShape(>matrix.shape[0]), yShape(>matrix.shape[1]) and a color,\n",
    "    this function extends the matrix using the dimensions given by xShape and\n",
    "    yShape by coloring the extra pixels with the given color.\n",
    "    If xShape or yShape are not given, then they will be equal to matrix.shape[0]\n",
    "    of matrix.shape[1], respectively.\n",
    "    The position of the input matrix in the output matrix can be given by\n",
    "    specifying \"tl\", \"tr\", \"bl\" or \"br\" (top-left/top-right/bot-left/bot-right).\n",
    "    The default is top-left.\n",
    "    \"\"\"\n",
    "    if isSquare:\n",
    "        if goodDimension=='x':\n",
    "            xShape = matrix.shape[0]\n",
    "            yShape = matrix.shape[0]\n",
    "        if goodDimension=='y':\n",
    "            xShape = matrix.shape[1]\n",
    "            yShape = matrix.shape[1]\n",
    "    if xShape==None:\n",
    "        xShape = matrix.shape[0]\n",
    "    if yShape==None:\n",
    "        yShape = matrix.shape[1]\n",
    "    m = np.full((xShape, yShape), color, dtype=np.uint8)\n",
    "    if position==\"tl\":\n",
    "        m[0:matrix.shape[0], 0:matrix.shape[1]] = matrix.m.copy()\n",
    "    elif position==\"tr\":\n",
    "        m[0:matrix.shape[0], yShape-matrix.shape[1]:yShape] = matrix.m.copy()\n",
    "    elif position==\"bl\":\n",
    "        m[xShape-matrix.shape[0]:xShape, 0:matrix.shape[1]] = matrix.m.copy()\n",
    "    else:\n",
    "        m[xShape-matrix.shape[0]:xShape, yShape-matrix.shape[1]:yShape] = matrix.m.copy()\n",
    "    \n",
    "    return m\n",
    "\n",
    "def getBestExtendMatrix(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"extendMatrix\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    if t.backgroundColor==-1:\n",
    "        totalColorCount = Counter()\n",
    "        for s in t.trainSamples:\n",
    "            totalColorCount += s.inMatrix.colorCount\n",
    "        background = max(totalColorCount.items(), key=operator.itemgetter(1))[0]\n",
    "    else:\n",
    "        background=t.backgroundColor\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    \n",
    "    # Define xShape and yShape:\n",
    "    xShape = None\n",
    "    yShape = None\n",
    "    isSquare=False\n",
    "    goodDimension=None\n",
    "    \n",
    "    # If the outShape is given, easy\n",
    "    if t.sameOutShape:\n",
    "        xShape=t.outShape[0]\n",
    "        yShape=t.outShape[1]\n",
    "    # If the output xShape is always the same and the yShape keeps constant, pass the common xShape\n",
    "    elif len(set([s.outMatrix.shape[0] for s in t.trainSamples]))==1:\n",
    "        if all([s.outMatrix.shape[1]==s.inMatrix.shape[1] for s in t.trainSamples]):\n",
    "            xShape=t.trainSamples[0].outMatrix.shape[0]  \n",
    "    # If the output yShape is always the same and the xShape keeps constant, pass the common yShape\n",
    "    elif len(set([s.outMatrix.shape[1] for s in t.trainSamples]))==1:\n",
    "        if all([s.outMatrix.shape[0]==s.inMatrix.shape[0] for s in t.trainSamples]):\n",
    "            yShape=t.trainSamples[0].outMatrix.shape[1] \n",
    "    # If the matrix is always squared, and one dimension (x or y) is fixed, do this:\n",
    "    elif all([s.outMatrix.shape[0]==s.outMatrix.shape[1] for s in t.trainSamples]):\n",
    "        isSquare=True\n",
    "        if all([s.outMatrix.shape[1]==s.inMatrix.shape[1] for s in t.trainSamples]):\n",
    "            goodDimension='y'     \n",
    "        elif all([s.outMatrix.shape[0]==s.inMatrix.shape[0] for s in t.trainSamples]):\n",
    "            goodDimension='x'      \n",
    "    \n",
    "    for position in [\"tl\", \"tr\", \"bl\", \"br\"]:\n",
    "        f = partial(extendMatrix, color=background, xShape=xShape, yShape=yShape,\\\n",
    "                    position=position, isSquare=isSquare, goodDimension=goodDimension)\n",
    "        bestFunction, bestScore = updateBestFunction(t, f, bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        \n",
    "    return bestFunction\n",
    "    \n",
    "def getFactor(matrix, factor):\n",
    "    \"\"\"\n",
    "    Given a Task.inShapeFactor (that can be a string), this function\n",
    "    returns its corresponding tuple for the given matrix.\n",
    "    \"\"\"\n",
    "    if factor == \"squared\":\n",
    "        f = (matrix.shape[0], matrix.shape[1])\n",
    "    elif factor == \"xSquared\":\n",
    "        f = (matrix.shape[0], 1)\n",
    "    elif factor == \"ySquared\":\n",
    "        f = (1, matrix.shape[1])\n",
    "    elif factor == \"nColors\":\n",
    "        f = (matrix.nColors, matrix.nColors)\n",
    "    elif factor == \"nColors-1\":\n",
    "        f = (matrix.nColors-1, matrix.nColors-1)\n",
    "    else:\n",
    "        f = factor\n",
    "    return f\n",
    "\n",
    "def multiplyPixels(matrix, factor):\n",
    "    \"\"\"\n",
    "    Factor is a 2-dimensional tuple.\n",
    "    The output matrix has shape matrix.shape*factor. Each pixel of the input\n",
    "    matrix is expanded by factor.\n",
    "    \"\"\"\n",
    "    factor = getFactor(matrix, factor)\n",
    "    m = np.zeros(tuple(s * f for s, f in zip(matrix.shape, factor)), dtype=np.uint8)\n",
    "    for i,j in np.ndindex(matrix.m.shape):\n",
    "        for k,l in np.ndindex(factor):\n",
    "            m[i*factor[0]+k, j*factor[1]+l] = matrix.m[i,j]\n",
    "    return m\n",
    "\n",
    "def multiplyMatrix(matrix, factor):\n",
    "    \"\"\"\n",
    "    Copy the matrix \"matrix\" into every submatrix of the output, which has\n",
    "    shape matrix.shape * factor.\n",
    "    \"\"\"\n",
    "    factor = getFactor(matrix, factor)\n",
    "    m = np.zeros(tuple(s * f for s, f in zip(matrix.shape, factor)), dtype=np.uint8)\n",
    "    for i,j in np.ndindex(factor):\n",
    "        m[i*matrix.shape[0]:(i+1)*matrix.shape[0], j*matrix.shape[1]:(j+1)*matrix.shape[1]] = matrix.m.copy()\n",
    "    return m\n",
    "\n",
    "def matrixTopLeft(matrix, factor, background=0):\n",
    "    \"\"\"\n",
    "    Copy the matrix into the top left corner of the multiplied matrix\n",
    "    \"\"\"\n",
    "    factor = getFactor(matrix, factor)\n",
    "    m = np.full(tuple(s * f for s, f in zip(matrix.shape, factor)), background, dtype=np.uint8)\n",
    "    m[0:matrix.shape[0], 0:matrix.shape[1]] = matrix.m.copy()\n",
    "    return m\n",
    "    \n",
    "def matrixBotRight(matrix, factor, background=0):\n",
    "    \"\"\"\n",
    "    Copy the matrix into the bottom right corner of the multiplied matrix\n",
    "    \"\"\"\n",
    "    factor = getFactor(matrix, factor)\n",
    "    m = np.full(tuple(s * f for s, f in zip(matrix.shape, factor)), background, dtype=np.uint8)\n",
    "    m[(factor[0]-1)*matrix.shape[0]:factor[0]*matrix.shape[0], \\\n",
    "      (factor[1]-1)*matrix.shape[1]:factor[1]*matrix.shape[1]]\n",
    "    return m\n",
    "\n",
    "def getBestMosaic(t):\n",
    "    \"\"\"\n",
    "    Given a task t, this function tries to find the best way to generate a\n",
    "    mosaic, given that the output shape is always bigger than the input shape\n",
    "    with a shape factor that makes sense.\n",
    "    A mosaic is a matrix that takes an input matrix as reference, and then\n",
    "    copies it many times. The copies can include rotations or mirrorings.\n",
    "    \"\"\"\n",
    "    factor = t.inShapeFactor\n",
    "    ops = []\n",
    "    ops.append(partial(identityM))\n",
    "    ops.append(partial(mirror, axis=\"lr\"))\n",
    "    ops.append(partial(mirror, axis=\"ud\"))\n",
    "    ops.append(partial(rotate, angle=180))\n",
    "    if t.inMatricesSquared:\n",
    "        ops.append(partial(mirror, axis=\"d1\"))\n",
    "        ops.append(partial(mirror, axis=\"d2\"))\n",
    "        ops.append(partial(rotate, angle=90))\n",
    "        ops.append(partial(rotate, angle=270))\n",
    "    bestOps = []\n",
    "    for i in range(factor[0]):\n",
    "        bestOps.append([])\n",
    "        for j in range(factor[1]):\n",
    "            bestScore = 1000\n",
    "            bestOp = partial(identityM)\n",
    "            for op in ops:\n",
    "                score = 0\n",
    "                for s in t.trainSamples:\n",
    "                    inM = s.inMatrix.m.copy()\n",
    "                    outM = s.outMatrix.m[i*inM.shape[0]:(i+1)*inM.shape[0], j*inM.shape[1]:(j+1)*inM.shape[1]]\n",
    "                    score += incorrectPixels(op(inM),outM)\n",
    "                if score < bestScore:\n",
    "                    bestScore = score\n",
    "                    bestOp = op\n",
    "                    if score==0:\n",
    "                        break\n",
    "            bestOps[i].append(bestOp)\n",
    "    return bestOps\n",
    "\n",
    "def generateMosaic(matrix, ops, factor):\n",
    "    \"\"\"\n",
    "    Generates a mosaic from the given matrix using the operations given in the\n",
    "    list ops. The output matrix has shape matrix.shape*factor.\n",
    "    \"\"\"\n",
    "    m = np.zeros(tuple(s * f for s, f in zip(matrix.shape, factor)), dtype=np.uint8)\n",
    "    for i in range(factor[0]):\n",
    "        for j in range(factor[1]):\n",
    "            m[i*matrix.shape[0]:(i+1)*matrix.shape[0], j*matrix.shape[1]:(j+1)*matrix.shape[1]] = \\\n",
    "            ops[i][j](matrix)\n",
    "    return m\n",
    "\n",
    "# Only if the factor is squared\n",
    "def getBestMultiplyMatrix(t, falseColor): \n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"multiplyMatrix\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    def getFullMatrix(matrix, color):\n",
    "        return np.full(matrix.shape, color, dtype=np.uint8)\n",
    "    # Possible operations on the matrix\n",
    "    ops = []\n",
    "    ops.append(partial(identityM))\n",
    "    ops.append(partial(mirror, axis=\"lr\"))\n",
    "    ops.append(partial(mirror, axis=\"ud\"))\n",
    "    ops.append(partial(rotate, angle=180))\n",
    "    if t.inMatricesSquared:\n",
    "        ops.append(partial(mirror, axis=\"d1\"))\n",
    "        ops.append(partial(mirror, axis=\"d2\"))\n",
    "        ops.append(partial(rotate, angle=90))\n",
    "        ops.append(partial(rotate, angle=270))\n",
    "    if all([n==2 for n in t.nInColors]):\n",
    "        ops.append(partial(switchColors))\n",
    "    \n",
    "    # Conditions\n",
    "    def trueCondition(matrix, pixel):\n",
    "        return True\n",
    "    def maxColor(matrix, pixel):\n",
    "        x = [k for k, v in sorted(matrix.colorCount.items(), key=lambda item: item[1], reverse=True)]\n",
    "        if len(x)<2 or matrix.colorCount[x[0]]!=matrix.colorCount[x[1]]:\n",
    "            return pixel==max(matrix.colorCount, key=matrix.colorCount.get)\n",
    "        else:\n",
    "            return False\n",
    "    def minColor(matrix,pixel):\n",
    "        x = [k for k, v in sorted(matrix.colorCount.items(), key=lambda item: item[1])]\n",
    "        if len(x)<2 or matrix.colorCount[x[-1]]!=matrix.colorCount[x[-2]]:\n",
    "            return pixel==min(matrix.colorCount, key=matrix.colorCount.get)\n",
    "        else:\n",
    "            return False\n",
    "    def isColor(matrix, pixel, color):\n",
    "        return pixel==color\n",
    "    def nonZero(matrix, pixel):\n",
    "        return pixel!=0\n",
    "    def zero(matrix, pixel):\n",
    "        return pixel==0\n",
    "    conditions = []\n",
    "    conditions.append(partial(trueCondition))\n",
    "    conditions.append(partial(maxColor))\n",
    "    conditions.append(partial(minColor))\n",
    "    conditions.append(partial(nonZero))\n",
    "    conditions.append(partial(zero))\n",
    "    for c in t.colors:\n",
    "        conditions.append(partial(isColor, color=c))\n",
    "\n",
    "    bestScore = 1000\n",
    "    for op, cond in product(ops, conditions):\n",
    "        score = 0\n",
    "        for s in t.trainSamples:\n",
    "            factor = getFactor(s.inMatrix, t.inShapeFactor)\n",
    "            for i,j in np.ndindex(factor):\n",
    "                inM = s.inMatrix.m.copy()\n",
    "                outM = s.outMatrix.m[i*inM.shape[0]:(i+1)*inM.shape[0], j*inM.shape[1]:(j+1)*inM.shape[1]]\n",
    "                if cond(s.inMatrix, inM[i,j]):\n",
    "                    score += incorrectPixels(op(inM),outM)\n",
    "                else:\n",
    "                    score += incorrectPixels(getFullMatrix(inM, falseColor), outM)\n",
    "        if score < bestScore:\n",
    "            bestScore = score\n",
    "            opCond = (op, cond)\n",
    "            if score==0:\n",
    "                return opCond\n",
    "    return opCond\n",
    "\n",
    "def doBestMultiplyMatrix(matrix, opCond, falseColor):\n",
    "    factor = matrix.shape\n",
    "    m = np.full(tuple(s * f for s, f in zip(matrix.shape, factor)), falseColor, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(factor):\n",
    "        if opCond[1](matrix, matrix.m[i,j]):\n",
    "            m[i*matrix.shape[0]:(i+1)*matrix.shape[0], j*matrix.shape[1]:(j+1)*matrix.shape[1]] = \\\n",
    "            opCond[0](matrix)\n",
    "    return m\n",
    "\n",
    "# %% Multiply pixels\n",
    "\n",
    "def multiplyPixelsAndAnd(matrix, factor, falseColor):\n",
    "    \"\"\"\n",
    "    This function basically is the same as executing the functions\n",
    "    multiplyPixels, multiplyMatrix, and executing pixelwiseAnd with these two\n",
    "    matrices as inputs\n",
    "    \"\"\"\n",
    "    factor = getFactor(matrix, factor)\n",
    "    m = matrix.m.copy()\n",
    "    multipliedM = multiplyPixels(matrix, factor)\n",
    "    for i,j in np.ndindex(factor):\n",
    "        newM = multipliedM[i*m.shape[0]:(i+1)*m.shape[0], j*m.shape[1]:(j+1)*m.shape[1]]\n",
    "        multipliedM[i*m.shape[0]:(i+1)*m.shape[0], j*m.shape[1]:(j+1)*m.shape[1]] = pixelwiseAnd([m, newM], falseColor)\n",
    "    return multipliedM\n",
    "\n",
    "def multiplyPixelsAndOr(matrix, factor, falseColor):\n",
    "    \"\"\"\n",
    "    This function basically is the same as executing the functions\n",
    "    multiplyPixels, multiplyMatrix, and executing pixelwiseOr with these two\n",
    "    matrices as inputs\n",
    "    \"\"\"\n",
    "    factor = getFactor(matrix, factor)\n",
    "    m = matrix.m.copy()\n",
    "    multipliedM = multiplyPixels(matrix, factor)\n",
    "    for i,j in np.ndindex(factor):\n",
    "        newM = multipliedM[i*m.shape[0]:(i+1)*m.shape[0], j*m.shape[1]:(j+1)*m.shape[1]]\n",
    "        multipliedM[i*m.shape[0]:(i+1)*m.shape[0], j*m.shape[1]:(j+1)*m.shape[1]] = pixelwiseOr([m, newM], falseColor)\n",
    "    return multipliedM\n",
    "\n",
    "def multiplyPixelsAndXor(matrix, factor, falseColor):\n",
    "    \"\"\"\n",
    "    This function basically is the same as executing the functions\n",
    "    multiplyPixels, multiplyMatrix, and executing pixelwiseXor with these two\n",
    "    matrices as inputs\n",
    "    \"\"\"\n",
    "    factor = getFactor(matrix, factor)\n",
    "    m = matrix.m.copy()\n",
    "    multipliedM = multiplyPixels(matrix, factor)\n",
    "    for i,j in np.ndindex(factor):\n",
    "        newM = multipliedM[i*m.shape[0]:(i+1)*m.shape[0], j*m.shape[1]:(j+1)*m.shape[1]]\n",
    "        multipliedM[i*m.shape[0]:(i+1)*m.shape[0], j*m.shape[1]:(j+1)*m.shape[1]] = pixelwiseXor(m, newM, falseColor)\n",
    "    return multipliedM\n",
    "\n",
    "# %% Operations considering all submatrices of task with outShapeFactor\n",
    "    \n",
    "def getSubmatrices(m, factor):\n",
    "    \"\"\"\n",
    "    Given a matrix m and a factor, this function returns a list of all the\n",
    "    submatrices with shape determined by the factor.\n",
    "    \"\"\"\n",
    "    matrices = []\n",
    "    nRows = int(m.shape[0] / factor[0])\n",
    "    nCols = int(m.shape[1] / factor[1])\n",
    "    for i,j in np.ndindex(factor):\n",
    "        matrices.append(m[i*nRows:(i+1)*nRows, j*nCols:(j+1)*nCols])\n",
    "    return matrices\n",
    "\n",
    "def outputIsSubmatrix(t, isGrid=False):\n",
    "    \"\"\"\n",
    "    Given a task t that has outShapeFactor, this function returns true if any\n",
    "    of the submatrices is equal to the output matrix for every sample.\n",
    "    \"\"\"\n",
    "    for sample in t.trainSamples:\n",
    "        if isGrid:\n",
    "            matrices = [c[0].m for c in sample.inMatrix.grid.cellList]\n",
    "        else:\n",
    "            matrices = getSubmatrices(sample.inMatrix.m, sample.outShapeFactor)\n",
    "        anyIsSubmatrix = False\n",
    "        for m in matrices:\n",
    "            if np.array_equal(m, sample.outMatrix.m):\n",
    "                anyIsSubmatrix = True\n",
    "                break\n",
    "        if not anyIsSubmatrix:\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "def selectSubmatrixWithMaxColor(matrix, color, outShapeFactor=None, isGrid=False):\n",
    "    \"\"\"\n",
    "    Given a matrix, this function returns the submatrix with most appearances\n",
    "    of the color given. If the matrix is not a grid, an outShapeFactor must be\n",
    "    specified.\n",
    "    \"\"\"\n",
    "    if isGrid:\n",
    "        matrices = [c[0].m for c in matrix.grid.cellList]\n",
    "    else:\n",
    "        matrices = getSubmatrices(matrix.m, outShapeFactor)\n",
    "        \n",
    "    maxCount = 0\n",
    "    matricesWithProperty = 0\n",
    "    bestMatrix = None\n",
    "    for mat in matrices:\n",
    "        m = Matrix(mat)\n",
    "        if color in m.colors:\n",
    "            if m.colorCount[color]>maxCount:\n",
    "                bestMatrix = mat.copy()\n",
    "                maxCount = m.colorCount[color]\n",
    "                matricesWithProperty = 1\n",
    "            if m.colorCount[color]==maxCount:\n",
    "                matricesWithProperty += 1\n",
    "    if matricesWithProperty!=1:\n",
    "        return matrix.m.copy()\n",
    "    else:\n",
    "        return bestMatrix\n",
    "    \n",
    "def selectSubmatrixWithMinColor(matrix, color, outShapeFactor=None, isGrid=False):\n",
    "    \"\"\"\n",
    "    Given a matrix, this function returns the submatrix with least appearances\n",
    "    of the color given. If the matrix is not a grid, an outShapeFactor must be\n",
    "    specified.\n",
    "    \"\"\"\n",
    "    if isGrid:\n",
    "        matrices = [c[0].m for c in matrix.grid.cellList]\n",
    "    else:\n",
    "        matrices = getSubmatrices(matrix.m, outShapeFactor)\n",
    "        \n",
    "    minCount = 1000\n",
    "    matricesWithProperty = 0\n",
    "    bestMatrix = None\n",
    "    for mat in matrices:\n",
    "        m = Matrix(mat)\n",
    "        if color in m.colors:\n",
    "            if m.colorCount[color]<minCount:\n",
    "                bestMatrix = mat.copy()\n",
    "                minCount = m.colorCount[color]\n",
    "                matricesWithProperty = 1\n",
    "            elif m.colorCount[color]==minCount:\n",
    "                matricesWithProperty += 1\n",
    "    if matricesWithProperty!=1:\n",
    "        return matrix.m.copy()\n",
    "    else:\n",
    "        return bestMatrix\n",
    "    \n",
    "def selectSubmatrixWithMostColors(matrix, outShapeFactor=None, isGrid=False):\n",
    "    \"\"\"\n",
    "    Given a matrix, this function returns the submatrix with the most number of\n",
    "    colors. If the matrix is not a grid, an outShapeFactor must be specified.\n",
    "    \"\"\"\n",
    "    if isGrid:\n",
    "        matrices = [c[0].m for c in matrix.grid.cellList]\n",
    "    else:\n",
    "        matrices = getSubmatrices(matrix.m, outShapeFactor)\n",
    "        \n",
    "    maxNColors = 0\n",
    "    matricesWithProperty = 0\n",
    "    bestMatrix = None\n",
    "    for mat in matrices:\n",
    "        m = Matrix(mat)\n",
    "        if len(m.colorCount)>maxNColors:\n",
    "            bestMatrix = mat.copy()\n",
    "            maxNColors = len(m.colorCount)\n",
    "            matricesWithProperty = 1\n",
    "        elif len(m.colorCount)==maxNColors:\n",
    "            matricesWithProperty += 1\n",
    "    if matricesWithProperty!=1:\n",
    "        return matrix.m.copy()\n",
    "    else:\n",
    "        return bestMatrix\n",
    "    \n",
    "def selectSubmatrixWithLeastColors(matrix, outShapeFactor=None, isGrid=False):\n",
    "    \"\"\"\n",
    "    Given a matrix, this function returns the submatrix with the least number\n",
    "    of colors. If the matrix is not a grid, an outShapeFactor must be\n",
    "    specified.\n",
    "    \"\"\"\n",
    "    if isGrid:\n",
    "        matrices = [c[0].m for c in matrix.grid.cellList]\n",
    "    else:\n",
    "        matrices = getSubmatrices(matrix.m, outShapeFactor)\n",
    "        \n",
    "    minNColors = 1000\n",
    "    matricesWithProperty = 0\n",
    "    bestMatrix = None\n",
    "    for mat in matrices:\n",
    "        m = Matrix(mat)\n",
    "        if len(m.colorCount)<minNColors:\n",
    "            bestMatrix = mat.copy()\n",
    "            minNColors = len(m.colorCount)\n",
    "            matricesWithProperty = 1\n",
    "        elif len(m.colorCount)==minNColors:\n",
    "            matricesWithProperty += 1\n",
    "    if matricesWithProperty!=1:\n",
    "        return matrix.m.copy()\n",
    "    else:\n",
    "        return bestMatrix\n",
    "        \n",
    "def getBestSubmatrixPosition(t, outShapeFactor=None, isGrid=False):\n",
    "    \"\"\"\n",
    "    Given a task t, and assuming that all the input matrices have the same\n",
    "    shape and all the ouptut matrices have the same shape too, this function\n",
    "    tries to check whether the output matrix is just the submatrix in a given\n",
    "    position. If that's the case, it returns the position. Otherwise, it\n",
    "    returns 0.\n",
    "    \"\"\"\n",
    "    iteration = 0\n",
    "    possiblePositions = []\n",
    "    for sample in t.trainSamples:\n",
    "        if isGrid:\n",
    "            matrices = [c[0].m for c in sample.inMatrix.grid.cellList]\n",
    "        else:\n",
    "            matrices = getSubmatrices(sample.inMatrix.m, outShapeFactor)\n",
    "            \n",
    "        possiblePositions.append(set())\n",
    "        for m in range(len(matrices)):\n",
    "            if np.array_equal(matrices[m], sample.outMatrix.m):\n",
    "                possiblePositions[iteration].add(m)\n",
    "        \n",
    "        iteration += 1\n",
    "    positions = set.intersection(*possiblePositions)\n",
    "    if len(positions)==1:\n",
    "        return next(iter(positions))\n",
    "    else:\n",
    "        return 0\n",
    "                \n",
    "def selectSubmatrixInPosition(matrix, position, outShapeFactor=None, isGrid=False):\n",
    "    \"\"\"\n",
    "    Given a matrix and a position, this function returns the submatrix that\n",
    "    appears in the given position (submatrices are either defined by\n",
    "    outShapeFactor or by the shape of the grid cells).\n",
    "    \"\"\"\n",
    "    if isGrid:\n",
    "        matrices = [c[0].m for c in matrix.grid.cellList]\n",
    "    else:\n",
    "        matrices = getSubmatrices(matrix.m, outShapeFactor)\n",
    "        \n",
    "    return matrices[position].copy()\n",
    "\n",
    "def maxColorFromCell(matrix):\n",
    "    \"\"\"\n",
    "    Only to be called if matrix.isGrid.\n",
    "    Given a matrix with a grid, this function returns a matrix with the same\n",
    "    shape as the grid. Every pixel of the matrix will be colored with the \n",
    "    color that appears the most in the corresponding cell of the grid.\n",
    "    \"\"\"\n",
    "    m = np.zeros(matrix.grid.shape, dtype=np.uint8)\n",
    "    for i,j  in np.ndindex(matrix.grid.shape):\n",
    "        color = max(matrix.grid.cells[i][j][0].colorCount.items(), key=operator.itemgetter(1))[0]\n",
    "        m[i,j] = color\n",
    "    return m\n",
    "\n",
    "def colorAppearingXTimes(matrix, times):\n",
    "    m = np.zeros(matrix.grid.shape, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(matrix.grid.shape):\n",
    "        for k,v in matrix.grid.cells[i][j][0].colorCount.items():\n",
    "            if v==times:\n",
    "                m[i,j] = k\n",
    "    return m\n",
    "        \n",
    "def pixelwiseAndInSubmatrices(matrix, factor, falseColor, targetColor=None, trueColor=None):\n",
    "    matrices = getSubmatrices(matrix.m.copy(), factor)\n",
    "    return pixelwiseAnd(matrices, falseColor, targetColor, trueColor)\n",
    "\n",
    "def pixelwiseOrInSubmatrices(matrix, factor, falseColor, targetColor=None, trueColor=None, \\\n",
    "                             trueValues=None):\n",
    "    matrices = getSubmatrices(matrix.m.copy(), factor)\n",
    "    return pixelwiseOr(matrices, falseColor, targetColor, trueColor, trueValues)\n",
    "\n",
    "def pixelwiseXorInSubmatrices(matrix, factor, falseColor, targetColor=None, trueColor=None, \\\n",
    "                              firstTrue=None, secondTrue=None):\n",
    "    matrices = getSubmatrices(matrix.m.copy(), factor)\n",
    "    return pixelwiseXor(matrices[0], matrices[1], falseColor, targetColor, trueColor, firstTrue, secondTrue)\n",
    "\n",
    "# %% Operations considering all submatrices of a grid\n",
    "\n",
    "def pixelwiseAndInGridSubmatrices(matrix, falseColor, targetColor=None, trueColor=None):\n",
    "    matrices = [c[0].m for c in matrix.grid.cellList]\n",
    "    return pixelwiseAnd(matrices, falseColor, targetColor, trueColor)\n",
    "\n",
    "def pixelwiseOrInGridSubmatrices(matrix, falseColor, targetColor=None, trueColor=None, \\\n",
    "                                 trueValues=None):\n",
    "    matrices = [c[0].m for c in matrix.grid.cellList]\n",
    "    return pixelwiseOr(matrices, falseColor, targetColor, trueColor, trueValues)\n",
    "\n",
    "def pixelwiseXorInGridSubmatrices(matrix, falseColor, targetColor=None, trueColor=None, \\\n",
    "                                  firstTrue=None, secondTrue=None):\n",
    "    m1 = matrix.grid.cellList[0][0].m.copy()\n",
    "    m2 = matrix.grid.cellList[1][0].m.copy()\n",
    "    return pixelwiseXor(m1, m2, falseColor, targetColor, trueColor, firstTrue, secondTrue)\n",
    "\n",
    "# %% crop all shapes\n",
    "    \n",
    "def cropAllShapes(matrix, background, diagonal=False):\n",
    "    if diagonal:\n",
    "        shapes = [shape for shape in matrix.dShapes if shape.color!=background]\n",
    "    else:\n",
    "        shapes = [shape for shape in matrix.shapes if shape.color!=background]\n",
    "    shapes = sorted(shapes, key=lambda x: x.nPixels, reverse=True)\n",
    "    \n",
    "    if len(shapes)==0:\n",
    "        return matrix.m.copy()\n",
    "    \n",
    "    m = shapes[0].m.copy()\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        if m[i,j]==255:\n",
    "            m[i,j] = background\n",
    "    \n",
    "    outMatrix = Matrix(m)\n",
    "    if diagonal:\n",
    "        outShapes = [shape for shape in outMatrix.dShapes if shape.color==background]\n",
    "    else:\n",
    "        outShapes = [shape for shape in outMatrix.shapes if shape.color==background]\n",
    "    \n",
    "    for s in outShapes:\n",
    "        if s.color==background:\n",
    "            for shape in shapes:\n",
    "                if shape.hasSameShape(s):\n",
    "                    m = changeColorShapes(m, [s], shape.color)\n",
    "                    break\n",
    "    return m\n",
    "\n",
    "def getLayerDict(t):\n",
    "    \"\"\"\n",
    "    Only to be called when there is a potential bijective dictionary between submatrices of the input and output. This\n",
    "    function returns such dictionary. \n",
    "    \"\"\"\n",
    "    shList = []\n",
    "    for s in t.trainSamples:\n",
    "        shList += [sh.shape for sh in s.outMatrix.shapes]\n",
    "    sc = Counter(shList)\n",
    "    shape = sc.most_common(1)[0][0]\n",
    "    layerDict = dict()\n",
    "    for s in t.trainSamples:\n",
    "        inM = s.inMatrix.m\n",
    "        outM = s.outMatrix.m\n",
    "        xR = inM.shape[0]//shape[0]\n",
    "        yR = inM.shape[1]//shape[1]\n",
    "        for x in range(xR):\n",
    "            for y in range(yR):\n",
    "                if inM[x*shape[0]:(x+1)*shape[0],y*shape[1]:(y+1)*shape[1]].tobytes() in layerDict:\n",
    "                    if np.all(layerDict[inM[x*shape[0]:(x+1)*shape[0],y*shape[1]:(y+1)*shape[1]].tobytes()] !=\\\n",
    "                                                outM[x*shape[0]:(x+1)*shape[0],y*shape[1]:(y+1)*shape[1]]):\n",
    "                        return None, dict()\n",
    "                layerDict[inM[x*shape[0]:(x+1)*shape[0],y*shape[1]:(y+1)*shape[1]].tobytes()]\\\n",
    "                                    = outM[x*shape[0]:(x+1)*shape[0],y*shape[1]:(y+1)*shape[1]] \n",
    "    return shape, layerDict\n",
    "\n",
    "def subMatToLayer(matrix, shapeAndDict):\n",
    "    \"\"\"\n",
    "    Transforms each submatrix according to the given dictionary shapeAndDict. If the input submatrix is not \n",
    "    in the dictionary, returns the original matrix. \n",
    "    \"\"\"\n",
    "    shape = shapeAndDict[0]\n",
    "    layerDict = shapeAndDict[1]\n",
    "    m = matrix.m.copy()\n",
    "    if shape is None:\n",
    "        return m\n",
    "    xR = m.shape[0]//shape[0]\n",
    "    yR = m.shape[1]//shape[1]\n",
    "    for x in range(xR):\n",
    "        for y in range(yR):\n",
    "            if m[x*shape[0]:(x+1)*shape[0],y*shape[1]:(y+1)*shape[1]].tobytes() in layerDict:\n",
    "                m[x*shape[0]:(x+1)*shape[0],y*shape[1]:(y+1)*shape[1]] =\\\n",
    "                        layerDict[m[x*shape[0]:(x+1)*shape[0],y*shape[1]:(y+1)*shape[1]].tobytes()].copy()\n",
    "            else:\n",
    "                return m\n",
    "    return m\n",
    "    \n",
    "def getBestFitToFrame(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"fitToFrame\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    crop = False\n",
    "    if t.outSmallerThanIn:\n",
    "        crop = True\n",
    "    for s in [True, False]:\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(fitToFrame, crop=crop, includeFrame=True, scale=s), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(fitToFrame, crop=crop, includeFrame=False, scale=s), bestScore, bestFunction)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(fitToFrame, crop=crop, includeFrame=True, colorMatch=True), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def fitToFrame(matrix, crop=False, scale=False, includeFrame=True, colorMatch=False):\n",
    "    \"\"\"\n",
    "    To be called only if the task has a partial or full frame. Attempts to fit or adjust\n",
    "    a shape or list of shapes inside the frame.\n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    if len(matrix.partialFrames) != 1:\n",
    "        return m\n",
    "    frame = matrix.partialFrames[0]\n",
    "    found = False\n",
    "    maux = Matrix(deleteShape(m, frame, matrix.backgroundColor))\n",
    "    for sh in maux.multicolorDShapes:\n",
    "        if sh != frame and sh.shape[0]>1 and sh.shape[1]>1:\n",
    "            m = deleteShape(m, sh, matrix.backgroundColor)\n",
    "            found = True\n",
    "            break\n",
    "    if not found or sh.shape[0] > frame.shape[0] or sh.shape[1] > frame.shape[1]:\n",
    "        return m\n",
    "    if scale:\n",
    "        r = min((frame.shape[0]-2)//sh.shape[0], (frame.shape[1]-2)//sh.shape[1])\n",
    "        newSh = copy.deepcopy(sh)\n",
    "        newSh.m = np.repeat(np.repeat(sh.m, r, axis=1), r, axis=0)\n",
    "        newSh.shape = newSh.m.shape\n",
    "        newSh.pixels = set([(i,j) for i,j in np.ndindex(newSh.m.shape) if newSh.m[i,j]!=255])\n",
    "    else:\n",
    "        newSh = copy.deepcopy(sh)\n",
    "    newSh.position = (frame.position[0] + (frame.shape[0]-newSh.shape[0])//2, frame.position[1] + (frame.shape[1]-newSh.shape[1])//2)\n",
    "    if colorMatch:\n",
    "        if len(set(sh.m[:,0]).intersection(matrix.m[frame.position[0]:frame.position[0]+\\\n",
    "               frame.shape[0],frame.position[1]])-set([255])) == 0:\n",
    "            newSh.m = np.fliplr(newSh.m)\n",
    "    m = insertShape(m, newSh)\n",
    "    if crop:\n",
    "        bC = matrix.backgroundColor\n",
    "        if np.all(m == bC):\n",
    "            return m\n",
    "        x1, x2, y1, y2 = 0, m.shape[0]-1, 0, m.shape[1]-1\n",
    "        while x1 <= x2 and np.all(m[x1,:] == bC):\n",
    "            x1 += 1\n",
    "        while x2 >= x1 and np.all(m[x2,:] == bC):\n",
    "            x2 -= 1\n",
    "        while y1 <= y2 and np.all(m[:,y1] == bC):\n",
    "            y1 += 1\n",
    "        while y2 >= y1 and np.all(m[:,y2] == bC):\n",
    "            y2 -= 1\n",
    "        if includeFrame:\n",
    "            return(m[x1:x2+1,y1:y2+1])\n",
    "        elif x1+1<x2 and y1+1<y2:\n",
    "            return(m[x1+1:x2,y1+1:y2])\n",
    "    return m\n",
    "    \n",
    "def getBestCountColors(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"countColors\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    outShape = [None]\n",
    "    sliced = [False]\n",
    "    if t.sameOutShape:\n",
    "        outShape.append(t.trainSamples[0].outMatrix.shape)\n",
    "    if t.outSmallerThanIn or not t.sameIOShapes:\n",
    "        sliced += [True]\n",
    "    if t.sameIOShapes:\n",
    "        outShape=['inShape']\n",
    "    for byShape in [True,False]:\n",
    "        for outCol in [t.backgroundColor, t.trainSamples[0].outMatrix.backgroundColor, 0]:\n",
    "            for rotate in range(-1,2):\n",
    "                for flip in [True, False]:\n",
    "                    for ignore in ['max', 'min', None]:\n",
    "                        for sl in sliced:\n",
    "                            for oSh in outShape:\n",
    "                                bestFunction, bestScore = updateBestFunction(t, partial(countColors,\\\n",
    "                                            rotate=rotate,outBackgroundColor=outCol, flip=flip, sliced=sl,\\\n",
    "                                            ignore=ignore, outShape=oSh, byShape=byShape), bestScore, bestFunction)\n",
    "                                bestFunction, bestScore = updateBestFunction(t, partial(countColors,\\\n",
    "                                            rotate=rotate,outBackgroundColor=outCol, flip=flip, sliced=sl,\\\n",
    "                                            ignore=ignore, outShape=oSh, byShape=byShape, sortByColor=1), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def countColors(matrix, outBackgroundColor=-1, outShape=None,ignoreBackground=True,\\\n",
    "                ignore=False, sliced=False, rotate=0, flip=False, byShape=False, sortByColor=False):\n",
    "    \"\"\"\n",
    "    Counts the colors of a given input matrix as pixels. The parameters flip, rotate and sliced determine the layout of the\n",
    "    output pixels. The argument sortByColor sorts the colors by their number, otherwise they are sorted by occurrence. The \n",
    "    parameter ignore allows to ignore either the most or least common colors.\n",
    "    \"\"\"\n",
    "    if byShape:\n",
    "        cc = Counter([sh.color for sh in matrix.shapes if (sh.color != matrix.backgroundColor or not ignoreBackground)])\n",
    "        cc = sorted(cc.items(), key=lambda x: x[1], reverse=True)\n",
    "    else:\n",
    "        cc = matrix.colorCount\n",
    "        #add line to sort geographically\n",
    "        cc = sorted(cc.items(), key=lambda x: x[1], reverse=True)\n",
    "    if sortByColor:\n",
    "        cc = sorted(cc, key=lambda x: x[0])\n",
    "    if outBackgroundColor == -1:\n",
    "        bC = matrix.backgroundColor\n",
    "    else:\n",
    "        bC = outBackgroundColor\n",
    "    if ignoreBackground:\n",
    "        cc = [c for c in cc if c[0] != bC]\n",
    "    if ignore == 'max':\n",
    "        cc = [cc[i] for i in range(1,len(cc))]\n",
    "    elif ignore == 'min':\n",
    "        cc = [c for c in cc if c[1] == max([c[1] for c in cc]) ]\n",
    "    if len(cc) == 0:\n",
    "        return matrix.m.copy()\n",
    "    if outShape == None:\n",
    "        m = np.full((max([c[1] for c in cc]),len(cc)), fill_value=bC)\n",
    "        for j in range(len(cc)):\n",
    "            for i in range(cc[j][1]):\n",
    "                if i >= m.shape[0] or j >= m.shape[1]:\n",
    "                    break\n",
    "                m[i,j] = cc[j][0]\n",
    "    else:\n",
    "        if outShape == 'inShape':\n",
    "            m = np.full(matrix.shape, fill_value=bC)\n",
    "            if m.shape[0] > m.shape[1]:\n",
    "                m = np.rot90(m)\n",
    "            for j in range(len(cc)):\n",
    "                for i in range(cc[j][1]):\n",
    "                    if i >= m.shape[0] or j >= m.shape[1]:\n",
    "                        break\n",
    "                    m[i,j] = cc[j][0]\n",
    "        else:\n",
    "            m = np.full(outShape, fill_value=bC)\n",
    "            cc = [c[0] for c in cc for j in range(c[1])]\n",
    "            i = 0\n",
    "            while i < m.shape[0]:\n",
    "                j = 0\n",
    "                while j < m.shape[1]:\n",
    "                    if i*m.shape[1]+j >= len(cc):\n",
    "                        break\n",
    "                    m[i,j] = cc[i*m.shape[1]+j]\n",
    "                    j += 1 \n",
    "                i += 1 \n",
    "    if sliced:\n",
    "        m = [m[0,:]]\n",
    "    m = np.rot90(m, rotate)\n",
    "    if flip:\n",
    "        m = np.flipud(m)\n",
    "    if outShape == 'inShape' and m.shape !=  matrix.shape:\n",
    "        m = np.rot90(m)\n",
    "    return m\n",
    "\n",
    "def getBestCountShapes(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"countShapes\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    if all([len(s.inMatrix.shapes)>15 for s in t.trainSamples]):\n",
    "            return bestFunction\n",
    "    if t.sameIOShapes:\n",
    "        oSh = 'inShape'\n",
    "    elif t.sameOutShape:\n",
    "        oSh = t.trainSamples[0].outMatrix.shape\n",
    "    else:\n",
    "        oSh = None\n",
    "        \n",
    "    for outC in set([None]).union(set.intersection(*t.outColors)):\n",
    "        for inC in set([-1]).union(set.intersection(*t.inColors)):\n",
    "            for sh in [None] + t.commonInShapes:\n",
    "                for lay in ['h','d']:\n",
    "                    for skip in [True, False]:\n",
    "                        bestFunction, bestScore = updateBestFunction(t, partial(countShapes,color=inC,\\\n",
    "                                        outShape=oSh, lay=lay, outColor=outC, shape=sh, skip=skip), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def countShapes(matrix, color=-1, shape=None, outColor=None, outShape=None, lay='d', skip=False):\n",
    "    \"\"\"\n",
    "    This function returns the count of shapes of a given color or shape, and lays that many pixels. \n",
    "    The pixels layout is specified by the arguments outShape, skip and lay. \n",
    "    \"\"\"\n",
    "    if color < 0:\n",
    "        shc = [sh for sh in matrix.shapes]\n",
    "    else:\n",
    "        shc = [sh for sh in matrix.shapes if sh.color == color]\n",
    "    if shape != None:\n",
    "        shc = [sh for sh in shc if sh == shape]\n",
    "    if outShape == None:\n",
    "        m = np.full((len(shc),len(shc)), fill_value=matrix.backgroundColor)\n",
    "    elif outShape == 'inShape':\n",
    "        m = np.full(matrix.shape, fill_value=matrix.backgroundColor)\n",
    "    else:\n",
    "        m = np.full(outShape, fill_value=matrix.backgroundColor)\n",
    "    if lay == 'd':\n",
    "        for d in range(min(len(shc), m.shape[0], m.shape[1])):\n",
    "            if outColor == None:\n",
    "                m[d,d] = shc[d].color\n",
    "            else:\n",
    "                m[d,d] = outColor\n",
    "    elif lay == 'h':\n",
    "        shc = [sh.color for sh in shc]\n",
    "        if skip:\n",
    "            shc = [[c,matrix.backgroundColor] for c in shc]\n",
    "            shc = [c for p in shc for c in p]\n",
    "        i = 0\n",
    "        while i < m.shape[0]:\n",
    "            j = 0\n",
    "            while j < m.shape[1]:\n",
    "                if i*m.shape[1]+j >= len(shc):\n",
    "                    break\n",
    "                if outColor == None:\n",
    "                    m[i,j] = shc[i*m.shape[1]+j]\n",
    "                elif shc[i*m.shape[1] + j] != matrix.backgroundColor:\n",
    "                    m[i,j] = outColor\n",
    "                j += 1           \n",
    "            i += 1 \n",
    "    return m \n",
    "\n",
    "def getBestSymmetrizeAllShapes(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"symmetrizeAllShapes\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    for cc in set.intersection(*t.inColors).union(set([-1])):\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(symmetrizeAllShapes, targetColor = cc), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(symmetrizeAllShapes, targetColor = cc, byColor=True), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(symmetrizeAllShapes, targetColor = cc, context=True), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(symmetrizeAllShapes, targetColor = cc, context=True, byColor=True), bestScore, bestFunction)\n",
    "\n",
    "    return bestFunction\n",
    "\n",
    "def symmetrizeAllShapes(matrix, diagonal=True, multicolor=True, targetColor=-1,\\\n",
    "                        context=False, lr = True, ud = True, byColor=False):\n",
    "    \"\"\"\n",
    "    Symmetrize all shapes of a given type with respect to specified axes lr or ud. If tagetColor is specified, \n",
    "    then only shapes of that color will be symmetrized. \n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    bC = matrix.backgroundColor\n",
    "    if byColor:\n",
    "        shList = [sh for sh in matrix.shapesByColor if (sh.shape[0]<m.shape[0] and sh.shape[1]<m.shape[1])]\n",
    "    else:    \n",
    "        if not multicolor: \n",
    "            if diagonal:   \n",
    "                shList = [sh for sh in matrix.dShapes]\n",
    "            else:   \n",
    "                shList = [sh for sh in matrix.shapes]\n",
    "        else:\n",
    "            if diagonal: \n",
    "                shList = [sh for sh in matrix.multicolorDShapes]\n",
    "            else:\n",
    "                shList = [sh for sh in matrix.multicolorShapes]\n",
    "    if targetColor > -1:\n",
    "        shList = [sh for sh in shList if hasattr(sh, 'color') and sh.color == targetColor]\n",
    "    for sh in shList:\n",
    "        if context:\n",
    "            shM = m[sh.position[0]:sh.position[0]+sh.shape[0], sh.position[1]:sh.position[1]+sh.shape[1]]\n",
    "        else:\n",
    "            shM = sh.m.copy()\n",
    "        if lr:\n",
    "            shMlr = np.fliplr(shM)\n",
    "            for i,j in np.ndindex(sh.shape):\n",
    "                if shM[i,j] == bC or shM[i,j] == 255:\n",
    "                    shM[i,j] = shMlr[i,j]\n",
    "        if ud:\n",
    "            shMud = np.flipud(shM)\n",
    "            for i,j in np.ndindex(sh.shape):\n",
    "                if shM[i,j] == bC or shM[i,j] == 255:\n",
    "                    shM[i,j] = shMud[i,j]\n",
    "        if context:\n",
    "            m[sh.position[0]:sh.position[0]+sh.shape[0], sh.position[1]:sh.position[1]+sh.shape[1]] = shM\n",
    "        else:\n",
    "            newInsert = copy.deepcopy(sh)\n",
    "            newInsert.m = shM\n",
    "            m = insertShape(m, newInsert)\n",
    "    return m\n",
    "\n",
    "def paintGridLikeBackground(matrix):\n",
    "    \"\"\"\n",
    "    Ignores the grid by paiting it in the background color (most repeated color or second if first coincides with grid color).\n",
    "    In some cases cells have more than one color but it is still worth ignoring the grid. \n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    bC = max(matrix.colorCount,key=matrix.colorCount.get)\n",
    "    if matrix.isGrid:\n",
    "        m[m==matrix.grid.color] = bC\n",
    "    elif matrix.isAsymmetricGrid:\n",
    "        m[m==matrix.asymmetricGrid.color] = bC\n",
    "    return m  \n",
    "\n",
    "def downsizeMode(matrix, newShape, falseColor=None):\n",
    "    \"\"\"\n",
    "    Given a matrix and a shape, this function returns a new matrix with the\n",
    "    given shape. The elements of the return matrix are given by the colors of \n",
    "    each of the submatrices. If a submatrix has more than one color the mode is\n",
    "    chosen.\n",
    "    \"\"\"\n",
    "    if (matrix.shape[0]%newShape[0])!=0 or (matrix.shape[1]%newShape[1])!=0:\n",
    "        return matrix.m.copy()\n",
    "    xBlock = int(matrix.shape[0]/newShape[0])\n",
    "    yBlock = int(matrix.shape[1]/newShape[1])\n",
    "    m = np.full(newShape, matrix.backgroundColor, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(newShape[0], newShape[1]):\n",
    "        colorCount = Counter(matrix.m[i*xBlock: (i+1)*xBlock, j*yBlock: (j+1)*yBlock].flatten())\n",
    "        m[i,j] = max(colorCount, key=colorCount.get)\n",
    "    return m    \n",
    "\n",
    "def getBestColorByPixels(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"colorByPixels\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    delPix = False\n",
    "    if isDeleteTask(t) or t.outSmallerThanIn:\n",
    "        delPix = True\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(colorByPixels, deletePixels=delPix), bestScore, bestFunction)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(colorByPixels, colorMap=True, deletePixels=delPix), bestScore, bestFunction)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(colorByPixels, oneColor=True, deletePixels=delPix), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def colorByPixels(matrix, colorMap=False, oneColor=False, deletePixels=False):\n",
    "    \"\"\"\n",
    "    Attempts to find color changes dictated by pixels. Be it pixels determine the color of the closest shape,\\\n",
    "    be it adjacent pixels determine a color map. \n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    shList = [sh for sh in matrix.shapes if (sh.color != matrix.backgroundColor) and len(sh.pixels)>1]\n",
    "    pixList = [sh for sh in matrix.dShapes if (sh.color != matrix.backgroundColor) and len(sh.pixels)==1]\n",
    "    ogPixList = [p for p in pixList]\n",
    "    if len(shList)==0 or len(pixList)==0 or len(pixList)>15:\n",
    "        return m\n",
    "    if colorMap:\n",
    "        cMap = dict()\n",
    "        seenP = []\n",
    "        for p1 in pixList:\n",
    "            for p2 in pixList:\n",
    "                if abs(p1.position[1]-p2.position[1])==1 and p1.position[0]==p2.position[0] and p1.color not in cMap.keys():\n",
    "                    cMap[p1.color] = p2.color\n",
    "                    seenP.append(p1.position)                   \n",
    "        if deletePixels:\n",
    "            for pix in pixList:\n",
    "                m[pix.position[0], pix.position[1]] = matrix.backgroundColor\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            if m[i,j] in cMap.keys() and (i,j) not in seenP:\n",
    "                m[i,j] = cMap[m[i,j]]     \n",
    "    else:\n",
    "        if oneColor:\n",
    "            cc = Counter([sh.color for sh in pixList])\n",
    "            newC = max(cc, key=cc.get)\n",
    "            if deletePixels:\n",
    "                m[m==newC]=matrix.backgroundColor\n",
    "            m[m!=matrix.backgroundColor]=newC\n",
    "        else:\n",
    "            if len(pixList) < len(shList):\n",
    "                return m\n",
    "            c = 0\n",
    "            nSh = len(shList)\n",
    "            while c < nSh:\n",
    "                minD, newD = 1000, 1000\n",
    "                bestSh, bestPix = None, None\n",
    "                for pix in pixList:\n",
    "                    for i in range(len(pixList)):\n",
    "                        for sh in shList:\n",
    "                            newD = min(np.linalg.norm(np.subtract(pix.position,np.add(p, sh.position))) for p in sh.pixels) \n",
    "                            if newD < minD:\n",
    "                                minD = newD\n",
    "                                bestSh = sh\n",
    "                                bestPix = pix\n",
    "                if bestSh != None:\n",
    "                    for i,j in np.ndindex(bestSh.shape):\n",
    "                        if bestSh.m[i,j] != 255:\n",
    "                            m[bestSh.position[0]+i, bestSh.position[1]+j]=bestPix.color\n",
    "                    c += 1\n",
    "                    shList.remove(bestSh)\n",
    "                    pixList.remove(bestPix)\n",
    "    if deletePixels:\n",
    "        for pix in ogPixList:\n",
    "            m[pix.position] = matrix.backgroundColor\n",
    "    return m  \n",
    "\n",
    "def isDeleteTask(t):\n",
    "    \"\"\"\n",
    "    Check if the task involves replacing pixels by background color, and thus deleting content. \n",
    "    \"\"\"\n",
    "    if hasattr(t, 'colorChanges') and t.backgroundColor in [c[1] for c in t.colorChanges]:\n",
    "        return True\n",
    "    return False\n",
    "\n",
    "def getBestDeleteShapes(t, multicolor=False, diagonal=True):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"deleteShapes\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    attrs = set(['LaSh','SmSh','MoCl','MoCo','PiXl'])\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    for attr in attrs:\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(deleteShapes, diagonal=diagonal, multicolor=multicolor,attributes=set([attr])), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def getDeleteAttributes(t, diagonal=True):\n",
    "    \"\"\"\n",
    "    Given a Task t that involves deleting shapes, this function returns the identifying attributes of the\n",
    "    deleted shapes. \n",
    "    \"\"\"\n",
    "    bC = max(0, t.backgroundColor)\n",
    "    if diagonal:\n",
    "        if t.nCommonInOutDShapes == 0:\n",
    "            return set()\n",
    "        attrs = set.union(*[s.inMatrix.getShapeAttributes(backgroundColor=bC,\\\n",
    "                    singleColor=True, diagonals=True)[s.inMatrix.dShapes.index(sh[0])]\\\n",
    "                    for s in t.trainSamples for sh in s.commonDShapes])\n",
    "        nonAttrs = set()\n",
    "        c = 0\n",
    "        for s in t.trainSamples:\n",
    "            shAttrs = s.inMatrix.getShapeAttributes(backgroundColor=bC, singleColor=True, diagonals=True)\n",
    "            for shi in range(len(s.inMatrix.shapes)):\n",
    "                if any(s.inMatrix.dShapes[shi] == sh2[0] for sh2 in s.commonDShapes) or s.inMatrix.dShapes[shi].color == bC:\n",
    "                    continue\n",
    "                else:\n",
    "                    if c == 0:\n",
    "                        nonAttrs = shAttrs[shi]\n",
    "                        c += 1\n",
    "                    else:\n",
    "                        nonAttrs = nonAttrs.intersection(shAttrs[shi])\n",
    "                        c += 1\n",
    "    else:\n",
    "        if t.nCommonInOutShapes == 0:\n",
    "            return set()\n",
    "        attrs = set.union(*[s.inMatrix.getShapeAttributes(backgroundColor=bC,\\\n",
    "                    singleColor=True, diagonals=False)[s.inMatrix.shapes.index(sh[0])]\\\n",
    "                    for s in t.trainSamples for sh in s.commonShapes])\n",
    "        nonAttrs = set()\n",
    "        c = 0\n",
    "        for s in t.trainSamples:\n",
    "            shAttrs = s.inMatrix.getShapeAttributes(backgroundColor=bC, singleColor=True, diagonals=False)\n",
    "            for shi in range(len(s.inMatrix.shapes)):\n",
    "                if any(s.inMatrix.shapes[shi] == sh2[0] for sh2 in s.commonShapes) or s.inMatrix.shapes[shi].color == bC:\n",
    "                    continue\n",
    "                else:\n",
    "                    if c == 0:\n",
    "                        nonAttrs = shAttrs[shi]\n",
    "                        c += 1\n",
    "                    else:\n",
    "                        nonAttrs = nonAttrs.intersection(shAttrs[shi])\n",
    "                        c += 1\n",
    "    return set(nonAttrs - attrs)\n",
    "\n",
    "def deleteShapes(matrix, attributes, diagonal, multicolor):\n",
    "    \"\"\"\n",
    "    Deletes the shapes of matrix that have specified attributes. \n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    if not multicolor: \n",
    "        if diagonal:   \n",
    "            shList = [sh for sh in matrix.dShapes]\n",
    "        else:   \n",
    "            shList = [sh for sh in matrix.shapes]\n",
    "    else:\n",
    "        if diagonal: \n",
    "            shList = [sh for sh in matrix.multicolorDShapes]\n",
    "        else:\n",
    "            shList = [sh for sh in matrix.multicolorShapes]\n",
    "    attrList = matrix.getShapeAttributes(matrix.backgroundColor, not multicolor, diagonal)\n",
    "    for shi in range(len(shList)):\n",
    "        if len(attrList[shi].intersection(attributes)) > 0:\n",
    "            m = deleteShape(m, shList[shi], matrix.backgroundColor)\n",
    "    return m\n",
    "\n",
    "def getBestArrangeShapes(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"arrangeShapes\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    if hasattr(t, 'outShape'):\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(arrangeShapes, outShape=t.outShape,\\\n",
    "                                                            diagonal=True, multicolor=False), bestScore, bestFunction)\n",
    "    elif t.outIsInMulticolorShapeSize:\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(arrangeShapes, diagonal=True,\\\n",
    "                                                                multicolor=True,outShape='LaSh'), bestScore, bestFunction)\n",
    "    else:\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(arrangeShapes,shByColor=True,outShape='LaSh'), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(arrangeShapes,shByColor=True, fullFrames=True,outShape='LaSh'), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def arrangeShapes (matrix, outShape = None, multicolor=True, diagonal=True, shByColor=False,\\\n",
    "                   fullFrames=False, outDummyMatrix=None, outDummyColor=0):\n",
    "    \"\"\"\n",
    "    Given an output size outShape, attempt to arrange and fit all the shapes of a given typs inside it. Alternatively,\n",
    "    a template matrix outDummyMatrix can be passed, in which case attempts to reproduce its pattern.\n",
    "    \"\"\"\n",
    "    def completeFrames(shape,rotate=False,fill=False):\n",
    "        'version of symmetrize shape intended for frame-like shapes' \n",
    "        m = shape.m.copy()\n",
    "        if m.shape[0]>m.shape[1]:   \n",
    "            newm = np.full((m.shape[0], m.shape[0]), fill_value=255)\n",
    "            sideL = m.shape[0]\n",
    "        else:\n",
    "            newm = np.full((m.shape[1], m.shape[1]), fill_value=255)\n",
    "            sideL = m.shape[1]\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            if m[i,j] != 255:\n",
    "                if newm[i,j] == 255:\n",
    "                    newm[i,j] = m[i,j]\n",
    "                if newm[sideL - j - 1,i] == 255:\n",
    "                    newm[sideL - j - 1,i] = m[i,j]\n",
    "                if newm[sideL - i - 1,sideL - j - 1] == 255:\n",
    "                    newm[sideL - i - 1,sideL - j - 1] = m[i,j]\n",
    "                if newm[j,sideL - i - 1] == 255 and m[i,j]:\n",
    "                    newm[j,sideL - i - 1] = m[i,j]\n",
    "        newSh = copy.deepcopy(shape)\n",
    "        newSh.m = newm\n",
    "        newSh.shape = newm.shape\n",
    "        return newSh\n",
    "    \n",
    "    def tessellateShapes (mat, shL, n, bC, rotation=False):\n",
    "        m = mat.copy()\n",
    "        arrFound = False\n",
    "        rot = 1\n",
    "        \"\"\"\n",
    "        Attempts to tessellate matrix mat with background color bC shapes is list sh.\n",
    "        \"\"\"\n",
    "        if rotation:\n",
    "            rot = 4\n",
    "            if len(shL[n].pixels)==1:\n",
    "                rot = 1\n",
    "            if len(shL[n].pixels)==2:\n",
    "                rot = 2\n",
    "        for x in range(rot):\n",
    "            sh = copy.deepcopy(shL[n])\n",
    "            sh.m = np.rot90(sh.m,x).copy()\n",
    "            sh.shape = sh.m.shape\n",
    "            if mat.shape[0] < sh.shape[0] or mat.shape[1] < sh.shape[1]:\n",
    "                continue\n",
    "            for i, j in np.ndindex(tuple(mat.shape[k] - sh.shape[k] + 1 for k in (0,1))):\n",
    "                if np.all(np.logical_or(m[i: i+sh.shape[0], j: j+sh.shape[1]] == bC, sh.m == 255)):\n",
    "                    for k, l in np.ndindex(sh.shape):\n",
    "                        if sh.m[k,l] != 255:\n",
    "                            m[i+k,j+l] = sh.m[k,l]\n",
    "                    if n == len(shL) - 1:\n",
    "                        return m, True\n",
    "                    m, arrFound = tessellateShapes(m, shL, n+1, bC, rotation)\n",
    "                    if arrFound:\n",
    "                        return m, True\n",
    "                    if not arrFound:\n",
    "                        for k, l in np.ndindex(sh.shape):\n",
    "                            if sh.m[k,l] != 255:\n",
    "                                m[i+k,j+l] = bC\n",
    "        return m, False\n",
    "                                            \n",
    "    if shByColor:\n",
    "        shList = [sh for sh in matrix.shapesByColor]\n",
    "        if fullFrames:\n",
    "            shList = [completeFrames(sh) for sh in matrix.shapesByColor]\n",
    "    else:\n",
    "        if not multicolor: \n",
    "            if diagonal:   \n",
    "                shList = [sh for sh in matrix.dShapes if sh.color != matrix.backgroundColor]\n",
    "            else:   \n",
    "                shList = [sh for sh in matrix.shapes if sh.color != matrix.backgroundColor]\n",
    "        else:\n",
    "            if diagonal: \n",
    "                shList = [sh for sh in matrix.multicolorDShapes]\n",
    "            else:\n",
    "                shList = [sh for sh in matrix.multicolorShapes]\n",
    "    if len(shList) < 2 or len(shList)>7:\n",
    "        return matrix.m.copy()\n",
    "    if outDummyMatrix is None:\n",
    "        shList.sort(key=lambda x: x.shape[0]*x.shape[1], reverse=True)\n",
    "        if outShape == 'LaSh':\n",
    "            outShape = shList[0].shape    \n",
    "        if outShape == None:\n",
    "            outShape = matrix.shape\n",
    "        if all((sh.shape[0]<=outShape[0] and sh.shape[1]<=outShape[1]) for sh in shList) and\\\n",
    "                                sum(len(sh.pixels) for sh in shList) <= outShape[0]*outShape[1]:\n",
    "            m, tessellate = tessellateShapes(np.full(outShape, fill_value=matrix.backgroundColor),shList,\\\n",
    "                                             0,matrix.backgroundColor)\n",
    "            if tessellate:\n",
    "                return m\n",
    "            m, tessellate = tessellateShapes(np.full(outShape, fill_value=matrix.backgroundColor),shList,\\\n",
    "                                             0,matrix.backgroundColor,rotation=True)\n",
    "            if tessellate:\n",
    "                return m\n",
    "    else:\n",
    "        m = np.full(outDummyMatrix.shape,fill_value=outDummyColor)\n",
    "        shC = Counter([sh.shape for sh in shList])\n",
    "        pSh = shC.most_common(1)[0][0]\n",
    "        if pSh[0]>m.shape[0] or pSh[1]>m.shape[1] or m.shape[0]%pSh[0] != 0 or m.shape[1]%pSh[1] != 0:\n",
    "            return matrix.m.copy()\n",
    "        for i, j in np.ndindex(m.shape[0]//pSh[0], m.shape[1]//pSh[1]):\n",
    "            for k, l in np.ndindex(matrix.shape[0]-pSh[0]+1, matrix.shape[1]-pSh[1]+1):\n",
    "                if np.all((matrix.m[k:k+pSh[0],l:l+pSh[1]] == outDummyColor)==(outDummyMatrix[i*pSh[0]:(i+1)*pSh[0],j*pSh[1]:(j+1)*pSh[1]]==0)):\n",
    "                    m[i*pSh[0]:(i+1)*pSh[0],j*pSh[1]:(j+1)*pSh[1]] = matrix.m[k:k+pSh[0],l:l+pSh[1]]\n",
    "                    break\n",
    "        return m\n",
    "    return matrix.m.copy()\n",
    "\n",
    "def getBestLayShapes(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"layShapes\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    outShape = None\n",
    "    if all([len(s.inMatrix.shapes) > 15 for s in t.trainSamples]):\n",
    "        return bestFunction\n",
    "    if t.sameIOShapes:\n",
    "        outShape = 'inShape'\n",
    "    elif hasattr(t, 'outShape'):\n",
    "        outShape = t.outShape\n",
    "    if outShape != None:\n",
    "        for sortBy in ['grid','lr','ud','smallToLarge',set.intersection(*t.inColors)]:\n",
    "            for reverse in [False, True]:\n",
    "                for overlap in [(0,0), (1,1), (-1,-1)]:\n",
    "                    for multicolor in [True, False]:\n",
    "                        for direction in [(1,0), (0,1), (1,1)]:\n",
    "                            bestFunction, bestScore = updateBestFunction(t, partial(layShapes, firstPos=(0,0), diagonal=True, multicolor=multicolor,\\\n",
    "                                                            outShape=outShape, overlap=overlap, direction=direction, sortBy=sortBy, reverse=reverse), bestScore, bestFunction)\n",
    "                        bestFunction, bestScore = updateBestFunction(t, partial(layShapes, firstPos=(1,0), direction=(-1,0), diagonal=True, multicolor=multicolor,\\\n",
    "                                                            outShape=outShape, overlap=overlap, sortBy=sortBy, reverse=reverse), bestScore, bestFunction)\n",
    "                        bestFunction, bestScore = updateBestFunction(t, partial(layShapes, firstPos=(0,1), direction=(0,-1), diagonal=True, multicolor=multicolor,\\\n",
    "                                                            outShape=outShape, overlap=overlap, sortBy=sortBy, reverse=reverse), bestScore, bestFunction)\n",
    "                        bestFunction, bestScore = updateBestFunction(t, partial(layShapes, firstPos=(1,1), direction=(-1,-1), diagonal=True, multicolor=multicolor,\\\n",
    "                                                            outShape=outShape, overlap=overlap, sortBy=sortBy, reverse=reverse), bestScore, bestFunction)\n",
    "    elif t.outSmallerThanIn:\n",
    "         bestFunction, bestScore = updateBestFunction(t, partial(layShapes,completeRect=True,diagonal=True,\\\n",
    "                                         direction=(0,0),sortBy='smallToLarge', reverse=True, outShape='LaSh',multicolor=False), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def layShapes(matrix, firstPos=(0,0), direction=(0,1), overlap=(0,0), outShape='inShape', multicolor=True,\\\n",
    "              diagonal=True, sortBy='lrud', completeRect=False, reverse=True):\n",
    "    \"\"\"\n",
    "    Moves all shapes and lays them in an appropriate way in a matrix of size outShape. \n",
    "    The first shape is position at firstPos, and the succesive shapes are layed in the given direction,\n",
    "    with possible overlap. Shapes can be sorted before by position, size or color. \n",
    "    \"\"\"\n",
    "    def completeRectangles(shape):\n",
    "        \"\"\"\n",
    "        Version of complete rectangles shape intended for frame-like shapes.\n",
    "        \"\"\"\n",
    "        newSh = copy.deepcopy(shape)\n",
    "        newSh.m = np.full(shape.shape, fill_value=shape.color)\n",
    "        newSh.shape = newSh.m.shape\n",
    "        return newSh\n",
    "    m = matrix.m.copy()\n",
    "    if not multicolor: \n",
    "        if diagonal:   \n",
    "            shList = [sh for sh in matrix.dShapes if sh.color != matrix.backgroundColor]\n",
    "        else:   \n",
    "            shList = [sh for sh in matrix.shapes if sh.color != matrix.backgroundColor]\n",
    "    else:\n",
    "        if diagonal: \n",
    "            shList = [sh for sh in matrix.multicolorDShapes]\n",
    "        else:\n",
    "            shList = [sh for sh in matrix.multicolorShapes]\n",
    "    if completeRect and (not multicolor):\n",
    "        shList = [completeRectangles(sh) for sh in shList]\n",
    "    if sortBy == 'smallToLarge':\n",
    "        shList.sort(key=lambda x: len(x.pixels), reverse=reverse)\n",
    "    elif sortBy == 'lr':\n",
    "        shList.sort(key=lambda x: x.position[1], reverse=reverse)\n",
    "    elif sortBy == 'ud':\n",
    "        shList.sort(key=lambda x: x.position[0], reverse=reverse)\n",
    "    elif sortBy == 'grid':\n",
    "        shList.sort(key=lambda x: x.position[0])\n",
    "        newList = []\n",
    "        shList.sort(key=lambda x: x.position[0])\n",
    "        gridD = int(len(shList)**(1/2))\n",
    "        for i in range(gridD):\n",
    "            newList += sorted(shList[i*gridD: (i + 1)*gridD], key=lambda x: x.position[1])\n",
    "        shList = [sh for sh in newList]\n",
    "    elif type(sortBy) == int:\n",
    "        shList.sort(key=lambda x: x.colorCount[sortBy])\n",
    "    if len(shList) == 0:\n",
    "        return m\n",
    "    if outShape == 'inShape':\n",
    "        m = np.full(matrix.shape, fill_value=matrix.backgroundColor)\n",
    "    elif outShape == 'LaSh' and sortBy == 'smallToLarge' and reverse:\n",
    "        m = np.full(shList[0].m.shape, fill_value=matrix.backgroundColor)\n",
    "    else:\n",
    "        m = np.full(outShape, fill_value=matrix.backgroundColor)\n",
    "    shList = [sh for sh in shList if (sh.shape[0] <= m.shape[0] and sh.shape[1] <= m.shape[1])]\n",
    "    startPos = (firstPos[0]*(m.shape[0]), firstPos[1]*(m.shape[1]))\n",
    "    (currentX, currentY) = startPos\n",
    "    for sh in shList:\n",
    "        if currentX + sh.shape[0]*direction[0] > m.shape[0] or currentX + sh.shape[0]*direction[0] < 0:\n",
    "            (currentX, currentY) = (startPos[0], currentY + sh.shape[1] - overlap[1])\n",
    "            if currentY > m.shape[1] or currentY < 0:\n",
    "                return matrix.m.copy()\n",
    "        if currentY + sh.shape[1]*direction[1] > m.shape[1] or currentY + sh.shape[1]*direction[1] < 0:\n",
    "            (currentX, currentY) = (currentX + sh.shape[0] - overlap[0], startPos[1])\n",
    "            if currentX > m.shape[0] or currentX < 0:\n",
    "                return matrix.m.copy()\n",
    "        newInsert = copy.deepcopy(sh)\n",
    "        newInsert.position = (currentX, currentY)\n",
    "        if direction[0] < 0:\n",
    "            newInsert.position = (newInsert.position[0] - sh.shape[0], newInsert.position[1])\n",
    "        if direction[1] < 0:\n",
    "            newInsert.position = (newInsert.position[0], newInsert.position[1] - sh.shape[1])\n",
    "        m = insertShape(m, newInsert)\n",
    "        currentX, currentY = (currentX + (sh.shape[0]- overlap[0])*direction[0] , currentY + (sh.shape[1]- overlap[1])*direction[1])\n",
    "    return m\n",
    "\n",
    "def getBestAlignShapes(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"alignShapes\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    if t.sameIOShapes:\n",
    "        for cc in set.intersection(*t.inColors):\n",
    "            bestFunction, bestScore = updateBestFunction(t, partial(alignShapes, refColor=cc), bestScore, bestFunction)\n",
    "    elif t.outSmallerThanIn:\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(alignShapes, compress=True, crop=True), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(alignShapes, compress=False, crop=True), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def alignShapes(matrix, compress=True, diagonal=True, multicolor=False, refColor=None, crop=False):\n",
    "    \"\"\"\n",
    "    Attempts to align the shapes in matrix. If refColor is passed, then the shapes of that color remain unmoved.\n",
    "    The arguments compress and crop, are called if the output of the task is smaller than the input. \n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    if not multicolor: \n",
    "        if diagonal:   \n",
    "            shList = [sh for sh in matrix.dShapes if sh.color != matrix.backgroundColor]\n",
    "        else:   \n",
    "            shList = [sh for sh in matrix.shapes if sh.color != matrix.backgroundColor]\n",
    "    else:\n",
    "        if diagonal: \n",
    "            shList = [sh for sh in matrix.multicolorDShapes]\n",
    "        else:\n",
    "            shList = [sh for sh in matrix.multicolorShapes]\n",
    "    if len(shList) < 2:\n",
    "        return m\n",
    "    shList.sort(key=lambda x: x.position[1])\n",
    "    arrFound = False\n",
    "    if all(shList[i].position[1]+shList[i].shape[1] <= shList[i+1].position[1] for i in range(len(shList)-1)):\n",
    "        arrFound = (0,1)\n",
    "    if arrFound == False:\n",
    "        shList.sort(key=lambda x: x.position[0])\n",
    "        if all(shList[i].position[0]+shList[i].shape[0] <= shList[i+1].position[0] for i in range(len(shList)-1)):\n",
    "           arrFound = (1,0)\n",
    "    if arrFound == False:\n",
    "        return m  \n",
    "    for sh in shList:\n",
    "        m = deleteShape(m, sh, matrix.backgroundColor)\n",
    "    (currentX, currentY) = (0, 0)\n",
    "    if refColor != None and not crop:\n",
    "        refPos = -1\n",
    "        for sh in shList:\n",
    "            if sh.color == refColor:\n",
    "                refPos = sh.position\n",
    "                break\n",
    "        if refPos == -1:\n",
    "            return matrix.m.copy()\n",
    "        for sh in shList:\n",
    "            newInsert = copy.deepcopy(sh)\n",
    "            if arrFound == (0,1):\n",
    "                newInsert.position = (refPos[0], sh.position[1])\n",
    "            elif arrFound == (1,0):\n",
    "                newInsert.position = (sh.position[0], refPos[1])\n",
    "            m = insertShape(m, newInsert)\n",
    "    else:\n",
    "        for sh in shList:\n",
    "            newInsert = copy.deepcopy(sh)\n",
    "            if compress:\n",
    "                newInsert.position = (currentX, currentY)\n",
    "                (currentX, currentY) = (currentX + sh.shape[0]*arrFound[0], currentY + sh.shape[1]*arrFound[1])\n",
    "            else:\n",
    "                newInsert.position = (0,0)\n",
    "            m = insertShape(m, newInsert)\n",
    "        \n",
    "    if crop:\n",
    "        bC = matrix.backgroundColor\n",
    "        if np.all(m == bC):\n",
    "            return m\n",
    "        x1, x2, y1, y2 = 0, m.shape[0]-1, 0, m.shape[1]-1\n",
    "        while x1 <= x2 and np.all(m[x1,:] == bC):\n",
    "            x1 += 1\n",
    "        while x2 >= x1 and np.all(m[x2,:] == bC):\n",
    "            x2 -= 1\n",
    "        while y1 <= y2 and np.all(m[:,y1] == bC):\n",
    "            y1 += 1\n",
    "        while y2 >= y1 and np.all(m[:,y2] == bC):\n",
    "            y2 -= 1\n",
    "        return(m[x1:x2+1,y1:y2+1])\n",
    "    else:\n",
    "        return m\n",
    "\n",
    "def isReplicateTask(t):\n",
    "    \"\"\"\n",
    "    Identify if a given task involves the action of replicating shapes. In this case, return a list of \n",
    "    booleans indicating the types of shapes involved. \n",
    "    \"\"\"\n",
    "    if all(any(sh[2] > 1 for sh in s.commonMulticolorDShapes) for s in t.trainSamples):\n",
    "        return [True, True, True]\n",
    "    elif all(any(sh[2] > 1 for sh in s.commonMulticolorShapes) for s in t.trainSamples):\n",
    "        return [True, True, False]\n",
    "    elif all(any(sh[2] > 1 for sh in s.commonShapes) for s in t.trainSamples):\n",
    "        return [True, False, False]\n",
    "    elif all(any(sh[2] > 1 for sh in s.commonDShapes) for s in t.trainSamples):\n",
    "        return [True, False, True]\n",
    "    return [False]\n",
    "\n",
    "def getBestReplicateShapes(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"replicateShapes\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    isReplicateParams = isReplicateTask(t)\n",
    "    deleteOriginal = False\n",
    "    multicolor = True\n",
    "    diagonal = True\n",
    "    if isReplicateParams[0]:\n",
    "        multicolor = isReplicateParams[1]\n",
    "        diagonal = isReplicateParams[2]\n",
    "    if isDeleteTask(t):\n",
    "        deleteOriginal = True\n",
    "    \n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, diagonal=diagonal, multicolor=multicolor,deleteOriginal=deleteOriginal,\\\n",
    "                                                            anchorType='subframe'), bestScore, bestFunction)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, diagonal=diagonal, multicolor=multicolor,deleteOriginal=deleteOriginal,\\\n",
    "                                                            anchorType='subframe', allCombs=False,attributes=set(['MoCl'])), bestScore, bestFunction)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, diagonal=diagonal, multicolor=multicolor,deleteOriginal=deleteOriginal,\\\n",
    "                                                            anchorType='subframe', allCombs=False,scale=True,attributes=set(['MoCl'])), bestScore, bestFunction)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, diagonal=diagonal, multicolor=multicolor,deleteOriginal=deleteOriginal,\\\n",
    "                                                            anchorType='subframe', allCombs=True,attributes=set(['MoCl'])), bestScore, bestFunction)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, diagonal=diagonal, multicolor=multicolor,deleteOriginal=deleteOriginal,\\\n",
    "                                                            anchorType='subframe', allCombs=True,scale=True,attributes=set(['MoCl'])), bestScore, bestFunction)\n",
    "    \n",
    "    \"\"\"\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes,diagonal=diagonal, multicolor=False, anchorType='subframe', allCombs=False,\\\n",
    "                                                                adoptAnchorColor=True), bestScore, bestFunction)\n",
    "    if isReplicateParams[0]:\n",
    "        if bestScore == 0:\n",
    "            return bestFunction\n",
    "        for attributes in [set(['MoCl'])]:\n",
    "            cC = Counter([cc[0] for cc in t.colorChanges])\n",
    "            cc = max(cC, key=cC.get)\n",
    "            bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, attributes=attributes, diagonal=diagonal, multicolor=multicolor, anchorType='all', anchorColor=cc,\\\n",
    "                                    mirror=None, rotate=0, allCombs=True, scale=False, deleteOriginal=deleteOriginal), bestScore, bestFunction)\n",
    "            bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, attributes=attributes, diagonal=diagonal, multicolor=multicolor, anchorType='all', anchorColor=cc,\\\n",
    "                                    mirror=None, rotate=0, allCombs=True, scale=False, deleteOriginal=deleteOriginal), bestScore, bestFunction)\n",
    "            if bestScore == 0:\n",
    "                return bestFunction\n",
    "            for mirror in [None, 'lr', 'ud']:\n",
    "                for rotate in range(0, 4):\n",
    "                    bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, attributes=attributes, diagonal=diagonal, multicolor=multicolor, anchorType='all', anchorColor=cc,\\\n",
    "                                    mirror=mirror, rotate=rotate, allCombs=False, scale=False, deleteOriginal=deleteOriginal), bestScore, bestFunction)\n",
    "                    bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, attributes=attributes, diagonal=diagonal, multicolor=multicolor, anchorType='all', anchorColor=cc,\\\n",
    "                                    mirror=mirror, rotate=rotate, allCombs=False, scale=True, deleteOriginal=deleteOriginal), bestScore, bestFunction)\n",
    "                    if bestScore == 0:      \n",
    "                        return bestFunction\n",
    "    \"\"\"\n",
    "                                            \n",
    "    cC = Counter([cc[0] for cc in t.colorChanges])\n",
    "    cc = max(cC, key=cC.get)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, diagonal=False, multicolor=multicolor, anchorType='all', anchorColor=cc,\\\n",
    "                        allCombs=False, scale=False, deleteOriginal=deleteOriginal,perfectFit=True), bestScore, bestFunction)\n",
    "    for attributes in [set(['UnCo'])]:    \n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, attributes=attributes, diagonal=True, multicolor=False, anchorType='all', anchorColor=cc,\\\n",
    "                        allCombs=True, scale=False, deleteOriginal=deleteOriginal), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, attributes=attributes, diagonal=True, multicolor=False, anchorType='all', anchorColor=cc,\\\n",
    "                        allCombs=False, scale=False, deleteOriginal=deleteOriginal, perfectFit=True), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, attributes=attributes, diagonal=True, multicolor=False, anchorType='all', anchorColor=cc,\\\n",
    "                        allCombs=False, scale=True, deleteOriginal=deleteOriginal, perfectFit=False), bestScore, bestFunction)\n",
    "   \n",
    "    \"\"\"\n",
    "    if t.hasPartialFrame:\n",
    "        for attributes in [set(['IsRef'])]:    \n",
    "            bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, attributes=attributes, diagonal=True, multicolor=False, anchorType='all', anchorColor=cc,\\\n",
    "                            allCombs=True, scale=False, deleteOriginal=deleteOriginal), bestScore, bestFunction)\n",
    "            bestFunction, bestScore = updateBestFunction(t, partial(replicateShapes, attributes=attributes, diagonal=True, multicolor=False, anchorType='all', anchorColor=cc,\\\n",
    "                            allCombs=False, scale=False, deleteOriginal=deleteOriginal, perfectFit=True), bestScore, bestFunction)\n",
    "    \"\"\"        \n",
    "    return bestFunction\n",
    "\n",
    "def replicateShapes(matrix, attributes=None, diagonal=False, multicolor=True, anchorType=None, anchorColor=0,\\\n",
    "                    mirror=None, rotate=0, allCombs=False, scale=False, deleteOriginal=False, perfectFit=False,\n",
    "                    adoptAnchorColor=False):\n",
    "    \"\"\"\n",
    "    Replicates shapes at target locations. The arguments diagonal, multicolor, attributes determine the shapes to be\n",
    "    replicated. The arguments mirror, rotate, allCombs and scale modify the shape before it is replicated. \n",
    "    The arguments anchorType, anchorColor and perfectFit determine the places where the shapes should be replicated.\n",
    "    The options adoptAnchorColor, modify the shape once it has been replicated, either by changing its color. The option\n",
    "    deleteOriginal, deletes the original shape after replicating it. \n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    #first find the shape or shapes to replicate\n",
    "    if diagonal:\n",
    "        if multicolor:\n",
    "            shList = matrix.multicolorDShapes\n",
    "        else:\n",
    "            shList = matrix.dShapes\n",
    "    else:\n",
    "        if multicolor:\n",
    "            shList = matrix.multicolorShapes\n",
    "        else:\n",
    "            shList = matrix.shapes\n",
    "    if attributes != None:\n",
    "        repList = []\n",
    "        attrList = matrix.getShapeAttributes(backgroundColor=matrix.backgroundColor,\\\n",
    "                                             singleColor=not multicolor, diagonals=diagonal)\n",
    "        for shi in range(len(shList)):\n",
    "            if len(attrList[shi].intersection(attributes)) == 1:\n",
    "                repList.append(shList[shi])\n",
    "        if len(repList) == 0:\n",
    "            return m\n",
    "    else:\n",
    "        if multicolor:\n",
    "            repList = [sh for sh in shList if (len(sh.pixels)>1 and not sh.isSquare)]\n",
    "        else:\n",
    "            repList = [sh for sh in shList if (sh.color != matrix.backgroundColor and not sh.isSquare)]\n",
    "    delList = [sh for sh in repList]\n",
    "    if len(repList) > 10:\n",
    "        return m\n",
    "    #apply transformations to replicating shapes\n",
    "    if allCombs:\n",
    "        newList = []\n",
    "        for repShape in repList:\n",
    "            for r in range(0,4):\n",
    "                mr, mrM = np.rot90(repShape.m.copy(), r), np.rot90(repShape.m[::-1,::].copy(), r)\n",
    "                newRep, newRepM = copy.deepcopy(repShape), copy.deepcopy(repShape)\n",
    "                newRep.m, newRepM.m = mr, mrM\n",
    "                newRep.shape, newRepM.shape = mr.shape, mrM.shape\n",
    "                newList.append(newRep)\n",
    "                newList.append(newRepM)\n",
    "        repList = [sh for sh in newList]           \n",
    "    elif mirror == 'lr' and len(repList) == 1:\n",
    "        newRep = copy.deepcopy(repList[0])\n",
    "        newRep.m = repList[0].m[::,::-1]\n",
    "        repList = [newRep]\n",
    "    elif mirror == 'ud' and len(repList) == 1:\n",
    "        newRep = copy.deepcopy(repList[0])\n",
    "        newRep.m = repList[0].m[::-1,::]\n",
    "        repList = [newRep]\n",
    "    elif rotate > 0 and len(repList) == 1:\n",
    "        newRep = copy.deepcopy(repList[0])\n",
    "        newRep.m = np.rot90(repList[0].m,rotate)\n",
    "        newRep.shape = newRep.m.shape\n",
    "        repList = [newRep]\n",
    "    if scale == True:\n",
    "        newRepList=[]\n",
    "        for repShape in repList:\n",
    "            for sc in range(4,0,-1):\n",
    "                newRep = copy.deepcopy(repShape)\n",
    "                newRep.m = np.repeat(np.repeat(repShape.m, sc, axis=1), sc, axis=0)\n",
    "                newRep.shape = newRep.m.shape\n",
    "                newRep.pixels = set([(i,j) for i,j in np.ndindex(newRep.m.shape) if newRep.m[i,j]!=255])\n",
    "                newRepList.append(newRep)\n",
    "        repList = [sh for sh in newRepList]\n",
    "    repList.sort(key=lambda x: len(x.pixels), reverse=True)\n",
    "    if anchorType == 'subframe' and scale:\n",
    "        repList.sort(key=lambda x: len(x.pixels))\n",
    "    #then find places to replicate\n",
    "    if anchorType == 'all':\n",
    "        seenM = np.zeros(m.shape, dtype=int)\n",
    "        for repSh in repList:\n",
    "            if np.all(np.logical_or(repSh.m==255,repSh.m==anchorColor)):\n",
    "                continue\n",
    "            for j in range(matrix.shape[1] - repSh.shape[1]+1):\n",
    "                for i in range(matrix.shape[0] - repSh.shape[0]+1):\n",
    "                    if np.all(np.logical_or(m[i:i+repSh.shape[0],j:j+repSh.shape[1]]==anchorColor,repSh.m==255))\\\n",
    "                                                    and np.all(seenM[i:i+repSh.shape[0],j:j+repSh.shape[1]]==0):\n",
    "                        if perfectFit:\n",
    "                            surrPixList = set([(i+p[0]+1, j+p[1]) for p in repSh.pixels]+[(i+p[0], j+p[1]+1) for p in repSh.pixels]\\\n",
    "                                               +[(i+p[0]-1, j+p[1]) for p in repSh.pixels]+[(i+p[0], j+p[1]-1) for p in repSh.pixels])\n",
    "                            surrPixList = surrPixList - set([(i+p[0], j+p[1]) for p in repSh.pixels])\n",
    "                            surrPixList = set([p for p in surrPixList if (p[0]>=0 and p[1]>=0 and p[0]<m.shape[0] and p[1]<m.shape[1])])\n",
    "                            if len(set([m[p[0],p[1]] for p in surrPixList]))==1 and m[list(surrPixList)[0]]!=anchorColor: \n",
    "                                newInsert = copy.deepcopy(repSh)\n",
    "                                newInsert.position = (i, j)\n",
    "                                m = insertShape(m, newInsert)\n",
    "                                for l, k in np.ndindex(repSh.shape):\n",
    "                                    seenM[i+l,j+k]=1\n",
    "                        else:\n",
    "                            newInsert = copy.deepcopy(repSh)\n",
    "                            newInsert.position = (i, j)\n",
    "                            m = insertShape(m, newInsert)\n",
    "                            for l, k in np.ndindex(repSh.shape):\n",
    "                                seenM[i+l,j+k]=1\n",
    "    elif anchorType == 'subframe':\n",
    "        seenM = np.zeros(matrix.shape)\n",
    "        if attributes != None:\n",
    "            for repSh in delList:\n",
    "                for i,j in np.ndindex(repSh.shape):\n",
    "                    seenM[i+repSh.position[0],j+repSh.position[1]]=1\n",
    "        for sh2 in shList:\n",
    "            score, bestScore= 0, 0\n",
    "            bestSh = None\n",
    "            for repSh in repList:\n",
    "                if adoptAnchorColor:\n",
    "                    if sh2.isSubshape(repSh,sameColor=False,rotation=False,mirror=False) and len(sh2.pixels)<len(repSh.pixels):\n",
    "                        for x in range((repSh.shape[0]-sh2.shape[0])+1):\n",
    "                            for y in range((repSh.shape[1]-sh2.shape[1])+1):\n",
    "                                mAux = m[max(sh2.position[0]-x, 0):min(sh2.position[0]-x+repSh.shape[0], m.shape[0]), max(sh2.position[1]-y, 0):min(sh2.position[1]-y+repSh.shape[1], m.shape[1])]\n",
    "                                shAux = repSh.m[max(0, x-sh2.position[0]):min(repSh.shape[0],m.shape[0]+x-sh2.position[0]),max(0, y-sh2.position[1]):min(repSh.shape[1],m.shape[1]+y-sh2.position[1])]\n",
    "                                seenAux = seenM[max(sh2.position[0]-x, 0):min(sh2.position[0]-x+repSh.shape[0], m.shape[0]), max(sh2.position[1]-y, 0):min(sh2.position[1]-y+repSh.shape[1], m.shape[1])]\n",
    "                                if np.all(np.logical_or(shAux!=255, mAux == matrix.backgroundColor)) and np.all(seenAux==0):\n",
    "                                    score = np.count_nonzero(mAux!=matrix.backgroundColor)\n",
    "                                    if score > bestScore:\n",
    "                                        bestScore = score\n",
    "                                        bestX, bestY = sh2.position[0]-x, sh2.position[1]-y\n",
    "                                        bestSh = copy.deepcopy(repSh)\n",
    "                else:\n",
    "                    if sh2.isSubshape(repSh,sameColor=True,rotation=False,mirror=False) and len(sh2.pixels)<len(repSh.pixels):\n",
    "                        for x in range((repSh.shape[0]-sh2.shape[0])+1):\n",
    "                            for y in range((repSh.shape[1]-sh2.shape[1])+1):\n",
    "                                mAux = m[max(sh2.position[0]-x, 0):min(sh2.position[0]-x+repSh.shape[0], m.shape[0]), max(sh2.position[1]-y, 0):min(sh2.position[1]-y+repSh.shape[1], m.shape[1])]\n",
    "                                shAux = repSh.m[max(0, x-sh2.position[0]):min(repSh.shape[0],m.shape[0]+x-sh2.position[0]),max(0, y-sh2.position[1]):min(repSh.shape[1],m.shape[1]+y-sh2.position[1])]\n",
    "                                seenAux = seenM[max(sh2.position[0]-x, 0):min(sh2.position[0]-x+repSh.shape[0], m.shape[0]), max(sh2.position[1]-y, 0):min(sh2.position[1]-y+repSh.shape[1], m.shape[1])]\n",
    "                                if np.all(np.logical_or(mAux==shAux, mAux == matrix.backgroundColor)) and np.all(seenAux==0):\n",
    "                                    score = np.count_nonzero(mAux==shAux)\n",
    "                                    if score > bestScore:\n",
    "                                        bestScore = score\n",
    "                                        bestX, bestY = sh2.position[0]-x, sh2.position[1]-y\n",
    "                                        bestSh = copy.deepcopy(repSh)\n",
    "            if bestSh != None:\n",
    "                for i,j in np.ndindex(bestSh.shape):\n",
    "                    if i+bestX>=0 and i+bestX<seenM.shape[0] and j+bestY>=0 and j+bestY<seenM.shape[1]:\n",
    "                        seenM[i+bestX,j+bestY]=1\n",
    "                newInsert = copy.deepcopy(bestSh)\n",
    "                if adoptAnchorColor:\n",
    "                    newInsert.m[newInsert.m!=255]=sh2.color\n",
    "                newInsert.position = (bestX, bestY)\n",
    "                newInsert.shape = newInsert.m.shape\n",
    "                m = insertShape(m, newInsert)        \n",
    "    if deleteOriginal:\n",
    "        for sh in delList:\n",
    "            m = deleteShape(m, sh, matrix.backgroundColor)\n",
    "    return(m)\n",
    "    \n",
    "def getBestReplicateOneShape(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"replicateOneShape\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    if t.outSmallerThanIn or (not t.sameIOShapes):\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(replicateOneShape, lay='pixelwise', reshape=True,\\\n",
    "                                                                multicolor=True), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(replicateOneShape, lay='horizontal', reshape=True,\\\n",
    "                                                                multicolor=True), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(replicateOneShape, lay='vertical', reshape=True,\\\n",
    "                                                                multicolor=True), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(replicateOneShape, lay='pixelwise', reshape=True,\\\n",
    "                                                                multicolor=False, paintLikePix=True), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(replicateOneShape, lay='horizontal', reshape=True,\\\n",
    "                                                                multicolor=False, paintLikePix=True), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(replicateOneShape, lay='vertical', reshape=True,\\\n",
    "                                                                multicolor=False, paintLikePix=True), bestScore, bestFunction)\n",
    "    deleteO = [False]\n",
    "    deleteA = [False]\n",
    "    if isDeleteTask(t):\n",
    "        deleteO += [True]\n",
    "        deleteA += [True]\n",
    "    if t.sameIOShapes:\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(replicateOneShape, lay='pixelwise', multicolor=False,\\\n",
    "                                                                paintLikePix=True), bestScore, bestFunction)\n",
    "        for deleteOriginal in deleteO:\n",
    "            for deleteAnchor in deleteA:\n",
    "                bestFunction, bestScore = updateBestFunction(t, partial(replicateOneShape, deleteOriginal=deleteOriginal,\\\n",
    "                                                                deleteAnchor=deleteAnchor), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def replicateOneShape(matrix, diagonal=True, multicolor=True, deleteOriginal=False,\\\n",
    "                      deleteAnchor=False, lay=False, reshape=False,overlap=0, paintLikePix=False):\n",
    "    \"\"\"\n",
    "    To be called when the task involved replicating one shape. The target positions are specified by pixels. Pixels may\n",
    "    specify a color change with paitLikePix, or can be deleted after with deleteAnchor. The arguments reshape and lay modify\n",
    "    the output by moving the shapes or changing the output shape. \n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    #first find the shape or shapes to replicate\n",
    "    if diagonal:\n",
    "        if multicolor:\n",
    "            shList = [sh for sh in matrix.multicolorDShapes if len(sh.pixels)>1]\n",
    "        else:\n",
    "            shList = [sh for sh in matrix.dShapes if (len(sh.pixels)>1 and sh.color != matrix.backgroundColor)]\n",
    "    else:\n",
    "        if multicolor:\n",
    "            shList = [sh for sh in matrix.multicolorShapes if len(sh.pixels)>1]\n",
    "        else:\n",
    "            shList = [sh for sh in matrix.shapes if (len(sh.pixels)>1 and sh.color != matrix.backgroundColor)]\n",
    "    pixList = matrix.isolatedPixels#[pix for pix in matrix.dShapes if len(pix.pixels)==1]\n",
    "    if len(shList) != 1 or len(pixList) == 0:\n",
    "        return m\n",
    "    repSh = shList[0]\n",
    "    if lay != False:\n",
    "        if lay == 'pixelwise':\n",
    "            if len(pixList) < 2:\n",
    "                return m\n",
    "            if len(set([p.position[0] for p in pixList])) == 1:\n",
    "                pixList.sort(key=lambda x: x.position[1])\n",
    "                steps = [(0, pixList[i].position[1] - pixList[i-1].position[1] - 1) for i in range(1,len(pixList))]\n",
    "                direction = (0, (pixList[1].position[1] - pixList[0].position[1] - 1)//abs(pixList[1].position[1] - pixList[0].position[1] - 1))\n",
    "            elif len(set([p.position[1] for p in pixList])) == 1:\n",
    "                pixList.sort(key=lambda x: x.position[0])\n",
    "                steps = [(pixList[i].position[0] - pixList[i-1].position[0] - 1, 0) for i in range(1,len(pixList))]\n",
    "                direction = ((pixList[1].position[0] - pixList[0].position[0] - 1)//abs(pixList[1].position[0] - pixList[0].position[0] - 1), 0)\n",
    "            else:\n",
    "                return m\n",
    "            if paintLikePix and repSh.color == pixList[-1].color:\n",
    "                    steps = steps[::-1]\n",
    "                    steps = [(-st[0], -st[1]) for st in steps]\n",
    "                    direction = (-direction[0], -direction[1])\n",
    "                    pixList = pixList[::-1]\n",
    "            if reshape:\n",
    "                m = np.full((repSh.shape[0]*(1 + (len(pixList)-1)*abs(direction[0])), repSh.shape[1]*(1 + (len(pixList)-1)*abs(direction[1]))),\\\n",
    "                            fill_value = matrix.backgroundColor)\n",
    "                for i in range(len(pixList)):\n",
    "                    newInsert = copy.deepcopy(repSh)\n",
    "                    newInsert.position = (i*repSh.shape[0]*direction[0], i*repSh.shape[1]*direction[1])\n",
    "                    if paintLikePix:\n",
    "                        newInsert.m[repSh.m == repSh.color] = pixList[i].color\n",
    "                    m = insertShape(m, newInsert)\n",
    "                deleteOriginal = False\n",
    "            else:\n",
    "                pos = repSh.position\n",
    "                for (p,i) in zip(steps, [j for j in range(1,len(pixList))]):\n",
    "                    pos = (pos[0] + direction[0]*repSh.shape[0] + p[0], pos[1] + direction[1]*repSh.shape[1] + p[1])\n",
    "                    newInsert = copy.deepcopy(repSh)\n",
    "                    newInsert.position = pos\n",
    "                    if paintLikePix:\n",
    "                        newInsert.m[repSh.m == repSh.color] = pixList[i].color\n",
    "                    m = insertShape(m, newInsert)\n",
    "        elif lay == 'horizontal': \n",
    "            m = np.full((repSh.shape[0], len(pixList)*repSh.shape[1]), fill_value = matrix.backgroundColor)\n",
    "            deleteOriginal = False\n",
    "            for i in range(len(pixList)):\n",
    "                newInsert = copy.deepcopy(repSh)\n",
    "                newInsert.position = (0, i*repSh.shape[1])\n",
    "                m = insertShape(m, newInsert)\n",
    "        elif lay == 'vertical': \n",
    "            m = np.full((len(pixList)*repSh.shape[0], repSh.shape[1]), fill_value = matrix.backgroundColor)\n",
    "            deleteOriginal = False\n",
    "            for i in range(len(pixList)):\n",
    "                newInsert = copy.deepcopy(repSh)\n",
    "                newInsert.position = (i*repSh.shape[0], 0)\n",
    "                m = insertShape(m, newInsert)\n",
    "    else:\n",
    "        for pix in pixList:\n",
    "            if (pix.position[0] >= repSh.position[0]) and (pix.position[1] >= repSh.position[1]) \\\n",
    "                and (pix.position[0] < repSh.position[0]+repSh.shape[0]) and (pix.position[1] < repSh.position[1]+repSh.shape[1]):\n",
    "                continue\n",
    "            newInsert = copy.deepcopy(repSh)\n",
    "            if pix.m[0,0] in repSh.m:\n",
    "                newInsert = copy.deepcopy(repSh)\n",
    "                for i, j in np.ndindex(repSh.shape):\n",
    "                    if repSh.m[i,j] == pix.m[0,0]:\n",
    "                        newInsert.position = (pix.position[0]-i, pix.position[1]-j)\n",
    "                        break\n",
    "            else:\n",
    "                newInsert.position = (pix.position[0] - (repSh.shape[0]-1)//2, pix.position[1] - (repSh.shape[1]-1)//2)\n",
    "            m = insertShape(m, newInsert)\n",
    "            if deleteAnchor:\n",
    "                m = deleteShape(m, pix, matrix.backgroundColor)\n",
    "    if deleteOriginal:\n",
    "        m = deleteShape(m, repSh, matrix.backgroundColor)\n",
    "    return m\n",
    "\n",
    "def getBestMoveToPanel(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"moveToPanel\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    for fit in [True, False]:\n",
    "        for igPan in [True, False]:\n",
    "            for uniq in [True, False]:\n",
    "                bestFunction, bestScore = updateBestFunction(t, partial(moveToPanel,fit=fit,\\\n",
    "                                                ignorePanel=igPan), bestScore, bestFunction)\n",
    "    return bestFunction\n",
    "\n",
    "def moveToPanel(matrix, diagonal=True, fit=False, ignorePanel=False, cropPanel=True, uniq=True):\n",
    "    \"\"\"\n",
    "    Moves shapes and places them inside a larger square shape panel. The arguments ignorePanel and cropPanel\n",
    "    modify the output in the cases where the output matrix has a different shape. \n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    shList = [sh for sh in matrix.multicolorDShapes if len(sh.pixels)>1]\n",
    "    if len(shList) < 2 or len(shList) > 8:\n",
    "        return m\n",
    "    shList.sort(key=lambda x: x.shape[0]*x.shape[1],reverse=True)\n",
    "    panel = shList[0]\n",
    "    shList = shList[1:]\n",
    "    if fit and hasattr(panel, 'color'):\n",
    "        pC = panel.color\n",
    "        for sh in shList:\n",
    "            found = False\n",
    "            for x in range(4):\n",
    "                rotSh = np.rot90(sh.m, x).copy()\n",
    "                if panel.shape[0]-rotSh.shape[0]+1<0 or panel.shape[1]-rotSh.shape[1]+1<0:\n",
    "                    continue\n",
    "                for i, j in np.ndindex(panel.shape[0]-rotSh.shape[0]+1,panel.shape[1]-rotSh.shape[1]+1):\n",
    "                    if np.all((rotSh==pC)==(panel.m[i:i+rotSh.shape[0],j:j+rotSh.shape[1]]==255)):\n",
    "                        newInsert = copy.deepcopy(sh)\n",
    "                        newInsert.m = rotSh\n",
    "                        newInsert.shape = rotSh.shape\n",
    "                        newInsert.position = (panel.position[0]+i, panel.position[1]+j)\n",
    "                        m = insertShape(m, newInsert)\n",
    "                        found = True\n",
    "                        break\n",
    "                if found:\n",
    "                    break\n",
    "    else:\n",
    "        pixList = [pix for pix in matrix.dShapes if len(pix.pixels)==1]\n",
    "        pixList = [pix for pix in pixList if all(pix.position[i]>=panel.position[i]\\\n",
    "                                        and pix.position[i]<panel.position[i]+panel.shape[i] for i in [0,1])]\n",
    "        if len(pixList)==0:\n",
    "            return m\n",
    "        newInsList = []\n",
    "        for pix in pixList:\n",
    "            for sh in shList:\n",
    "                if pix.m[0,0] in sh.m:\n",
    "                    newInsert = copy.deepcopy(sh)\n",
    "                    if sh.nColors == 1:\n",
    "                        newInsert.position = (pix.position[0] - (sh.shape[0]-1)//2, pix.position[1] - (sh.shape[1]-1)//2)\n",
    "                    else:\n",
    "                        for i, j in np.ndindex(sh.shape):\n",
    "                            if sh.m[i,j] == pix.m[0,0]:\n",
    "                                newInsert.position = (pix.position[0]-i, pix.position[1]-j) \n",
    "                    newInsList.append(newInsert)\n",
    "                    if uniq:\n",
    "                        break\n",
    "            if ignorePanel:\n",
    "                m = deleteShape(m, panel, matrix.backgroundColor)\n",
    "            for sh in newInsList:\n",
    "                m = insertShape(m, sh)\n",
    "    if cropPanel:\n",
    "        return m[panel.position[0]:panel.position[0]+panel.shape[0],\\\n",
    "                 panel.position[1]:panel.position[1]+panel.shape[1]]\n",
    "    return m\n",
    "        \n",
    "def printShapes(matrices, base=0, backgroundColor=0):\n",
    "    \"\"\"\n",
    "    This function returns the result of printing one matrices on the other.\n",
    "    The matrices are shape matrices and may contain 255.\n",
    "    \"\"\"\n",
    "    if base == 1:\n",
    "        matrices = matrices[::-1]\n",
    "    m = np.zeros(matrices[0].shape, dtype=int)\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        if matrices[0][i,j] != 255 and matrices[0][i,j] != backgroundColor:\n",
    "            if matrices[1][i,j] != 255:\n",
    "                m[i,j] = matrices[1][i,j]\n",
    "            else:\n",
    "                m[i,j] = matrices[0][i,j]\n",
    "        else:\n",
    "            m[i,j] = backgroundColor\n",
    "    return m \n",
    "\n",
    "def multiplyMatrices(matrices, outShape=None, background=0, base=0, color=0):\n",
    "    \"\"\"\n",
    "    Copy m1 matrix into every pixel of m2 if it is not background. Output has shape\n",
    "    m1.shape*m2.shape. base and color are arguments to swap m1 and m2 and to choose\n",
    "    color appropriately\n",
    "    \"\"\"\n",
    "    if base == 1:\n",
    "        matrices = matrices[::-1]\n",
    "    m1 = matrices[0].copy()\n",
    "    m2 = matrices[1].copy()\n",
    "    s1 = m1.shape\n",
    "    s2 = m2.shape\n",
    "    m = np.full((m1.shape[0]*m2.shape[0],m1.shape[1]*m2.shape[1]), fill_value=background)\n",
    "    if color == 0:\n",
    "        for i,j in np.ndindex(s1):\n",
    "            if m1[i,j] != background:\n",
    "                m[i*s2[0]:(i+1)*s2[0], j*s2[1]:(j+1)*s2[1]] = m2\n",
    "    else:\n",
    "       for i,j in np.ndindex(s1):\n",
    "            if m1[i,j] != background:\n",
    "                for k,l in np.ndindex(m2.shape):\n",
    "                    if m2[k,l] != background:\n",
    "                        m[i*s2[0]+k, j*s2[1]+l] = m1[i,j] \n",
    "    return m   \n",
    "\n",
    "def overlapSubmatrices(matrix, colorHierarchy, shapeFactor=None):\n",
    "    \"\"\"\n",
    "    This function returns the result of overlapping all submatrices of a given\n",
    "    shape factor pixelswise with a given color hierarchy. Includes option to overlap\n",
    "    all grid cells.     \n",
    "    \"\"\"\n",
    "    if shapeFactor == None:\n",
    "       submat = [t[0].m for t in matrix.grid.cellList]\n",
    "\n",
    "    else:\n",
    "        matrix = matrix.m\n",
    "        sF = tuple(sin // sfact for sin, sfact in zip(matrix.shape, shapeFactor))\n",
    "        submat = [matrix[sF[0]*i:sF[0]*(i+1),sF[1]*j:sF[1]*(j+1)] for i,j in np.ndindex(shapeFactor)]\n",
    "    return overlapMatrices(submat, colorHierarchy)\n",
    "\n",
    "def overlapMatrices(matrices, colorHierarchy):\n",
    "    \"\"\"\n",
    "    Overlaps matrices of a given shape according to the color hierarchy\n",
    "    \"\"\"\n",
    "    m = np.zeros(matrices[0].shape, dtype=np.uint8)\n",
    "    for i,j in np.ndindex(m.shape):\n",
    "        m[i,j] = colorHierarchy[max([colorHierarchy.index(x[i,j]) for x in matrices])]\n",
    "    return m\n",
    "\n",
    "def overlapShapes(matrix, diagonal=True, multicolor=True, byColor=False, hierarchy=[0,1,2,3,4,5,6,7,8,9]):\n",
    "    \"\"\"\n",
    "    Overlaps shapes of a given type of the same size. If there is a color conflict, a color is chosen \n",
    "    acording to the given hierarchy. \n",
    "    \"\"\"\n",
    "    if not multicolor: \n",
    "        if diagonal:   \n",
    "            shList = [sh for sh in matrix.dShapes]\n",
    "        else:   \n",
    "            shList = [sh for sh in matrix.shapes]\n",
    "    else:\n",
    "        if diagonal: \n",
    "            shList = [sh for sh in matrix.multicolorDShapes]\n",
    "        else:\n",
    "            shList = [sh for sh in matrix.multicolorShapes]\n",
    "    if byColor:\n",
    "        shList = [sh for sh in matrix.shapesByColor]\n",
    "    shList = [sh for sh in shList if sh.isRectangle]\n",
    "    if len(set([sh.shape for sh in shList])) != 1:\n",
    "        return matrix.m.copy()\n",
    "    return overlapMatrices([sh.m for sh in shList],hierarchy)\n",
    "\n",
    "def getCropAttributes(t, diagonal, multicolor, sameColor=True):\n",
    "    \"\"\"\n",
    "    If the task involved cropping a shape, this function returns the unique identifying arguments of that shape. \n",
    "    \"\"\"\n",
    "    bC = max(0, t.backgroundColor)\n",
    "    if diagonal and not multicolor:\n",
    "        if t.nCommonInOutDShapes == 0:\n",
    "            return set()\n",
    "        attrs = set.intersection(*[s.inMatrix.getShapeAttributes(backgroundColor=bC,\\\n",
    "                    singleColor=True, diagonals=True)[s.inMatrix.dShapes.index(s.commonDShapes[0][0])] for s in t.trainSamples])\n",
    "        nonAttrs = set()\n",
    "        for s in t.trainSamples:\n",
    "            shAttrs = s.inMatrix.getShapeAttributes(backgroundColor=bC, singleColor=True, diagonals=True)\n",
    "            for shi in range(len(s.inMatrix.dShapes)):\n",
    "                if s.inMatrix.dShapes[shi] == s.commonDShapes[0][0]:\n",
    "                    continue\n",
    "                else:\n",
    "                    nonAttrs = nonAttrs.union(shAttrs[shi])  \n",
    "    if not diagonal and not multicolor:\n",
    "        if t.nCommonInOutShapes == 0:\n",
    "            return set()\n",
    "        attrs = set.intersection(*[s.inMatrix.getShapeAttributes(backgroundColor=bC,\\\n",
    "                    singleColor=True, diagonals=False)[s.inMatrix.shapes.index(s.commonShapes[0][0])] for s in t.trainSamples])\n",
    "        nonAttrs = set()\n",
    "        for s in t.trainSamples:\n",
    "            shAttrs = s.inMatrix.getShapeAttributes(backgroundColor=bC, singleColor=True, diagonals=False)\n",
    "            for shi in range(len(s.inMatrix.shapes)):\n",
    "                if s.inMatrix.shapes[shi] == s.commonShapes[0][0]:\n",
    "                    continue\n",
    "                else:\n",
    "                    nonAttrs = nonAttrs.union(shAttrs[shi]) \n",
    "    if not diagonal and multicolor:\n",
    "        if not t.outIsInMulticolorShapeSize:\n",
    "            return set()                               \n",
    "        attrs = set()\n",
    "        nonAttrs = set()\n",
    "        for s in t.trainSamples:\n",
    "            shAttrs = s.inMatrix.getShapeAttributes(backgroundColor=bC, singleColor=False, diagonals=False)\n",
    "            crop = False\n",
    "            for shi in range(len(s.inMatrix.multicolorShapes)):\n",
    "                if s.inMatrix.multicolorShapes[shi].shape == s.outMatrix.shape and\\\n",
    "                np.all(np.logical_or(s.inMatrix.multicolorShapes[shi].m == s.outMatrix.m, s.inMatrix.multicolorShapes[shi].m==255)):\n",
    "                    crop = True\n",
    "                    if len(attrs) == 0:\n",
    "                        attrs = shAttrs[shi]\n",
    "                    attrs = attrs.intersection(shAttrs[shi])\n",
    "                else:\n",
    "                    nonAttrs = nonAttrs.union(shAttrs[shi])    \n",
    "        if not crop:\n",
    "                return set()\n",
    "    if diagonal and multicolor:\n",
    "        if not t.outIsInMulticolorDShapeSize:\n",
    "            return set()                               \n",
    "        attrs = set()\n",
    "        nonAttrs = set()\n",
    "        for s in t.trainSamples:\n",
    "            shAttrs = s.inMatrix.getShapeAttributes(backgroundColor=bC, singleColor=False, diagonals=True)\n",
    "            crop = False\n",
    "            for shi in range(len(s.inMatrix.multicolorDShapes)):\n",
    "                if s.inMatrix.multicolorDShapes[shi].shape == s.outMatrix.shape and\\\n",
    "                np.all(np.logical_or(s.inMatrix.multicolorDShapes[shi].m == s.outMatrix.m, s.inMatrix.multicolorDShapes[shi].m==255)):\n",
    "                    crop = True\n",
    "                    if len(attrs) == 0:\n",
    "                        attrs = shAttrs[shi]\n",
    "                    attrs = attrs.intersection(shAttrs[shi])\n",
    "                else:\n",
    "                    nonAttrs = nonAttrs.union(shAttrs[shi])    \n",
    "        if not crop:\n",
    "                return set()\n",
    "    return(attrs - nonAttrs)\n",
    "        \n",
    "def getBestCropShape(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"cropShape\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    bC = max(0, t.backgroundColor)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(cropShape, attributes=getCropAttributes(t,True, False),\\\n",
    "                                                           backgroundColor=bC, singleColor=True, diagonals=True), bestScore, bestFunction)\n",
    "    if bestScore==0:\n",
    "        return bestFunction\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(cropShape, attributes=getCropAttributes(t,False, False),\\\n",
    "                                                           backgroundColor=bC, singleColor=True, diagonals=False), bestScore, bestFunction)\n",
    "    if bestScore==0:\n",
    "        return bestFunction\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(cropShape, attributes=getCropAttributes(t,True, True),\\\n",
    "                                                           backgroundColor=bC, singleColor=False, diagonals=True), bestScore, bestFunction)\n",
    "    if bestScore==0:\n",
    "        return bestFunction\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(cropShape, attributes=getCropAttributes(t,False, True),\\\n",
    "                                                           backgroundColor=bC, singleColor=False, diagonals=False), bestScore, bestFunction)\n",
    "    if bestScore==0:\n",
    "        return bestFunction\n",
    "    for attr in ['LaSh', 'MoCo', 'MoCl', 'UnSh', 'UnSi']:\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(cropShape, attributes=set([attr]),\\\n",
    "                                                           backgroundColor=bC, singleColor=True, diagonals=True), bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(cropShape, attributes=set([attr]),\\\n",
    "                                                           backgroundColor=bC, singleColor=True, diagonals=False), bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(cropShape, attributes=set([attr]),\\\n",
    "                                                           backgroundColor=bC, singleColor=False, diagonals=True), bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(cropShape, attributes=set([attr]),\\\n",
    "                                                           backgroundColor=bC, singleColor=True, diagonals=False), bestScore, bestFunction)\n",
    "        if bestScore==0:\n",
    "            return bestFunction\n",
    "        \n",
    "    return bestFunction\n",
    "    \n",
    "def cropShape(matrix, attributes, backgroundColor=0, singleColor=True, diagonals=True, context=False):\n",
    "    \"\"\"\n",
    "    This function crops the shape out of a matrix with the maximum common attributes. If the argument context is passed,\n",
    "    the pixels lying inside the shape matrix are included. \n",
    "    \"\"\"\n",
    "    if singleColor: \n",
    "        if diagonals:   \n",
    "            shapeList = [sh for sh in matrix.dShapes]\n",
    "        else:   \n",
    "            shapeList = [sh for sh in matrix.shapes]\n",
    "    else:\n",
    "        if diagonals: \n",
    "            shapeList = [sh for sh in matrix.multicolorDShapes]\n",
    "        else:\n",
    "            shapeList = [sh for sh in matrix.multicolorShapes]\n",
    "    bestShapes = []\n",
    "    score = 0\n",
    "    attrList = matrix.getShapeAttributes(backgroundColor, singleColor, diagonals)\n",
    "    for i in range(len(shapeList)):\n",
    "        shscore = len(attributes.intersection(attrList[i]))\n",
    "        if shscore > score:\n",
    "            score = shscore\n",
    "            bestShapes = [i]\n",
    "        elif shscore == score:\n",
    "            bestShapes += [i]\n",
    "    if len(bestShapes) == 0:\n",
    "        return matrix\n",
    "    if context:\n",
    "        bestShape = shapeList[bestShapes[0]]\n",
    "        m = matrix.m[bestShape.position[0]:bestShape.position[0]+bestShape.shape[0], bestShape.position[1]:bestShape.position[1]+bestShape.shape[1]]\n",
    "        return m.copy() \n",
    "    else:\n",
    "        bestShape = shapeList[bestShapes[0]].m.copy()\n",
    "        bestShape[bestShape==255]=backgroundColor\n",
    "    return bestShape\n",
    "    \n",
    "def getBestCropReference(t):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function returns the partial function that uses the\n",
    "    function \"cropReference\" and works best for the training samples.\n",
    "    \"\"\"\n",
    "    bestFunction = partial(identityM)\n",
    "    bestScore = 1000\n",
    "    for sh in t.commonInShapes:\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(cropShapeReference, refShape=sh,\\\n",
    "                                refType='subshape', multicolor=True, diagonal=False), bestScore, bestFunction)\n",
    "    if len(t.commonInShapes) == 1:\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(cropShapeReference, refShape=t.commonInShapes[0],\\\n",
    "                                refType='mark', multicolor=False, diagonal=True), bestScore, bestFunction)\n",
    "        bestFunction, bestScore = updateBestFunction(t, partial(cropShapeReference, refShape=t.commonInShapes[0],\\\n",
    "                                refType='mark', multicolor=False, diagonal=False), bestScore, bestFunction)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(cropShapeReference,\\\n",
    "                                refType='pixels', maxOrMin='min',multicolor=True, diagonal=False), bestScore, bestFunction)\n",
    "    bestFunction, bestScore = updateBestFunction(t, partial(cropShapeReference,\\\n",
    "                                refType='pixels', maxOrMin='max',multicolor=True, diagonal=False), bestScore, bestFunction)\n",
    "    #add referenced by frames\n",
    "    return bestFunction\n",
    "\n",
    "def cropShapeReference(matrix, refShape=None, refType='subshape', maxOrMin='max', sameColor=True, multicolor=True, diagonal=False):\n",
    "    \"\"\"\n",
    "    Given a reference type refType, this function returns a referenced shape of the given type. \n",
    "    \"\"\"\n",
    "    if multicolor:\n",
    "        if diagonal:\n",
    "            shList = matrix.multicolorDShapes\n",
    "        else:\n",
    "            shList = matrix.multicolorShapes\n",
    "    else:\n",
    "        if diagonal:\n",
    "            shList = matrix.dShapes\n",
    "        else:\n",
    "            shList = matrix.shapes\n",
    "    if refType == 'subshape':\n",
    "        bestSh, bestScore = None, 0\n",
    "        for sh in shList:\n",
    "            if hasattr(sh, 'subshapes') and refShape in sh.subshapes:\n",
    "                score = np.count_nonzero([refShape == sh2 for sh2 in sh.subshapes])\n",
    "                if score > bestScore:\n",
    "                    bestSh = sh\n",
    "                    bestScore = score\n",
    "        if bestSh == None:\n",
    "            return matrix.m.copy()\n",
    "        return bestSh.m\n",
    "    elif refType == 'pixels':\n",
    "        if maxOrMin == 'max':\n",
    "            bestSh, bestScore = None, 0\n",
    "        else:\n",
    "            bestSh, bestScore = None, 1000\n",
    "        for sh in shList:\n",
    "            if hasattr(sh, 'subshapes'):\n",
    "                score = len([p for p in sh.subshapes if len(p.pixels) == 1])\n",
    "                if maxOrMin == 'max' and score > bestScore:\n",
    "                    bestSh = sh\n",
    "                    bestScore = score\n",
    "                if maxOrMin == 'min' and score < bestScore:\n",
    "                    bestSh = sh\n",
    "                    bestScore = score\n",
    "        if bestSh == None:\n",
    "            return matrix.m.copy()\n",
    "        return bestSh.m              \n",
    "    elif refType == 'frame':\n",
    "        return matrix.m.copy()\n",
    "    elif refType == 'mark':\n",
    "        foundRef = False\n",
    "        for sh in shList:\n",
    "            if sh == refShape:\n",
    "                refShape = sh\n",
    "                foundRef = True\n",
    "                break\n",
    "        if not foundRef:\n",
    "            return matrix.m\n",
    "        #otherwise return closest to reference\n",
    "        bestShape = None\n",
    "        dist = 1000\n",
    "        refPos = refShape.position\n",
    "        refPixels = [(p[0]+refPos[0], p[1]+refPos[1]) for p in refShape.pixels]\n",
    "        for sh in shList:\n",
    "            if sh == refShape or sh.color == matrix.backgroundColor:\n",
    "                continue\n",
    "            for p2 in [(p[0]+sh.position[0], p[1]+sh.position[1]) for p in sh.pixels]:\n",
    "                if min(abs((p[0]-p2[0]))+abs((p[1]-p2[1])) for p in refPixels) < dist:\n",
    "                    bestShape = sh\n",
    "                    dist = min(abs((p[0]-p2[0])+(p[1]-p2[1])) for p in refPixels)\n",
    "        if bestShape == None:\n",
    "            return matrix.m.copy()\n",
    "        bestShape=bestShape.m\n",
    "        bestShape[bestShape==255]=matrix.backgroundColor\n",
    "        return bestShape         \n",
    "\n",
    "def cropAllBackground(matrix):\n",
    "    \"\"\"\n",
    "    Deletes external rows and columns as long as they are colored fully by the background color. \n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    bC = matrix.backgroundColor\n",
    "    if np.all(m == bC):\n",
    "        return m\n",
    "    x1, x2, y1, y2 = 0, m.shape[0]-1, 0, m.shape[1]-1\n",
    "    while x1 <= x2 and np.all(m[x1,:] == bC):\n",
    "        x1 += 1\n",
    "    while x2 >= x1 and np.all(m[x2,:] == bC):\n",
    "        x2 -= 1\n",
    "    while y1 <= y2 and np.all(m[:,y1] == bC):\n",
    "        y1 += 1\n",
    "    while y2 >= y1 and np.all(m[:,y2] == bC):\n",
    "        y2 -= 1\n",
    "    return(m[x1:x2+1,y1:y2+1])\n",
    "\n",
    "def cropOnlyMulticolorShape(matrix, diagonals=False):\n",
    "    \"\"\"\n",
    "    This function is supposed to be called if there is one and only one \n",
    "    multicolor shape in all the input samples. This function just returns it.\n",
    "    \"\"\"\n",
    "    if diagonals:\n",
    "        m = matrix.multicolorDShapes[0].m.copy()\n",
    "        m[m==255] = matrix.multicolorDShapes[0].background\n",
    "    else:\n",
    "        m = matrix.multicolorShapes[0].m.copy()\n",
    "        m[m==255] = matrix.multicolorShapes[0].background\n",
    "    return m\n",
    "\n",
    "def cropFullFrame(matrix, includeBorder=True, bigOrSmall = None):\n",
    "    m = matrix.m.copy()\n",
    "    if bigOrSmall == None and len(matrix.fullFrames) != 1:\n",
    "        return m\n",
    "    if bigOrSmall == \"small\":\n",
    "        frame = matrix.fullFrames[-1]\n",
    "    else:\n",
    "        frame = matrix.fullFrames[0]\n",
    "    if includeBorder:\n",
    "        return m[frame.position[0]:frame.position[0]+frame.shape[0], \\\n",
    "                 frame.position[1]:frame.position[1]+frame.shape[1]]\n",
    "    else:\n",
    "        return m[frame.position[0]+1:frame.position[0]+frame.shape[0]-1, \\\n",
    "                 frame.position[1]+1:frame.position[1]+frame.shape[1]-1]\n",
    "        \n",
    "def cropPartialFrame(matrix, includeBorder=True, bigOrSmall = None):\n",
    "    \"\"\"\n",
    "    Crop the unique partial frame of a matrix. Options to include the border of the frame,\n",
    "    or to choose the frame according to its size. \n",
    "    \"\"\"\n",
    "    m = matrix.m.copy()\n",
    "    if len(matrix.partialFrames) == 0:\n",
    "        return m\n",
    "    if bigOrSmall == \"small\":\n",
    "        frame = matrix.partialFrames[-1]\n",
    "    else:\n",
    "        frame = matrix.partialFrames[0]\n",
    "    if includeBorder:\n",
    "        return m[frame.position[0]:frame.position[0]+frame.shape[0], \\\n",
    "                 frame.position[1]:frame.position[1]+frame.shape[1]]\n",
    "    else:\n",
    "        return m[frame.position[0]+1:frame.position[0]+frame.shape[0]-1, \\\n",
    "                 frame.position[1]+1:frame.position[1]+frame.shape[1]-1]\n",
    "        \n",
    "def getBestTwoShapeFunction(t):\n",
    "    \"\"\"\n",
    "    Inputs a task with two input shapes and tries a series of operations that make sense.\n",
    "    \"\"\"\n",
    "    bestScore = 1000\n",
    "    bestFunction = partial(identityM)\n",
    "    multicolor = t.twoShapeTask[1]\n",
    "    diagonal = t.twoShapeTask[2]\n",
    "    typ = t.twoShapeTask[-1]\n",
    "    cropAfter = [0] \n",
    "    if t.outSmallerThanIn:\n",
    "        cropAfter += [1,2]\n",
    "    #try possible operations\n",
    "    for crop in cropAfter:\n",
    "        for flip in [True,False]:\n",
    "            if t.twoShapeTask[3]:\n",
    "                #pixelwise and/or\n",
    "                for c in permutations(t.totalOutColors,2):\n",
    "                    bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(pixelwiseAnd, falseColor=c[0],\\\n",
    "                                targetColor=None,trueColor=c[1]), diagonal=diagonal,typ=typ, multicolor=multicolor, crop=crop, flip=flip), bestScore, bestFunction)\n",
    "                    bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(pixelwiseOr, falseColor=c[0],\\\n",
    "                                targetColor=None,trueColor=c[1]), diagonal=diagonal,typ=typ, multicolor=multicolor, crop=crop, flip=flip), bestScore, bestFunction)\n",
    "                #print shapes            \n",
    "                for base in [0,1]:                \n",
    "                    for bC in t.commonInColors:          \n",
    "                        bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(printShapes, base=base,\\\n",
    "                                backgroundColor=bC),diagonal=diagonal,typ=typ, multicolor=multicolor, crop=crop, flip=flip), bestScore, bestFunction)\n",
    "                #overlap matrices\n",
    "                if typ > 1:\n",
    "                    bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(overlapMatrices,\\\n",
    "                                colorHierarchy=[0,1,2,3,4,5,6,7,8,9]),diagonal=diagonal,typ=typ, multicolor=multicolor, crop=crop, flip=flip), bestScore, bestFunction)\n",
    "            else:\n",
    "                for c in permutations(t.totalOutColors,2):\n",
    "                    for target in [None]:\n",
    "                        bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(pixelwiseAnd, falseColor=c[0],\\\n",
    "                                    targetColor=target,trueColor=c[1]), diagonal=diagonal, multicolor=multicolor,typ=typ, crop=crop, flip=flip, downsizeToSmallest=True), bestScore, bestFunction)\n",
    "                        bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(pixelwiseOr, falseColor=c[0],\\\n",
    "                                    targetColor=target,trueColor=c[1]), diagonal=diagonal, multicolor=multicolor,typ=typ, crop=crop, flip=flip, downsizeToSmallest=True), bestScore, bestFunction)\n",
    "                        bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(pixelwiseAnd, falseColor=c[0],\\\n",
    "                                    targetColor=target,trueColor=c[1]), diagonal=diagonal, multicolor=multicolor,typ=typ, crop=crop, flip=flip, scaleToLargest=True), bestScore, bestFunction)\n",
    "                        bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(pixelwiseOr, falseColor=c[0],\\\n",
    "                                    targetColor=target,trueColor=c[1]), diagonal=diagonal, multicolor=multicolor,typ=typ, crop=crop, flip=flip, scaleToLargest=True), bestScore, bestFunction)  \n",
    "                for base in [0,1]:\n",
    "                    for bC in t.commonInColors:\n",
    "                        bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(printShapes, base=base,\\\n",
    "                                backgroundColor=bC),diagonal=diagonal, multicolor=multicolor, crop=crop,typ=typ, flip=flip, downsizeToSmallest=True), bestScore, bestFunction)\n",
    "                        bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(printShapes, base=base,\\\n",
    "                                backgroundColor=bC),diagonal=diagonal, multicolor=multicolor, crop=crop,typ=typ, flip=flip, scaleToLargest=True), bestScore, bestFunction)\n",
    "            #multiply matrices\n",
    "            for c in [0,1]:\n",
    "                for b in [0,1]:\n",
    "                    bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(multiplyMatrices, base=b,\\\n",
    "                                background=0, color=c), typ=typ, crop=crop, flip=flip), bestScore, bestFunction)\n",
    "                    bestFunction, bestScore = updateBestFunction(t, partial(twoShapeFun, f=partial(multiplyMatrices, base=b,\\\n",
    "                                background=0, color=c), typ=typ, crop=crop, flip=flip, downsizeToSmallest=True), bestScore, bestFunction)\n",
    "               \n",
    "    return bestFunction\n",
    "\n",
    "def twoShapeFun(matrix, f=partial(identityM), typ=1, diagonal=True, multicolor=True,\\\n",
    "                flip=False, rotate=False, downsizeToSmallest=False, scaleToLargest=False, crop=False):\n",
    "    \"\"\"\n",
    "    Apply function f to the shapes of matrix. By default, this function should only be called when matrix has two shapes.\n",
    "    The arguments downsizeToSmallest of scaleToLargest match the size of the shapes if possible. The argument crop returns\n",
    "    the result of the operation without context. The arguments flip and rotate modify the shapes before the operations are \n",
    "    performed. \n",
    "    \"\"\"\n",
    "    def cropAllBackgroundM(m, background):\n",
    "        if np.all(m == background):\n",
    "            return m\n",
    "        x1, x2, y1, y2 = 0, m.shape[0]-1, 0, m.shape[1]-1\n",
    "        while x1 <= x2 and np.all(m[x1,:] == background):\n",
    "            x1 += 1\n",
    "        while x2 >= x1 and np.all(m[x2,:] == background):\n",
    "            x2 -= 1\n",
    "        while y1 <= y2 and np.all(m[:,y1] == background):\n",
    "            y1 += 1\n",
    "        while y2 >= y1 and np.all(m[:,y2] == background):\n",
    "            y2 -= 1\n",
    "        return(m[x1:x2+1,y1:y2+1])\n",
    "                                            \n",
    "    def minimizeM(m):\n",
    "        x = 1\n",
    "        for i in range(1, m.shape[0]):\n",
    "            if np.array_equal(m[x,:],m[x-1,:]):\n",
    "                m = np.delete(m, (x), axis=0)\n",
    "            else:\n",
    "                x+=1\n",
    "        x = 1\n",
    "        for i in range(1, m.shape[1]):\n",
    "            if np.array_equal(m[:,x],m[:,x-1]):\n",
    "                m = np.delete(m, (x), axis=1)\n",
    "            else:\n",
    "                x+=1\n",
    "        return m\n",
    "                                            \n",
    "    def scaleM(m, s):\n",
    "        sc = min(s[0]//m.shape[0], s[1]//m.shape[1])\n",
    "        sm = np.repeat(np.repeat(m, sc, axis=1), sc, axis=0)\n",
    "        return sm\n",
    "                                            \n",
    "    m = matrix.m.copy()\n",
    "    if typ == 1:\n",
    "        if multicolor:\n",
    "            if diagonal: \n",
    "                shList = [sh for sh in matrix.multicolorDShapes]\n",
    "            else:\n",
    "                shList = [sh for sh in matrix.multicolorShapes]\n",
    "        else: \n",
    "            if diagonal:   \n",
    "                shList = [sh for sh in matrix.dShapes]\n",
    "            else:   \n",
    "                shList = [sh for sh in matrix.shapes]\n",
    "        posList = [sh.position for sh in shList]\n",
    "        mList = [sh.m.copy() for sh in shList]\n",
    "    elif typ == 2:\n",
    "        if not hasattr(matrix,'grid'):\n",
    "            return m\n",
    "        mList = [c[0].m for c in matrix.grid.cellList]\n",
    "        posList = [c[1] for c in matrix.grid.cellList]\n",
    "    elif typ == 3:\n",
    "        if matrix.shape[0] == 2*matrix.shape[1]:\n",
    "            mList = [matrix.m[:matrix.shape[0]//2,:].copy(), matrix.m[matrix.shape[0]//2:,:].copy()]\n",
    "            posList = [(0,0),(matrix.shape[1],0)]\n",
    "        elif matrix.shape[1] == 2*matrix.shape[0]:\n",
    "            mList = [matrix.m[:,:matrix.shape[1]//2].copy(), matrix.m[:,matrix.shape[1]//2:].copy()]\n",
    "            posList = [(0,0),(0,matrix.shape[0])]\n",
    "        else:\n",
    "            return m\n",
    "        mList.sort(key=lambda x: len(np.unique(x))) \n",
    "    if len(mList)  != 2:\n",
    "        return m    \n",
    "    if flip:\n",
    "        mList[1] = np.fliplr(mList[1])\n",
    "    #sort list largest first\n",
    "    if mList[0].shape[0]*mList[0].shape[1]<mList[1].shape[0]*mList[1].shape[1]:\n",
    "        mList = mList[::-1]\n",
    "        posList = posList[::-1]\n",
    "    if downsizeToSmallest:\n",
    "        mList[0] = minimizeM(mList[0])\n",
    "    if scaleToLargest:\n",
    "        mList[1] = scaleM(mList[1], mList[0].shape)\n",
    "    if mList[0].shape != mList[1].shape and f.func.__name__ != 'multiplyMatrices':\n",
    "        return m\n",
    "    if  f.func.__name__ == 'multiplyMatrices':\n",
    "        return f(mList)\n",
    "    elif crop == 1:\n",
    "        return(f(mList))\n",
    "    elif crop == 2:\n",
    "        return cropAllBackgroundM(f(mList),matrix.backgroundColor)\n",
    "    else:\n",
    "        if scaleToLargest:\n",
    "            pos = posList[0]\n",
    "            sh = mList[0]\n",
    "            m[pos[0]:pos[0] + sh.shape[0], pos[1]:pos[1] + sh.shape[1]] = f(mList)                \n",
    "        elif downsizeToSmallest:\n",
    "            pos = posList[1]\n",
    "            sh = mList[1]\n",
    "            m[pos[0]:pos[0] + sh.shape[0], pos[1]:pos[1] + sh.shape[1]] = f(mList)   \n",
    "        return m\n",
    "    \n",
    "###############################################################################\n",
    "###############################################################################\n",
    "# %% Main function: getPossibleOperations\n",
    "def getPossibleOperations(t, c):\n",
    "    \"\"\"\n",
    "    Given a Task t and a Candidate c, this function returns a list of all the\n",
    "    possible operations that make sense applying to the input matrices of c.\n",
    "    The elements of the list to be returned are partial functions, whose input\n",
    "    is a Matrix and whose output is a numpy.ndarray (2-dim matrix).\n",
    "    \"\"\" \n",
    "    candTask = c.t\n",
    "    x = [] # List to be returned\n",
    "    \n",
    "    ###########################################################################\n",
    "    # Fill the blanks\n",
    "    if t.fillTheBlank:\n",
    "        params = fillTheBlankParameters(t)\n",
    "        x.append(partial(fillTheBlank, params=params))\n",
    "        \n",
    "    # switchColors\n",
    "    if all([n==2 for n in candTask.nInColors]):\n",
    "        x.append(partial(switchColors))\n",
    "        \n",
    "    # downsize\n",
    "    if candTask.sameOutShape:\n",
    "        outShape = candTask.outShape\n",
    "        x.append(partial(downsize, newShape=outShape))\n",
    "        x.append(partial(downsizeMode,newShape=outShape))\n",
    "        if t.backgroundColor!=-1:\n",
    "            x.append(partial(downsize, newShape=outShape, falseColor=t.backgroundColor))\n",
    "        #if candTask.sameOutDummyMatrix and candTask.backgroundColor != -1:\n",
    "        #    x.append(partial(arrangeShapes,outDummyMatrix=candTask.trainSamples[0].outMatrix.dummyMatrix,\\\n",
    "        #                     outDummyColor=candTask.trainSamples[0].outMatrix.backgroundColor))\n",
    "    x.append(getBestCountColors(candTask))  \n",
    "    x.append(getBestCountShapes(candTask))  \n",
    "    ###########################################################################\n",
    "    # sameIOShapes\n",
    "    if candTask.sameIOShapes:\n",
    "        \n",
    "        #######################################################################\n",
    "        \n",
    "        # ColorMap\n",
    "        ncc = len(candTask.colorChanges)\n",
    "        if len(set([cc[0] for cc in candTask.colorChanges])) == ncc and ncc != 0:\n",
    "            x.append(partial(colorMap, cMap=dict(candTask.colorChanges)))\n",
    "            \n",
    "        x.append(partial(revertColorOrder))\n",
    "            \n",
    "        # Symmetrize\n",
    "        if all([len(x)==1 for x in candTask.changedInColors]):\n",
    "            color = next(iter(candTask.changedInColors[0]))\n",
    "            axis = []\n",
    "            if candTask.lrSymmetric:\n",
    "                axis.append(\"lr\")\n",
    "            if candTask.udSymmetric:\n",
    "                axis.append(\"ud\")\n",
    "            if all([sample.inMatrix.shape[0]==sample.inMatrix.shape[1] for sample in candTask.testSamples]):\n",
    "                if candTask.d1Symmetric:\n",
    "                    axis.append(\"d1\")\n",
    "                if candTask.d2Symmetric:\n",
    "                    axis.append(\"d2\")\n",
    "            x.append(partial(symmetrize, axis=axis, color=color))\n",
    "            if candTask.totalOutColors==1:\n",
    "                for fc in candTask.fixedColors:\n",
    "                    x.append(partial(symmetrize, axis=axis, refColor=fc,\\\n",
    "                                     outColor=next(iter(candTask.totalOutColors))))\n",
    "    \n",
    "        # Color symmetric pixels\n",
    "        x.append(getBestColorSymmetricPixels(candTask))\n",
    "    \n",
    "        # Complete rectangles\n",
    "        if candTask.backgroundColor!=-1 and len(candTask.fixedColors)==1 and \\\n",
    "        len(candTask.colorChanges)==1:\n",
    "            sc = next(iter(candTask.fixedColors))\n",
    "            nc = next(iter(candTask.colorChanges))[1]\n",
    "            x.append(partial(completeRectangles, sourceColor=sc, newColor=nc))\n",
    "        \n",
    "        x.append(partial(deletePixels, diagonals=True))\n",
    "        x.append(partial(deletePixels, diagonals=False))\n",
    "        \n",
    "        #######################################################################\n",
    "        # For LinearShapeModel we need to have the same shapes in the input\n",
    "        # and in the output, and in the exact same positions.\n",
    "        # This model predicts the color of the shape in the output.\n",
    "        \n",
    "        if candTask.onlyShapeColorChanges:\n",
    "            ccwf = getColorChangesWithFeatures(candTask)\n",
    "            fsf = candTask.fixedShapeFeatures\n",
    "            x.append(partial(changeShapesWithFeatures, ccwf=ccwf, fixedColors=candTask.fixedColors,\\\n",
    "                             fixedShapeFeatures=fsf))\n",
    "                \n",
    "            if all([\"getBestLSTM\" not in str(op.func) for op in c.ops]):        \n",
    "                x.append(getBestLSTM(candTask))\n",
    "            \n",
    "            # Other deterministic functions that change the color of shapes.\n",
    "            for cc in candTask.commonColorChanges:\n",
    "                for border, bs in product([True, False, None], [\"big\", \"small\", None]):\n",
    "                    x.append(partial(changeShapes, inColor=cc[0], outColor=cc[1],\\\n",
    "                                     bigOrSmall=bs, isBorder=border))\n",
    "                    \n",
    "            if isReplicateTask(candTask)[0]:\n",
    "                x.append(getBestReplicateShapes(candTask))\n",
    "            x.append(getBestColorByPixels(candTask))\n",
    "            \n",
    "            return x\n",
    "        \n",
    "        #######################################################################\n",
    "        # Complete row/col patterns\n",
    "        colStep=None\n",
    "        rowStep=None\n",
    "        if candTask.followsRowPattern:\n",
    "            if candTask.allEqual(candTask.rowPatterns):\n",
    "                rowStep=candTask.rowPatterns[0]\n",
    "        if candTask.followsColPattern:\n",
    "            if candTask.allEqual(candTask.colPatterns):\n",
    "                colStep=candTask.colPatterns[0]\n",
    "        if candTask.allEqual(candTask.changedInColors) and len(candTask.changedInColors[0])==1:\n",
    "            c2c = next(iter(candTask.changedInColors[0]))\n",
    "        else:\n",
    "            c2c=None\n",
    "                \n",
    "        if candTask.followsRowPattern and candTask.followsColPattern:\n",
    "            x.append(partial(followPattern, rc=\"both\", colorToChange=c2c,\\\n",
    "                             rowStep=rowStep, colStep=colStep))\n",
    "        elif candTask.followsRowPattern:\n",
    "            x.append(partial(followPattern, rc=\"row\", colorToChange=c2c,\\\n",
    "                             rowStep=rowStep, colStep=colStep))\n",
    "        elif candTask.followsColPattern:\n",
    "            x.append(partial(followPattern, rc=\"col\", colorToChange=c2c,\\\n",
    "                             rowStep=rowStep, colStep=colStep))\n",
    "\n",
    "        #######################################################################\n",
    "        # CNNs\n",
    "        \n",
    "        #x.append(getBestCNN(candTask))\n",
    "        #if candTask.sameNSampleColors and all([\"predictCNN\" not in str(op.func) for op in c.ops]):\n",
    "        #    x.append(getBestSameNSampleColorsCNN(candTask))\n",
    "\n",
    "        \"\"\"\n",
    "        if t.backgroundColor != -1:\n",
    "            model = trainCNNDummyColor(candTask, 5, -1)\n",
    "            x.append(partial(predictCNNDummyColor, model=model))\n",
    "            model = trainCNNDummyColor(candTask, 3, 0)\n",
    "            x.append(partial(predictCNNDummyColor, model=model))\n",
    "            #model = trainOneConvModelDummyColor(candTask, 7, -1)\n",
    "            #x.append(partial(predictConvModelDummyColor, model=model))\n",
    "        \"\"\"\n",
    "            \n",
    "        #cc = list(t.commonSampleColors)\n",
    "        #model = trainCNNDummyCommonColors(t, cc, 3, -1)\n",
    "        #x.append(partial(predictCNNDummyCommonColors, model=model,\\\n",
    "        #                commonColors=cc))\n",
    "        \n",
    "        #######################################################################\n",
    "        # Transformations if the color count is always the same:\n",
    "        # Rotations, Mirroring, Move Shapes, Mirror Shapes, ...\n",
    "        if candTask.sameColorCount:\n",
    "            for axis in [\"lr\", \"ud\"]:\n",
    "                x.append(partial(mirror, axis = axis))\n",
    "            # You can only mirror d1/d2 or rotate if the matrix is squared.\n",
    "            if candTask.inMatricesSquared:\n",
    "                for axis in [\"d1\", \"d2\"]:\n",
    "                    x.append(partial(mirror, axis = axis))\n",
    "                for angle in [90, 180, 270]:\n",
    "                    x.append(partial(rotate, angle = angle))\n",
    "                \n",
    "                                                         \n",
    "            # Mirror shapes\n",
    "            x.append(getBestFlipAllShapes(candTask))\n",
    "                                \n",
    "        #######################################################################\n",
    "        # Other sameIOShapes functions\n",
    "        # Move shapes\n",
    "        x.append(getBestMoveShapes(candTask, candidate=c))\n",
    "        \n",
    "        pr = pixelRecolor(candTask)\n",
    "        if len(pr)!=1:\n",
    "            x.append(partial(executePixelRecolor, Best_Dict=pr[0], Best_v=pr[1], Best_Q1=pr[2], Best_Q2=pr[3]))\n",
    "        \n",
    "        fun = getPixelChangeCriteria(candTask)\n",
    "        if fun != 0:\n",
    "            x.append(fun)\n",
    "            \n",
    "        # extendColor\n",
    "        x.append(getBestExtendColor(candTask))\n",
    "        \n",
    "        # surround shapes\n",
    "        x.append(getBestSurroundShapes(candTask))\n",
    "        \n",
    "        # Paint shapes in half\n",
    "        x.append(getBestPaintShapesInHalf(candTask))\n",
    "        \n",
    "        # Paint shapes from border color\n",
    "        x.append(getBestPaintShapeFromBorderColor(candTask))\n",
    "            \n",
    "        # fillRectangleInside\n",
    "        for cic in candTask.commonChangedInColors:\n",
    "            for coc in candTask.commonChangedOutColors:\n",
    "                x.append(partial(fillRectangleInside, rectangleColor=cic, fillColor=coc))\n",
    "        \n",
    "        # Color longest lines\n",
    "        if len(candTask.colorChanges)==1:\n",
    "            change = next(iter(candTask.colorChanges))\n",
    "            x.append(partial(colorLongestLines, cic=change[0], coc=change[1], direction='h'))\n",
    "            x.append(partial(colorLongestLines, cic=change[0], coc=change[1], direction='v'))\n",
    "            x.append(partial(colorLongestLines, cic=change[0], coc=change[1], direction='hv'))\n",
    "            x.append(partial(colorLongestLines, cic=change[0], coc=change[1], direction='d'))\n",
    "\n",
    "        # Connect Pixels\n",
    "        x.append(partial(connectAnyPixels))\n",
    "        x.append(partial(connectAnyPixels, diagonal=True))\n",
    "        if all([len(x)==1 for x in candTask.changedInColors]):\n",
    "            x.append(partial(connectAnyPixels, connColor=next(iter(candTask.changedOutColors[0]))))\n",
    "            x.append(partial(connectAnyPixels, connColor=next(iter(candTask.changedOutColors[0])), \\\n",
    "                             diagonal=True))\n",
    "\n",
    "        fc = candTask.fixedColors\n",
    "        #if hasattr(t, \"fixedColors\"):\n",
    "        #    tfc = candTask.fixedColors\n",
    "        #else:\n",
    "        #    tfc = set()\n",
    "        for pc in candTask.colors - candTask.commonChangedInColors:\n",
    "            for cc in candTask.commonChangedOutColors:\n",
    "                x.append(partial(connectAnyPixels, pixelColor=pc, \\\n",
    "                                 connColor=cc, fixedColors=fc))\n",
    "                x.append(partial(connectAnyPixels, pixelColor=pc, \\\n",
    "                                 connColor=cc, fixedColors=fc, diagonal=True))\n",
    "        for pc in candTask.colors - candTask.commonChangedInColors:\n",
    "            x.append(partial(connectAnyPixels, pixelColor=pc, allowedChanges=dict(candTask.colorChanges)))\n",
    "            x.append(partial(connectAnyPixels, pixelColor=pc, allowedChanges=dict(candTask.colorChanges), \\\n",
    "                             diagonal=True))\n",
    "            x.append(partial(connectAnyPixels, pixelColor=pc, allowedChanges=dict(candTask.colorChanges),\\\n",
    "                             lineExclusive=True))\n",
    "            x.append(partial(connectAnyPixels, pixelColor=pc, allowedChanges=dict(candTask.colorChanges),\\\n",
    "                             lineExclusive=True, diagonal=True))\n",
    "                \n",
    "        for cc in candTask.commonColorChanges:\n",
    "            for cc in candTask.commonColorChanges:\n",
    "                for border, bs in product([True, False, None], [\"big\", \"small\", None]):\n",
    "                    x.append(partial(changeShapes, inColor=cc[0], outColor=cc[1],\\\n",
    "                                     bigOrSmall=bs, isBorder=border))\n",
    "        \n",
    "        if all(len(set([sh.shape for sh in s.outMatrix.shapes]))==1 for s in candTask.trainSamples):\n",
    "            x.append(partial(subMatToLayer,shapeAndDict=getLayerDict(candTask)))\n",
    "        #replicate/symmterize/other shape related tasks\n",
    "        #x.append(getBestAlignShapes(candTask))\n",
    "        x.append(getBestSymmetrizeSubmatrix(candTask))\n",
    "        x.append(partial(replicateShapes,diagonal=True, multicolor=True, allCombs=False,anchorType='subframe', scale=False))\n",
    "        x.append(partial(replicateShapes,diagonal=True, multicolor=True, allCombs=True,anchorType='subframe', scale=False, deleteOriginal=True))\n",
    "        if isReplicateTask(candTask)[0]:\n",
    "            x.append(getBestReplicateShapes(candTask))\n",
    "        #x.append(getBestReplicateOneShape(candTask))\n",
    "        x.append(getBestSymmetrizeAllShapes(candTask))\n",
    "        x.append(getBestColorByPixels(candTask))\n",
    "        \n",
    "        #if len(candTask.colorChanges) == 1:\n",
    "        #    x.append(partial(replicateShapes,diagonal=True, multicolor=False, allCombs=True,\\\n",
    "        #                     anchorColor = list(candTask.colorChanges)[0][0], anchorType='all', attributes=set(['UnCo'])))\n",
    "\n",
    "        #delete shapes\n",
    "        if isDeleteTask(candTask) and all(t.backgroundColor == c[1] for c in candTask.colorChanges):\n",
    "            x.append(getBestDeleteShapes(candTask, True, True))\n",
    "                    \n",
    "    ###########################################################################\n",
    "    # Cases in which the input has always the same shape, and the output too\n",
    "    if candTask.sameInShape and candTask.sameOutShape and \\\n",
    "    all(candTask.trainSamples[0].inMatrix.shape == s.inMatrix.shape for s in candTask.testSamples):\n",
    "            \n",
    "        if candTask.sameNSampleColors:\n",
    "            #Cases where output colors are a subset of input colors\n",
    "            if candTask.sameNInColors and all(s.inHasOutColors for s in candTask.trainSamples):\n",
    "                if hasattr(candTask, 'outShapeFactor') or (hasattr(candTask,\\\n",
    "                              'gridCellIsOutputShape') and candTask.gridCellIsOutputShape):\n",
    "                    ch = dict(sum([Counter(s.outMatrix.colorCount) for s in candTask.trainSamples],Counter()))\n",
    "                    ch = sorted(ch, key=ch.get)\n",
    "                    if candTask.backgroundColor in ch:\n",
    "                        ch.remove(candTask.backgroundColor)\n",
    "                    ch = list(set([0,1,2,3,4,5,6,7,8,9]).difference(set(ch))) + ch\n",
    "                    if hasattr(candTask, 'outShapeFactor'):\n",
    "                        x.append(partial(overlapSubmatrices, colorHierarchy=ch, shapeFactor=candTask.outShapeFactor))\n",
    "                    else:\n",
    "                        x.append(partial(overlapSubmatrices, colorHierarchy=ch))\n",
    "        \n",
    "        pixelMap = pixelCorrespondence(candTask)\n",
    "        if len(pixelMap) != 0:\n",
    "            x.append(partial(mapPixels, pixelMap=pixelMap, outShape=candTask.outShape))\n",
    "    \n",
    "    ###########################################################################\n",
    "    # Evolve\n",
    "    if candTask.sameIOShapes and all([len(x)==1 for x in candTask.changedInColors]) and\\\n",
    "    len(candTask.commonChangedInColors)==1 and candTask.sameNSampleColors:\n",
    "        x.append(getBestEvolve(candTask))\n",
    "        #cfn = evolve(candTask)\n",
    "        #x.append(getBestEvolve(candTask, cfn))\n",
    "    #    x.append(partial(applyEvolve, cfn=cfn, nColors=candTask.trainSamples[0].nColors,\\\n",
    "    #                     kernel=5, nIterations=1))\n",
    "    \n",
    "    if candTask.sameIOShapes and all([len(x)==1 for x in candTask.changedInColors]) and\\\n",
    "    len(candTask.commonChangedInColors)==1:\n",
    "        x.append(getBestEvolvingLines(candTask))\n",
    "        \n",
    "    ###########################################################################\n",
    "    # Other cases\n",
    "    \n",
    "    if candTask.sameIOShapes and len(candTask.fixedColors)!=0:\n",
    "        for color in candTask.fixedColors:\n",
    "            for outColor in candTask.commonOutColors:\n",
    "                x.append(partial(paintCrossedCoordinates, refColor=color,\\\n",
    "                                 outColor=outColor, fixedColors=candTask.fixedColors))\n",
    "    \n",
    "    if candTask.inSmallerThanOut and t.inSmallerThanOut:\n",
    "        x.append(getBestExtendMatrix(candTask))\n",
    "        \n",
    "    if hasattr(candTask, 'inShapeFactor'):\n",
    "        x.append(partial(multiplyPixels, factor=candTask.inShapeFactor))\n",
    "        x.append(partial(multiplyMatrix, factor=candTask.inShapeFactor))\n",
    "        \n",
    "        for c in candTask.commonSampleColors:\n",
    "            x.append(partial(multiplyPixelsAndAnd, factor=candTask.inShapeFactor,\\\n",
    "                             falseColor=c))\n",
    "            x.append(partial(multiplyPixelsAndOr, factor=candTask.inShapeFactor,\\\n",
    "                             falseColor=c))\n",
    "            x.append(partial(multiplyPixelsAndXor, factor=candTask.inShapeFactor,\\\n",
    "                             falseColor=c))\n",
    "            \n",
    "        if type(candTask.inShapeFactor)==tuple:\n",
    "            ops = getBestMosaic(candTask)\n",
    "            x.append(partial(generateMosaic, ops=ops, factor=candTask.inShapeFactor))\n",
    "            \n",
    "        if all([s.inMatrix.shape[0]**2 == s.outMatrix.shape[0] and \\\n",
    "                s.inMatrix.shape[1]**2 == s.outMatrix.shape[1] for s in candTask.trainSamples]):\n",
    "            totalColorCount = Counter()\n",
    "            for sample in t.trainSamples:\n",
    "                for color in sample.outMatrix.colorCount.keys():\n",
    "                    totalColorCount[color] += sample.outMatrix.colorCount[color]\n",
    "            falseColor = max(totalColorCount.items(), key=operator.itemgetter(1))[0]\n",
    "            opCond = getBestMultiplyMatrix(candTask, falseColor)\n",
    "            x.append(partial(doBestMultiplyMatrix, opCond=opCond, falseColor=falseColor))\n",
    "            \n",
    "            \n",
    "    if hasattr(candTask, 'outShapeFactor'):\n",
    "        if outputIsSubmatrix(candTask):\n",
    "            for color in range(10):\n",
    "                x.append(partial(selectSubmatrixWithMaxColor, color=color, outShapeFactor=candTask.outShapeFactor))\n",
    "                x.append(partial(selectSubmatrixWithMinColor, color=color, outShapeFactor=candTask.outShapeFactor))\n",
    "            x.append(partial(selectSubmatrixWithMostColors, outShapeFactor=candTask.outShapeFactor))\n",
    "            x.append(partial(selectSubmatrixWithLeastColors, outShapeFactor=candTask.outShapeFactor))\n",
    "            position = getBestSubmatrixPosition(candTask, outShapeFactor=candTask.outShapeFactor)\n",
    "            x.append(partial(selectSubmatrixInPosition, position=position, outShapeFactor=candTask.outShapeFactor))\n",
    "        \n",
    "        # Pixelwise And\n",
    "        for c in candTask.commonOutColors:\n",
    "            x.append(partial(pixelwiseAndInSubmatrices, factor=candTask.outShapeFactor,\\\n",
    "                             falseColor=c))\n",
    "        if len(candTask.totalOutColors) == 2:\n",
    "            for target in candTask.totalInColors:\n",
    "                for c in permutations(candTask.totalOutColors, 2):\n",
    "                    x.append(partial(pixelwiseAndInSubmatrices, \\\n",
    "                                     factor=candTask.outShapeFactor, falseColor=c[0],\\\n",
    "                                     targetColor=target, trueColor=c[1]))\n",
    "        \n",
    "        # Pixelwise Or\n",
    "        for c in candTask.commonOutColors:\n",
    "            x.append(partial(pixelwiseOrInSubmatrices, factor=candTask.outShapeFactor,\\\n",
    "                             falseColor=c))\n",
    "        if len(candTask.totalOutColors) == 2:\n",
    "            for target in candTask.totalInColors:\n",
    "                for c in permutations(candTask.totalOutColors, 2):\n",
    "                    x.append(partial(pixelwiseOrInSubmatrices, \\\n",
    "                                     factor=candTask.outShapeFactor, falseColor=c[0],\\\n",
    "                                     targetColor=target, trueColor=c[1]))\n",
    "        if candTask.backgroundColor!=-1:\n",
    "            colors = candTask.commonOutColors - candTask.commonInColors\n",
    "            if candTask.outShapeFactor[0]*candTask.outShapeFactor[1]==len(colors):\n",
    "                for c in permutations(colors, len(colors)):\n",
    "                    x.append(partial(pixelwiseOrInSubmatrices, factor=candTask.outShapeFactor,\\\n",
    "                                     falseColor=candTask.backgroundColor,\\\n",
    "                                     trueValues = c))\n",
    "        \n",
    "        # Pixelwise Xor\n",
    "        if candTask.outShapeFactor in [(2,1), (1,2)]:\n",
    "            for c in candTask.commonOutColors:\n",
    "                x.append(partial(pixelwiseXorInSubmatrices, factor=candTask.outShapeFactor,\\\n",
    "                                 falseColor=c))\n",
    "            if len(candTask.commonOutColors - candTask.commonInColors)==2 and\\\n",
    "            candTask.backgroundColor!=-1:\n",
    "                colors = candTask.commonOutColors - candTask.commonInColors\n",
    "                for c in permutations(colors, 2):\n",
    "                    x.append(partial(pixelwiseXorInSubmatrices, falseColor=candTask.backgroundColor,\\\n",
    "                                     firstTrue=c[0], secondTrue=c[1]))\n",
    "            if len(candTask.totalOutColors) == 2:\n",
    "                for target in candTask.totalInColors:\n",
    "                    for c in permutations(candTask.totalOutColors, 2):\n",
    "                        x.append(partial(pixelwiseXorInSubmatrices, \\\n",
    "                                         factor=candTask.outShapeFactor, falseColor=c[0],\\\n",
    "                                         targetColor=target, trueColor=c[1]))\n",
    "    \n",
    "    if hasattr(candTask, 'gridCellIsOutputShape') and candTask.gridCellIsOutputShape:\n",
    "        if outputIsSubmatrix(candTask, isGrid=True):\n",
    "            for color in range(10):\n",
    "                x.append(partial(selectSubmatrixWithMaxColor, color=color, isGrid=True))\n",
    "                x.append(partial(selectSubmatrixWithMinColor, color=color, isGrid=True))\n",
    "            x.append(partial(selectSubmatrixWithMostColors, isGrid=True))\n",
    "            x.append(partial(selectSubmatrixWithLeastColors, isGrid=True))\n",
    "            position = getBestSubmatrixPosition(candTask, isGrid=True)\n",
    "            x.append(partial(selectSubmatrixInPosition, position=position, isGrid=True))\n",
    "        \n",
    "        # Pixelwise And\n",
    "        for c in candTask.commonOutColors:\n",
    "            x.append(partial(pixelwiseAndInGridSubmatrices, falseColor=c))\n",
    "        if len(candTask.totalOutColors) == 2:\n",
    "            for target in candTask.totalInColors:\n",
    "                for c in permutations(candTask.totalOutColors, 2):\n",
    "                    x.append(partial(pixelwiseAndInGridSubmatrices, falseColor=c[0],\\\n",
    "                                     targetColor=target, trueColor=c[1]))\n",
    "                        \n",
    "        # Pixelwise Or\n",
    "        for c in candTask.commonOutColors:\n",
    "            x.append(partial(pixelwiseOrInGridSubmatrices, falseColor=c))\n",
    "        if len(candTask.totalOutColors) == 2:\n",
    "            for target in candTask.totalInColors:\n",
    "                for c in permutations(candTask.totalOutColors, 2):\n",
    "                    x.append(partial(pixelwiseOrInGridSubmatrices, falseColor=c[0],\\\n",
    "                                     targetColor=target, trueColor=c[1]))\n",
    "        if candTask.backgroundColor!=-1:\n",
    "            colors = candTask.commonOutColors - candTask.commonInColors\n",
    "            if candTask.trainSamples[0].inMatrix.grid.nCells==len(colors):\n",
    "                for c in permutations(colors, len(colors)):\n",
    "                    x.append(partial(pixelwiseOrInGridSubmatrices,\\\n",
    "                                     falseColor=candTask.backgroundColor,\\\n",
    "                                     trueValues = c))\n",
    "        \n",
    "        # Pixelwise Xor\n",
    "        if all([s.inMatrix.grid.nCells == 2 for s in candTask.trainSamples]) \\\n",
    "        and all([s.inMatrix.grid.nCells == 2 for s in candTask.testSamples]):\n",
    "            for c in candTask.commonOutColors:\n",
    "                x.append(partial(pixelwiseXorInGridSubmatrices, falseColor=c))\n",
    "            if len(candTask.commonOutColors - candTask.commonInColors)==2 and\\\n",
    "            candTask.backgroundColor!=-1:\n",
    "                colors = candTask.commonOutColors - candTask.commonInColors\n",
    "                for c in permutations(colors, 2):\n",
    "                    x.append(partial(pixelwiseXorInGridSubmatrices, falseColor=candTask.backgroundColor,\\\n",
    "                                     firstTrue=c[0], secondTrue=c[1]))\n",
    "            if len(candTask.totalOutColors) == 2:\n",
    "                for target in candTask.totalInColors:\n",
    "                    for c in permutations(candTask.totalOutColors, 2):\n",
    "                        x.append(partial(pixelwiseXorInGridSubmatrices, falseColor=c[0],\\\n",
    "                                         targetColor=target, trueColor=c[1]))\n",
    "       \n",
    "    if candTask.inputIsGrid:\n",
    "        if all([s.inMatrix.grid.shape==s.outMatrix.shape for s in candTask.trainSamples]):\n",
    "            for times in range(1, 6):\n",
    "                x.append(partial(colorAppearingXTimes, times=times))\n",
    "            x.append(partial(maxColorFromCell))\n",
    "               \n",
    "    x.append(getBestLayShapes(candTask)) \n",
    "    #x.append(getBestReplicateOneShape(candTask))\n",
    "    #tasks with two shapes\n",
    "    if candTask.twoShapeTask[0]:\n",
    "        x.append(getBestTwoShapeFunction(t))\n",
    "    # Cropshape\n",
    "    #x.append(partial(colorByPixels))\n",
    "    #x.append(partial(colorByPixels, colorMap=True))\n",
    "    #x.append(partial(colorByPixels, oneColor=True))\n",
    "    x.append(partial(overlapShapes))\n",
    "    if candTask.outSmallerThanIn:\n",
    "        #x.append(getBestAlignShapes(candTask))\n",
    "        #x.append(partial(deleteShapes, attributes = getDeleteAttributes(candTask, diagonal = False), diagonal = False, multicolor=False))\n",
    "        #x.append(partial(replicateShapes, allCombs=True, scale=False,attributes=set(['MoCl']),anchorType='subframe',deleteOriginal=True))\n",
    "        #x.append(partial(replicateShapes, allCombs=False, scale=True,attributes=set(['MoCl']),anchorType='subframe',deleteOriginal=True))\n",
    "        #x.append(getBestArrangeShapes(candTask))\n",
    "        if candTask.backgroundColor!=-1:\n",
    "            x.append(partial(cropAllShapes, background=candTask.backgroundColor, diagonal=True))\n",
    "            x.append(partial(cropAllShapes, background=candTask.backgroundColor, diagonal=False))\n",
    "        \n",
    "        bestCrop = getBestCropShape(candTask)\n",
    "        if 'attributes' in bestCrop.keywords.keys():\n",
    "            for attr in bestCrop.keywords['attributes']:\n",
    "                if isinstance(attr, str) and attr[:2] == 'mo' and len(bestCrop.keywords['attributes']) > 1:\n",
    "                    continue\n",
    "                newCrop = copy.deepcopy(bestCrop)\n",
    "                newCrop.keywords['attributes'] = set([attr])\n",
    "                x.append(newCrop)\n",
    "                \n",
    "        x.append(getBestCropReference(candTask))  \n",
    "    \n",
    "        for attrs in [set(['LaSh']),set(['UnCo'])]:\n",
    "            x.append(partial(cropShape, attributes=attrs, backgroundColor=max(0,candTask.backgroundColor),\\\n",
    "                                 singleColor=True, diagonals=True)) \n",
    "            x.append(partial(cropShape, attributes=attrs, backgroundColor=max(0,candTask.backgroundColor),\\\n",
    "                                 singleColor=True, diagonals=True, context=True)) \n",
    "        if candTask.outIsInMulticolorShapeSize:\n",
    "            x.append(getBestMoveToPanel(candTask))\n",
    "    if all([len(s.inMatrix.multicolorShapes)==1 for s in candTask.trainSamples+candTask.testSamples]):\n",
    "        x.append(partial(cropOnlyMulticolorShape, diagonals=False))\n",
    "    if all([len(s.inMatrix.multicolorDShapes)==1 for s in candTask.trainSamples+candTask.testSamples]):\n",
    "        x.append(partial(cropOnlyMulticolorShape, diagonals=True))\n",
    "    if all([len(sample.inMatrix.fullFrames)==1 for sample in candTask.trainSamples+candTask.testSamples]):\n",
    "        x.append(partial(cropFullFrame))\n",
    "        x.append(partial(cropFullFrame, includeBorder=False))\n",
    "    if all([len(sample.inMatrix.fullFrames)>1 for sample in candTask.trainSamples+candTask.testSamples]):\n",
    "        x.append(partial(cropFullFrame, bigOrSmall=\"big\"))\n",
    "        x.append(partial(cropFullFrame, bigOrSmall=\"small\"))\n",
    "        x.append(partial(cropFullFrame, bigOrSmall=\"big\", includeBorder=False))\n",
    "        x.append(partial(cropFullFrame, bigOrSmall=\"small\", includeBorder=False))\n",
    "    \n",
    "    #more frames\n",
    "    if candTask.hasPartialFrame:\n",
    "        x.append(getBestFitToFrame(candTask))\n",
    "        if candTask.outSmallerThanIn:\n",
    "            x.append(partial(cropPartialFrame, includeBorder=False))\n",
    "            x.append(partial(cropPartialFrame, includeBorder=True))\n",
    "    \n",
    "    if candTask.sameIOShapes:        \n",
    "        if candTask.sameNSampleColors and all([\"predictCNN\" not in str(op.func) for op in c.ops]):\n",
    "            x.append(getBestSameNSampleColorsCNN(candTask))\n",
    "    \n",
    "    # startOps\n",
    "\n",
    "    x.append(partial(paintGridLikeBackground))\n",
    "    x.append(partial(cropAllBackground))\n",
    "    \n",
    "    # minimize\n",
    "    if not candTask.sameIOShapes:\n",
    "        x.append(partial(minimize))\n",
    "    \n",
    "    return x\n",
    "\n",
    "###############################################################################\n",
    "###############################################################################\n",
    "# Submission Setup\n",
    "    \n",
    "class Candidate():\n",
    "    \"\"\"\n",
    "    Objects of the class Candidate store the information about a possible\n",
    "    candidate for the solution.\n",
    "\n",
    "    ...\n",
    "    Attributes\n",
    "    ----------\n",
    "    ops: list\n",
    "        A list containing the operations to be performed to the input matrix\n",
    "        in order to get to the solution. The elements of the list are partial\n",
    "        functions (from functools.partial).\n",
    "    score: int\n",
    "        The score of the candidate. The score is defined as the sum of the\n",
    "        number incorrect pixels when applying ops to the input matrices of the\n",
    "        train samples of the task.\n",
    "    tasks: list\n",
    "        A list containing the tasks (in its original format) after performing\n",
    "        each of the operations in ops, starting from the original inputs.\n",
    "    t: Task\n",
    "        The Task object corresponding to the current status of the task.\n",
    "        This is, the status after applying all the operations of ops to the\n",
    "        input matrices of the task.\n",
    "    seen: bool\n",
    "        True if we have already checked and executed the possible operations\n",
    "        of the candidate. False otherwise.\n",
    "    \"\"\"\n",
    "    def __init__(self, ops, tasks, score=1000, predictions=np.zeros((2,2))):\n",
    "        self.ops = ops\n",
    "        self.score = score\n",
    "        self.tasks = tasks\n",
    "        self.t = None\n",
    "        self.predictions = predictions\n",
    "        self.seen = False\n",
    "\n",
    "    def __lt__(self, other):\n",
    "        \"\"\"\n",
    "        A candidate is better than another one if its score is lower.\n",
    "        \"\"\"\n",
    "        if self.score == other.score:\n",
    "            return len(self.ops) < len(other.ops)\n",
    "        return self.score < other.score\n",
    "\n",
    "    def generateTask(self):\n",
    "        \"\"\"\n",
    "        Assign to the attribute t the Task object corresponding to the\n",
    "        current task status.\n",
    "        \"\"\"\n",
    "        self.t = Task(self.tasks[-1], 'dummyIndex', submission=True)\n",
    "\n",
    "class Best3Candidates():\n",
    "    \"\"\"\n",
    "    An object of this class stores the three best candidates of a task.\n",
    "\n",
    "    ...\n",
    "    Attributes\n",
    "    ----------\n",
    "    candidates: list\n",
    "        A list of three elements, each one of them being an object of the class\n",
    "        Candidate.\n",
    "    \"\"\"\n",
    "    def __init__(self, Candidate1, Candidate2, Candidate3):\n",
    "        self.candidates = [Candidate1, Candidate2, Candidate3]\n",
    "\n",
    "    def maxCandidate(self):\n",
    "        \"\"\"\n",
    "        Returns the index of the candidate with highest score.\n",
    "        \"\"\"\n",
    "        x = 0\n",
    "        if self.candidates[1] > self.candidates[0]:\n",
    "            x = 1\n",
    "        if self.candidates[2] > self.candidates[x]:\n",
    "            x = 2\n",
    "        return x\n",
    "\n",
    "    def addCandidate(self, c):\n",
    "        \"\"\"\n",
    "        Given a candidate c, this function substitutes c with the worst\n",
    "        candidate in self.candidates only if it's a better candidate (its score\n",
    "        is lower).\n",
    "        \"\"\"\n",
    "        if all([self.candidates[i].score < c.score for i in range(3)]):\n",
    "            return\n",
    "        \n",
    "        for i in range(3):\n",
    "            if all([np.array_equal(self.candidates[i].predictions[x], c.predictions[x]) \\\n",
    "                    for x in range(len(c.predictions))]):\n",
    "                return\n",
    "        iMaxCand = self.maxCandidate()\n",
    "        for i in range(3):\n",
    "            if c < self.candidates[iMaxCand]:\n",
    "                c.generateTask()\n",
    "                self.candidates[iMaxCand] = c\n",
    "                break\n",
    "\n",
    "    def allPerfect(self):\n",
    "        return all([c.score==0 for c in self.candidates])\n",
    "\n",
    "    def getOrderedIndices(self):\n",
    "        \"\"\"\n",
    "        Returns a list of 3 indices (from 0 to 2) with the candidates ordered\n",
    "        from best to worst.\n",
    "        \"\"\"\n",
    "        orderedList = [0]\n",
    "        if self.candidates[1] < self.candidates[0]:\n",
    "            orderedList.insert(0, 1)\n",
    "        else:\n",
    "            orderedList.append(1)\n",
    "        if self.candidates[2] < self.candidates[orderedList[0]]:\n",
    "            orderedList.insert(0, 2)\n",
    "        elif self.candidates[2] < self.candidates[orderedList[1]]:\n",
    "            orderedList.insert(1, 2)\n",
    "        else:\n",
    "            orderedList.append(2)\n",
    "        return orderedList\n",
    "    \n",
    "# Separate task by shapes\n",
    "class TaskSeparatedByShapes():\n",
    "    \"\"\"\n",
    "    An object of this class stores a Task that has been separated by Shapes, \n",
    "    as well as the necessary information to restore it into the original the \n",
    "    matrices according to the original inputs.\n",
    "    ...\n",
    "    Attributes\n",
    "    ----------\n",
    "    originalTask: dict\n",
    "        The original task in its original format.\n",
    "    separatedTask: Task\n",
    "        The task separated by shapes, according to the criteria defined in\n",
    "        needsSeparationByShapes.\n",
    "    nShapes: dict\n",
    "        Stores the number of Shapes used to separate each train and test\n",
    "        sample.\n",
    "    background: int\n",
    "        The background color of the original task, necessary for the restoring\n",
    "        step.\n",
    "    mergeColor: int\n",
    "        Color to use when merging matrices if there is a conflict.\n",
    "    \"\"\"\n",
    "    def __init__(self, task, background, diagonal=False):\n",
    "        self.originalTask = task\n",
    "        self.separatedTask = None\n",
    "        self.nShapes = {'train': [], 'test': []}\n",
    "        self.background = background\n",
    "        self.mergeColor = None\n",
    "\n",
    "    def getRange(self, trainOrTest, index):\n",
    "        \"\"\"\n",
    "        Returns the range of matrices in the separated task to be merged in\n",
    "        order to obtain the final trainOrTest matrix on the given index.\n",
    "        \"\"\"\n",
    "        i, position = 0, 0\n",
    "        while i < index:\n",
    "            position += self.nShapes[trainOrTest][i]\n",
    "            i += 1\n",
    "        return (position, position+self.nShapes[trainOrTest][index])\n",
    "            \n",
    "def needsSeparationByShapes(t):\n",
    "    \"\"\"\n",
    "    This function checks whether the Task t needs to be separated by Shapes.\n",
    "    If that's the case, it returns an object of the class TaskSeparatedByShapes, \n",
    "    and otherwise it returns False.\n",
    "    A Task can be separated by Shapes if all the Samples can. A Sample can be\n",
    "    separated by Shapes if it has a background color, and the same number of\n",
    "    Shapes in the input and in the output, in a way that a Shape in the input\n",
    "    clearly corresponds to a Shape in the output. For the task to need\n",
    "    separation by Shapes, at least one of the Samples has to have two separable\n",
    "    Shapes.\n",
    "    Separability by Shapes is checked for single color, multicolor and diagonal\n",
    "    Shapes.\n",
    "    \"\"\"\n",
    "    def getOverlap(inShape, inPos, outShape, outPos):\n",
    "        x1a, y1a, x1b, y1b = inPos[0], inPos[1], outPos[0], outPos[1]\n",
    "        x2a, y2a = inPos[0]+inShape[0]-1, inPos[1]+inShape[1]-1\n",
    "        x2b, y2b = outPos[0]+outShape[0]-1, outPos[1]+outShape[1]-1\n",
    "        if x1a<=x1b:\n",
    "            if x2a<=x1b:\n",
    "                return 0\n",
    "            x = x2a-x1b+1\n",
    "        elif x1b<=x1a:\n",
    "            if x2b<=x1a:\n",
    "                return 0\n",
    "            x = x2b-x1a+1\n",
    "        if y1a<=y1b:\n",
    "            if y2a<=y1b:\n",
    "                return 0\n",
    "            y = y2a-y1b+1\n",
    "        elif y1b<=y1a:\n",
    "            if y2b<=y1a:\n",
    "                return 0\n",
    "            y = y2b-y1a+1\n",
    "\n",
    "        return x*y\n",
    "    \n",
    "    def generateNewTask(inShapes, outShapes, testShapes):\n",
    "        # Assign every input shape to the output shape with maximum overlap\n",
    "        separatedTask = TaskSeparatedByShapes(t.task.copy(), t.backgroundColor)\n",
    "        task = {'train': [], 'test': []}\n",
    "        for s in range(t.nTrain):\n",
    "            seenIndices = set()\n",
    "            for inShape in inShapes[s]:\n",
    "                shapeIndex = 0\n",
    "                maxOverlap = 0\n",
    "                bestIndex = -1\n",
    "                for outShape in outShapes[s]:\n",
    "                    overlap = getOverlap(inShape.shape, inShape.position, outShape.shape, outShape.position)\n",
    "                    if overlap > maxOverlap:\n",
    "                        maxOverlap = overlap\n",
    "                        bestIndex = shapeIndex\n",
    "                    shapeIndex += 1\n",
    "                if bestIndex!=-1 and bestIndex not in seenIndices:\n",
    "                    seenIndices.add(bestIndex)\n",
    "                    # Generate the new input and output matrices\n",
    "                    inM = np.full(t.trainSamples[s].inMatrix.shape, t.backgroundColor ,dtype=np.uint8)\n",
    "                    outM = inM.copy()\n",
    "                    inM = insertShape(inM, inShape)\n",
    "                    outM = insertShape(outM, outShapes[s][bestIndex])\n",
    "                    task['train'].append({'input': inM.tolist(), 'output': outM.tolist()})\n",
    "            # If we haven't dealt with all the shapes successfully, then return\n",
    "            if len(seenIndices) != len(inShapes[s]):\n",
    "                return False\n",
    "            # Record the number of new samples generated by sample s\n",
    "            separatedTask.nShapes['train'].append(len(inShapes[s]))\n",
    "        for s in range(t.nTest):\n",
    "            for testShape in testShapes[s]:\n",
    "                inM = np.full(t.testSamples[s].inMatrix.shape, t.backgroundColor ,dtype=np.uint8)\n",
    "                inM = insertShape(inM, testShape)\n",
    "                if t.submission:\n",
    "                    task['test'].append({'input': inM.tolist()})\n",
    "                else:\n",
    "                    task['test'].append({'input': inM.tolist(), 'output': t.testSamples[s].outMatrix.m.tolist()})\n",
    "            # Record the number of new samples generated by sample s\n",
    "            separatedTask.nShapes['test'].append(len(testShapes[s]))\n",
    "                \n",
    "        \n",
    "        # Complete and return the TaskSeparatedByShapes object\n",
    "        separatedTask.separatedTask = task.copy()\n",
    "        return separatedTask\n",
    "        \n",
    "\n",
    "    # I need to have a background color to generate the new task object\n",
    "    if t.backgroundColor==-1 or not t.sameIOShapes:\n",
    "        return False\n",
    "    # Only consider tasks without small matrices\n",
    "    if any([s.inMatrix.shape[0]*s.inMatrix.shape[1]<43 for s in t.trainSamples+t.testSamples]):\n",
    "        return False\n",
    "\n",
    "    mergeColors = t.commonOutColors - t.totalInColors\n",
    "    if len(mergeColors) == 1:\n",
    "        mergeColor = next(iter(mergeColors))\n",
    "    else:\n",
    "        mergeColor = None\n",
    "    \n",
    "    # First, consider normal shapes (not background, not diagonal, not multicolor) (Task 84 as example)\n",
    "    inShapes = [[shape for shape in s.inMatrix.shapes if shape.color!=t.backgroundColor] for s in t.trainSamples]\n",
    "    outShapes = [[shape for shape in s.outMatrix.shapes if shape.color!=t.backgroundColor] for s in t.trainSamples]\n",
    "    testShapes = [[shape for shape in s.inMatrix.shapes if shape.color!=t.backgroundColor] for s in t.testSamples]\n",
    "    if all([len(inShapes[s])<=7 and len(inShapes[s])==len(outShapes[s]) for s in range(t.nTrain)]):\n",
    "        newTask = generateNewTask(inShapes, outShapes, testShapes)\n",
    "        if newTask != False:\n",
    "            if len(mergeColors) == 1:\n",
    "                newTask.mergeColor = mergeColor\n",
    "            return newTask\n",
    "        \n",
    "    # Now, consider diagonal shapes (Task 681 as example)\n",
    "    inShapes = [[shape for shape in s.inMatrix.dShapes if shape.color!=t.backgroundColor] for s in t.trainSamples]\n",
    "    outShapes = [[shape for shape in s.outMatrix.dShapes if shape.color!=t.backgroundColor] for s in t.trainSamples]\n",
    "    testShapes = [[shape for shape in s.inMatrix.dShapes if shape.color!=t.backgroundColor] for s in t.testSamples]\n",
    "    if all([len(inShapes[s])<=5 and len(inShapes[s])==len(outShapes[s]) for s in range(t.nTrain)]):\n",
    "        newTask = generateNewTask(inShapes, outShapes, testShapes)\n",
    "        if newTask != False:\n",
    "            if len(mergeColors) == 1:\n",
    "                newTask.mergeColor = mergeColor\n",
    "            return newTask\n",
    "    \n",
    "    # Now, multicolor non-diagonal shapes (Task 611 as example)\n",
    "    inShapes = [[shape for shape in s.inMatrix.multicolorShapes] for s in t.trainSamples]\n",
    "    outShapes = [[shape for shape in s.outMatrix.multicolorShapes] for s in t.trainSamples]\n",
    "    testShapes = [[shape for shape in s.inMatrix.multicolorShapes] for s in t.testSamples]\n",
    "    if all([len(inShapes[s])<=7 and len(inShapes[s])==len(outShapes[s]) for s in range(t.nTrain)]):\n",
    "        newTask = generateNewTask(inShapes, outShapes, testShapes)\n",
    "        if newTask != False:\n",
    "            if len(mergeColors) == 1:\n",
    "                newTask.mergeColor = mergeColor\n",
    "            return newTask\n",
    "    \n",
    "    # Finally, multicolor diagonal (Task 610 as example)\n",
    "    inShapes = [[shape for shape in s.inMatrix.multicolorDShapes] for s in t.trainSamples]\n",
    "    outShapes = [[shape for shape in s.outMatrix.multicolorDShapes] for s in t.trainSamples]\n",
    "    testShapes = [[shape for shape in s.inMatrix.multicolorDShapes] for s in t.testSamples]\n",
    "    if all([len(inShapes[s])<=5 and len(inShapes[s])==len(outShapes[s]) for s in range(t.nTrain)]):\n",
    "        newTask = generateNewTask(inShapes, outShapes, testShapes)\n",
    "        if newTask != False:\n",
    "            if len(mergeColors) == 1:\n",
    "                newTask.mergeColor = mergeColor\n",
    "            return newTask\n",
    "\n",
    "    return False\n",
    "\n",
    "# Separate task by colors\n",
    "class TaskSeparatedByColors():\n",
    "    \"\"\"\n",
    "    An object of this class stores a Task that has been separated by colors, \n",
    "    as well as the necessary information to restore it into the original the \n",
    "    matrices according to the original inputs.\n",
    "    ...\n",
    "    Attributes\n",
    "    ----------\n",
    "    originalTask: dict\n",
    "        The original task in its original format.\n",
    "    separatedTask: Task\n",
    "        The task separated by shapes, according to the criteria defined in\n",
    "        needsSeparationByShapes.\n",
    "    commonColors: set\n",
    "        The colors that appear in every sample.\n",
    "    extraColors: dict\n",
    "        The number of colors that appear in every sample and are different from\n",
    "        the commonColors.\n",
    "    \"\"\"\n",
    "    def __init__(self, task):\n",
    "        self.originalTask = task\n",
    "        self.separatedTask = None\n",
    "        self.commonColors = None\n",
    "        self.extraColors = {'train': [], 'test': []}\n",
    "\n",
    "    def getRange(self, trainOrTest, index):\n",
    "        \"\"\"\n",
    "        Returns the range of matrices in the separated task to be merged in\n",
    "        order to obtain the final trainOrTest matrix on the given index.\n",
    "        \"\"\"\n",
    "        i, position = 0, 0\n",
    "        while i < index:\n",
    "            position += len(self.extraColors[trainOrTest][i])\n",
    "            i += 1\n",
    "        return (position, position+len(self.extraColors[trainOrTest][index]))\n",
    "\n",
    "\n",
    "def needsSeparationByColors(t):\n",
    "    \"\"\"\n",
    "    This function checks whether the Task t needs to be separated by colors.\n",
    "    If that's the case, it returns an object of the class TaskSeparatedByColors, \n",
    "    and otherwise it returns False.\n",
    "    A Task can be separated by colors the number of colors in each Sample is\n",
    "    different. If that's the case, every Sample will be converted into as many\n",
    "    samples as non-common colors it has (common color refers to color that \n",
    "    appears in every Sample).\n",
    "    \"\"\"\n",
    "    def generateMatrix(matrix, colorsToKeep, backgroundColor):\n",
    "        m = matrix.copy()\n",
    "        for i,j in np.ndindex(matrix.shape):\n",
    "            if m[i,j] not in colorsToKeep:\n",
    "                m[i,j] = backgroundColor\n",
    "\n",
    "        return m\n",
    "\n",
    "    def generateNewTask(commonColors, backgroundColor):\n",
    "        # Assign every input shape to the output shape with maximum overlap\n",
    "        separatedTask = TaskSeparatedByColors(t.task.copy())\n",
    "        task = {'train': [], 'test': []}\n",
    "        for s in range(t.nTrain):\n",
    "            separatedTask.extraColors['train'].append([])\n",
    "            colorsToConsider = (t.trainSamples[s].inMatrix.colors | t.trainSamples[s].outMatrix.colors)\\\n",
    "                                - commonColors\n",
    "            if len(colorsToConsider)==0:\n",
    "                return False\n",
    "            for color in colorsToConsider:\n",
    "                separatedTask.extraColors['train'][s].append(color)\n",
    "                inM = generateMatrix(t.trainSamples[s].inMatrix.m, commonColors|set([color]), backgroundColor)\n",
    "                outM = generateMatrix(t.trainSamples[s].outMatrix.m, commonColors|set([color]), backgroundColor)\n",
    "                task['train'].append({'input': inM.tolist(), 'output': outM.tolist()})\n",
    "\n",
    "        for s in range(t.nTest):\n",
    "            separatedTask.extraColors['test'].append([])\n",
    "            if t.submission:\n",
    "                colorsToConsider = t.testSamples[s].inMatrix.colors - commonColors\n",
    "                if len(colorsToConsider)==0:\n",
    "                    return False\n",
    "                for color in colorsToConsider:\n",
    "                    separatedTask.extraColors['test'][s].append(color)\n",
    "                    inM = generateMatrix(t.testSamples[s].inMatrix.m, commonColors|set([color]), backgroundColor)\n",
    "                    task['test'].append({'input': inM.tolist()})\n",
    "            else:\n",
    "                colorsToConsider = (t.testSamples[s].inMatrix.colors | t.testSamples[s].outMatrix.colors)\\\n",
    "                                    - commonColors\n",
    "                if len(colorsToConsider)==0:\n",
    "                    return False\n",
    "                for color in colorsToConsider:\n",
    "                    separatedTask.extraColors['test'][s].append(color)\n",
    "                    inM = generateMatrix(t.testSamples[s].inMatrix.m, commonColors|set([color]), backgroundColor)\n",
    "                    outM = generateMatrix(t.testSamples[s].outMatrix.m, commonColors|set([color]), backgroundColor)\n",
    "                    task['test'].append({'input': inM.tolist(), 'output': t.testSamples[s].outMatrix.m.tolist()})\n",
    "\n",
    "        # Complete and return the TaskSeparatedByShapes object\n",
    "        separatedTask.separatedTask = task.copy()\n",
    "        return separatedTask\n",
    "\n",
    "\n",
    "    # I need to have a background color to generate the new task object\n",
    "    if t.backgroundColor==-1 or not t.sameIOShapes:\n",
    "        return False\n",
    "    # Only consider tasks without small matrices\n",
    "    if any([s.inMatrix.shape[0]*s.inMatrix.shape[1]<43 for s in t.trainSamples+t.testSamples]):\n",
    "        return False\n",
    "\n",
    "    commonColors = t.commonInColors | t.commonOutColors\n",
    "\n",
    "    if all([sample.nColors == len(commonColors) for sample in t.trainSamples]):\n",
    "        return False\n",
    "    if any([sample.nColors < len(commonColors) for sample in t.trainSamples]):\n",
    "        return False\n",
    "\n",
    "    newTask = generateNewTask(commonColors, t.backgroundColor)\n",
    "\n",
    "    return newTask\n",
    "\n",
    "# Crop task if necessary\n",
    "def getCroppingPosition(matrix):\n",
    "    \"\"\"\n",
    "    Function to be used only if the the Task is to be cropped. If that's the\n",
    "    case, given a matrix, this function returns the cropping position.\n",
    "    \"\"\"\n",
    "    bC = matrix.backgroundColor\n",
    "    x, xMax, y, yMax = 0, matrix.m.shape[0]-1, 0, matrix.m.shape[1]-1\n",
    "    while x <= xMax and np.all(matrix.m[x,:] == bC):\n",
    "        x += 1\n",
    "    while y <= yMax and np.all(matrix.m[:,y] == bC):\n",
    "        y += 1\n",
    "    return [x,y]\n",
    "    \n",
    "def needsCropping(t):\n",
    "    \"\"\"\n",
    "    This function checks whether a Task needs to be cropped. A Task needs to be\n",
    "    cropped if any of the background is considered irrelevant, this is, if a\n",
    "    part of the background is not modified.\n",
    "    \"\"\"\n",
    "    # Only to be used if t.sameIOShapes\n",
    "    for sample in t.trainSamples:\n",
    "        if sample.inMatrix.backgroundColor != sample.outMatrix.backgroundColor:\n",
    "            return False\n",
    "        if getCroppingPosition(sample.inMatrix) != getCroppingPosition(sample.outMatrix):\n",
    "            return False\n",
    "        inMatrix = cropAllBackground(sample.inMatrix)\n",
    "        outMatrix = cropAllBackground(sample.outMatrix)\n",
    "        if inMatrix.shape!=outMatrix.shape or sample.inMatrix.shape==inMatrix.shape:\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "def cropTask(t, task):\n",
    "    \"\"\"\n",
    "    This function crops the background of all the matrices of the Task t. The\n",
    "    result of cropping the matrices is performed in-place in for the given\n",
    "    task, and the positions and background colors of every matrix are returned\n",
    "    in order to recover them in the revert process.\n",
    "    \"\"\"\n",
    "    positions = {\"train\": [], \"test\": []}\n",
    "    backgrounds = {\"train\": [], \"test\": []}\n",
    "    for s in range(t.nTrain):\n",
    "        task[\"train\"][s][\"input\"] = cropAllBackground(t.trainSamples[s].inMatrix).tolist()\n",
    "        task[\"train\"][s][\"output\"] = cropAllBackground(t.trainSamples[s].outMatrix).tolist()\n",
    "        backgrounds[\"train\"].append(t.trainSamples[s].inMatrix.backgroundColor)\n",
    "        positions[\"train\"].append(getCroppingPosition(t.trainSamples[s].inMatrix))\n",
    "    for s in range(t.nTest):\n",
    "        task[\"test\"][s][\"input\"] = cropAllBackground(t.testSamples[s].inMatrix).tolist()\n",
    "        backgrounds[\"test\"].append(t.testSamples[s].inMatrix.backgroundColor)\n",
    "        positions[\"test\"].append(getCroppingPosition(t.testSamples[s].inMatrix))\n",
    "        if not t.submission:\n",
    "            task[\"test\"][s][\"output\"] = cropAllBackground(t.testSamples[s].outMatrix).tolist()\n",
    "    return positions, backgrounds\n",
    "\n",
    "def recoverCroppedMatrix(matrix, outShape, position, backgroundColor):\n",
    "    \"\"\"\n",
    "    Function to revert the cropping of the matrices that had been cropped with\n",
    "    the function cropTask.\n",
    "    \"\"\"\n",
    "    m = np.full(outShape, backgroundColor, dtype=np.uint8)\n",
    "    m[position[0]:position[0]+matrix.shape[0], position[1]:position[1]+matrix.shape[1]] = matrix.copy()\n",
    "    return m\n",
    "    \n",
    "def needsRecoloring(t):\n",
    "    \"\"\"\n",
    "    This method determines whether the task t needs recoloring or not.\n",
    "    It needs recoloring if every color in an output matrix appears either\n",
    "    in the input or in every output matrix.\n",
    "    Otherwise a recoloring doesn't make sense.\n",
    "    If this function returns True, then orderTaskColors should be executed\n",
    "    as the first part of the preprocessing of t.\n",
    "    \"\"\"\n",
    "    for sample in t.trainSamples:\n",
    "        for color in sample.outMatrix.colors:\n",
    "            if (color not in sample.inMatrix.colors) and (color not in t.commonOutColors):\n",
    "                return False\n",
    "    return True\n",
    "\n",
    "def orderTaskColors(t):\n",
    "    \"\"\"\n",
    "    Given a task t, this function generates a new task (as a dictionary) by\n",
    "    recoloring all the matrices in a specific way.\n",
    "    The goal of this function is to impose that if two different colors\n",
    "    represent the exact same thing in two different samples, then they have the\n",
    "    same color in both of the samples.\n",
    "    Right now, the criterium to order colors is:\n",
    "        1. Common colors ordered according to Task.orderColors\n",
    "        2. Colors that appear both in the input and the output\n",
    "        3. Colors that only appear in the input\n",
    "        4. Colors that only appear in the output\n",
    "    In steps 2-4, if there is more that one color satisfying that condition, \n",
    "    the ordering will happen according to the colorCount.\n",
    "    \"\"\"\n",
    "    def orderColors(trainOrTest):\n",
    "        if trainOrTest==\"train\":\n",
    "            samples = t.trainSamples\n",
    "        else:\n",
    "            samples = t.testSamples\n",
    "        for sample in samples:\n",
    "            sampleColors = t.orderedColors.copy()\n",
    "            sortedColors = [k for k, v in sorted(sample.inMatrix.colorCount.items(), key=lambda item: item[1])]\n",
    "            for c in sortedColors:\n",
    "                if c not in sampleColors:\n",
    "                    sampleColors.append(c)\n",
    "            if trainOrTest==\"train\" or t.submission==False:\n",
    "                sortedColors = [k for k, v in sorted(sample.outMatrix.colorCount.items(), key=lambda item: item[1])]\n",
    "                for c in sortedColors:\n",
    "                    if c not in sampleColors:\n",
    "                        sampleColors.append(c)\n",
    "                    \n",
    "            rel, invRel = relDicts(sampleColors)\n",
    "            if trainOrTest==\"train\":\n",
    "                trainRels.append(rel)\n",
    "                trainInvRels.append(invRel)\n",
    "            else:\n",
    "                testRels.append(rel)\n",
    "                testInvRels.append(invRel)\n",
    "                \n",
    "            inMatrix = np.zeros(sample.inMatrix.shape, dtype=np.uint8)\n",
    "            for c in sample.inMatrix.colors:\n",
    "                inMatrix[sample.inMatrix.m==c] = invRel[c]\n",
    "            if trainOrTest=='train' or t.submission==False:\n",
    "                outMatrix = np.zeros(sample.outMatrix.shape, dtype=np.uint8)\n",
    "                for c in sample.outMatrix.colors:\n",
    "                    outMatrix[sample.outMatrix.m==c] = invRel[c]\n",
    "                if trainOrTest=='train':\n",
    "                    task['train'].append({'input': inMatrix.tolist(), 'output': outMatrix.tolist()})\n",
    "                else:\n",
    "                    task['test'].append({'input': inMatrix.tolist(), 'output': outMatrix.tolist()})\n",
    "            else:\n",
    "                task['test'].append({'input': inMatrix.tolist()})\n",
    "        \n",
    "    task = {'train': [], 'test': []}\n",
    "    trainRels = []\n",
    "    trainInvRels = []\n",
    "    testRels = []\n",
    "    testInvRels = []\n",
    "    \n",
    "    orderColors(\"train\")\n",
    "    orderColors(\"test\")\n",
    "    \n",
    "    return task, trainRels, trainInvRels, testRels, testInvRels\n",
    "\n",
    "def recoverOriginalColors(matrix, rel):\n",
    "    \"\"\"\n",
    "    Given a matrix, this function is intended to recover the original colors\n",
    "    before being modified in the orderTaskColors function.\n",
    "    rel is supposed to be either one of the trainRels or testRels outputs of\n",
    "    that function.\n",
    "    \"\"\"\n",
    "    m = matrix.copy()\n",
    "    for i,j in np.ndindex(matrix.shape):\n",
    "        if matrix[i,j] in rel.keys(): # TODO Task 162 fails. Delete this when fixed\n",
    "            m[i,j] = rel[matrix[i,j]][0]\n",
    "    return m\n",
    "\n",
    "def ignoreGrid(t, task, inMatrix=True, outMatrix=True):\n",
    "    \"\"\"\n",
    "    Given a Task t with a grid and its corresponding task dictionary, this\n",
    "    function modifies task in-place by considering each cell of the grid as a\n",
    "    pixel in the new matrix. For doing this, all the cells in the grid need to\n",
    "    have only one color.\n",
    "    \"\"\"\n",
    "    for s in range(t.nTrain):\n",
    "        if inMatrix:\n",
    "            m = np.zeros(t.trainSamples[s].inMatrix.grid.shape, dtype=np.uint8)\n",
    "            for i,j in np.ndindex(m.shape):\n",
    "                m[i,j] = next(iter(t.trainSamples[s].inMatrix.grid.cells[i][j][0].colors))\n",
    "            task[\"train\"][s][\"input\"] = m.tolist()\n",
    "        if outMatrix:\n",
    "            m = np.zeros(t.trainSamples[s].outMatrix.grid.shape, dtype=np.uint8)\n",
    "            for i,j in np.ndindex(m.shape):\n",
    "                m[i,j] = next(iter(t.trainSamples[s].outMatrix.grid.cells[i][j][0].colors))\n",
    "            task[\"train\"][s][\"output\"] = m.tolist()\n",
    "    for s in range(t.nTest):\n",
    "        if inMatrix:\n",
    "            m = np.zeros(t.testSamples[s].inMatrix.grid.shape, dtype=np.uint8)\n",
    "            for i,j in np.ndindex(m.shape):\n",
    "                m[i,j] = next(iter(t.testSamples[s].inMatrix.grid.cells[i][j][0].colors))\n",
    "            task[\"test\"][s][\"input\"] = m.tolist()\n",
    "        if outMatrix and not t.submission:\n",
    "            m = np.zeros(t.testSamples[s].outMatrix.grid.shape, dtype=np.uint8)\n",
    "            for i,j in np.ndindex(m.shape):\n",
    "                m[i,j] = next(iter(t.testSamples[s].outMatrix.grid.cells[i][j][0].colors))\n",
    "            task[\"test\"][s][\"output\"] = m.tolist()\n",
    "\n",
    "def recoverGrid(t, x, s):\n",
    "    \"\"\"\n",
    "    Given a matrix x, this function recovers the grid that had been removed\n",
    "    with ignoreGrid.\n",
    "    \"\"\"\n",
    "    realX = t.testSamples[s].inMatrix.m.copy()\n",
    "    cells = t.testSamples[s].inMatrix.grid.cells\n",
    "    for cellI in range(len(cells)):\n",
    "        for cellJ in range(len(cells[0])):\n",
    "            cellShape = cells[cellI][cellJ][0].shape\n",
    "            position = cells[cellI][cellJ][1]\n",
    "            for k,l in np.ndindex(cellShape):\n",
    "                realX[position[0]+k, position[1]+l] = x[cellI,cellJ]\n",
    "    return realX\n",
    "\n",
    "def ignoreAsymmetricGrid(t, task):\n",
    "    \"\"\"\n",
    "    Given a Task t with a grid and its corresponding task dictionary, this\n",
    "    function modifies task in-place by considering each cell of the asymmetric\n",
    "    grid as a pixel in the new matrix. For doing this, all the cells in the\n",
    "    grid need to have only one color.\n",
    "    \"\"\"\n",
    "    for s in range(t.nTrain):\n",
    "        m = np.zeros(t.trainSamples[s].inMatrix.asymmetricGrid.shape, dtype=np.uint8)\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            m[i,j] = next(iter(t.trainSamples[s].inMatrix.asymmetricGrid.cells[i][j][0].colors))\n",
    "        task[\"train\"][s][\"input\"] = m.tolist()\n",
    "        m = np.zeros(t.trainSamples[s].outMatrix.asymmetricGrid.shape, dtype=np.uint8)\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            m[i,j] = next(iter(t.trainSamples[s].outMatrix.asymmetricGrid.cells[i][j][0].colors))\n",
    "        task[\"train\"][s][\"output\"] = m.tolist()\n",
    "    for s in range(t.nTest):\n",
    "        m = np.zeros(t.testSamples[s].inMatrix.asymmetricGrid.shape, dtype=np.uint8)\n",
    "        for i,j in np.ndindex(m.shape):\n",
    "            m[i,j] = next(iter(t.testSamples[s].inMatrix.asymmetricGrid.cells[i][j][0].colors))\n",
    "        task[\"test\"][s][\"input\"] = m.tolist()\n",
    "        if not t.submission:\n",
    "            m = np.zeros(t.testSamples[s].outMatrix.asymmetricGrid.shape, dtype=np.uint8)\n",
    "            for i,j in np.ndindex(m.shape):\n",
    "                m[i,j] = next(iter(t.testSamples[s].outMatrix.asymmetricGrid.cells[i][j][0].colors))\n",
    "            task[\"test\"][s][\"output\"] = m.tolist()\n",
    "\n",
    "def recoverAsymmetricGrid(t, x, s):\n",
    "    \"\"\"\n",
    "    Given a matrix x, this function recovers the asymmetric grid that had been\n",
    "    removed with ignoreAsymmetricGrid.\n",
    "    \"\"\"\n",
    "    realX = t.testSamples[s].inMatrix.m.copy()\n",
    "    cells = t.testSamples[s].inMatrix.asymmetricGrid.cells\n",
    "    for cellI in range(len(cells)):\n",
    "        for cellJ in range(len(cells[0])):\n",
    "            cellShape = cells[cellI][cellJ][0].shape\n",
    "            position = cells[cellI][cellJ][1]\n",
    "            for k,l in np.ndindex(cellShape):\n",
    "                realX[position[0]+k, position[1]+l] = x[cellI,cellJ]\n",
    "    return realX\n",
    "\n",
    "def rotateTaskWithOneBorder(t, task):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function determines whether any of the Samples needs\n",
    "    have its Matrices rotated and, if that's the case, the task with the\n",
    "    rotated Matrices is returned. Otherwise, it returns false.\n",
    "    Matrices will be rotated if all the Samples contain a fixed Frontier that\n",
    "    is in one of the borders of the Matrix. If that's the case, it will be made\n",
    "    sure that the Frontier will be vertical and along the first column.\n",
    "    \"\"\"\n",
    "    rotTask = copy.deepcopy(task)\n",
    "    rotations = {'train': [], 'test': []}\n",
    "    for s in range(t.nTrain):\n",
    "        border = t.trainSamples[s].commonFullBorders[0]\n",
    "        if border.direction=='h' and border.position==0:\n",
    "            rotations['train'].append(1)\n",
    "            rotTask['train'][s]['input'] = np.rot90(t.trainSamples[s].inMatrix.m, 1).tolist()\n",
    "            rotTask['train'][s]['output'] = np.rot90(t.trainSamples[s].outMatrix.m, 1).tolist()\n",
    "        elif border.direction=='v' and border.position==t.trainSamples[s].inMatrix.shape[1]-1:\n",
    "            rotations['train'].append(2)\n",
    "            rotTask['train'][s]['input'] = np.rot90(t.trainSamples[s].inMatrix.m, 2).tolist()\n",
    "            rotTask['train'][s]['output'] = np.rot90(t.trainSamples[s].outMatrix.m, 2).tolist()\n",
    "        elif border.direction=='h' and border.position==t.trainSamples[s].inMatrix.shape[0]-1:\n",
    "            rotations['train'].append(3)\n",
    "            rotTask['train'][s]['input'] = np.rot90(t.trainSamples[s].inMatrix.m, 3).tolist()\n",
    "            rotTask['train'][s]['output'] = np.rot90(t.trainSamples[s].outMatrix.m, 3).tolist()\n",
    "        else:\n",
    "            rotations['train'].append(0)\n",
    "    \n",
    "    for s in range(t.nTest):\n",
    "        if t.submission:\n",
    "            hasBorder=False\n",
    "            for border in t.testSamples[s].inMatrix.fullBorders:\n",
    "                if border.color!=t.testSamples[s].inMatrix.backgroundColor:\n",
    "                    if border.direction=='h' and border.position==0:\n",
    "                        rotations['test'].append(1)\n",
    "                        rotTask['test'][s]['input'] = np.rot90(t.testSamples[s].inMatrix.m, 1).tolist()\n",
    "                    elif border.direction=='v' and border.position==t.testSamples[s].inMatrix.shape[1]-1:\n",
    "                        rotations['test'].append(2)\n",
    "                        rotTask['test'][s]['input'] = np.rot90(t.testSamples[s].inMatrix.m, 2).tolist()\n",
    "                    elif border.direction=='h' and border.position==t.testSamples[s].inMatrix.shape[0]-1:\n",
    "                        rotations['test'].append(3)\n",
    "                        rotTask['test'][s]['input'] = np.rot90(t.testSamples[s].inMatrix.m, 3).tolist()\n",
    "                    else:\n",
    "                        rotations['test'].append(0)\n",
    "                    hasBorder=True\n",
    "                    break\n",
    "            if not hasBorder:\n",
    "                return False\n",
    "        else:\n",
    "            border = t.testSamples[s].commonFullBorders[0]\n",
    "            if border.direction=='h' and border.position==0:\n",
    "                rotations['test'].append(1)\n",
    "                rotTask['test'][s]['input'] = np.rot90(t.testSamples[s].inMatrix.m, 1).tolist()\n",
    "                rotTask['test'][s]['output'] = np.rot90(t.testSamples[s].outMatrix.m, 1).tolist()\n",
    "            elif border.direction=='v' and border.position==t.testSamples[s].inMatrix.shape[1]-1:\n",
    "                rotations['test'].append(2)\n",
    "                rotTask['test'][s]['input'] = np.rot90(t.testSamples[s].inMatrix.m, 2).tolist()\n",
    "                rotTask['test'][s]['output'] = np.rot90(t.testSamples[s].outMatrix.m, 2).tolist()\n",
    "            elif border.direction=='h' and border.position==t.testSamples[s].inMatrix.shape[0]-1:\n",
    "                rotations['test'].append(3)\n",
    "                rotTask['test'][s]['input'] = np.rot90(t.testSamples[s].inMatrix.m, 3).tolist()\n",
    "                rotTask['test'][s]['output'] = np.rot90(t.testSamples[s].outMatrix.m, 3).tolist()\n",
    "            else:\n",
    "                rotations['test'].append(0)\n",
    "        \n",
    "    return rotTask, rotations\n",
    "\n",
    "def rotateHVTask(t, task):\n",
    "    \"\"\"\n",
    "    Given a Task t, this function determines whether any of the Samples needs\n",
    "    have its Matrices rotated and, if that's the case, the task with the\n",
    "    rotated Matrices is returned. Otherwise, it returns false.\n",
    "    Matrices will be rotated if all the Samples are either \"Horizontal\" or\n",
    "    \"Vertical\". In that case, it will be made sure that all the Matrices of the\n",
    "    rotated task are \"Horizontal\".\n",
    "    \"\"\"\n",
    "    rotTask = copy.deepcopy(task)\n",
    "    rotations = {'train': [], 'test': []}\n",
    "    \n",
    "    for s in range(t.nTrain):\n",
    "        if t.trainSamples[s].isVertical:\n",
    "            rotations['train'].append(1)\n",
    "            rotTask['train'][s]['input'] = np.rot90(t.trainSamples[s].inMatrix.m, 1).tolist()\n",
    "            rotTask['train'][s]['output'] = np.rot90(t.trainSamples[s].outMatrix.m, 1).tolist()\n",
    "        else:\n",
    "            rotations['train'].append(0)\n",
    "    \n",
    "    for s in range(t.nTest):\n",
    "        if t.submission:\n",
    "            if t.testSamples[s].inMatrix.isHorizontal:\n",
    "                rotations['test'].append(0)\n",
    "            elif t.testSamples[s].inMatrix.isVertical:\n",
    "                rotations['test'].append(1)\n",
    "                rotTask['test'][s]['input'] = np.rot90(t.testSamples[s].inMatrix.m, 1).tolist()\n",
    "            else:\n",
    "                return False, False\n",
    "        else:\n",
    "            if t.testSamples[s].isHorizontal:\n",
    "                rotations['test'].append(0)\n",
    "            elif t.testSamples[s].isVertical:\n",
    "                rotations['test'].append(1)\n",
    "                rotTask['test'][s]['input'] = np.rot90(t.testSamples[s].inMatrix.m, 1).tolist()\n",
    "                rotTask['test'][s]['output'] = np.rot90(t.testSamples[s].outMatrix.m, 1).tolist()\n",
    "            else:\n",
    "                return False, False\n",
    "            \n",
    "    return rotTask, rotations\n",
    "\n",
    "def recoverRotations(matrix, trainOrTest, s, rotations):\n",
    "    \"\"\"\n",
    "    Revert the rotation executed the given matrix during the preprocessing.\n",
    "    \"\"\"\n",
    "    if rotations[trainOrTest][s] == 1:\n",
    "        m = np.rot90(matrix, 3)\n",
    "    elif rotations[trainOrTest][s] == 2:\n",
    "        m = np.rot90(matrix, 2)\n",
    "    elif rotations[trainOrTest][s] == 3:\n",
    "        m = np.rot90(matrix, 1)\n",
    "    else:\n",
    "        m = matrix.copy()        \n",
    "    return m\n",
    "\n",
    "def tryOperations(t, c, cTask, b3c, firstIt=False):\n",
    "    \"\"\"\n",
    "    Given a Task t and a Candidate c, this function applies all the\n",
    "    operations that make sense to the input matrices of c. After a certain\n",
    "    operation is performed to all the input matrices, a new candidate is\n",
    "    generated from the resulting output matrices. If the score of the candidate\n",
    "    improves the score of any of the 3 best candidates, it will be saved in the\n",
    "    variable b3c, which is an object of the class Best3Candidates.\n",
    "    \"\"\"\n",
    "    if c.seen or c.score==0 or b3c.allPerfect():\n",
    "        return\n",
    "    startOps = (\"switchColors\", \"cropShape\", \"cropAllBackground\", \"minimize\", \\\n",
    "                \"maxColorFromCell\", \"deleteShapes\", \"replicateShapes\",\"colorByPixels\",\\\n",
    "                \"paintGridLikeBackground\")\n",
    "    repeatIfPerfect = (\"extendColor\", \"moveAllShapes\")\n",
    "    \n",
    "    # Get the operations for the given candidate and try all of them.\n",
    "    possibleOps = getPossibleOperations(t, c)\n",
    "    for op in possibleOps:\n",
    "        for s in range(t.nTrain):\n",
    "            cTask[\"train\"][s][\"input\"] = op(c.t.trainSamples[s].inMatrix).tolist()\n",
    "            if c.t.sameIOShapes and len(c.t.fixedColors) != 0:\n",
    "                cTask[\"train\"][s][\"input\"] = correctFixedColors(\\\n",
    "                     c.t.trainSamples[s].inMatrix.m,\\\n",
    "                     np.array(cTask[\"train\"][s][\"input\"]),\\\n",
    "                     c.t.fixedColors, c.t.commonOnlyChangedInColors).tolist()\n",
    "        newPredictions = []\n",
    "        for s in range(t.nTest):\n",
    "            newOutput = op(c.t.testSamples[s].inMatrix)\n",
    "            newPredictions.append(newOutput)\n",
    "            cTask[\"test\"][s][\"input\"] = newOutput.tolist()\n",
    "            if c.t.sameIOShapes and len(c.t.fixedColors) != 0:\n",
    "                cTask[\"test\"][s][\"input\"] = correctFixedColors(\\\n",
    "                     c.t.testSamples[s].inMatrix.m,\\\n",
    "                     np.array(cTask[\"test\"][s][\"input\"]),\\\n",
    "                     c.t.fixedColors, c.t.commonOnlyChangedInColors).tolist()\n",
    "        cScore = sum([incorrectPixels(np.array(cTask[\"train\"][s][\"input\"]), \\\n",
    "                                            t.trainSamples[s].outMatrix.m) for s in range(t.nTrain)])\n",
    "        changedPixels = sum([incorrectPixels(c.t.trainSamples[s].inMatrix.m, \\\n",
    "                                                  np.array(cTask[\"train\"][s][\"input\"])) for s in range(t.nTrain)])\n",
    "        \n",
    "        # Generate a new Candidate after applying the operation, and update\n",
    "        # the best 3 candidates.\n",
    "        newCandidate = Candidate(c.ops+[op], c.tasks+[copy.deepcopy(cTask)], cScore,\\\n",
    "                                 predictions=newPredictions)\n",
    "        b3c.addCandidate(newCandidate)\n",
    "        if firstIt and str(op)[28:60].startswith(startOps):\n",
    "            if all([np.array_equal(np.array(cTask[\"train\"][s][\"input\"]), \\\n",
    "                                   t.trainSamples[s].inMatrix.m) for s in range(t.nTrain)]):\n",
    "                continue\n",
    "            newCandidate.generateTask()\n",
    "            tryOperations(t, newCandidate, cTask, b3c)\n",
    "        elif str(op)[28:60].startswith(repeatIfPerfect) and c.score - changedPixels == cScore and changedPixels != 0:\n",
    "            newCandidate.generateTask()\n",
    "            tryOperations(t, newCandidate, cTask, b3c)\n",
    "    c.seen = True\n",
    "            \n",
    "def getPredictionsFromTask(originalT, task):\n",
    "    \"\"\"\n",
    "    Given a task in its dict format, and originalT, an object of the class Task\n",
    "    storing information of the original task, this function returns the\n",
    "    predictions and the corresponding Best3Candidates object computed for\n",
    "    the task.\n",
    "    \"\"\"\n",
    "    # Preprocessing\n",
    "    taskNeedsRecoloring = needsRecoloring(originalT)\n",
    "    if taskNeedsRecoloring:\n",
    "        task, trainRels, trainInvRels, testRels, testInvRels = orderTaskColors(originalT)\n",
    "        t = Task(task, task_id, submission=True)\n",
    "    else:\n",
    "        t = originalT\n",
    "\n",
    "    cTask = copy.deepcopy(task)\n",
    "\n",
    "    if t.sameIOShapes:\n",
    "        taskNeedsCropping = needsCropping(t)\n",
    "    else:\n",
    "        taskNeedsCropping = False\n",
    "    if taskNeedsCropping:\n",
    "        cropPositions, backgrounds = cropTask(t, cTask)\n",
    "        t2 = Task(cTask, task_id, submission=True, backgrounds=backgrounds)\n",
    "    elif t.hasUnchangedGrid:\n",
    "        if t.gridCellsHaveOneColor:\n",
    "            ignoreGrid(t, cTask) # This modifies cTask, ignoring the grid\n",
    "            t2 = Task(cTask, task_id, submission=True)\n",
    "        elif t.outGridCellsHaveOneColor:\n",
    "            ignoreGrid(t, cTask, inMatrix=False)\n",
    "            t2 = Task(cTask, task_id, submission=True)\n",
    "        else:\n",
    "            t2 = t\n",
    "    elif t.hasUnchangedAsymmetricGrid and t.assymmetricGridCellsHaveOneColor:\n",
    "        ignoreAsymmetricGrid(t, cTask)\n",
    "        t2 = Task(cTask, task_id, submission=True)\n",
    "    else:\n",
    "        t2 = t\n",
    "        \n",
    "    if t2.sameIOShapes:\n",
    "        hasRotated = False\n",
    "        if t2.hasOneFullBorder:\n",
    "            hasRotated, rotateParams = rotateTaskWithOneBorder(t2, cTask)\n",
    "        elif t2.requiresHVRotation:\n",
    "            hasRotated, rotateParams = rotateHVTask(t2, cTask)\n",
    "        if hasRotated!=False:\n",
    "            cTask = hasRotated.copy()\n",
    "            t2 = Task(cTask, task_id, submission=True)\n",
    "\n",
    "     # Generate the three candidates with best possible score\n",
    "    \n",
    "    cScore = sum([incorrectPixels(np.array(cTask[\"train\"][s][\"input\"]), \\\n",
    "                                         t2.trainSamples[s].outMatrix.m) for s in range(t.nTrain)])\n",
    "    dummyPredictions = [sample.inMatrix.m for sample in t2.testSamples]\n",
    "    c = Candidate([], [task], score=cScore, predictions=dummyPredictions)\n",
    "    c.t = t2\n",
    "    b3c = Best3Candidates(c, c, c)\n",
    "\n",
    "    prevScore = sum([c.score for c in b3c.candidates])\n",
    "    firstIt = True\n",
    "    while True:\n",
    "        copyB3C = copy.deepcopy(b3c)\n",
    "        for c in copyB3C.candidates:\n",
    "            if c.score == 0:\n",
    "                continue\n",
    "            tryOperations(t2, c, cTask, b3c, firstIt)\n",
    "            if firstIt:\n",
    "                firstIt = False\n",
    "                break\n",
    "        score = sum([c.score for c in b3c.candidates])\n",
    "        if score >= prevScore:\n",
    "            break\n",
    "        else:\n",
    "            prevScore = score\n",
    "            \n",
    "    taskPredictions = []\n",
    "    \n",
    "    # Once the best 3 candidates have been found, make the predictions and\n",
    "    # revert the preprocessing.\n",
    "    for s in range(t.nTest):\n",
    "        taskPredictions.append([])\n",
    "        for c in b3c.candidates:\n",
    "            x = t2.testSamples[s].inMatrix.m.copy()\n",
    "            for opI in range(len(c.ops)):\n",
    "                newX = c.ops[opI](Matrix(x))\n",
    "                if t2.sameIOShapes and len(t2.fixedColors) != 0:\n",
    "                    x = correctFixedColors(x, newX, t2.fixedColors, t2.commonOnlyChangedInColors)\n",
    "                else:\n",
    "                    x = newX.copy()\n",
    "            if t2.sameIOShapes and hasRotated!=False:\n",
    "                x = recoverRotations(x, \"test\", s, rotateParams)\n",
    "            if taskNeedsCropping:\n",
    "                x = recoverCroppedMatrix(x, originalT.testSamples[s].inMatrix.shape, \\\n",
    "                                         cropPositions[\"test\"][s], t.testSamples[s].inMatrix.backgroundColor)\n",
    "            elif t.hasUnchangedGrid and (t.gridCellsHaveOneColor or t.outGridCellsHaveOneColor):\n",
    "                x = recoverGrid(t, x, s)\n",
    "            elif t.hasUnchangedAsymmetricGrid and t.assymmetricGridCellsHaveOneColor:\n",
    "                x = recoverAsymmetricGrid(t, x, s)\n",
    "            if taskNeedsRecoloring:\n",
    "                x = recoverOriginalColors(x, testRels[s])\n",
    "            \n",
    "            taskPredictions[s].append(x)\n",
    "            \n",
    "    return taskPredictions, b3c\n",
    "        \n",
    "###############################################################################\n",
    "# %% Main Loop and submission\n",
    "\n",
    "submission = pd.DataFrame(columns=['output'])\n",
    "submission.index.name = 'output_id'\n",
    "    \n",
    "for task_id, task in data.items():    \n",
    "    bestScores = []\n",
    "                    \n",
    "    originalT = Task(task, task_id, submission=True)\n",
    "        \n",
    "    predictions, b3c = getPredictionsFromTask(originalT, task.copy())\n",
    "    \n",
    "    if any([c.score==0 for c in b3c.candidates]):\n",
    "        perfectScore = True\n",
    "    \n",
    "    separationByShapes = needsSeparationByShapes(originalT)\n",
    "    separationByColors = needsSeparationByColors(originalT)\n",
    "\n",
    "    if separationByShapes != False:\n",
    "        separatedT = Task(separationByShapes.separatedTask, task_id, submission=True)\n",
    "        sepPredictions, sepB3c = getPredictionsFromTask(separatedT, separationByShapes.separatedTask.copy())\n",
    "\n",
    "        if any([c.score==0 for c in sepB3c.candidates]):\n",
    "            perfectScore = True\n",
    "        \n",
    "        mergedPredictions = []\n",
    "        for s in range(originalT.nTest):\n",
    "            mergedPredictions.append([])\n",
    "            matrixRange = separationByShapes.getRange(\"test\", s)\n",
    "            matrices = [[sepPredictions[i][cand] for i in range(matrixRange[0], matrixRange[1])] \\\n",
    "                         for cand in range(3)]\n",
    "            for cand in range(3):\n",
    "                pred = mergeMatrices(matrices[cand], originalT.backgroundColor, separationByShapes.mergeColor)\n",
    "                mergedPredictions[s].append(pred)\n",
    "        \n",
    "        finalPredictions = []\n",
    "        for s in range(originalT.nTest):\n",
    "            finalPredictions.append([[], [], []])\n",
    "        \n",
    "        b3cIndices = b3c.getOrderedIndices()\n",
    "        sepB3cIndices = sepB3c.getOrderedIndices()\n",
    "\n",
    "        b3cIndex, sepB3cIndex = 0, 0\n",
    "        i = 0\n",
    "        if b3c.candidates[b3cIndices[0]].score==0:\n",
    "            bestScores.append(0)\n",
    "            for s in range(originalT.nTest):\n",
    "                finalPredictions[s][0] = predictions[s][b3cIndices[0]]\n",
    "            i += 1\n",
    "        if sepB3c.candidates[sepB3cIndices[0]].score==0:\n",
    "            bestScores.append(0)\n",
    "            for s in range(originalT.nTest):\n",
    "                finalPredictions[s][i] = mergedPredictions[s][sepB3cIndices[0]]\n",
    "            i += 1\n",
    "        while i < 3:\n",
    "            if b3c.candidates[b3cIndices[b3cIndex]] < sepB3c.candidates[sepB3cIndices[sepB3cIndex]]:\n",
    "                bestScores.append(b3c.candidates[b3cIndices[b3cIndex]].score)\n",
    "                for s in range(originalT.nTest):\n",
    "                    finalPredictions[s][i] = predictions[s][b3cIndices[b3cIndex]]\n",
    "                b3cIndex += 1\n",
    "            else:\n",
    "                bestScores.append(sepB3c.candidates[sepB3cIndices[sepB3cIndex]].score)\n",
    "                for s in range(originalT.nTest):\n",
    "                    finalPredictions[s][i] = mergedPredictions[s][sepB3cIndices[sepB3cIndex]]\n",
    "                sepB3cIndex += 1\n",
    "            i += 1\n",
    "\n",
    "    elif separationByColors != False:\n",
    "        separatedT = Task(separationByColors.separatedTask, task_id, submission=True)\n",
    "        sepPredictions, sepB3c = getPredictionsFromTask(separatedT, separationByColors.separatedTask.copy())\n",
    "\n",
    "        if any([c.score==0 for c in sepB3c.candidates]):\n",
    "            perfectScore = True\n",
    "        \n",
    "        mergedPredictions = []\n",
    "        for s in range(originalT.nTest):\n",
    "            mergedPredictions.append([])\n",
    "            matrixRange = separationByColors.getRange(\"test\", s)\n",
    "            matrices = [[sepPredictions[i][cand] for i in range(matrixRange[0], matrixRange[1])] \\\n",
    "                         for cand in range(3)]\n",
    "            for cand in range(3):\n",
    "                pred = mergeMatrices(matrices[cand], originalT.backgroundColor)\n",
    "                mergedPredictions[s].append(pred)\n",
    "\n",
    "        finalPredictions = []\n",
    "        for s in range(originalT.nTest):\n",
    "            finalPredictions.append([[], [], []])\n",
    "        \n",
    "        b3cIndices = b3c.getOrderedIndices()\n",
    "        sepB3cIndices = sepB3c.getOrderedIndices()\n",
    "\n",
    "        b3cIndex, sepB3cIndex = 0, 0\n",
    "        i = 0\n",
    "        if b3c.candidates[b3cIndices[0]].score==0:\n",
    "            bestScores.append(0)\n",
    "            for s in range(originalT.nTest):\n",
    "                finalPredictions[s][0] = predictions[s][b3cIndices[0]]\n",
    "            i += 1\n",
    "        if sepB3c.candidates[sepB3cIndices[0]].score==0:\n",
    "            bestScores.append(0)\n",
    "            for s in range(originalT.nTest):\n",
    "                finalPredictions[s][i] = mergedPredictions[s][sepB3cIndices[0]]\n",
    "            i += 1\n",
    "        while i < 3:\n",
    "            if b3c.candidates[b3cIndices[b3cIndex]] < sepB3c.candidates[sepB3cIndices[sepB3cIndex]]:\n",
    "                bestScores.append(b3c.candidates[b3cIndices[b3cIndex]].score)\n",
    "                for s in range(originalT.nTest):\n",
    "                    finalPredictions[s][i] = predictions[s][b3cIndices[b3cIndex]]\n",
    "                b3cIndex += 1\n",
    "            else:\n",
    "                bestScores.append(sepB3c.candidates[sepB3cIndices[sepB3cIndex]].score)\n",
    "                for s in range(originalT.nTest):\n",
    "                    finalPredictions[s][i] = mergedPredictions[s][sepB3cIndices[sepB3cIndex]]\n",
    "                sepB3cIndex += 1\n",
    "            i += 1\n",
    "    else:\n",
    "        for c in b3c.candidates:\n",
    "            bestScores.append(c.score)\n",
    "        finalPredictions = predictions\n",
    "\n",
    "    pred = []\n",
    "    for pair_id in range(originalT.nTest):\n",
    "        for i in range(len(finalPredictions[pair_id])):\n",
    "            pred.append(flattener(finalPredictions[pair_id][i].astype(int).tolist()))\n",
    "    predictions = pred\n",
    "\n",
    "    if len(predictions) == 0:\n",
    "        pred = '|0| |0| |0|'\n",
    "    elif len(predictions) == 1:\n",
    "        pred = predictions[0] + ' ' + predictions[0] + ' ' + predictions[0]\n",
    "    elif len(predictions) == 2:\n",
    "        pred =  predictions[0] + ' ' + predictions[1] + ' ' + predictions[0]\n",
    "    elif len(predictions) == 3:\n",
    "        pred = predictions[0] + ' ' + predictions[1] + ' ' + predictions[2]\n",
    "    \n",
    "    for pair_id in range(originalT.nTest):\n",
    "        submission.loc[task_id+'_'+str(pair_id), 'output'] = pred\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "import os\n",
    "import math\n",
    "import inspect\n",
    "import collections\n",
    "from os.path import join as path_join\n",
    "import torch\n",
    "from torch.nn import CrossEntropyLoss\n",
    "from torch.optim import Adam\n",
    "from torch.nn import Conv2d\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch import FloatTensor, LongTensor\n",
    "import random\n",
    "from time import time\n",
    "from random import randint\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import colors\n",
    "from itertools import combinations\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from time import sleep\n",
    "import copy\n",
    "import gc\n",
    "from pdb import set_trace as st\n",
    "import timeit\n",
    "import itertools\n",
    "from numpy.lib.stride_tricks import as_strided\n",
    "from scipy.spatial import distance\n",
    "from collections import defaultdict\n",
    "import warnings\n",
    "from skimage import measure\n",
    "warnings.filterwarnings('ignore')\n",
    "pd.set_option('display.max_columns', 100)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [],
   "source": [
    "def seed_everything(seed=1234):\n",
    "    random.seed(seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "seed_everything(0)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neighbors import KNeighborsClassifier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# helper funcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_symmetry_lr(inp):\n",
    "    return np.array(inp).tolist() == np.fliplr(inp).tolist()\n",
    "def is_symmetry_ud(inp):\n",
    "    return np.array(inp).tolist() == np.flipud(inp).tolist()\n",
    "\n",
    "def input_output_shape_is_same(task):\n",
    "    return all([np.array(el['input']).shape == np.array(el['output']).shape for el in task['train']])\n",
    "\n",
    "def sliding_window_search(large, small):\n",
    "    m1, n1 = np.array(large).shape\n",
    "    m2, n2 = np.array(small).shape\n",
    "    for m in range(m1 - m2 + 1):\n",
    "        for n in range(n1 - n2 + 1):\n",
    "            if np.array(small).tolist() == np.array(large)[m:m+m2, n:n+n2].tolist():\n",
    "                return True, m, n\n",
    "    return False, -1, -1\n",
    "\n",
    "def matrix_use_color(m):\n",
    "    return list(set(list(itertools.chain.from_iterable(m))))\n",
    "\n",
    "def is_square(task):\n",
    "    return all([np.array(el['input']).shape[0] == np.array(el['input']).shape[1] for el in task['train']])\n",
    "\n",
    "def inouts_flip(task_train):\n",
    "    for inout in task_train:\n",
    "        inout['input'] = np.flip(inout['input'])\n",
    "    return task_train\n",
    "\n",
    "def inouts_flipud(task_train):\n",
    "    for inout in task_train:\n",
    "        inout['input'] = np.flipud(inout['input'])\n",
    "    return task_train\n",
    "\n",
    "def inouts_fliplr(task_train):\n",
    "    for inout in task_train:\n",
    "        inout['input'] = np.fliplr(inout['input'])\n",
    "    return task_train\n",
    "\n",
    "def match_rate(mat1, mat2):\n",
    "    m1 = np.array(mat1)\n",
    "    m2 = np.array(mat2)\n",
    "    if m1.shape != m2.shape:\n",
    "        return -1\n",
    "    v1 = list(itertools.chain.from_iterable(mat1))\n",
    "    v2 = list(itertools.chain.from_iterable(mat2))\n",
    "    score1 = np.sum(m1==m2) / len(v1)\n",
    "    score2 = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))\n",
    "    return (score1 + score2) / 2\n",
    "\n",
    "def flattener(pred):\n",
    "    str_pred = str([row for row in pred])\n",
    "    str_pred = str_pred.replace(', ', '')\n",
    "    str_pred = str_pred.replace('[[', '|')\n",
    "    str_pred = str_pred.replace('][', '|')\n",
    "    str_pred = str_pred.replace(']]', '|')\n",
    "    return str_pred\n",
    "\n",
    "def getDefaultPred(inp):\n",
    "    pred_1 = flattener(inp)\n",
    "    # for the second guess, change all 0s to 5s\n",
    "    inp = [[5 if i==0 else i for i in j] for j in inp]\n",
    "    pred_2 = flattener(inp)\n",
    "    # for the last gues, change everything to 0\n",
    "    inp = [[0 for i in j] for j in inp]\n",
    "    pred_3 = flattener(inp)\n",
    "    # concatenate and add to the submission output\n",
    "    pred = pred_1 + ' ' + pred_2 + ' ' + pred_3 + ' ' \n",
    "    return pred\n",
    "\n",
    "def preds_to_str(preds_list, idx):\n",
    "    pred_strs = []\n",
    "#     st()\n",
    "    for i in range(len(preds_list[0])):\n",
    "        pred_str = ''\n",
    "        for j, preds in enumerate(reversed(preds_list)):\n",
    "            if j == 3:\n",
    "                break\n",
    "            pred_str += flattener(np.array(preds[i]).tolist()) + ' '\n",
    "        pred_strs.append(pred_str)\n",
    "    return pred_strs\n",
    "\n",
    "def preds_to_str_only0(preds0, idx):\n",
    "    preds = []\n",
    "    for i in range(len(preds0)):\n",
    "        pred0 = flattener(np.array(preds0[i]).tolist())\n",
    "        pred1 = flattener(np.array([[0]]).tolist())\n",
    "        preds.append(pred0 + ' ' + pred1 + ' ' + pred1 + ' ')\n",
    "    return preds\n",
    "\n",
    "def get_not_use_num(matrix1, matrix2):\n",
    "    v1 = list(itertools.chain.from_iterable(matrix1))\n",
    "    v2 = list(itertools.chain.from_iterable(matrix2))\n",
    "    for j in range(1,10):\n",
    "        if (j not in v1) & (j not in v2):\n",
    "            return j\n",
    "    return 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# test NN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_nn(train, test):\n",
    "    task={'train':train, 'test':test}\n",
    "\n",
    "    train_dataset = ArcDataset(task, mode=\"train\", augment=False)\n",
    "    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True, collate_fn=device_collate)\n",
    "    valid_dataset = ArcDataset(task, mode=\"test\", augment=False)\n",
    "    valid_dataloader = torch.utils.data.DataLoader(valid_dataset, batch_size=1, shuffle=False, collate_fn=device_collate)\n",
    "\n",
    "    net = Task330Net().to(device)\n",
    "    criterion = hinge_loss\n",
    "    #optimizer = torch.optim.SGD(net.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4)\n",
    "    optimizer = torch.optim.Adam(net.parameters(), lr=1e-3, weight_decay=1e-4)\n",
    "    t0 = time()\n",
    "    tmp_train_loss = 1\n",
    "\n",
    "    for param in net.named_parameters():\n",
    "        print(f\"{param[0]:>15} {list(param[1].shape)}\")\n",
    "    count = 0\n",
    "    for epoch in range(5000):\n",
    "        train_loss = valid_loss = 0.0\n",
    "        train_loss_denom = valid_loss_denom = 0\n",
    "\n",
    "        ####################\n",
    "        # train\n",
    "        ####################\n",
    "        net.train()\n",
    "        for i, (feature, target) in enumerate(train_dataloader):\n",
    "            outputs = net(feature)\n",
    "            loss = criterion(outputs, target)\n",
    "\n",
    "            # backprop\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            # record\n",
    "            train_loss += loss.item()\n",
    "            train_loss_denom += feature.shape[0]\n",
    "\n",
    "        train_loss /= train_loss_denom\n",
    "\n",
    "        ####################\n",
    "        # eval\n",
    "        ####################\n",
    "#         net.eval()\n",
    "#         with torch.no_grad():\n",
    "#             for i, (feature, target) in enumerate(valid_dataloader):\n",
    "#                 feature = feature.to(device)\n",
    "#                 target = target.to(device)\n",
    "\n",
    "#                 outputs = net(feature)\n",
    "#                 loss = criterion(outputs, target)\n",
    "\n",
    "#                 # record\n",
    "#                 valid_loss += loss.item()\n",
    "#                 valid_loss_denom += feature.shape[0]\n",
    "\n",
    "#         valid_loss /= valid_loss_denom\n",
    "    \n",
    "        if epoch%100==0:\n",
    "#             print(f\"epoch {epoch:4d}  |  train_loss: {train_loss:5.6f}  valid_loss: {valid_loss:5.6f}  |  time: {time()-t0:7.1f} sec\")\n",
    "            print(f\"epoch {epoch:4d}  |  train_loss: {train_loss:5.6f}  |  time: {time()-t0:7.1f} sec\")\n",
    "            if tmp_train_loss <= train_loss:\n",
    "                count += 1\n",
    "            if count >= 4:\n",
    "                break\n",
    "            tmp_train_loss = train_loss\n",
    "\n",
    "#             if best_loss > valid_loss:\n",
    "#                 best_loss = valid_loss\n",
    "#                 filename = f\"./work/trained_weight/{MODEL_NAME}_epoch{epoch:03d}_loss{valid_loss:.3f}.pth\"\n",
    "#                 torch.save(net.state_dict(), filename)\n",
    "\n",
    "    return check(task, lambda x: task_train330(x, net))\n",
    "\n",
    "class ArcDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, task=None, mode=\"train\", augment=False):\n",
    "        if task is not None:\n",
    "            assert mode in [\"train\", \"test\"]\n",
    "            self.mode = mode\n",
    "            self.task = task[mode]\n",
    "        self.augment = augment\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.task)\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        t = self.task[index]\n",
    "        t_in = torch.tensor(t[\"input\"])\n",
    "        t_out = torch.tensor(t[\"output\"])\n",
    "        t_in, t_out = self.preprocess(t_in, t_out)\n",
    "        return t_in, t_out\n",
    "    \n",
    "    def preprocess(self, t_in, t_out):\n",
    "        if self.augment:\n",
    "            t_in, t_out = self._random_rotate(t_in, t_out)\n",
    "        t_in = self._one_hot_encode(t_in)\n",
    "        t_out = self._one_hot_encode(t_out)\n",
    "        return t_in, t_out\n",
    "    \n",
    "    def _one_hot_encode(self, x):\n",
    "        return torch.eye(10)[x].permute(2, 0, 1)\n",
    "    \n",
    "    def _random_rotate(self, t_in, t_out):\n",
    "        t_in_shape = t_in.shape\n",
    "        t_out_shape = t_out.shape\n",
    "        t_in = t_in.reshape(-1, *t_in_shape[-2:])\n",
    "        t_out = t_out.reshape(-1, *t_out_shape[-2:])\n",
    "        r = randint(0, 7)\n",
    "        if r%2 == 0:\n",
    "            t_in = t_in.permute(0, 2, 1)\n",
    "            t_out = t_out.permute(0, 2, 1)\n",
    "        r //= 2\n",
    "        if r%2 == 0:\n",
    "            t_in = t_in[:, :, torch.arange(t_in.shape[-1]-1, -1, -1)]\n",
    "            t_out = t_out[:, :, torch.arange(t_out.shape[-1]-1, -1, -1)]\n",
    "        r //= 2\n",
    "        if r%2 == 0:\n",
    "            t_in = t_in[:, torch.arange(t_in.shape[-2]-1, -1, -1), :]\n",
    "            t_out = t_out[:, torch.arange(t_out.shape[-2]-1, -1, -1), :]\n",
    "        t_in = t_in.reshape(*t_in_shape[:-2], *t_in.shape[-2:])\n",
    "        t_out = t_out.reshape(*t_out_shape[:-2], *t_out.shape[-2:])\n",
    "        return t_in, t_out\n",
    "    \n",
    "def device_collate(batch):\n",
    "    return tuple(map(lambda x: torch.stack(x).to(device), zip(*batch)))\n",
    "\n",
    "def hinge_loss(y_pred, y_true):\n",
    "    loss = y_pred.clone()\n",
    "    loss[y_true>0.5] = 1-loss[y_true>0.5]\n",
    "    loss[loss<0] = 0\n",
    "    return loss.sum(0).mean()\n",
    "\n",
    "class Task330Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Task330Net, self).__init__()\n",
    "        siz = 16\n",
    "        self.conv1 = nn.Conv2d(10, siz, 3, padding=1)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.dropout = torch.nn.Dropout2d(p=0.1)\n",
    "        self.conv1x1 = nn.Conv2d(siz, 10, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x2 = self.conv1(x)\n",
    "        x2 = self.relu(x2)\n",
    "        x2 = self.dropout(x2)\n",
    "        x2 = self.conv1x1(x2)\n",
    "        x = x + x2  # skip connection\n",
    "        return x\n",
    "    \n",
    "def check(task, pred_func):\n",
    "    preds = []\n",
    "    for i, t in enumerate(task[\"test\"]):\n",
    "        t_in = np.array(t[\"input\"])\n",
    "        preds.append(pred_func(t_in))\n",
    "    return preds\n",
    "\n",
    "def task_train330(x, net):\n",
    "    def one_hot_decode(x):\n",
    "        return x.argmax(0)\n",
    "    net.eval()\n",
    "    with torch.no_grad():\n",
    "        x = torch.tensor(x).to(device)\n",
    "        y_dummy = x.clone()\n",
    "        dataset = ArcDataset(augment=False)\n",
    "        x = dataset.preprocess(x, y_dummy)[0].unsqueeze(0)\n",
    "        y = net(x).detach()\n",
    "        y = one_hot_decode(y.squeeze(0))\n",
    "    y = y.to(\"cpu\").numpy()\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CAModel(nn.Module):\n",
    "    def __init__(self, num_states):\n",
    "        super(CAModel, self).__init__()\n",
    "        self.transition = nn.Sequential(\n",
    "            nn.Conv2d(num_states, 128, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(128, num_states, kernel_size=1)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x, steps=1):\n",
    "        for _ in range(steps):\n",
    "            x = self.transition(torch.softmax(x, dim=1))\n",
    "        return x\n",
    "\n",
    "def solve_task(task, max_steps=10):\n",
    "    model = CAModel(10).to(device)\n",
    "    num_epochs = 100\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    losses = np.zeros((max_steps - 1) * num_epochs)\n",
    "\n",
    "    for num_steps in range(1, max_steps):\n",
    "        optimizer = torch.optim.Adam(model.parameters(), lr=(0.1 / (num_steps * 2)))\n",
    "        \n",
    "        for e in range(num_epochs):\n",
    "            optimizer.zero_grad()\n",
    "            loss = 0.0\n",
    "\n",
    "            for sample in task:\n",
    "                # predict output from input\n",
    "#                 st()\n",
    "                x = torch.from_numpy(inp2img(sample[\"input\"])).unsqueeze(0).float().to(device)\n",
    "                y = torch.tensor(sample[\"output\"]).long().unsqueeze(0).to(device)\n",
    "                y_pred = model(x, num_steps)\n",
    "                loss += criterion(y_pred, y)\n",
    "                \n",
    "                # predit output from output\n",
    "                # enforces stability after solution is reached\n",
    "                y_in = torch.from_numpy(inp2img(sample[\"output\"])).unsqueeze(0).float().to(device)\n",
    "                y_pred = model(y_in, 1) \n",
    "                loss += criterion(y_pred, y)\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            losses[(num_steps - 1) * num_epochs + e] = loss.item()\n",
    "    return model, num_steps, losses\n",
    "                \n",
    "@torch.no_grad()\n",
    "def predict(model, task):\n",
    "    predictions = []\n",
    "    for sample in task:\n",
    "        x = torch.from_numpy(inp2img(sample[\"input\"])).unsqueeze(0).float().to(device)\n",
    "        pred = model(x, 100).argmax(1).squeeze().cpu().numpy()\n",
    "        predictions.append(pred)\n",
    "    return predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inp2img(inp):\n",
    "    inp = np.array(inp)\n",
    "    img = np.full((10, inp.shape[0], inp.shape[1]), 0, dtype=np.uint8)\n",
    "    for i in range(10):\n",
    "        img[i] = (inp==i)\n",
    "    return img\n",
    "\n",
    "class TaskSolver:\n",
    "    def train(self, task_train, n_epoch=30, preprocess_funcs=[],final=False):\n",
    "        \"\"\"basic pytorch train loop\"\"\"\n",
    "        self.net = Conv2d(in_channels=10, out_channels=10, kernel_size=5, padding=2)\n",
    "\n",
    "        criterion = CrossEntropyLoss()\n",
    "        optimizer = Adam(self.net.parameters(), lr = 0.1)\n",
    "        for epoch in range(n_epoch):\n",
    "            for sample in task_train:\n",
    "                inputs = copy.deepcopy(sample['input'])\n",
    "                for preprocess_func in preprocess_funcs:\n",
    "                    inputs = preprocess_func(inputs)\n",
    "                inputs = FloatTensor(inp2img(inputs)).unsqueeze(dim=0)\n",
    "                labels = LongTensor(sample['output']).unsqueeze(dim=0)\n",
    "                optimizer.zero_grad()\n",
    "                outputs = self.net(inputs)\n",
    "#                 import pdb; pdb.set_trace()\n",
    "                loss = criterion(outputs, labels)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "#                 st()\n",
    "        return self\n",
    "\n",
    "    def predict(self, task_test, preprocess_funcs=[], success_map={}, idx='searching', final_score_map={}, final=False):\n",
    "        predictions = []\n",
    "        with torch.no_grad():\n",
    "            for i, in_out in enumerate(task_test):\n",
    "                inputs = copy.deepcopy(in_out['input'])\n",
    "                # input変更\n",
    "                for preprocess_func in preprocess_funcs:\n",
    "                    inputs = preprocess_func(inputs)\n",
    "                inputs = FloatTensor(inp2img(inputs)).unsqueeze(dim=0)\n",
    "                outputs = self.net(inputs)\n",
    "                pred =  outputs.squeeze(dim=0).cpu().numpy().argmax(0)\n",
    "                if ('output' in in_out) & (idx != 'searching') & (not final):\n",
    "                    similarity = match_rate(in_out['output'], pred)\n",
    "                    if (idx not in final_score_map) or (final_score_map.get(idx, 101) < similarity):\n",
    "                        final_score_map[idx] = similarity\n",
    "                predictions.append(pred)\n",
    "\n",
    "        return predictions\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# post process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [],
   "source": [
    "def post_processes(preds, origin_task):\n",
    "    processed_preds = []\n",
    "    for pred in preds:\n",
    "        all_input_same_colors, all_output_same_colors, all_input_and_output_same_colors, each_input_and_output_same_colors = search_inout_used_colors(origin_task['train'])\n",
    "        colors_are_black_and_one_flag, only_used_color_num = colors_are_black_and_one(origin_task['train'][0]['output'])\n",
    "        if all_output_same_colors & colors_are_black_and_one_flag:\n",
    "            pred = np.where(pred != 0, only_used_color_num, pred)\n",
    "        processed_preds.append(pred)\n",
    "    return processed_preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [],
   "source": [
    "def colors_are_black_and_one(matrix):\n",
    "    unique_nums = np.unique(matrix).tolist()\n",
    "    if 0 not in unique_nums:\n",
    "        return False, None\n",
    "        \n",
    "    unique_nums.remove(0)\n",
    "    if len(unique_nums) == 1:\n",
    "        return True, unique_nums[0]\n",
    "    else:\n",
    "        return False, None\n",
    "\n",
    "def search_inout_used_colors(task_train):\n",
    "    all_input_same_colors = True\n",
    "    all_output_same_colors = True\n",
    "    all_input_and_output_same_colors = True\n",
    "    each_input_and_output_same_colors = True\n",
    "    input_unique_nums=[]\n",
    "    for i, inout in enumerate(task_train):\n",
    "        if each_input_and_output_same_colors:\n",
    "            if np.unique(inout['input']).tolist() != np.unique(inout['output']).tolist():\n",
    "                each_input_and_output_same_colors = False\n",
    "                all_input_and_output_same_colors = False\n",
    "        if i == 0:\n",
    "            input_unique_nums = np.unique(inout['input']).tolist()\n",
    "            output_unique_nums = np.unique(inout['output']).tolist()\n",
    "            continue\n",
    "\n",
    "        if input_unique_nums != np.unique(inout['input']).tolist():\n",
    "            all_input_same_colors = False\n",
    "            all_input_and_output_same_colors = False\n",
    "        if output_unique_nums != np.unique(inout['output']).tolist():\n",
    "            all_output_same_colors = False\n",
    "            all_input_and_output_same_colors = False\n",
    "    return all_input_same_colors, all_output_same_colors, all_input_and_output_same_colors, each_input_and_output_same_colors            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# mirror aug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fliplr_aug(train):\n",
    "    return mirror_aug(train, np.fliplr)\n",
    "\n",
    "def flipud_aug(train):\n",
    "    return mirror_aug(train, np.flipud)\n",
    "\n",
    "def flip_aug(train):\n",
    "    return mirror_aug(train, np.flip)\n",
    "\n",
    "def transpose_aug(train):\n",
    "    return mirror_aug(train, np.transpose)\n",
    "\n",
    "def mirror_aug(train, aug_func):\n",
    "    inouts = []\n",
    "    for j, inout_origin in enumerate(train):\n",
    "        inout = copy.deepcopy(inout_origin)\n",
    "        same_flag = False\n",
    "        inout['input'] = aug_func(inout['input']).tolist()\n",
    "        inout['output'] = aug_func(inout['output']).tolist()\n",
    "        if inout['input'] != inout_origin['input'].tolist():\n",
    "            for io in inouts:\n",
    "                if io['input'] == inout['input']:\n",
    "                    same_flag = True\n",
    "                    break\n",
    "            if not same_flag:\n",
    "                for inout_other in train:\n",
    "                    if inout_other['input'].tolist() == inout['input']:\n",
    "                        same_flag = True\n",
    "                        break\n",
    "            if not same_flag:\n",
    "                inouts.append({'input': inout['input'], 'output': inout['output']})\n",
    "    return inouts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# color aug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "metadata": {},
   "outputs": [],
   "source": [
    "def has_duplicates(seq):\n",
    "    return len(seq) != len(set(seq))\n",
    "\n",
    "def color_augument(tasks):\n",
    "    for idx, task in tasks.iteritems():\n",
    "        task['train'] += one_train_color_augument(task['train'])\n",
    "\n",
    "def one_train_color_augument(train, verbose=False):\n",
    "    two_colors = []\n",
    "    for pair in combinations([0,1,2,3,4,5,6,7,8,9], 2):\n",
    "        two_colors.append(pair)\n",
    "    color_pairs_list = []\n",
    "    for i in range(1, 6):\n",
    "        for color_pairs in combinations(two_colors, i):\n",
    "            if has_duplicates(list(itertools.chain.from_iterable(list(color_pairs)))):\n",
    "                continue\n",
    "            color_pairs_list.append(color_pairs)\n",
    "    inouts = []\n",
    "    for color_pairs in color_pairs_list:\n",
    "        for j, inout_origin in enumerate(train):\n",
    "            inout = copy.deepcopy(inout_origin)\n",
    "            same_flag = False\n",
    "            inout['input'] = np.array(inout['input'])\n",
    "            inout['output'] = np.array(inout['output'])\n",
    "            tmp_inp = copy.deepcopy(inout['input'])\n",
    "            tmp_out = copy.deepcopy(inout['output'])\n",
    "\n",
    "            for color_pair in color_pairs:\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[0], color_pair[1], inout['input'])\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[1], color_pair[0], inout['input'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[0], color_pair[1], inout['output'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[1], color_pair[0], inout['output'])\n",
    "\n",
    "            inout['input'] = inout['input'].tolist()\n",
    "            inout['output'] = inout['output'].tolist()\n",
    "            if inout['input'] != tmp_inp.tolist():\n",
    "                for io in inouts:\n",
    "                    if io['input'] == inout['input']:\n",
    "                        same_flag = True\n",
    "                        break\n",
    "                if not same_flag:\n",
    "                    for inout_other in train:\n",
    "                        if inout_other['input'].tolist() == inout['input']:\n",
    "                            same_flag = True\n",
    "                            break\n",
    "                if not same_flag:\n",
    "                    inouts.append({'input': inout['input'], 'output': inout['output']})\n",
    "    if verbose:\n",
    "        print(len(inouts))\n",
    "\n",
    "    return inouts\n",
    "\n",
    "\n",
    "def super_heavy_color_augument(tasks):\n",
    "    for idx, task in tasks.iteritems():\n",
    "        task['train'] += one_train_super_heavy_color_augument(task['train'])\n",
    "\n",
    "def one_train_super_heavy_color_augument(train, verbose=False):\n",
    "    two_colors = []\n",
    "    for pair in combinations([1,2,3,4,5,6,7,8,9], 2):\n",
    "        two_colors.append(pair)\n",
    "    color_pairs_list = []\n",
    "    for i in [2, 3, 4]:\n",
    "        for color_pairs in combinations(two_colors, i):\n",
    "            if has_duplicates(list(itertools.chain.from_iterable(list(color_pairs)))):\n",
    "                continue\n",
    "            color_pairs_list.append(color_pairs)\n",
    "    inouts = []\n",
    "    for color_pairs in color_pairs_list:\n",
    "        for j, inout_origin in enumerate(train):\n",
    "            inout = copy.deepcopy(inout_origin)\n",
    "            same_flag = False\n",
    "            inout['input'] = np.array(inout['input'])\n",
    "            inout['output'] = np.array(inout['output'])\n",
    "            tmp_inp = copy.deepcopy(inout['input'])\n",
    "            tmp_out = copy.deepcopy(inout['output'])\n",
    "\n",
    "            for color_pair in color_pairs:\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[0], color_pair[1], inout['input'])\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[1], color_pair[0], inout['input'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[0], color_pair[1], inout['output'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[1], color_pair[0], inout['output'])\n",
    "\n",
    "            inout['input'] = inout['input'].tolist()\n",
    "            inout['output'] = inout['output'].tolist()\n",
    "            if inout['input'] != tmp_inp.tolist():\n",
    "                for io in inouts:\n",
    "                    if io['input'] == inout['input']:\n",
    "                        same_flag = True\n",
    "                        break\n",
    "                if not same_flag:\n",
    "                    for inout_other in train:\n",
    "                        if inout_other['input'].tolist() == inout['input']:\n",
    "                            same_flag = True\n",
    "                            break\n",
    "                if not same_flag:\n",
    "                    inouts.append({'input': inout['input'], 'output': inout['output']})\n",
    "    if verbose:\n",
    "        print(len(inouts))\n",
    "\n",
    "    return inouts\n",
    "\n",
    "\n",
    "def heavy_color_augument(tasks):\n",
    "    for idx, task in tasks.iteritems():\n",
    "        task['train'] += one_train_heavy_color_augument(task['train'])\n",
    "\n",
    "def one_train_heavy_color_augument(train, verbose=False):\n",
    "    two_colors = []\n",
    "    for pair in combinations([1,2,3,4,5,6,7,8,9], 2):\n",
    "        two_colors.append(pair)\n",
    "    color_pairs_list = []\n",
    "    for i in [3, 4]:\n",
    "        for color_pairs in combinations(two_colors, i):\n",
    "            if has_duplicates(list(itertools.chain.from_iterable(list(color_pairs)))):\n",
    "                continue\n",
    "            color_pairs_list.append(color_pairs)\n",
    "\n",
    "    inouts = []\n",
    "    for color_pairs in color_pairs_list:\n",
    "        for j, inout_origin in enumerate(train):\n",
    "            inout = copy.deepcopy(inout_origin)\n",
    "            same_flag = False\n",
    "            inout['input'] = np.array(inout['input'])\n",
    "            inout['output'] = np.array(inout['output'])\n",
    "            tmp_inp = copy.deepcopy(inout['input'])\n",
    "            tmp_out = copy.deepcopy(inout['output'])\n",
    "\n",
    "            for color_pair in color_pairs:\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[0], color_pair[1], inout['input'])\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[1], color_pair[0], inout['input'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[0], color_pair[1], inout['output'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[1], color_pair[0], inout['output'])\n",
    "\n",
    "            inout['input'] = inout['input'].tolist()\n",
    "            inout['output'] = inout['output'].tolist()\n",
    "            if inout['input'] != tmp_inp.tolist():\n",
    "                for io in inouts:\n",
    "                    if io['input'] == inout['input']:\n",
    "                        same_flag = True\n",
    "                        break\n",
    "                if not same_flag:\n",
    "                    for inout_other in train:\n",
    "                        if inout_other['input'].tolist() == inout['input']:\n",
    "                            same_flag = True\n",
    "                            break\n",
    "                if not same_flag:\n",
    "                    inouts.append({'input': inout['input'], 'output': inout['output']})\n",
    "    if verbose:\n",
    "        print(len(inouts))\n",
    "\n",
    "    return inouts\n",
    "\n",
    "\n",
    "def medium_heavy_color_augument(tasks):\n",
    "    for idx, task in tasks.iteritems():\n",
    "        task['train'] += one_train_medium_heavy_color_augument(task['train'])\n",
    "\n",
    "def one_train_medium_heavy_color_augument(train, verbose=False):\n",
    "    two_colors = []\n",
    "    for pair in combinations([1,2,3,4,5,6,7,8,9], 2):\n",
    "        two_colors.append(pair)\n",
    "    color_pairs_list = []\n",
    "    for color_pairs in combinations(two_colors, 2):\n",
    "        if has_duplicates(list(itertools.chain.from_iterable(list(color_pairs)))):\n",
    "            continue\n",
    "        color_pairs_list.append(color_pairs)\n",
    "    inouts = []\n",
    "    for color_pairs in color_pairs_list:\n",
    "        for j, inout_origin in enumerate(train):\n",
    "            inout = copy.deepcopy(inout_origin)\n",
    "            same_flag = False\n",
    "            inout['input'] = np.array(inout['input'])\n",
    "            inout['output'] = np.array(inout['output'])\n",
    "            tmp_inp = copy.deepcopy(inout['input'])\n",
    "            tmp_out = copy.deepcopy(inout['output'])\n",
    "\n",
    "            for color_pair in color_pairs:\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[0], color_pair[1], inout['input'])\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[1], color_pair[0], inout['input'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[0], color_pair[1], inout['output'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[1], color_pair[0], inout['output'])\n",
    "\n",
    "            inout['input'] = inout['input'].tolist()\n",
    "            inout['output'] = inout['output'].tolist()\n",
    "            if inout['input'] != tmp_inp.tolist():\n",
    "                for io in inouts:\n",
    "                    if io['input'] == inout['input']:\n",
    "                        same_flag = True\n",
    "                        break\n",
    "                if not same_flag:\n",
    "                    for inout_other in train:\n",
    "                        if inout_other['input'].tolist() == inout['input']:\n",
    "                            same_flag = True\n",
    "                            break\n",
    "                if not same_flag:\n",
    "                    inouts.append({'input': inout['input'], 'output': inout['output']})\n",
    "    if verbose:\n",
    "        print(len(inouts))\n",
    "\n",
    "    return inouts\n",
    "\n",
    "\n",
    "def medium_color_augument(tasks):\n",
    "    for idx, task in tasks.iteritems():\n",
    "        task['train'] += one_train_medium_color_augument(task['train'])\n",
    "\n",
    "def one_train_medium_color_augument(train, verbose=False):\n",
    "    color_pairs_list = [\n",
    "        [[2, 3], [4, 5], [6, 7], [8, 9]],\n",
    "        [[1, 3], [4, 6], [8, 7], [2, 9]],\n",
    "        [[6, 3], [2, 5], [4, 7], [1, 9]],\n",
    "        [[2, 4], [7, 5], [6, 9], [8, 1]],\n",
    "        [[1, 4], [5, 6], [8, 3], [2, 7]],\n",
    "        [[7, 3], [6, 1], [8, 4], [5, 9]],\n",
    "    ]\n",
    "    inouts = []\n",
    "    for color_pairs in color_pairs_list:\n",
    "        if has_duplicates(list(itertools.chain.from_iterable(list(color_pairs)))):\n",
    "            continue\n",
    "        for j, inout_origin in enumerate(train):\n",
    "            inout = copy.deepcopy(inout_origin)\n",
    "            same_flag = False\n",
    "            inout['input'] = np.array(inout['input'])\n",
    "            inout['output'] = np.array(inout['output'])\n",
    "            tmp_inp = copy.deepcopy(inout['input'])\n",
    "            tmp_out = copy.deepcopy(inout['output'])\n",
    "\n",
    "            for color_pair in color_pairs:\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[0], color_pair[1], inout['input'])\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[1], color_pair[0], inout['input'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[0], color_pair[1], inout['output'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[1], color_pair[0], inout['output'])\n",
    "\n",
    "            inout['input'] = inout['input'].tolist()\n",
    "            inout['output'] = inout['output'].tolist()\n",
    "            if inout['input'] != tmp_inp.tolist():\n",
    "                for io in inouts:\n",
    "                    if io['input'] == inout['input']:\n",
    "                        same_flag = True\n",
    "                        break\n",
    "                if not same_flag:\n",
    "                    for inout_other in train:\n",
    "                        if inout_other['input'].tolist() == inout['input']:\n",
    "                            same_flag = True\n",
    "                            break\n",
    "                if not same_flag:\n",
    "                    inouts.append({'input': inout['input'], 'output': inout['output']})\n",
    "    if verbose:\n",
    "        print(len(inouts))\n",
    "\n",
    "    return inouts\n",
    "\n",
    "\n",
    "def medium_light_color_augument(tasks):\n",
    "    for idx, task in tasks.iteritems():\n",
    "        task['train'] += one_train_medium_light_color_augument(task['train'])\n",
    "\n",
    "def one_train_medium_light_color_augument(train, verbose=False):\n",
    "    color_pairs_list = [\n",
    "        [[2, 3], [4, 5], [6, 7], [8, 9]],\n",
    "        [[1, 3], [4, 6], [8, 7], [2, 9]],\n",
    "        [[6, 3], [2, 5], [4, 7], [1, 9]],\n",
    "    ]\n",
    "    inouts = []\n",
    "    for color_pairs in color_pairs_list:\n",
    "        if has_duplicates(list(itertools.chain.from_iterable(list(color_pairs)))):\n",
    "            continue\n",
    "        for j, inout_origin in enumerate(train):\n",
    "            inout = copy.deepcopy(inout_origin)\n",
    "            same_flag = False\n",
    "            inout['input'] = np.array(inout['input'])\n",
    "            inout['output'] = np.array(inout['output'])\n",
    "            tmp_inp = copy.deepcopy(inout['input'])\n",
    "            tmp_out = copy.deepcopy(inout['output'])\n",
    "\n",
    "            for color_pair in color_pairs:\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[0], color_pair[1], inout['input'])\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[1], color_pair[0], inout['input'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[0], color_pair[1], inout['output'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[1], color_pair[0], inout['output'])\n",
    "\n",
    "            inout['input'] = inout['input'].tolist()\n",
    "            inout['output'] = inout['output'].tolist()\n",
    "            if inout['input'] != tmp_inp.tolist():\n",
    "                for io in inouts:\n",
    "                    if io['input'] == inout['input']:\n",
    "                        same_flag = True\n",
    "                        break\n",
    "                if not same_flag:\n",
    "                    for inout_other in train:\n",
    "                        if inout_other['input'].tolist() == inout['input']:\n",
    "                            same_flag = True\n",
    "                            break\n",
    "                if not same_flag:\n",
    "                    inouts.append({'input': inout['input'], 'output': inout['output']})\n",
    "    if verbose:\n",
    "        print(len(inouts))\n",
    "\n",
    "    return inouts\n",
    "\n",
    "\n",
    "def light_color_augument(tasks):\n",
    "    for idx, task in tasks.iteritems():\n",
    "        task['train'] += one_train_light_color_augument(task['train'])\n",
    "\n",
    "def one_train_light_color_augument(train, verbose=False):\n",
    "    color_pairs_list = [\n",
    "        [[2, 3], [4, 5], [6, 7], [8, 9]],\n",
    "    ]\n",
    "    inouts = []\n",
    "    for color_pairs in color_pairs_list:\n",
    "        if has_duplicates(list(itertools.chain.from_iterable(list(color_pairs)))):\n",
    "            continue\n",
    "        for j, inout_origin in enumerate(train):\n",
    "            inout = copy.deepcopy(inout_origin)\n",
    "            same_flag = False\n",
    "            inout['input'] = np.array(inout['input'])\n",
    "            inout['output'] = np.array(inout['output'])\n",
    "            tmp_inp = copy.deepcopy(inout['input'])\n",
    "            tmp_out = copy.deepcopy(inout['output'])\n",
    "\n",
    "            for color_pair in color_pairs:\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[0], color_pair[1], inout['input'])\n",
    "                inout['input'] = np.where(tmp_inp == color_pair[1], color_pair[0], inout['input'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[0], color_pair[1], inout['output'])\n",
    "                inout['output'] = np.where(tmp_out == color_pair[1], color_pair[0], inout['output'])\n",
    "\n",
    "            inout['input'] = inout['input'].tolist()\n",
    "            inout['output'] = inout['output'].tolist()\n",
    "            if inout['input'] != tmp_inp.tolist():\n",
    "                for io in inouts:\n",
    "                    if io['input'] == inout['input']:\n",
    "                        same_flag = True\n",
    "                        break\n",
    "                if not same_flag:\n",
    "                    for inout_other in train:\n",
    "                        if inout_other['input'].tolist() == inout['input']:\n",
    "                            same_flag = True\n",
    "                            break\n",
    "                if not same_flag:\n",
    "                    inouts.append({'input': inout['input'], 'output': inout['output']})\n",
    "    if verbose:\n",
    "        print(len(inouts))\n",
    "\n",
    "    return inouts\n",
    "\n",
    "def mirror_augument(tasks):\n",
    "    for idx, task in tasks.iteritems():\n",
    "        task['train'] += one_train_mirror_augument(task['train'])\n",
    "\n",
    "def one_train_mirror_augument(train):\n",
    "    aug_func_petterns = [np.transpose, mirror_x, mirror_y]\n",
    "    inouts = []\n",
    "    for i in range(len(aug_func_petterns)):\n",
    "        for func_combi in combinations(aug_func_petterns, i+1):\n",
    "            for j, inout_origin in enumerate(train):\n",
    "                inout = copy.deepcopy(inout_origin)\n",
    "                same_flag = False\n",
    "                for func in list(func_combi):\n",
    "                    inout['input'] = func(inout['input']).tolist()\n",
    "                    inout['output'] = func(inout['output']).tolist()\n",
    "                if inout['input'] != inout_origin['input']:\n",
    "                    for io in inouts:\n",
    "                        if io['input'] == inout['input']:\n",
    "                            same_flag = True\n",
    "                            break\n",
    "                    if not same_flag:\n",
    "                        for inout_other in train:\n",
    "                            if inout_other['input'].tolist() == inout['input']:\n",
    "                                same_flag = True\n",
    "                                break\n",
    "                    if not same_flag:\n",
    "                        inouts.append({'input': inout['input'], 'output': inout['output']})\n",
    "\n",
    "    return inouts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# get_mode_preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_mode_preds(preds_list):\n",
    "    preds = []\n",
    "    origin_shape_map = {}\n",
    "    each_question_preds = defaultdict(list)\n",
    "    for i in range(len(preds_list[0])):\n",
    "        origin_shape_map[i] = np.array(preds_list[0][i]).shape\n",
    "        for preds in preds_list:\n",
    "            each_question_preds[i].append(np.array(preds[i]).reshape(-1))\n",
    "    mode_preds = []\n",
    "    \n",
    "    for i in range(len(preds_list[0])):\n",
    "        ans = []\n",
    "        for j in range(len(each_question_preds[i][0])):\n",
    "            ms = [m[j] for m in each_question_preds[i]]\n",
    "            ans.append(np.argmax(np.bincount(ms)))\n",
    "        mode_preds.append(np.array(ans).reshape(origin_shape_map[i]).tolist())\n",
    "    return mode_preds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# final score update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "metadata": {},
   "outputs": [],
   "source": [
    "def final_score_update(test_tasks, preds_list, final_score_map, idx, success_map):\n",
    "    print(f'{idx}, 順番: mode, (aug CA), aug func0')\n",
    "    for i in range(len(preds_list[0])):\n",
    "        pred_str = ''\n",
    "        for j, preds in enumerate(reversed(preds_list)):\n",
    "            pred = np.array(preds[i])\n",
    "            if test_tasks[i]['output'] == pred.tolist():\n",
    "                success_map[f'{idx}_{i}'] = True            \n",
    "            similarity = match_rate(pred, test_tasks[i]['output'])\n",
    "            print(f'similarity: {similarity}')\n",
    "            if (idx not in final_score_map) or (final_score_map.get(idx, 101) < similarity):\n",
    "                final_score_map[idx] = similarity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# final_train_and_predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {},
   "outputs": [],
   "source": [
    "each_preds = defaultdict(lambda: defaultdict(list))\n",
    "ca_skip = False\n",
    "def final_train_and_predict(task_train, task_train2, task_train_aug, task_test, task_test2, idx, success_map, final_score_map, final=False, promising=False, origin_task=None):\n",
    "    funcs0 = []\n",
    "    funcs1 = []\n",
    "    preds_list = []\n",
    "    mode_preds_list = []\n",
    "    if not final:\n",
    "#         st()\n",
    "        ts = TaskSolver()\n",
    "        ts.train(task_train, preprocess_funcs=funcs0)\n",
    "        preds = ts.predict(task_test, preprocess_funcs=funcs0, success_map=success_map, idx=idx, final_score_map=final_score_map)\n",
    "        return preds \n",
    "    \n",
    "    # 別のNNに通す\n",
    "#     if promising:\n",
    "#         train = copy.deepcopy(task_train_aug)\n",
    "#         test = copy.deepcopy(task_test)\n",
    "#         for func in funcs0:\n",
    "#             for inout in train:\n",
    "#                 inout['input'] = func(inout['input'])\n",
    "#             for inout in test:\n",
    "#                 inout['input'] = func(inout['input'])\n",
    "#         preds = test_nn(train, test)\n",
    "#         preds = post_processes(preds, origin_task)\n",
    "#         each_preds[idx]['another NN'] = preds\n",
    "#         preds_list.append(preds)\n",
    "#         mode_preds_list.append(preds)\n",
    "#         preds = preds_to_str_only0(preds, idx)\n",
    "\n",
    "    # not aug, funcs0\n",
    "#     ts = TaskSolver()\n",
    "#     ts.train(task_train, preprocess_funcs=funcs0, final=final)\n",
    "#     preds3 = ts.predict(task_test, preprocess_funcs=funcs0, success_map=success_map, idx=idx, final_score_map=final_score_map, final=final)\n",
    "#     preds3 = post_processes(preds3, origin_task)\n",
    "#     preds_list.append(preds3)\n",
    "#     mode_preds_list.append(preds3)\n",
    "#     each_preds[idx]['not aug, normal NN, funcs0'] = preds3\n",
    "#     st()\n",
    "    \n",
    "    # not aug, funcs1\n",
    "#     ts = TaskSolver()\n",
    "#     ts.train(task_train2, preprocess_funcs=funcs1, final=final)\n",
    "#     preds4 = ts.predict(task_test2, preprocess_funcs=funcs1, success_map=success_map, idx=idx, final_score_map=final_score_map, final=final)\n",
    "#     preds4 = post_processes(preds4, origin_task)\n",
    "#     each_preds[idx]['not aug, normal NN, funcs1'] = preds4\n",
    "#     preds_list.append(preds4)\n",
    "#     mode_preds_list.append(preds4)\n",
    "\n",
    "    # not aug, func0\n",
    "    ts = TaskSolver()\n",
    "    ts.train(task_train, preprocess_funcs=funcs1, final=final)\n",
    "    preds4 = ts.predict(task_test, preprocess_funcs=funcs1, success_map=success_map, idx=idx, final_score_map=final_score_map, final=final)\n",
    "    preds4 = post_processes(preds4, origin_task)\n",
    "    each_preds[idx]['not aug, normal NN, funcs1'] = preds4\n",
    "    preds_list.append(preds4)\n",
    "    mode_preds_list.append(preds4)    \n",
    "    \n",
    "    # aug, funcs0\n",
    "    ts = TaskSolver()\n",
    "    ts.train(task_train_aug, preprocess_funcs=funcs0, final=final)\n",
    "    preds2 = ts.predict(task_test, preprocess_funcs=funcs0, success_map=success_map, idx=idx, final_score_map=final_score_map, final=final)\n",
    "    preds2 = post_processes(preds2, origin_task)\n",
    "    preds_list.append(preds2)\n",
    "    mode_preds_list.append(preds2)\n",
    "    mode_preds_list.append(preds2) # weight2\n",
    "    each_preds[idx]['aug, normal NN, funcs0'] = preds2\n",
    "\n",
    "    if (len(task_train_aug) < 200) & (not ca_skip):\n",
    "#     if False:\n",
    "#         print('CA1 start')\n",
    "        # not aug CA\n",
    "        # TODO: 消すか検証する\n",
    "#         train_changed = []\n",
    "#         test_changed = []\n",
    "#         for inout in task_train:\n",
    "#             io = copy.deepcopy(inout)\n",
    "#             for f in funcs0:\n",
    "#                 io['input'] = f(io['input'])\n",
    "#             train_changed.append(io)\n",
    "#         for inout in task_test:\n",
    "#             io = copy.deepcopy(inout)\n",
    "#             for f in funcs0:\n",
    "#                 io['input'] = f(io['input'])\n",
    "#             test_changed.append(io)\n",
    "#         model, _, _ = solve_task(train_changed)\n",
    "#         preds0 = predict(model, test_changed)\n",
    "#         preds0 = post_processes(preds0, origin_task)\n",
    "#         each_preds[idx]['not aug CA'] = preds0\n",
    "#         preds_list.append(preds0)\n",
    "#         preds_list.append(preds0)\n",
    "\n",
    "        # aug, CA\n",
    "        print('CA1 start')\n",
    "#         train_changed = []\n",
    "#         test_changed = []\n",
    "#         for inout in task_train_aug:\n",
    "#             io = copy.deepcopy(inout)\n",
    "#             for f in funcs0:\n",
    "#                 io['input'] = f(io['input'])\n",
    "#             train_changed.append(io)\n",
    "#         for inout in task_test:\n",
    "#             io = copy.deepcopy(inout)\n",
    "#             for f in funcs0:\n",
    "#                 io['input'] = f(io['input'])\n",
    "#             test_changed.append(io)\n",
    "        model, _, _ = solve_task(task_train_aug)\n",
    "        preds1 = predict(model, task_test)\n",
    "        preds1 = post_processes(preds1, origin_task)\n",
    "        preds_list.append(preds1)\n",
    "        mode_preds_list.append(preds1)\n",
    "        mode_preds_list.append(preds1) # weight2\n",
    "        each_preds[idx]['aug CA'] = preds1\n",
    "\n",
    "    preds_mode = get_mode_preds(mode_preds_list)\n",
    "    each_preds[idx]['mode matrix'] = preds_mode\n",
    "    preds_list.append(preds_mode)\n",
    "    if 'output' in task_test[0]:\n",
    "#         st()\n",
    "        final_score_update(task_test, preds_list, final_score_map, idx, success_map)\n",
    "\n",
    "    preds = preds_to_str(preds_list, idx)\n",
    "#     st()\n",
    "\n",
    "    return preds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# apply_aug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_mirror_aug(train, preprocess_best_score, idx, use_transpose_flag, promising_map):\n",
    "    use_mirror_augs = []\n",
    "    for aug_func in [flipud_aug, fliplr_aug, flip_aug]:\n",
    "        inouts = aug_func(train)\n",
    "#         st()\n",
    "        similarity = get_similarity(train+inouts, [], 'searching_mirror_aug_'+idx, {})\n",
    "        print(f'{aug_func}: ------> {similarity}')\n",
    "        if similarity > 0.99:\n",
    "            promising_map[idx] = True\n",
    "        if (similarity > preprocess_best_score) or ((similarity==1) & (preprocess_best_score != 1)):\n",
    "            use_mirror_augs.append(aug_func)\n",
    "    if use_transpose_flag:\n",
    "        print(transpose_aug)\n",
    "        inouts = transpose_aug(train)\n",
    "        similarity = get_similarity(train+inouts, [], 'searching_transpose_aug_'+idx, {})\n",
    "        print(similarity, preprocess_best_score)\n",
    "        if (similarity > preprocess_best_score) or ((similarity==1) & (preprocess_best_score != 1)):\n",
    "            use_mirror_augs.append(transpose_aug)\n",
    "        if similarity > 0.99:\n",
    "            promising_map[idx] = True\n",
    "    return use_mirror_augs\n",
    "\n",
    "def apply_transpose_aug(train):\n",
    "    for inout_origin in train:\n",
    "        inout = copy.deepcopy(inout_origin)\n",
    "        m, n = np.array(inout['input']).shape\n",
    "        if m != n:\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "def apply_color_aug(train, preprocess_best_score, best_aug_score_map, idx, promising_map):\n",
    "    best_aug_score_map[idx] = 0\n",
    "    use_inouts = []\n",
    "    use_aug_func = return_arg\n",
    "    skip_heavy_flag = False\n",
    "    heavy_funcs = [one_train_medium_heavy_color_augument, one_train_heavy_color_augument, one_train_super_heavy_color_augument]\n",
    "#     for aug_func in [one_train_light_color_augument, one_train_medium_light_color_augument, one_train_medium_color_augument, one_train_medium_heavy_color_augument, one_train_heavy_color_augument, one_train_super_heavy_color_augument]:\n",
    "    for aug_func in [one_train_light_color_augument, one_train_medium_light_color_augument, one_train_medium_color_augument, one_train_medium_heavy_color_augument, one_train_heavy_color_augument]:\n",
    "#     for aug_func in [one_train_medium_heavy_color_augument, one_train_light_color_augument, one_train_medium_light_color_augument, one_train_medium_color_augument]:\n",
    "        if aug_func in heavy_funcs:\n",
    "            if skip_heavy_flag == True:\n",
    "                continue\n",
    "            if (best_aug_score_map[idx] < 0.997) or (best_aug_score_map[idx] < preprocess_best_score+0.04):\n",
    "                skip_heavy_flag = True\n",
    "                continue\n",
    "        inouts = aug_func(train)\n",
    "        scores = []\n",
    "        # 重い場合時間がかかるので学習は一度にする\n",
    "        if aug_func in heavy_funcs:\n",
    "            ts = TaskSolver()\n",
    "            tmp_train = train + inouts\n",
    "            if len(tmp_train) < 10:\n",
    "                continue\n",
    "            val_train, tmp_train = tmp_train[:3], tmp_train[3:]\n",
    "            ts.train(tmp_train, preprocess_funcs=[])\n",
    "            for i in range(3):\n",
    "                preds = ts.predict([val_train[i]], preprocess_funcs=[])\n",
    "                similarity = match_rate(preds[0], val_train[i]['output'])\n",
    "                scores.append(similarity)\n",
    "        # 軽い場合はなるべくpreと条件を揃えたいので都度学習\n",
    "        else:\n",
    "            for i in range(3):\n",
    "                similarity = train_and_evaluate(train+inouts, [], seed=i, idx='searching_aug', success_map={})\n",
    "                scores.append(similarity)\n",
    "        score = np.mean(scores)\n",
    "        print(f'{aug_func}: ------> {score}')\n",
    "        if score > 0.9999:\n",
    "            promising_map[idx] = True\n",
    "            return use_inouts, aug_func\n",
    "        if score < 0.8:\n",
    "            return use_inouts, use_aug_func\n",
    "        if (score > best_aug_score_map[idx]) & (score > preprocess_best_score):\n",
    "            best_aug_score_map[idx] = score\n",
    "            use_inouts = inouts\n",
    "            use_aug_func = aug_func\n",
    "        # 常に更新、かつ直前で0.99以上じゃない場合、重い処理をスキップ\n",
    "        if score < best_aug_score_map[idx]:\n",
    "            skip_heavy_flag = True\n",
    "        if (aug_func == one_train_medium_heavy_color_augument) & ((score < 0.997)):\n",
    "            skip_heavy_flag = True\n",
    "\n",
    "    if best_aug_score_map[idx] > 0.98:\n",
    "        promising_map[idx] = True\n",
    "    return use_inouts, use_aug_func\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# train_and_evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_and_evaluate(train, func_combi, seed, idx, success_map, search_func=False):\n",
    "    ts = TaskSolver()\n",
    "    tmp_train = copy.deepcopy(train)\n",
    "    val_train = [tmp_train.pop(seed % len(tmp_train))]\n",
    "    ts.train(tmp_train, preprocess_funcs=func_combi)\n",
    "    preds = ts.predict(val_train, preprocess_funcs=func_combi) # idxを渡すとsimilarityがprintされる\n",
    "    return match_rate(preds[0], val_train[0]['output'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add fill closed area"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_fill_closed_area(train):\n",
    "    inp_origin = train[0]['input']\n",
    "    out = train[0]['output']\n",
    "    apply_flag = False\n",
    "    for func in [np.array, np.flip, np.fliplr, np.flipud]:\n",
    "        inp = np.array(inp_origin.copy())\n",
    "        inp = func(inp)\n",
    "        if len(set([ele for ele in np.array(out).reshape(-1) - np.array(inp).reshape(-1) if ele != 0])) == 1:\n",
    "            fill_color = [ele for ele in np.array(out).reshape(-1) - np.array(inp).reshape(-1) if ele != 0][0]\n",
    "            apply_flag = True\n",
    "            break\n",
    "    if not apply_flag:\n",
    "        return [inouts_array]\n",
    "    best_score = 0\n",
    "    best_enclose_color = 0\n",
    "\n",
    "    for enclose_color in range(1,10):\n",
    "        inp_copy = inp.copy()\n",
    "        if enclose_color == fill_color:\n",
    "            continue\n",
    "        H, W = inp_copy.shape\n",
    "        Dy = [0, -1, 0, 1]\n",
    "        Dx = [1, 0, -1, 0]\n",
    "        arr_padded = np.pad(inp_copy, ((1,1),(1,1)), \"constant\", constant_values=0)\n",
    "        searched = np.zeros(arr_padded.shape, dtype=bool)\n",
    "        searched[0, 0] = True\n",
    "        q = [(0, 0)]\n",
    "        while q:\n",
    "            y, x = q.pop()\n",
    "            for dy, dx in zip(Dy, Dx):\n",
    "                y_, x_ = y+dy, x+dx\n",
    "                if not 0 <= y_ < H+2 or not 0 <= x_ < W+2:\n",
    "                    continue\n",
    "                if not searched[y_][x_] and arr_padded[y_][x_]==0:\n",
    "                    q.append((y_, x_))\n",
    "                    searched[y_, x_] = True\n",
    "        res = searched[1:-1, 1:-1]\n",
    "        res |= inp_copy==enclose_color\n",
    "        inp_copy[~res]=fill_color\n",
    "        similarity = match_rate(inp_copy, out)\n",
    "        if similarity > best_score:\n",
    "            best_score = similarity\n",
    "            best_enclose_color = enclose_color\n",
    "    def fill_closed_area(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            inp = np.array(inp)\n",
    "            H, W = inp.shape\n",
    "            Dy = [0, -1, 0, 1]\n",
    "            Dx = [1, 0, -1, 0]\n",
    "            arr_padded = np.pad(inp, ((1,1),(1,1)), \"constant\", constant_values=0)\n",
    "            searched = np.zeros(arr_padded.shape, dtype=bool)\n",
    "            searched[0, 0] = True\n",
    "            q = [(0, 0)]\n",
    "            while q:\n",
    "                y, x = q.pop()\n",
    "                for dy, dx in zip(Dy, Dx):\n",
    "                    y_, x_ = y+dy, x+dx\n",
    "                    if not 0 <= y_ < H+2 or not 0 <= x_ < W+2:\n",
    "                        continue\n",
    "                    if not searched[y_][x_] and arr_padded[y_][x_]==0:\n",
    "                        q.append((y_, x_))\n",
    "                        searched[y_, x_] = True\n",
    "            res = searched[1:-1, 1:-1]\n",
    "            res |= inp==best_enclose_color\n",
    "\n",
    "            inp[~res] = fill_color\n",
    "    #         st()\n",
    "            inout['input'] = inp\n",
    "        return task_train\n",
    "\n",
    "    return [inouts_array, fill_closed_area]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_train0_double(task_train, m1,n1,m2,n2):\n",
    "    if not((m2 >= m1) & (n2 >= n1) & (m2 % m1 == 0) & (n2 % n1 == 0)):\n",
    "        return []\n",
    "    for inout in task_train:\n",
    "        m1_, n1_ = np.array(inout['input']).shape\n",
    "        m2_, n2_ = np.array(inout['output']).shape\n",
    "        if (m2 / m1 != m2_ / m1_) or (n2 / n1 != n2_ / n1_):\n",
    "            return []    \n",
    "    def train0_double(task_train):\n",
    "        for inout in task_train:\n",
    "            x = inout['input']\n",
    "            x = np.array(x)\n",
    "            m, n = m2//m1, n2//n1\n",
    "            x_upsampled = x.repeat(m, axis=0).repeat(n, axis=1)\n",
    "            x_tiled = np.tile(x, (m, n))\n",
    "            y = x_upsampled & x_tiled\n",
    "            inout['input'] = y    \n",
    "        return task_train\n",
    "    return [train0_double]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add double func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 前提：same_shape_inputs, same_shape_outputs\n",
    "def add_gdc_double_funcs_with_same_shape_inputs(task_train, m1, n1, m2, n2):\n",
    "    if (m1==m2) & (n1==n2):\n",
    "        return []\n",
    "    m_gdc = math.gcd(m1, m2)\n",
    "    n_gdc = math.gcd(n1, n2)\n",
    "    if ((m_gdc==1) or (n_gdc==1)):\n",
    "        return []\n",
    "    if not ((m2 >= m1) & (n2 >= n1)):\n",
    "        return []\n",
    "    transpose_funcs = [np.array]\n",
    "    if m1 == n1:\n",
    "        transpose_funcs.append(np.transpose)\n",
    "    transpose_flip_map = {}\n",
    "    flip_cols_map = {}\n",
    "    flip_rows_map = {}\n",
    "    inp_mn_map = {}\n",
    "    for m in range(m2 // m_gdc):\n",
    "        for n in range(n2 // n_gdc):\n",
    "            transpose_flip_map[str(m)+','+str(n)] = [np.array] # 初期値\n",
    "            flip_cols_map[str(m)+','+str(n)] = []\n",
    "            flip_rows_map[str(m)+','+str(n)] = []\n",
    "            inp_mn_map[str(m)+','+str(n)] = [0,0]\n",
    "            correct_flag = False\n",
    "            pickup_output = np.array(task_train[0]['output'])[m*m_gdc:(m+1)*m_gdc, n*n_gdc:(n+1)*n_gdc]\n",
    "            best_score = 0\n",
    "            for transpose_func in transpose_funcs:\n",
    "                if correct_flag:\n",
    "                    break\n",
    "                for inp_m in range(m1 // m_gdc):\n",
    "                    if correct_flag:\n",
    "                        break\n",
    "                    for inp_n in range(n1 // n_gdc):\n",
    "                        if correct_flag:\n",
    "                            break\n",
    "                        inp_copy = np.array(copy.deepcopy(task_train[0]['input']))\n",
    "                        inp_copy = inp_copy[inp_m*m_gdc:(inp_m+1)*m_gdc, inp_n*n_gdc:(inp_n+1)*n_gdc]\n",
    "                        inp_copy = transpose_func(inp_copy)\n",
    "                        for flip_func in [np.flip, np.flipud, np.fliplr, np.array]:\n",
    "                            if correct_flag:\n",
    "                                break\n",
    "                            inp_copy_copy = copy.deepcopy(inp_copy)\n",
    "                            inp_copy_copy = flip_func(inp_copy_copy)\n",
    "                            if pickup_output.tolist() == inp_copy_copy.tolist():\n",
    "                                correct_flag = True\n",
    "                                transpose_flip_map[str(m)+','+str(n)] = [transpose_func, flip_func]\n",
    "                                inp_mn_map[str(m)+','+str(n)] = [inp_n, inp_m]\n",
    "                                flip_cols_map[str(m)+','+str(n)] = []\n",
    "                                flip_rows_map[str(m)+','+str(n)] = []\n",
    "                                break\n",
    "                            similarity = match_rate(pickup_output, inp_copy_copy)\n",
    "                            if best_score < similarity:\n",
    "                                best_score = similarity\n",
    "                                transpose_flip_map[str(m)+','+str(n)] = [transpose_func, flip_func]\n",
    "                                inp_mn_map[str(m)+','+str(n)] = [inp_n, inp_m]\n",
    "                                flip_cols_map[str(m)+','+str(n)] = []\n",
    "                                flip_rows_map[str(m)+','+str(n)] = []\n",
    "#                                 st()\n",
    "                        for i in range(m_gdc+1):\n",
    "                            if correct_flag:\n",
    "                                break\n",
    "                            for change_rows in combinations(range(m_gdc), i):\n",
    "                                if correct_flag:\n",
    "                                    break\n",
    "                                change_rows = list(change_rows)\n",
    "                                for j in range(n_gdc+1):\n",
    "                                    if correct_flag:\n",
    "                                        break\n",
    "                                    for change_cols in combinations(range(n_gdc), j):\n",
    "                                        change_cols = list(change_cols)\n",
    "                                        inp_copy_copy = copy.deepcopy(inp_copy)\n",
    "                                        inp_copy_copy[change_rows, :] = np.fliplr(inp_copy_copy[change_rows, :])\n",
    "                                        inp_copy_copy[:, change_cols] = np.flipud(inp_copy_copy[:, change_cols])\n",
    "                                        if pickup_output.tolist() == inp_copy_copy.tolist():\n",
    "                                            correct_flag = True\n",
    "                                            transpose_flip_map[str(m)+','+str(n)] = [transpose_func, flip_func]\n",
    "                                            inp_mn_map[str(m)+','+str(n)] = [inp_n, inp_m]\n",
    "                                            flip_cols_map[str(m)+','+str(n)] = change_cols\n",
    "                                            flip_rows_map[str(m)+','+str(n)] = change_rows  \n",
    "                                            break\n",
    "                                        \n",
    "                                        similarity = match_rate(pickup_output, inp_copy_copy)\n",
    "                                        if best_score < similarity:\n",
    "                                            best_score = similarity\n",
    "                                            transpose_flip_map[str(m)+','+str(n)] = [transpose_func, flip_func]\n",
    "                                            inp_mn_map[str(m)+','+str(n)] = [inp_n, inp_m]\n",
    "                                            flip_cols_map[str(m)+','+str(n)] = change_cols\n",
    "                                            flip_rows_map[str(m)+','+str(n)] = change_rows\n",
    "\n",
    "    def double(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            ans = np.zeros((m2, n2)).astype(int)\n",
    "            for coordinate, transpose_funcs in transpose_flip_map.items():\n",
    "                m, n = coordinate.split(',')\n",
    "                m, n = int(m), int(n)\n",
    "                inp_copy = np.array(copy.deepcopy(inp))\n",
    "                inp_n, inp_m = inp_mn_map[coordinate]\n",
    "                inp_copy = inp_copy[inp_m*m_gdc:(inp_m+1)*m_gdc, inp_n*n_gdc:(inp_n+1)*n_gdc]\n",
    "                for transpose_func in transpose_funcs:\n",
    "                    inp_copy = transpose_func(inp_copy)\n",
    "                change_cols = flip_cols_map[coordinate]\n",
    "                change_rows = flip_rows_map[coordinate]\n",
    "                inp_copy[:, change_cols] = np.flipud(inp_copy[:, change_cols])\n",
    "                inp_copy[change_rows, :] = np.fliplr(inp_copy[change_rows, :])\n",
    "                ans[m*m_gdc:(m+1)*m_gdc, n*n_gdc:(n+1)*n_gdc] = inp_copy\n",
    "            inout['input'] = ans\n",
    "        return task_train\n",
    "\n",
    "    return [double]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "metadata": {},
   "outputs": [],
   "source": [
    "# each shapeが違って割合が一緒の場合\n",
    "def add_gcd_double_funcs(task_train, m1, n1, m2, n2):\n",
    "    if (m2==m1) & (n2==n1):\n",
    "        return []\n",
    "    m = math.gcd(m1, m2)\n",
    "    n = math.gcd(n1, n2)    \n",
    "    if not ((m2 >= m1) & (n2 >= n1) & (m2 % m1 == 0) & (n2 % n1 == 0)):\n",
    "        return []\n",
    "    for inout in task_train:\n",
    "        m1_, n1_ = np.array(inout['input']).shape\n",
    "        m2_, n2_ = np.array(inout['output']).shape\n",
    "        if (m2 / m1 != m2_ / m1_) or (n2 / n1 != n2_ / n1_):\n",
    "            return []\n",
    "\n",
    "    transpose_funcs = [np.array]\n",
    "    if m1 == n1:\n",
    "        transpose_funcs.append(np.transpose)\n",
    "    transpose_flip_map = {}\n",
    "    flip_cols_map = {}\n",
    "    flip_rows_map = {}\n",
    "    for m in range(m2 // m1):\n",
    "        for n in range(n2 // n1):\n",
    "            transpose_flip_map[str(m)+','+str(n)] = [np.array]\n",
    "            correct_flag = False\n",
    "            pickup_output = np.array(task_train[0]['output'])[m*m1:(m+1)*m1, n*n1:(n+1)*n1]\n",
    "            best_score = 0\n",
    "            for flip_func in [np.flip, np.flipud, np.fliplr, np.array]:\n",
    "                for transpose_func in transpose_funcs:\n",
    "                    if correct_flag:\n",
    "                        break\n",
    "                    inp_copy = copy.deepcopy(task_train[0]['input'])\n",
    "                    inp_copy = transpose_func(inp_copy)\n",
    "                    inp_copy = flip_func(inp_copy)\n",
    "                    if pickup_output.tolist() == inp_copy.tolist():\n",
    "                        correct_flag = True\n",
    "                        transpose_flip_map[str(m)+','+str(n)] = [flip_func, transpose_func]\n",
    "                        flip_cols_map[str(m)+','+str(n)] = []\n",
    "                        flip_rows_map[str(m)+','+str(n)] = []\n",
    "                    similarity = match_rate(pickup_output, inp_copy)\n",
    "                    if best_score < similarity:\n",
    "                        best_score = similarity\n",
    "                        transpose_flip_map[str(m)+','+str(n)] = [flip_func, transpose_func]\n",
    "                        flip_cols_map[str(m)+','+str(n)] = []\n",
    "                        flip_rows_map[str(m)+','+str(n)] = []\n",
    "    def double(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            inp_m, inp_n = np.array(inp).shape\n",
    "            m, n = m2//m1 - 1, n2 // n1 - 1\n",
    "            ans = np.zeros((inp_m*(int(m)+1), inp_n*(int(n)+1))).astype(int)\n",
    "            for coordinate, transpose_funcs in transpose_flip_map.items():\n",
    "                m, n = coordinate.split(',')\n",
    "                m, n = int(m), int(n)\n",
    "                inp_copy = copy.deepcopy(inp)\n",
    "                for transpose_func in transpose_funcs:\n",
    "                    inp_copy = transpose_func(inp_copy)\n",
    "\n",
    "                ans[m*inp_m:(m+1)*inp_m, n*inp_n:(n+1)*inp_n] = inp_copy\n",
    "            inout['input'] = ans\n",
    "        return task_train\n",
    "\n",
    "    return [double]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {},
   "outputs": [],
   "source": [
    "# all_output_shape are same\n",
    "# outputsの形は違うけど倍率が同じ場合は別の関数で\n",
    "def add_double_funcs_with_same_shape_inputs(train, m1, n1, m2, n2):\n",
    "    inp = np.array(train[0]['input'].copy())\n",
    "    out = np.array(train[0]['output'].copy())\n",
    "    tmp_ans = np.zeros(out.shape)\n",
    "    ans = np.zeros(out.shape)\n",
    "    best_ans_ms = []\n",
    "    best_ans_ns = []\n",
    "    best_m_flips= []\n",
    "    best_n_flips = []\n",
    "    if (m2==m1) & (n2==n1):\n",
    "        return []\n",
    "    # 縦横どちらかでもoutputが小さい場合は別関数で\n",
    "    if not ((m2 >= m1) & (n2 >= n1)):\n",
    "        return []\n",
    "    for inout in train:\n",
    "        m1_, n1_ = np.array(inout['input']).shape\n",
    "        m2_, n2_ = np.array(inout['output']).shape\n",
    "        if (m1 != m1_) or (n1 != n1_) or (m2 != m2_) or (n2 != n2_):\n",
    "            return []\n",
    "\n",
    "    for ans_m in range(m2):\n",
    "        o = out[ans_m:ans_m+1, :n1]\n",
    "        best_score = 0\n",
    "        for inp_m in range(m1):\n",
    "            i = inp[inp_m:inp_m+1, :]\n",
    "#             st()\n",
    "            for flip in [np.array, np.fliplr]:\n",
    "                similarity = match_rate(flip(i), flip(o))\n",
    "                if best_score < similarity:\n",
    "                    best_score = similarity\n",
    "                    best_ans_m = inp_m\n",
    "                    best_flip = flip\n",
    "\n",
    "        best_ans_ms.append(best_ans_m)\n",
    "        best_m_flips.append(best_flip)\n",
    "\n",
    "    for i, (flip, m) in enumerate(zip(best_m_flips, best_ans_ms)):\n",
    "        tmp_ans[i:i+1, :n1] = flip(inp[m:m+1, :])\n",
    "    for ans_n in range(n2):\n",
    "        o = out[:, ans_n:ans_n+1]\n",
    "        best_score = 0\n",
    "        for inp_n in range(n1):\n",
    "            i = tmp_ans[:, inp_n:inp_n+1]\n",
    "            for flip in [np.array, np.fliplr]:\n",
    "                similarity = match_rate(flip(i), flip(o))\n",
    "                if best_score < similarity:\n",
    "                    best_score = similarity\n",
    "                    best_ans_n = inp_n\n",
    "                    best_flip = flip\n",
    "        best_ans_ns.append(best_ans_n)    \n",
    "        best_n_flips.append(best_flip)\n",
    "    def double(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            inp = np.array(inp)\n",
    "            tmp_ans = np.zeros(out.shape)\n",
    "            ans = np.zeros(out.shape)        \n",
    "            for i, (flip, m) in enumerate(zip(best_m_flips, best_ans_ms)):\n",
    "                tmp_ans[i:i+1, :n1] = flip(inp[m:m+1, :])\n",
    "\n",
    "            for i, (flip, n) in enumerate(zip(best_n_flips, best_ans_ns)):\n",
    "                ans[:, i:i+1] = flip(tmp_ans[:, n:n+1])\n",
    "            inout['input'] = ans\n",
    "        return task_train\n",
    "    \n",
    "    return [double]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_period_length_vertical(arr):\n",
    "    arr = np.array(arr)\n",
    "    H, W = arr.shape\n",
    "    period = 1\n",
    "    while True:\n",
    "        cycled = np.pad(arr[:period, :], ((0,H-period),(0,0)), 'wrap')\n",
    "        if (cycled==arr).all():\n",
    "            return period\n",
    "        period += 1\n",
    "        \n",
    "def add_train2_double_vertical(task_train, m1, n1, m2, n2):\n",
    "    if not((n1 == n2) & (m2 > m1)):\n",
    "        return []\n",
    "    \n",
    "    def train2_double(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            inp = np.array(inp)\n",
    "            period = get_period_length_vertical(inp)\n",
    "            y = inp[:period, :]\n",
    "            y = np.pad(y, ((0,m2-period),(0,0)), 'wrap')\n",
    "            inout['input'] = y\n",
    "        return task_train\n",
    "    return [train2_double]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_period_length_horizontal(arr):\n",
    "    arr = np.array(arr)\n",
    "    H, W = arr.shape\n",
    "    period = 1\n",
    "    while True:\n",
    "        cycled = np.pad(arr[:, :period], ((0,0),(0,W-period)), 'wrap')\n",
    "        if (cycled==arr).all():\n",
    "            return period\n",
    "        period += 1\n",
    "        \n",
    "def add_train2_double_horizontal(task_train, m1, n1, m2, n2):\n",
    "    if not((m1 == m2) & (n2 > n1)):\n",
    "        return []\n",
    "    \n",
    "    def train2_double(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            inp = np.array(inp)\n",
    "            period = get_period_length_horizontal(inp)\n",
    "            y = inp[:, :period]\n",
    "            y = np.pad(y, ((0,0),(0,n2-period)), 'wrap')\n",
    "            inout['input'] = y\n",
    "        return task_train    \n",
    "    return [train2_double]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add crop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "metadata": {},
   "outputs": [],
   "source": [
    "def crop_width(task_train_origin):\n",
    "    for inout in task_train_origin:\n",
    "        inout['input'] = np.array(inout['input'])\n",
    "    try:\n",
    "        task_train = copy.deepcopy(task_train_origin)\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            max_width = 0\n",
    "            max_width_color = 0\n",
    "            for c in [color for color in np.unique(inp) if color != 0]:\n",
    "                coords = np.argwhere(inp==c)\n",
    "                x_min, y_min = coords.min(axis=0)\n",
    "                x_max, y_max = coords.max(axis=0)        \n",
    "                if max_width < y_max - y_min:\n",
    "                    max_width = y_max - y_min\n",
    "                    max_width_color = c\n",
    "            coords = np.argwhere(inp==max_width_color)\n",
    "            x_min, y_min = coords.min(axis=0)\n",
    "            x_max, y_max = coords.max(axis=0)        \n",
    "            inout['input'] = inp[x_min:x_max+1, y_min:y_max+1]\n",
    "        return task_train\n",
    "    except:\n",
    "        return task_train_origin\n",
    "\n",
    "def add_crop_width(task_train):\n",
    "    for inout in task_train:\n",
    "        try:\n",
    "            inp = np.array(inout['input'])\n",
    "            out = np.array(inout['output'])\n",
    "            max_width = 0\n",
    "            max_width_color = 0\n",
    "            for c in [color for color in np.unique(inp) if color != 0]:\n",
    "                coords = np.argwhere(inp==c)\n",
    "                x_min, y_min = coords.min(axis=0)\n",
    "                x_max, y_max = coords.max(axis=0)   \n",
    "                if max_width < y_max - y_min:\n",
    "                    max_width = y_max - y_min\n",
    "                    max_width_color = c\n",
    "            coords = np.argwhere(inp==max_width_color)\n",
    "            x_min, y_min = coords.min(axis=0)\n",
    "            x_max, y_max = coords.max(axis=0)               \n",
    "            if (inp[x_min:x_max+1, y_min:y_max+1].shape != out.shape) & (inp[x_min:x_max+1, y_min:y_max+1].shape != out.T.shape):\n",
    "                return []   \n",
    "        except:\n",
    "            return []\n",
    "\n",
    "    return [crop_height]\n",
    "def crop_height(task_train_origin):\n",
    "    for inout in task_train_origin:\n",
    "        inout['input'] = np.array(inout['input'])\n",
    "    try:\n",
    "        task_train = copy.deepcopy(task_train_origin)\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            max_height = 0\n",
    "            max_height_color = 0\n",
    "            for c in [color for color in np.unique(inp) if color != 0]:\n",
    "                coords = np.argwhere(inp==c)\n",
    "                x_min, y_min = coords.min(axis=0)\n",
    "                x_max, y_max = coords.max(axis=0)        \n",
    "                if max_height < x_max - x_min:\n",
    "                    max_height = x_max - x_min\n",
    "                    max_height_color = c\n",
    "            coords = np.argwhere(inp==max_height_color)\n",
    "            x_min, y_min = coords.min(axis=0)\n",
    "            x_max, y_max = coords.max(axis=0)        \n",
    "            inout['input'] = inp[x_min:x_max+1, y_min:y_max+1]\n",
    "        return task_train\n",
    "    except:\n",
    "        return task_train_origin\n",
    "\n",
    "def add_crop_height(task_train):\n",
    "    for inout in task_train:\n",
    "        try:\n",
    "            inp = np.array(inout['input'])\n",
    "            out = np.array(inout['output'])\n",
    "            max_height = 0\n",
    "            max_height_color = 0\n",
    "            for c in [color for color in np.unique(inp) if color != 0]:\n",
    "                coords = np.argwhere(inp==c)\n",
    "                x_min, y_min = coords.min(axis=0)\n",
    "                x_max, y_max = coords.max(axis=0)   \n",
    "                if max_height < x_max - x_min:\n",
    "                    max_height = x_max - x_min\n",
    "                    max_height_color = c\n",
    "            coords = np.argwhere(inp==max_height_color)\n",
    "            x_min, y_min = coords.min(axis=0)\n",
    "            x_max, y_max = coords.max(axis=0)               \n",
    "            if (inp[x_min:x_max+1, y_min:y_max+1].shape != out.shape) & (inp[x_min:x_max+1, y_min:y_max+1].shape != out.T.shape):\n",
    "                return []   \n",
    "        except:\n",
    "            return []\n",
    "\n",
    "    return [crop_height]\n",
    "def crop_max(task_train_origin):\n",
    "    for inout in task_train_origin:\n",
    "        inout['input'] = np.array(inout['input'])\n",
    "    task_train = copy.deepcopy(task_train_origin)\n",
    "    try:\n",
    "        for inout in task_train:\n",
    "            a = inout['input']\n",
    "            b = np.bincount(a.flatten(),minlength=10)\n",
    "            b[0] = 255\n",
    "            c = np.argsort(b)[-2]\n",
    "            coords = np.argwhere(a==c)\n",
    "            x_min, y_min = coords.min(axis=0)\n",
    "            x_max, y_max = coords.max(axis=0)\n",
    "            inout['input'] = a[x_min:x_max+1, y_min:y_max+1]\n",
    "    except:\n",
    "        return task_train_origin\n",
    "    return task_train\n",
    "    \n",
    "def crop_min(task_train_origin):\n",
    "    for inout in task_train_origin:\n",
    "        inout['input'] = np.array(inout['input'])\n",
    "    task_train = copy.deepcopy(task_train_origin)\n",
    "    try:\n",
    "        for inout in task_train:\n",
    "            a = inout['input']\n",
    "            b = np.bincount(a.flatten(),minlength=10)\n",
    "            c = int(np.where(b==np.min(b[np.nonzero(b)]))[0])\n",
    "            coords = np.argwhere(a==c)\n",
    "            x_min, y_min = coords.min(axis=0)\n",
    "            x_max, y_max = coords.max(axis=0)\n",
    "            inout['input'] = a[x_min:x_max+1, y_min:y_max+1]\n",
    "    except:\n",
    "        return task_train_origin\n",
    "    return task_train\n",
    "\n",
    "def add_crop_max(task_train):\n",
    "    for inout in task_train:\n",
    "        try:\n",
    "            inp = np.array(inout['input'])\n",
    "            out = np.array(inout['output'])\n",
    "            bin_c = np.bincount(inp.flatten(), minlength=10)\n",
    "\n",
    "            bin_c[0] = 255\n",
    "            c = np.argsort(bin_c)[-2]\n",
    "            coords = np.argwhere(inp==c)\n",
    "            x_min, y_min = coords.min(axis=0)\n",
    "            x_max, y_max = coords.max(axis=0)\n",
    "            inp = inp[x_min:x_max+1, y_min:y_max+1]  \n",
    "            if (inp.shape != out.shape) & (inp.T.shape != out.shape):\n",
    "                return []\n",
    "        except:\n",
    "            return []\n",
    "    return [crop_max]\n",
    "\n",
    "def add_crop_min(task_train):\n",
    "    for inout in task_train:\n",
    "        try: \n",
    "            inp = np.array(inout['input'])\n",
    "            out = np.array(inout['output'])\n",
    "            bin_c = np.bincount(inp.flatten(), minlength=10)\n",
    "            c = int(np.where(bin_c==np.min(bin_c[np.nonzero(bin_c)]))[0])\n",
    "            coords = np.argwhere(inp==c)\n",
    "            x_min, y_min = coords.min(axis=0)\n",
    "            x_max, y_max = coords.max(axis=0)\n",
    "            inp = inp[x_min:x_max+1, y_min:y_max+1]\n",
    "            if (inp.shape != out.shape) & (inp.T.shape != out.shape):\n",
    "                return []\n",
    "        except:\n",
    "            return []\n",
    "    return [crop_min]    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# all_inputs_same_shape and all_outputs_same_shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "metadata": {},
   "outputs": [],
   "source": [
    "def all_inputs_same_shape_and_all_outputs_same_shape(task_train):\n",
    "    m1, n1 = np.array(task_train[0]['input']).shape\n",
    "    m2, n2 = np.array(task_train[0]['output']).shape\n",
    "    all_inputs_same_shape = True\n",
    "    all_outputs_same_shape = True\n",
    "    for inout in task_train:\n",
    "        m1_, n1_ = np.array(inout['input']).shape\n",
    "        m2_, n2_ = np.array(inout['output']).shape\n",
    "        if (m1_ != m1) or (n1_ != n1):\n",
    "            all_inputs_same_shape = False\n",
    "        if (m2_ != m2) or (n2_ != n2):\n",
    "            all_outputs_same_shape = False\n",
    "    return all_inputs_same_shape, all_outputs_same_shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add size change funcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_size_change_funcs(task_train, task_n):\n",
    "    size_change_funcs = [inouts_array]\n",
    "    m1, n1 = np.array(task_train[0]['input']).shape\n",
    "    m2, n2 = np.array(task_train[0]['output']).shape\n",
    "    all_inputs_same_shape = True\n",
    "    all_outputs_same_shape = True\n",
    "    for inout in task_train:\n",
    "        m1_, n1_ = np.array(inout['input']).shape\n",
    "        m2_, n2_ = np.array(inout['output']).shape\n",
    "        if (m1_ != m1) or (n1_ != n1):\n",
    "            all_inputs_same_shape = False       \n",
    "        if (m2_ != m2) or (n2_ != n2):\n",
    "            all_outputs_same_shape = False\n",
    "    if m1 == n1 == m2 == n2:\n",
    "        return size_change_funcs    \n",
    "\n",
    "    # grid\n",
    "    size_change_funcs += add_grid_funcs(m1, n1, m2, n2)\n",
    "    \n",
    "    # div\n",
    "    if (m1 >= m2*2) or (n1 > n2*2):\n",
    "        size_change_funcs += add_div_funcs(task_train, m1, n1, m2, n2)\n",
    "    else:\n",
    "        size_change_funcs += add_div_funcs2(task_train, m1, n1, m2, n2, vertically=True)\n",
    "        size_change_funcs += add_div_funcs2(task_train, m1, n1, m2, n2, vertically=False)\n",
    "    if (m1 > m2) & (n1 > n2) & (m1 < 20) & (n1 < 20):\n",
    "        size_change_funcs += add_object_detect2(task_train)\n",
    "        \n",
    "    # double\n",
    "    if all_inputs_same_shape & all_outputs_same_shape:\n",
    "        size_change_funcs += add_train2_double_horizontal(task_train, m1, n1, m2, n2)\n",
    "        size_change_funcs += add_train2_double_vertical(task_train, m1, n1, m2, n2)\n",
    "        size_change_funcs += add_gdc_double_funcs_with_same_shape_inputs(task_train, m1, n1, m2, n2)\n",
    "        size_change_funcs += add_recolor(task_train, task_n)\n",
    "    else:\n",
    "        size_change_funcs += add_gcd_double_funcs(task_train, m1, n1, m2, n2)\n",
    "        \n",
    "    size_change_funcs += add_train0_double(task_train, m1, n1, m2, n2)\n",
    "    if (m1 >= m2) & (n1 >= n2):\n",
    "        size_change_funcs += add_crop_max(task_train)\n",
    "        size_change_funcs += add_crop_min(task_train)\n",
    "        size_change_funcs += add_crop_height(task_train)\n",
    "        size_change_funcs += add_crop_width(task_train)\n",
    "        size_change_funcs += add_crop_by_line(task_train)\n",
    "    # 他にもたくさん足す\n",
    "    return size_change_funcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for interface...\n",
    "def return_arg(inp):\n",
    "    return inp\n",
    "\n",
    "def inouts_transpose(task_train):\n",
    "    for inout in task_train:\n",
    "        inout['input'] = np.transpose(inout['input'])\n",
    "    return task_train\n",
    "\n",
    "def inouts_array(task_train):\n",
    "    for inout in task_train:\n",
    "        inout['input'] = np.array(inout['input'])\n",
    "    return task_train\n",
    "\n",
    "def add_transpose(task_train):\n",
    "    m1, n1 = np.array(task_train[0]['input']).shape\n",
    "    m2, n2 = np.array(task_train[0]['output']).shape\n",
    "    if (m1==n2) & (n1==m2):\n",
    "        return [inouts_array, inouts_transpose]\n",
    "    else:\n",
    "        return [inouts_array]\n",
    "\n",
    "def add_grid_funcs(m1, n1, m2, n2):\n",
    "    grid_funcs = []\n",
    "    if (m1 <= m2) & (n1 <= n2):\n",
    "        return grid_funcs\n",
    "    if (m2%m1 == 0) & (n2%n1 == 0):\n",
    "        def griding(task_train):\n",
    "            for inout in enumerate(task_train):\n",
    "                inp = copy.deepcopy(inout['input'])\n",
    "                m_grid = m2 // m1\n",
    "                n_grid = n2 // n1\n",
    "                inp = np.array(inp)\n",
    "                m, n = inp.shape\n",
    "                ans_tmp = np.zeros((m*m_grid, n), dtype='int')\n",
    "                for i in range(m):\n",
    "                    for j in range(m_grid):\n",
    "                        ans_tmp[i*m_grid+j, :] = inp[i, :]\n",
    "                ans = copy.deepcopy(ans_tmp)\n",
    "                for stack_n in range(n_grid-1):\n",
    "                    ans = np.hstack([ans, ans_tmp])\n",
    "                for i in range(n):\n",
    "                    for j in range(n_grid):\n",
    "                        ans[:, i*n_grid+j] = ans_tmp[:, i]\n",
    "                inout['input'] = ans\n",
    "            return task_train\n",
    "        grid_funcs.append(griding)\n",
    "\n",
    "    if (m1 != n1) & (m2%n1 == 0) & (n2%m1 == 0):\n",
    "        def transpose_griding(task_train):\n",
    "            for inout in task_train:\n",
    "                inp = copy.deepcopy(inp_o)\n",
    "                m_grid = m2 // n1\n",
    "                n_grid = n2 // m1\n",
    "                inp = np.transpose(inp)\n",
    "                m, n = inp.shape\n",
    "                ans_tmp = np.zeros((m*m_grid, n), dtype='int')\n",
    "                for i in range(m):\n",
    "                    for j in range(m_grid):\n",
    "                        ans_tmp[i*m_grid+j, :] = inp[i, :]\n",
    "                ans = copy.deepcopy(ans_tmp)\n",
    "                for stack_n in range(n_grid-1):\n",
    "                    ans = np.hstack([ans, ans_tmp])\n",
    "                for i in range(n):\n",
    "                    for j in range(n_grid):\n",
    "                        ans[:, i*n_grid+j] = ans_tmp[:, i]\n",
    "                inout['input'] = ans\n",
    "            return task_train\n",
    "        grid_funcs.append(transpose_griding)\n",
    "    return grid_funcs\n",
    "\n",
    "def div_two_inputs(inp, long):\n",
    "    inp = np.array(inp)\n",
    "    m, n = inp.shape\n",
    "    if m == n:\n",
    "        return inp, inp\n",
    "    horizontal = False\n",
    "    # 縦長にする\n",
    "    if n > m:\n",
    "        horizontal = True\n",
    "        inp = inp.T\n",
    "        m, n = inp.shape\n",
    "\n",
    "    a = inp[:long, :]\n",
    "    b = inp[m-long:, :]\n",
    "    # 元が横長だったら戻す\n",
    "    if horizontal:\n",
    "        a = a.T\n",
    "        b = b.T\n",
    "    return a, b\n",
    "\n",
    "def add_div_funcs(train, m1, n1, m2, n2):\n",
    "    for inout in train:\n",
    "        m1_0, n1_0 = np.array(inout['input']).shape\n",
    "        m2_0, n2_0 = np.array(inout['output']).shape\n",
    "        if (m1_0 != m1) or (n1_0 != n1) or (m2_0 != m2) or (n2_0 != n2):\n",
    "            return []\n",
    "    if (m1 == n1) or (np.min([m1, n1]) != np.min([m2, n2])) or(np.max([m1,n1]) <= np.max([m2,n2])):\n",
    "        return []\n",
    "    long = np.max([m2,n2])\n",
    "    def div_and(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs(inp, long)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            inout['input'] = a_align_num&b_align_num\n",
    "        return task_train\n",
    "\n",
    "    def div_or(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs(inp, long)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            inout['input'] = a_align_num|b_align_num\n",
    "        return task_train\n",
    "\n",
    "    def div_xor(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs(inp, long)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            inout['input'] = a_align_num^b_align_num\n",
    "        return task_train\n",
    "    def div_not_and(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs(inp, long)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            c = a_align_num&b_align_num\n",
    "            inout['input'] = np.where(c == 0, not_use_num, 0)\n",
    "        return task_train\n",
    "\n",
    "    def div_not_or(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs(inp, long)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            c = a_align_num|b_align_num\n",
    "            inout['input'] = np.where(c == 0, not_use_num, 0)\n",
    "        return task_train\n",
    "\n",
    "    def div_not_xor(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs(inp, long)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            c = a_align_num^b_align_num\n",
    "            inout['input'] = np.where(c == 0, not_use_num, 0)\n",
    "        return task_train\n",
    "\n",
    "    return [div_and, div_or, div_xor, div_not_and, div_not_or, div_not_xor]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "metadata": {},
   "outputs": [],
   "source": [
    "def div_two_inputs2(inp, belt_length, vertically):\n",
    "    inp = np.array(inp)\n",
    "    if vertically:\n",
    "        inp = inp.T\n",
    "    m, n = inp.shape\n",
    "    after_n = (n-belt_length) // 2\n",
    "\n",
    "    a = inp[:, :after_n]\n",
    "    b = inp[:, n-after_n:]\n",
    "    if vertically:\n",
    "        a, b = a.T, b.T\n",
    "    return a, b\n",
    "\n",
    "# 真ん中のベルトで分けるタイプ。ベルトの長さは各inputで一定の場合のみ\n",
    "def add_div_funcs2(train, m1, n1, m2, n2, vertically):\n",
    "    if vertically:\n",
    "        if (n1 != n2) or (m1 < m2*2):\n",
    "            return []\n",
    "        belt_length = m1 - m2*2\n",
    "    else:\n",
    "        if (m1 != m2) or (n1 < n2*2):\n",
    "            return []\n",
    "        belt_length = n1 - n2*2\n",
    "\n",
    "    for inout in train:\n",
    "        m1_0, n1_0 = np.array(inout['input']).shape\n",
    "        m2_0, n2_0 = np.array(inout['output']).shape\n",
    "        if vertically:\n",
    "            if (n1_0 != n2_0) or (m1_0 != m2_0*2 + belt_length):\n",
    "                return []\n",
    "        else:\n",
    "            if (m1_0 != m2_0) or (n1_0 != n2_0*2 + belt_length):\n",
    "                return []\n",
    "\n",
    "    def div_and(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs2(inp, belt_length, vertically)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            inout['input'] = a_align_num&b_align_num\n",
    "        return task_train\n",
    "\n",
    "    def div_or(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs2(inp, belt_length, vertically)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            inout['input'] = a_align_num|b_align_num\n",
    "        return task_train\n",
    "\n",
    "    def div_xor(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs2(inp, belt_length, vertically)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            inout['input'] = a_align_num^b_align_num\n",
    "        return task_train\n",
    "\n",
    "    def div_not_and(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs2(inp, belt_length, vertically)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            c = a_align_num&b_align_num\n",
    "            inout['input'] = np.where(c == 0, not_use_num, 0)\n",
    "        return task_train\n",
    "\n",
    "    def div_not_or(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs2(inp, belt_length, vertically)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            c = a_align_num|b_align_num\n",
    "            inout['input'] = np.where(c == 0, not_use_num, 0)\n",
    "        return task_train\n",
    "\n",
    "    def div_not_xor(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            a, b = div_two_inputs2(inp, belt_length, vertically)\n",
    "            not_use_num = get_not_use_num(a, b)\n",
    "            a_align_num = np.where(a==0, 0, not_use_num)\n",
    "            b_align_num = np.where(b==0, 0, not_use_num)\n",
    "            c = a_align_num^b_align_num\n",
    "            inout['input'] = np.where(c == 0, not_use_num, 0)\n",
    "        return task_train\n",
    "\n",
    "    return [div_and, div_or, div_xor, div_not_and, div_not_or, div_not_xor]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add patch funcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {},
   "outputs": [],
   "source": [
    "patch_skip = False\n",
    "def add_patch_funcs(task_train, idx):\n",
    "    if patch_skip:\n",
    "        return [inouts_array]\n",
    "    \n",
    "    if correct_task_train(task_train):\n",
    "        return [inouts_array]\n",
    "    \n",
    "    if idx in ['15696249', 'e345f17b']:\n",
    "        return [inouts_array]\n",
    "    inp, out = np.array(task_train[-1][\"input\"]), np.array(task_train[-1][\"output\"])\n",
    "    for inout in task_train:\n",
    "        if np.array(inout[\"input\"]).shape != inp.shape:\n",
    "            return [inouts_array]\n",
    "\n",
    "    flip_funcs = [np.array, np.flip, np.flipud, np.fliplr]\n",
    "    transpose_funcs = [np.array, np.transpose]\n",
    "    best_score = match_rate(inp, out)\n",
    "    best_feat = None\n",
    "    for flip_func in flip_funcs:\n",
    "        for transpose_func in transpose_funcs:\n",
    "            inp_copy = copy.deepcopy(inp)\n",
    "            inp_copy = flip_func(inp_copy)\n",
    "            inp_copy = transpose_func(inp_copy)\n",
    "            pred, feat = call_pred_train(inp_copy, out, patch_image)\n",
    "            similarity = match_rate(out, pred)\n",
    "            if best_score < similarity:\n",
    "                best_score = similarity\n",
    "                best_flip_func = flip_func\n",
    "                best_transpose_func = transpose_func\n",
    "                best_feat = feat\n",
    "\n",
    "    def this_time_patch_image(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            if (best_feat is not None) & (best_feat != {}):\n",
    "                inp = best_flip_func(inp)\n",
    "                inp = best_transpose_func(inp)\n",
    "    #             print(best_feat)\n",
    "                pred = call_pred_test(inp, patch_image, best_feat)\n",
    "    #             if np.array(pred).shape != task['test'][0]''\n",
    "                if pred.shape != np.array(inp).shape:\n",
    "                    inout['input'] = np.array(inp)\n",
    "                    continue\n",
    "                inout['input'] = pred\n",
    "            else:\n",
    "                inout['input'] = np.array(inp)\n",
    "        return task_train\n",
    "    \n",
    "    return [this_time_patch_image, inouts_array]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [],
   "source": [
    "def in_out_diff(t_in, t_out):\n",
    "    x_in, y_in = t_in.shape\n",
    "    x_out, y_out = t_out.shape\n",
    "    diff = np.zeros((max(x_in, x_out), max(y_in, y_out)))\n",
    "    diff[:x_in, :y_in] -= t_in\n",
    "    diff[:x_out, :y_out] += t_out\n",
    "    return diff\n",
    "\n",
    "def check_symmetric(a):\n",
    "    try:\n",
    "        sym = 1\n",
    "        if np.array_equal(a, a.T):\n",
    "            sym *= 2 #Check main diagonal symmetric (top left to bottom right)\n",
    "        if np.array_equal(a, np.flip(a).T):\n",
    "            sym *= 3 #Check antidiagonal symmetric (top right to bottom left)\n",
    "        if np.array_equal(a, np.flipud(a)):\n",
    "            sym *= 5 # Check horizontal symmetric of array\n",
    "        if np.array_equal(a, np.fliplr(a)):\n",
    "            sym *= 7 # Check vertical symmetric of array\n",
    "        return sym\n",
    "    except:\n",
    "        return 0\n",
    "    \n",
    "def bbox(a):\n",
    "    try:\n",
    "        r = np.any(a, axis=1)\n",
    "        c = np.any(a, axis=0)\n",
    "        rmin, rmax = np.where(r)[0][[0, -1]]\n",
    "        cmin, cmax = np.where(c)[0][[0, -1]]\n",
    "        return rmin, rmax, cmin, cmax\n",
    "    except:\n",
    "        return 0,a.shape[0],0,a.shape[1]\n",
    "\n",
    "def cmask(t_in):\n",
    "    cmin = 999\n",
    "    cm = 0\n",
    "    for c in range(10):\n",
    "        t = t_in.copy().astype('int8')\n",
    "        t[t==c],t[t>0],t[t<0]=-1,0,1\n",
    "        b = bbox(t)\n",
    "        a = (b[1]-b[0])*(b[3]-b[2])\n",
    "        s = (t[b[0]:b[1],b[2]:b[3]]).sum()\n",
    "        if a>2 and a<cmin and s==a:\n",
    "            cmin=a\n",
    "            cm=c\n",
    "    return cm\n",
    "\n",
    "def mask_rect(a):\n",
    "    r,c = a.shape\n",
    "    m = a.copy().astype('uint8')\n",
    "    for i in range(r-1):\n",
    "        for j in range(c-1):\n",
    "            if m[i,j]==m[i+1,j]==m[i,j+1]==m[i+1,j+1]>=1:m[i,j]=2\n",
    "            if m[i,j]==m[i+1,j]==1 and m[i,j-1]==2:m[i,j]=2\n",
    "            if m[i,j]==m[i,j+1]==1 and m[i-1,j]==2:m[i,j]=2\n",
    "            if m[i,j]==1 and m[i-1,j]==m[i,j-1]==2:m[i,j]=2\n",
    "    m[m==1]=0\n",
    "    return (m==2)\n",
    "    \n",
    "def call_pred_train(t_in, t_out, pred_func):\n",
    "    try:\n",
    "        feat = {}\n",
    "        feat['s_out'] = t_out.shape\n",
    "        if t_out.shape==t_in.shape:\n",
    "            diff = in_out_diff(t_in,t_out)\n",
    "            feat['diff'] = diff\n",
    "            feat['cm'] = t_in[diff!=0].max()\n",
    "        else:\n",
    "            feat['diff'] = (t_in.shape[0]-t_out.shape[0],t_in.shape[1]-t_out.shape[1])\n",
    "            feat['cm'] = cmask(t_in)\n",
    "        feat['sym'] = check_symmetric(t_out)\n",
    "        args = inspect.getargspec(pred_func).args\n",
    "        if len(args)==1:\n",
    "            return pred_func(t_in), feat\n",
    "        elif len(args)==2:\n",
    "            t_pred = pred_func(t_in,feat[args[1]])    \n",
    "        elif len(args)==3:\n",
    "            t_pred = pred_func(t_in,feat[args[1]],feat[args[2]])\n",
    "        feat['sizeok'] = len(t_out)==len(t_pred)\n",
    "        t_pred = np.resize(t_pred,t_out.shape)\n",
    "        acc = (t_pred==t_out).sum()/t_out.size\n",
    "        return t_pred, feat\n",
    "    except:\n",
    "        return t_in, {}\n",
    "\n",
    "def call_pred_test(t_in, pred_func, feat):\n",
    "    args = inspect.getargspec(pred_func).args\n",
    "    if len(args)==1:\n",
    "        return pred_func(t_in)\n",
    "    elif len(args)==2:\n",
    "        t_pred = pred_func(t_in,feat[args[1]]) \n",
    "    elif len(args)==3:\n",
    "        t_pred = pred_func(t_in,feat[args[1]],feat[args[2]])\n",
    "    return t_pred\n",
    "\n",
    "num2color = [\"black\", \"blue\", \"red\", \"green\", \"yellow\", \"gray\", \"magenta\", \"orange\", \"sky\", \"brown\"]\n",
    "color2num = {c: n for n, c in enumerate(num2color)}\n",
    "\n",
    "def get_tile(img ,mask):\n",
    "    try:\n",
    "        m,n = img.shape\n",
    "        a = img.copy().astype('int8')\n",
    "        a[mask] = -1\n",
    "        r=c=0\n",
    "        for x in range(n):\n",
    "            if np.count_nonzero(a[0:m,x]<0):continue\n",
    "            for r in range(2,m):\n",
    "                if 2*r<m and (a[0:r,x]==a[r:2*r,x]).all():break\n",
    "            if r<m:break\n",
    "            else: r=0\n",
    "        for y in range(m):\n",
    "            if np.count_nonzero(a[y,0:n]<0):continue\n",
    "            for c in range(2,n):\n",
    "                if 2*c<n and (a[y,0:c]==a[y,c:2*c]).all():break\n",
    "            if c<n:break\n",
    "            else: c=0\n",
    "        if c>0:\n",
    "            for x in range(n-c):\n",
    "                if np.count_nonzero(a[:,x]<0)==0:\n",
    "                    a[:,x+c]=a[:,x]\n",
    "                elif np.count_nonzero(a[:,x+c]<0)==0:\n",
    "                    a[:,x]=a[:,x+c]\n",
    "        if r>0:\n",
    "            for y in range(m-r):\n",
    "                if np.count_nonzero(a[y,:]<0)==0:\n",
    "                    a[y+r,:]=a[y,:]\n",
    "                elif np.count_nonzero(a[y+r,:]<0)==0:\n",
    "                    a[y,:]=a[y+r,:]\n",
    "        return a[r:2*r,c:2*c]\n",
    "    except:\n",
    "        return a[0:1,0:1]\n",
    "    \n",
    "def patch_image(t_in,s_out,cm=0):\n",
    "    t_in = np.array(t_in)\n",
    "    try:\n",
    "        t = t_in.copy()\n",
    "        ty,tx=t.shape\n",
    "        if cm>0:\n",
    "            m = mask_rect(t==cm)\n",
    "        else:\n",
    "            m = (t==cm)   \n",
    "        tile = get_tile(t ,m)\n",
    "        if tile.size>2 and s_out==t.shape:\n",
    "            rt = np.tile(tile,(1+ty//tile.shape[0],1+tx//tile.shape[1]))[0:ty,0:tx]\n",
    "            if (rt[~m]==t[~m]).all():\n",
    "                return rt\n",
    "        for i in range(6):\n",
    "            m = (t==cm)\n",
    "            t -= cm\n",
    "            if tx==ty:\n",
    "                a = np.maximum(t,t.T)\n",
    "                if (a[~m]==t[~m]).all():t=a.copy()\n",
    "                a = np.maximum(t,np.flip(t).T)\n",
    "                if (a[~m]==t[~m]).all():t=a.copy()\n",
    "            a = np.maximum(t,np.flipud(t))\n",
    "            if (a[~m]==t[~m]).all():t=a.copy()\n",
    "            a = np.maximum(t,np.fliplr(t))\n",
    "            if (a[~m]==t[~m]).all():t=a.copy()\n",
    "            t += cm\n",
    "            m = (t==cm)\n",
    "            lms = measure.label(m.astype('uint8'))\n",
    "            for l in range(1,lms.max()+1):\n",
    "                lm = np.argwhere(lms==l)\n",
    "                lm = np.argwhere(lms==l)\n",
    "                x_min = max(0,lm[:,1].min()-1)\n",
    "                x_max = min(lm[:,1].max()+2,t.shape[0])\n",
    "                y_min = max(0,lm[:,0].min()-1)\n",
    "                y_max = min(lm[:,0].max()+2,t.shape[1])\n",
    "                gap = t[y_min:y_max,x_min:x_max]\n",
    "                sy,sx=gap.shape\n",
    "                if i==1:\n",
    "                    sy//=2\n",
    "                    y_max=y_min+sx\n",
    "                gap = t[y_min:y_max,x_min:x_max]\n",
    "                sy,sx=gap.shape\n",
    "                allst = as_strided(t, shape=(ty,tx,sy,sx),strides=2*t.strides)    \n",
    "                allst = allst.reshape(-1,sy,sx)\n",
    "                allst = np.array([a for a in allst if np.count_nonzero(a==cm)==0])\n",
    "                gm = (gap!=cm)\n",
    "                for a in allst:\n",
    "                    if sx==sy:\n",
    "                        fpd = a.T\n",
    "                        fad = np.flip(a).T\n",
    "                        if i==1:gm[sy-1,0]=gm[0,sx-1]=False\n",
    "                        if (fpd[gm]==gap[gm]).all():\n",
    "                            gm = (gap!=cm)\n",
    "                            np.putmask(gap,~gm,fpd)\n",
    "                            t[y_min:y_max,x_min:x_max] = gap\n",
    "                            break\n",
    "                        if i==1:gm[0,0]=gm[sy-1,sx-1]=False\n",
    "                        if (fad[gm]==gap[gm]).all():\n",
    "                            gm = (gap!=cm)\n",
    "                            np.putmask(gap,~gm,fad)\n",
    "                            t[y_min:y_max,x_min:x_max] = gap\n",
    "                            break \n",
    "                    fud = np.flipud(a)\n",
    "                    flr = np.fliplr(a)\n",
    "                    if i==1:gm[sy-1,0]=gm[0,sx-1]=gm[0,0]=gm[sy-1,sx-1]=False\n",
    "                    if (a[gm]==gap[gm]).all():\n",
    "                        gm = (gap!=cm)\n",
    "                        np.putmask(gap,~gm,a)\n",
    "                        t[y_min:y_max,x_min:x_max] = gap\n",
    "                        break\n",
    "                    elif (fud[gm]==gap[gm]).all():\n",
    "                        gm = (gap!=cm)\n",
    "                        np.putmask(gap,~gm,fud)\n",
    "                        t[y_min:y_max,x_min:x_max] = gap\n",
    "                        break\n",
    "                    elif (flr[gm]==gap[gm]).all():\n",
    "                        gm = (gap!=cm)\n",
    "                        np.putmask(gap,~gm,flr)\n",
    "                        t[y_min:y_max,x_min:x_max] = gap\n",
    "                        break\n",
    "        if s_out==t.shape:\n",
    "            return t\n",
    "        else:\n",
    "            m = (t_in==cm)\n",
    "            return np.resize(t[m],crop_min(m).shape)\n",
    "    except:\n",
    "        return t_in\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add train004 growth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_train4_growth(shaped_train):\n",
    "    if correct_task_train(shaped_train):\n",
    "        return [inouts_array]\n",
    "    x = shaped_train[0]['input'].copy()\n",
    "    out = shaped_train[0]['output'].copy()\n",
    "    x = np.array(x)\n",
    "    def get_base_pattern(arr, w, h):\n",
    "        # find maximum number of unique color tiles in 3x3 field\n",
    "        H, W = arr.shape\n",
    "        arr_onehot = 1<<arr\n",
    "        arr_bool = arr.astype(bool).astype(np.int32)\n",
    "        counts = np.zeros(arr.shape, dtype=np.int32)\n",
    "        colors = np.zeros(arr.shape, dtype=np.int32)\n",
    "        for y in range(H-2):\n",
    "            for x in range(W-2):\n",
    "                counts[y, x] = arr_bool[y:y+2, x:x+2].sum()\n",
    "                colors[y, x] = np.bitwise_or.reduce(arr_onehot[y:y+2, x:x+2].reshape(-1))\n",
    "        n_colors = np.zeros(arr.shape, dtype=np.int32)\n",
    "        for c in range(1, 10):\n",
    "            n_colors += colors>>c & 1\n",
    "        counts[n_colors>=2] = 0\n",
    "        res_y, res_x = np.unravel_index(np.argmax(counts), counts.shape)\n",
    "        pattern = arr[res_y:res_y+h, res_x:res_x+w].astype(bool).astype(np.int32)\n",
    "        return (res_y, res_x), pattern\n",
    "    repeat_num = 10\n",
    "    best_wh = (2, 2)\n",
    "    best_score = 0\n",
    "    correct_flag = False\n",
    "    for w in range(2,7):\n",
    "        if correct_flag:\n",
    "            break\n",
    "        for h in range(2,7):\n",
    "            (base_y, base_x), pattern = get_base_pattern(x, w, h)\n",
    "#             p(pattern)\n",
    "            try:\n",
    "                pad_size = repeat_num * np.max([w, h])\n",
    "                x_padded = np.pad(x, ((pad_size,pad_size),(pad_size,pad_size)), \"constant\", constant_values=0)\n",
    "                base_y += pad_size\n",
    "                base_x += pad_size\n",
    "                y = x_padded.copy()\n",
    "                for dy in [-(h+1), 0, h+1]:\n",
    "                    for dx in [-(w+1), 0, w+1]:\n",
    "                        y_, x_ = base_y+dy, base_x+dx\n",
    "                        if dy==dx==0:\n",
    "                            continue\n",
    "                        count = np.bincount(x_padded[y_:y_+h+1, x_:x_+w+1].reshape(-1))\n",
    "                        if count[0]==9:\n",
    "                            continue\n",
    "                        count[0] = 0\n",
    "                        color = count.argmax()\n",
    "                        for i in range(1, repeat_num):\n",
    "                            y[base_y+dy*i:base_y+dy*i+h, base_x+dx*i:base_x+dx*i+w] = color * pattern\n",
    "                y = y[pad_size:-pad_size, pad_size:-pad_size]\n",
    "                score = match_rate(y, out)\n",
    "                if best_score < score:\n",
    "                    best_score = score\n",
    "                    best_wh = (w, h)\n",
    "                    if score == 1:\n",
    "                        correct_flag = True\n",
    "                        break\n",
    "            except:\n",
    "                pass\n",
    "    def train4_growth(task_train):\n",
    "        for inout in task_train:\n",
    "            inp = inout['input']\n",
    "            x = np.array(inp)\n",
    "            try:\n",
    "                w, h = best_wh\n",
    "                (base_y, base_x), pattern = get_base_pattern(x, w, h)\n",
    "                pad_size = repeat_num * np.max([w, h])\n",
    "                x_padded = np.pad(x, ((pad_size,pad_size),(pad_size,pad_size)), \"constant\", constant_values=0)\n",
    "                base_y += pad_size\n",
    "                base_x += pad_size\n",
    "                y = x_padded.copy()\n",
    "                for dy in [-(h+1), 0, h+1]:\n",
    "                    for dx in [-(w+1), 0, w+1]:\n",
    "                        y_, x_ = base_y+dy, base_x+dx\n",
    "                        if dy==dx==0:\n",
    "                            continue\n",
    "                        count = np.bincount(x_padded[y_:y_+h+1, x_:x_+w+1].reshape(-1))\n",
    "                        if count[0]==9:\n",
    "                            continue\n",
    "                        count[0] = 0\n",
    "                        color = count.argmax()\n",
    "                        for i in range(1, repeat_num):\n",
    "                            y[base_y+dy*i:base_y+dy*i+h, base_x+dx*i:base_x+dx*i+w] = color * pattern\n",
    "                inout['input'] = y[pad_size:-pad_size, pad_size:-pad_size]\n",
    "            except:\n",
    "                inout['input'] = x\n",
    "        return task_train\n",
    "        \n",
    "    return [inouts_array, train4_growth]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add change color funcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_change_color_funcs(task_train):\n",
    "    if correct_task_train(task_train):\n",
    "        return [inouts_array]\n",
    "    in_use_colors, out_use_colors, color_changed = about_color(task_train)\n",
    "    if not color_changed:\n",
    "        return [inouts_array]\n",
    "    inout_map = {}\n",
    "    for in_color in in_use_colors:\n",
    "        for out_color in out_use_colors:\n",
    "            scores = []\n",
    "            best_score = 0\n",
    "            for inout in task_train:\n",
    "                inp = inout['input'].copy()\n",
    "                out = inout['output'].copy()\n",
    "                in_vec = list(itertools.chain.from_iterable(inp))\n",
    "                out_vec = list(itertools.chain.from_iterable(out))\n",
    "                if (in_color not in in_vec) or (out_color not in out_vec):\n",
    "                    continue\n",
    "                inp = np.where(np.array(inp) == in_color, out_color, inp)\n",
    "                scores.append(match_rate(inp, out))\n",
    "            if np.mean(scores) > best_score:\n",
    "                best_score = np.mean(scores)\n",
    "                inout_map[in_color] = out_color\n",
    "    def change_color(task_train):\n",
    "        for inout in task_train:\n",
    "            inp_origin = inout['input']\n",
    "            inp = np.array(inp_origin.copy())\n",
    "            vec = list(itertools.chain.from_iterable(inp_origin))\n",
    "            for in_color, out_color in inout_map.items():\n",
    "                if in_color in vec:\n",
    "                    inp = np.where(np.array(inp_origin) == in_color, out_color, inp)\n",
    "            inout['input'] = inp\n",
    "        return task_train\n",
    "    \n",
    "    return [inouts_array, change_color]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [],
   "source": [
    "def about_color(task_train):\n",
    "    in_colors = []\n",
    "    out_colors = []\n",
    "    color_changed = False\n",
    "    for inout in task_train:\n",
    "        in_vec = list(itertools.chain.from_iterable(inout['input']))\n",
    "        in_colors += list(set(in_vec))\n",
    "        out_vec = list(itertools.chain.from_iterable(inout['output']))\n",
    "        out_colors += list(set(out_vec))\n",
    "        if set(in_vec) != set(out_vec):\n",
    "            color_changed = True\n",
    "    return list(set(in_colors)), list(set(out_colors)), color_changed\n",
    "\n",
    "def about_color_for_test(task_test):\n",
    "    in_colors = []\n",
    "    out_colors = []\n",
    "    color_changed = False\n",
    "    for inout in task_test:\n",
    "        in_vec = list(itertools.chain.from_iterable(inout['input']))\n",
    "        in_colors += list(set(in_vec))\n",
    "    return list(set(in_colors))\n",
    "\n",
    "def about_color_for_task(task):\n",
    "    in_colors = []\n",
    "    out_colors = []\n",
    "    color_changed = False\n",
    "    for inout in task['train']:\n",
    "        in_vec = list(itertools.chain.from_iterable(inout['input']))\n",
    "        in_colors += list(set(in_vec))\n",
    "    for inout in task['test']:\n",
    "        in_vec = list(itertools.chain.from_iterable(inout['input']))\n",
    "        in_colors += list(set(in_vec))\n",
    "    return list(set(in_colors))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add task train6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_task_train6(task_train):\n",
    "    if correct_task_train(task_train):\n",
    "        return [inouts_array]\n",
    "    use_same_color = True\n",
    "    for inout in task_train:\n",
    "        in_vec = list(itertools.chain.from_iterable(inout['input']))\n",
    "        in_use_colors = list(set(in_vec))\n",
    "        in_use_colors.remove(0) if 0 in in_use_colors else 0\n",
    "        out_vec = list(itertools.chain.from_iterable(inout['output']))\n",
    "        out_use_colors = list(set(out_vec))\n",
    "        out_use_colors.remove(0) if 0 in out_use_colors else 0\n",
    "        if sorted(in_use_colors) != sorted(out_use_colors):\n",
    "            use_same_color = False\n",
    "    if use_same_color:    \n",
    "        return [inouts_array, task_train6]\n",
    "    else:\n",
    "        return [inouts_array]\n",
    "    \n",
    "def task_train6(task_train):\n",
    "    for inout in task_train:\n",
    "        x = inout['input']\n",
    "        x = np.array(x)\n",
    "        H, W = x.shape\n",
    "        vec = list(itertools.chain.from_iterable(x))\n",
    "        use_colors = list(set(vec))\n",
    "        use_colors.remove(0) if 0 in use_colors else 0\n",
    "        colors = [0] * len(use_colors)\n",
    "        for yy in range(H):\n",
    "            for xx in range(W):\n",
    "                color = x[yy, xx]\n",
    "                if color != 0:\n",
    "                    colors[(yy+xx)%len(use_colors)] = color\n",
    "        y = x.copy()\n",
    "        for yy in range(H):\n",
    "            for xx in range(W):\n",
    "                y[yy, xx] = colors[(yy+xx)%len(use_colors)]\n",
    "        inout['input'] = y\n",
    "    return task_train        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add move object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [],
   "source": [
    "move_skip = False\n",
    "def add_move_object(task_train):\n",
    "    start = time()\n",
    "    if move_skip:\n",
    "        return [inouts_array]\n",
    "    if correct_task_train(task_train):\n",
    "        return [inouts_array]\n",
    "    inp = np.array(task_train[0]['input'])\n",
    "    out = np.array(task_train[0]['output'])\n",
    "    in_use_colors, _, _ = about_color(task_train)\n",
    "    in_use_colors = [c for c in in_use_colors if c != 0]\n",
    "    best_score = match_rate(inp, out)\n",
    "    best_goal_color = 0\n",
    "    best_move_color = 0\n",
    "    best_change_color = 0\n",
    "    best_move_num = 0\n",
    "    best_direction = [0, 0]\n",
    "    best_correction = 0\n",
    "    Dy = [0, 1, 0, -1]\n",
    "    Dx = [1, 0, -1, 0]\n",
    "    should_change = False\n",
    "    for use_color_n, goal_color in enumerate(in_use_colors):\n",
    "        for move_color in in_use_colors:\n",
    "            if (time() - start > 60*60) & (len(in_use_colors) / 2 > use_color_n):\n",
    "                return [inouts_array]\n",
    "\n",
    "            goal_idx_set = set(tuple(idx) for idx in np.array(np.where(inp==goal_color)).T)\n",
    "            move_idx_list = [tuple(idx) for idx in np.array(np.where(inp==move_color)).T]\n",
    "            for dy, dx in zip(Dy, Dx):\n",
    "                for move_num in range(1, 40):\n",
    "                    obj_idx = set((idx[0]+dy*move_num, idx[1]+dx*move_num) for idx in move_idx_list)\n",
    "                    if obj_idx & goal_idx_set:\n",
    "                        for correction in [-2, -1, 0, 1, 2]:\n",
    "                            for change_color in range(10):\n",
    "                                inp_copy = copy.deepcopy(inp)\n",
    "                                for idx in obj_idx:\n",
    "                                    idx = (idx[0]+(dy*correction), idx[1]+(dx*correction))\n",
    "                                    if (idx[0] < 0) or (idx[1] < 0) or (inp_copy.shape[0] <= idx[0]) or (inp_copy.shape[1] <= idx[1]):\n",
    "                                        break\n",
    "                                    inp_copy[idx] = change_color\n",
    "                                for origin_move_pad_color in range(10):\n",
    "                                    inp_copy2 = copy.deepcopy(inp_copy)\n",
    "                                    for move_idx in move_idx_list:\n",
    "                                        inp_copy2[move_idx] = origin_move_pad_color\n",
    "                                    score = match_rate(inp_copy2, out)\n",
    "                                    if best_score < score:\n",
    "                                        should_change = True\n",
    "                                        best_score = score\n",
    "                                        best_goal_color = goal_color\n",
    "                                        best_move_color = move_color\n",
    "                                        best_move_num = move_num\n",
    "                                        best_direction = [dy, dx]\n",
    "                                        best_correction = correction\n",
    "                                        best_change_color = change_color\n",
    "                                        best_origin_move_pad_color = origin_move_pad_color\n",
    "\n",
    "    def move_object(task_train_origin):\n",
    "        for inout in task_train_origin:\n",
    "            inout['input'] = np.array(inout['input'])\n",
    "        if not should_change:\n",
    "            return task_train_origin\n",
    "        task_train = copy.deepcopy(task_train_origin)\n",
    "        for i, inout in enumerate(task_train):\n",
    "            finished = False\n",
    "            inp = np.array(inout['input'])\n",
    "            directions = [[0, 1], [1, 0], [0, -1], [-1, 0]]\n",
    "            for direction in directions:\n",
    "                if finished:\n",
    "                    break\n",
    "                for move_num in range(1, 50):\n",
    "                    if finished:\n",
    "                        break\n",
    "                    goal_idx_set = set(tuple(idx) for idx in np.array(np.where(inp==best_goal_color)).T)\n",
    "                    move_idx_list = [tuple(idx) for idx in np.array(np.where(inp==best_move_color)).T]\n",
    "                    obj_idx = set((idx[0] + direction[0]*move_num, idx[1] + direction[1]*move_num) for idx in move_idx_list)\n",
    "                    if obj_idx & goal_idx_set:\n",
    "                        for idx in obj_idx:\n",
    "                            idx = (idx[0]+(direction[0]*best_correction), idx[1]+(direction[1]*best_correction))\n",
    "                            if (idx[0] < 0) or (idx[1] < 0) or (inp.shape[0] <= idx[0]) or (inp.shape[1] <= idx[1]):\n",
    "                                continue\n",
    "                            inp[idx] = best_change_color\n",
    "                        for move_idx in move_idx_list:\n",
    "                            inp[move_idx] = best_origin_move_pad_color\n",
    "                        task_train[i]['input'] = inp\n",
    "                        finished = True\n",
    "        # if recursion:\n",
    "        #     for i in range(5):\n",
    "        #         if 'output' in task_train[0]:\n",
    "        #             if correct_task_train(task_train):\n",
    "        #                 return task_train\n",
    "        #         funcs = add_move_object(task_train, False)\n",
    "        #         for func in funcs:\n",
    "        #             task_train = func(task_train)\n",
    "        return task_train\n",
    "    return [move_object, inouts_array]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# correct task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "metadata": {},
   "outputs": [],
   "source": [
    "def correct_task_train(task_train):\n",
    "    correct = True\n",
    "    for inout in task_train:\n",
    "        if np.array(inout['input']).tolist() != np.array(inout['output']).tolist():\n",
    "            correct = False\n",
    "    return correct"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# check_p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_p(task, pred_func):\n",
    "    n = len(task[\"train\"]) + len(task[\"test\"])\n",
    "    fig, axs = plt.subplots(3, n, figsize=(4*n,12), dpi=50)\n",
    "    plt.subplots_adjust(wspace=0.3, hspace=0.3)\n",
    "    fnum = 0\n",
    "    for i, t in enumerate(task[\"train\"]):\n",
    "        t_in, t_out = np.array(t[\"input\"]).astype('uint8'), np.array(t[\"output\"]).astype('uint8')\n",
    "        t_pred, feat = call_pred_train(t_in, t_out, pred_func)\n",
    "        plot_one(axs[0,fnum],t_in,f'train-{i} input')\n",
    "        plot_one(axs[1,fnum],t_out,f'train-{i} output')\n",
    "        plot_one(axs[2,fnum],t_pred,f'train-{i} pred')\n",
    "        fnum += 1\n",
    "    for i, t in enumerate(task[\"test\"]):\n",
    "        t_in, t_out = np.array(t[\"input\"]).astype('uint8'), np.array(t[\"output\"]).astype('uint8')\n",
    "        t_pred = call_pred_test(t_in, pred_func, feat)        \n",
    "        plot_one(axs[0,fnum],t_in,f'test-{i} input')\n",
    "        plot_one(axs[1,fnum],t_out,f'test-{i} output')\n",
    "        plot_one(axs[2,fnum],t_pred,f'test-{i} pred')\n",
    "#         t_pred = np.resize(t_pred,t_out.shape)\n",
    "        fnum += 1\n",
    "    plt.show()\n",
    "    return 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# recolor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Defensive_Copy(A): \n",
    "    n = len(A)\n",
    "    k = len(A[0])\n",
    "    L = np.zeros((n,k), dtype = int)\n",
    "    for i in range(n):\n",
    "        for j in range(k):\n",
    "            L[i,j] = 0 + A[i][j]\n",
    "    return L.tolist()\n",
    "\n",
    "def Create(task, task_id = 0):\n",
    "    n = len(task['train'])\n",
    "    Input = [Defensive_Copy(task['train'][i]['input']) for i in range(n)]\n",
    "    Output = [Defensive_Copy(task['train'][i]['output']) for i in range(n)]\n",
    "    Input.append(Defensive_Copy(task['test'][task_id]['input']))\n",
    "    return Input, Output\n",
    "\n",
    "def add_recolor(task_train, task_n):\n",
    "    return [inouts_array]\n",
    "    m1, n1 = np.array(task_train[0]['input']).shape\n",
    "    m2, n2 = np.array(task_train[0]['output']).shape\n",
    "    all_inputs_same_shape = True\n",
    "    all_outputs_same_shape = True\n",
    "    for inout in task_train:\n",
    "        m1_, n1_ = np.array(inout['input']).shape\n",
    "        m2_, n2_ = np.array(inout['output']).shape\n",
    "        if (m1_ != m1) or (n1_ != n1):\n",
    "            all_inputs_same_shape = False\n",
    "        if (m2_ != m2) or (n2_ != n2):\n",
    "            all_outputs_same_shape = False\n",
    "    if (not all_inputs_same_shape) or (not all_outputs_same_shape):\n",
    "        return [inouts_array]\n",
    "\n",
    "    inputs = []\n",
    "    outputs = []\n",
    "    for inout in task_train:\n",
    "        inputs.append(copy.deepcopy(inout['input']))\n",
    "        outputs.append(copy.deepcopy(inout['output']))\n",
    "\n",
    "    N = len(inputs)\n",
    "    x0 = inputs[0]\n",
    "    y0 = outputs[0]\n",
    "    n = len(x0)\n",
    "    k = len(x0[0])\n",
    "    a = len(y0)\n",
    "    b = len(y0[0])\n",
    "\n",
    "    List1 = {}\n",
    "    List2 = {}\n",
    "    for i in range(n):\n",
    "        for j in range(k):\n",
    "            seq = []\n",
    "            for x in inputs:\n",
    "                seq.append(x[i][j])\n",
    "            List1[(i,j)] = seq\n",
    "\n",
    "    for p in range(a):\n",
    "        for q in range(b):\n",
    "            seq1 = []\n",
    "            for y in outputs:\n",
    "                seq1.append(y[p][q])\n",
    "\n",
    "            places = []\n",
    "            for key in List1:\n",
    "                if List1[key] == seq1:\n",
    "                    places.append(key)\n",
    "\n",
    "            List2[(p,q)] = places\n",
    "            if len(places) == 0:\n",
    "                return [inouts_array]\n",
    "\n",
    "    def recolor(task_train):\n",
    "        for inout in task_train:\n",
    "            inout['input'] = np.array(inout['input'])\n",
    "        answer = np.zeros((a,b), dtype = int)\n",
    "        for inout_n, inout in enumerate(task_train):\n",
    "            for p in range(a):\n",
    "                for q in range(b):\n",
    "                    palette = [0,0,0,0,0,0,0,0,0,0]\n",
    "                    for i, j in List2[(p,q)]:\n",
    "                        color = inout['input'][i][j]\n",
    "                        palette[color]+=1\n",
    "                    answer[p,q] =  np.argmax(palette)\n",
    "\n",
    "            task_train[inout_n]['input'] = np.array(answer)\n",
    "        return task_train\n",
    "\n",
    "    return [inouts_array, recolor]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# get_similarity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_similarity(train, func_combi, idx, search_func=True):\n",
    "    similarities = []\n",
    "    for seed in range(3):\n",
    "        similarity = train_and_evaluate(train, func_combi, seed, idx, {}, search_func=search_func)\n",
    "        similarities.append(similarity)\n",
    "    return np.mean(similarities)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add kneighbors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "metadata": {},
   "outputs": [],
   "source": [
    "def comparematrixes(a,b):\n",
    "    out=0;\n",
    "    for i in range(min(len(a),len(b))):\n",
    "        for j in range(min(len(a[0]),len(b[0]))):\n",
    "            if a[i][j]==b[i][j]:\n",
    "                out+=1\n",
    "    out/=len(a)*len(a[0]);\n",
    "    return 1-out;\n",
    "def add_kneighbors(task_train):\n",
    "    all_inputs_same_shape, all_outputs_same_shape = all_inputs_same_shape_and_all_outputs_same_shape(task_train)\n",
    "    if (not all_inputs_same_shape) or (not all_outputs_same_shape):\n",
    "#         return [inouts_array]\n",
    "        pass\n",
    "    ines=[];\n",
    "    outes=[];\n",
    "    for i in range(len(task_train)):\n",
    "        vx=task_train[i][\"input\"].copy();\n",
    "        vi=task_train[i][\"output\"].copy();\n",
    "        if (len(vx) > 10) or (len(vi) > 10):\n",
    "            return [inouts_array]\n",
    "        for k1 in range(min(len(vx),len(vi))):\n",
    "            for k2 in range(min(len(vx[0]),len(vi[0]))):\n",
    "                dtm=[];\n",
    "                for k3 in range(-2,2+1,1):\n",
    "                    for k4 in range(-2,2+1,1):\n",
    "                        if(k1+k3<len(vx) and k1+k3>=0 and k2+k4<len(vx[0]) and k2+k4>=0 and k1+k3<len(vi) and k1+k3>=0 and k2+k4<len(vi[0]) and k2+k4>=0):\n",
    "                            td=[0,0,0,0,0,0,0,0,0,0,0];\n",
    "                            if (vx[k1+k3][k2+k4] > 10) or (vi[k1+k3][k2+k4]):\n",
    "                                return [inouts_array]\n",
    "                            td[vx[k1+k3][k2+k4]]=1\n",
    "                            dtm+=td.copy();\n",
    "                            td=[0,0,0,0,0,0,0,0,0,0,0];\n",
    "                            td[vi[k1+k3][k2+k4]]=1;\n",
    "                            dtm+=td.copy();\n",
    "                        else:\n",
    "                            dtm+=[0,0,0,0,0,0,0,0,0,0,0];\n",
    "                            dtm+=[0,0,0,0,0,0,0,0,0,0,0];\n",
    "                ines.append(dtm);\n",
    "                if(len(vi)>k1 and len(vi[0])>k2 and k1>=0 and k2>=0):\n",
    "                    outes.append(vi[k1][k2]);\n",
    "                else:\n",
    "                    outes.append(0);\n",
    "    knn = KNeighborsClassifier(n_neighbors = 1);\n",
    "    ines=json.loads(json.dumps(ines));\n",
    "    knn.fit(ines,outes);\n",
    "    outs=[]\n",
    "    def kneighbors(task_train_origin):\n",
    "        for inout in task_train_origin:\n",
    "            inout['input'] = np.array(inout['input'])\n",
    "        task_train = copy.deepcopy(task_train_origin)\n",
    "        for i in range(len(task_train)):\n",
    "            thisdone=False;\n",
    "            vx=task_train[i][\"input\"].copy();\n",
    "            vi=task_train[i][\"input\"].copy();\n",
    "            for U in range(20):\n",
    "                for k1 in range(len(vx)):\n",
    "                    for k2 in range(len(vx[0])):\n",
    "                        dtm=[];\n",
    "                        for k3 in range(-2,2+1,1):\n",
    "                            for k4 in range(-2,2+1,1):\n",
    "                                if(k1+k3<len(vx) and k1+k3>=0 and k2+k4<len(vx[0]) and k2+k4>=0 and k1+k3<len(vi) and k1+k3>=0 and k2+k4<len(vi[0]) and k2+k4>=0):\n",
    "                                    td = [0,0,0,0,0,0,0,0,0,0,0];\n",
    "                                    td[vx[k1+k3][k2+k4]]=1\n",
    "                                    dtm+=td.copy();\n",
    "                                    td = [0,0,0,0,0,0,0,0,0,0,0];\n",
    "                                    td[vi[k1+k3][k2+k4]]=1;\n",
    "                                    dtm+=td.copy();\n",
    "                                else:\n",
    "                                    dtm+=[0,0,0,0,0,0,0,0,0,0,0];\n",
    "                                    dtm+=[0,0,0,0,0,0,0,0,0,0,0];\n",
    "                        vi[k1][k2]=int(knn.predict([dtm])[0]);\n",
    "                vx=vi.copy();\n",
    "            task_train[i]['input'] = vx\n",
    "        return task_train\n",
    "    return [inouts_array, kneighbors]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ARC solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ARC_solver:\n",
    "    def __init__(self):\n",
    "        self.identified_objects = []\n",
    "        self.io_inx = [] # the original index of the identified objects (io)\n",
    "        self.io_height = [] # height of io\n",
    "        self.io_width = [] # width of io\n",
    "        self.io_pixel_count = [] # count of non-background pixels\n",
    "        self.io_size = [] # overall grid size\n",
    "        self.io_unique_colors = [] # number of unique colors\n",
    "        self.io_main_color = [] # the dominating color\n",
    "\n",
    "    def reset(self):\n",
    "        self.identified_objects = []\n",
    "        self.io_inx = []\n",
    "        self.io_height = []\n",
    "        self.io_width = []\n",
    "        self.io_pixel_count = []\n",
    "        self.io_size = []\n",
    "        self.io_unique_colors = []\n",
    "        self.io_main_color = []\n",
    "\n",
    "    def get_background(self, image):\n",
    "        # if image contains 0\n",
    "        if 0 in image:\n",
    "            background = 0\n",
    "        # else use the most frequent pixel color\n",
    "        else:\n",
    "            unique_colors, counts = np.unique(image, return_counts = True)\n",
    "            background = unique_colors[np.argmax(counts)]\n",
    "        return background\n",
    "\n",
    "    def check_pairs(self, inx_pairs, this_pair, return_inx = False):\n",
    "        # check if this_pair is in inx_pairs\n",
    "        match = []\n",
    "        for pair in inx_pairs:\n",
    "            if pair[0] == this_pair[0] and pair[1] == this_pair[1]:\n",
    "                match.append(True)\n",
    "            else:\n",
    "                match.append(False)\n",
    "        if return_inx:\n",
    "            return any(match), np.where(match)\n",
    "        else:\n",
    "            return any(match)\n",
    "\n",
    "    def check_neighbors(self, all_pairs, this_pair, objectness, this_object):\n",
    "        # all_pairs: an array of index pairs for all nonzero/colored pixels\n",
    "        # this_pair: the index pair whose neighbors will be checked\n",
    "        # objectness: an array with the shape of original image, storage for how much objectness has been identified\n",
    "        # this_object: the current object we are looking at\n",
    "        row_inx = this_pair[0]\n",
    "        col_inx = this_pair[1]\n",
    "        objectness[row_inx, col_inx] = this_object\n",
    "        # find if any neighboring pixels contain color\n",
    "        if self.check_pairs(all_pairs, [row_inx-1, col_inx-1]): # up-left\n",
    "            objectness[row_inx-1, col_inx-1] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx-1, col_inx]): # up\n",
    "            objectness[row_inx-1, col_inx] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx-1, col_inx+1]): # up-right\n",
    "            objectness[row_inx-1, col_inx+1] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx, col_inx-1]): # left\n",
    "            objectness[row_inx, col_inx-1] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx, col_inx+1]): # right\n",
    "            objectness[row_inx, col_inx+1] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx+1, col_inx-1]): # down-left\n",
    "            objectness[row_inx+1, col_inx-1] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx+1, col_inx]): # down\n",
    "            objectness[row_inx+1, col_inx] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx+1, col_inx+1]): # down-right\n",
    "            objectness[row_inx+1, col_inx+1] = this_object\n",
    "        return objectness\n",
    "\n",
    "    def identify_object_by_color(self, true_image, background = 0):\n",
    "        # identify obeject by the color only\n",
    "        unique_colors = np.unique(true_image)\n",
    "        for i, color in enumerate(unique_colors):\n",
    "            image = np.copy(true_image) # make a copy from original first\n",
    "            if color == background:\n",
    "                continue\n",
    "            image[image != color] = background\n",
    "            inx = np.where(image == color)\n",
    "            obj = image[np.min(inx[0]):np.max(inx[0])+1, np.min(inx[1]):np.max(inx[1])+1]\n",
    "            # append the object attributes\n",
    "            self.identified_objects.append(obj)\n",
    "            self.io_inx.append(inx)\n",
    "            self.io_height.append(obj.shape[0])\n",
    "            self.io_width.append(obj.shape[1])\n",
    "            self.io_pixel_count.append(obj[obj != background].shape[0])\n",
    "            self.io_size.append(obj.size)\n",
    "            nc, c = np.unique(obj, return_counts = True)\n",
    "            self.io_unique_colors.append(nc)\n",
    "            self.io_main_color.append(nc[np.argmax(c)])\n",
    "\n",
    "    def identify_object_by_isolation(self, image, background = 0):\n",
    "        # identify all objects by physical isolation on the given image\n",
    "        all_pairs = np.array(np.where(image != background)).T\n",
    "        objectness = np.zeros(image.shape)\n",
    "        this_object = 1\n",
    "        while len(all_pairs) >= 1:\n",
    "            init_pair = all_pairs[0] # start with the first pair\n",
    "            objectness = self.check_neighbors(all_pairs, init_pair, objectness, this_object)\n",
    "            # get a list of index pairs whose neghbors haven't been checked\n",
    "            unchecked_pairs = np.array(np.where(objectness == this_object)).T\n",
    "            checked_pairs = np.zeros((0,2))\n",
    "            # check all the index pairs in the expanding unchecked_pairs untill all have been checked\n",
    "            while len(unchecked_pairs) != 0:\n",
    "                this_pair = unchecked_pairs[0]\n",
    "                objectness = self.check_neighbors(all_pairs, this_pair, objectness, this_object)\n",
    "                # append the checked_pairs\n",
    "                checked_pairs = np.vstack((checked_pairs, this_pair))\n",
    "                # get all index pairs for the currently identified object\n",
    "                current_object_pairs = np.array(np.where(objectness == this_object)).T\n",
    "                # delete the checked pairs from current object pairs\n",
    "                checked_inx = []\n",
    "                for pair in checked_pairs:\n",
    "                    _, inx = self.check_pairs(current_object_pairs, pair, return_inx = True)\n",
    "                    checked_inx.append(inx[0][0])\n",
    "                unchecked_pairs = np.delete(current_object_pairs, checked_inx, axis = 0)\n",
    "\n",
    "            # store this object to identified_objects\n",
    "            current_object_pairs = np.array(np.where(objectness == this_object)).T\n",
    "            cop = current_object_pairs.T\n",
    "            obj = image[np.min(cop[0]):np.max(cop[0])+1, np.min(cop[1]):np.max(cop[1])+1]\n",
    "            # delete the current object pairs from all_pairs\n",
    "            cop_inx = []\n",
    "            for pair in current_object_pairs:\n",
    "                _, this_cop_inx = self.check_pairs(all_pairs, pair, return_inx = True)\n",
    "                cop_inx.append(this_cop_inx[0][0])\n",
    "            all_pairs = np.delete(all_pairs, cop_inx, axis = 0)\n",
    "            # append the object attribute\n",
    "            # p(obj)\n",
    "            if np.array(obj).shape[0] * np.array(obj).shape[0] >= 3:\n",
    "                self.identified_objects.append(obj)\n",
    "            self.io_inx.append(inx)\n",
    "            self.io_height.append(obj.shape[0])\n",
    "            self.io_width.append(obj.shape[1])\n",
    "            self.io_pixel_count.append(obj[obj != background].shape[0])\n",
    "            self.io_size.append(obj.size)\n",
    "            nc, c = np.unique(obj, return_counts = True)\n",
    "            self.io_unique_colors.append(nc)\n",
    "            self.io_main_color.append(nc[np.argmax(c)])\n",
    "            # start identifying a new object\n",
    "            this_object += 1\n",
    "        return objectness\n",
    "\n",
    "    def identify_object_by_color_isolation(self, true_image, background = 0):\n",
    "        # identify objects first by color then by physical isolation\n",
    "        unique_colors = np.unique(true_image)\n",
    "        for i, color in enumerate(unique_colors):\n",
    "            image = np.copy(true_image) # make a copy from the original first\n",
    "            if color == background:\n",
    "                continue\n",
    "            # identify objects by isolation in this color only\n",
    "            image[image != color] = background\n",
    "            self.identify_object_by_isolation(image, background = background)\n",
    "            \n",
    "\n",
    "    def sort(self, objs, inp):\n",
    "        xs = []\n",
    "        ys = []\n",
    "        for i, o in enumerate(objs):\n",
    "            _, m, n = sliding_window_search(inp, o)\n",
    "            xs.append(m)\n",
    "            ys.append(n)\n",
    "\n",
    "        ans = [[[]],[[]],[[]],[[]]]\n",
    "        left = np.array(ys).argsort()[0:2] # 1,3\n",
    "        right = np.array(ys).argsort()[2:4] # 1,3\n",
    "        if xs[left[0]] <= xs[left[1]]:\n",
    "            ans[0] = objs[left[0]]\n",
    "            ans[2] = objs[left[1]]\n",
    "        else:\n",
    "            ans[2] = objs[left[0]]\n",
    "            ans[0] = objs[left[1]]    \n",
    "        if xs[right[0]] <= xs[right[1]]:\n",
    "            ans[1] = objs[right[0]]\n",
    "            ans[3] = objs[right[1]]\n",
    "        else:\n",
    "            ans[3] = objs[right[0]]\n",
    "            ans[1] = objs[right[1]]\n",
    "        return ans        \n",
    "            \n",
    "    def merge(self, objects, belt, use_color):\n",
    "#         ans = objects\n",
    "        ans=[[[]],[[]],[[]],[[]]]\n",
    "        for o in objects:\n",
    "            o = np.array(o)\n",
    "            max_total = 0\n",
    "            for x in [0,1]:\n",
    "                for y in [0,1]:\n",
    "                    if max_total < o[x:x+len(o)-1, y:y+len(o[0])-1].sum():\n",
    "                        max_total = o[x:x+len(o)-1, y:y+len(o[0])-1].sum()\n",
    "                        max_xy = (x, y)\n",
    "            if max_xy == (0,0):\n",
    "                ans[3] = o\n",
    "            elif max_xy == (0,1):\n",
    "                ans[2] = o\n",
    "            elif max_xy == (1,0):\n",
    "                ans[1] = o\n",
    "            else:\n",
    "                ans[0] = o\n",
    "\n",
    "        if belt == 0:\n",
    "            belt_list = [[use_color]]*len(ans[0])\n",
    "            u=np.hstack([ans[0], ans[1]])\n",
    "            u\n",
    "            s=np.hstack([ans[2], ans[3]])\n",
    "            return np.vstack([u,s])\n",
    "        else:\n",
    "            belt_list = [[use_color]*belt]*len(ans[0])\n",
    "\n",
    "            u=np.hstack([ans[0], belt_list, ans[1]])\n",
    "            s=np.hstack([ans[2], belt_list, ans[3]])\n",
    "            belt_list = [[use_color]*len(s[0])]*belt\n",
    "            return np.vstack([u,belt_list,s])\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add block merge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "metadata": {},
   "outputs": [],
   "source": [
    "def divide_block_and_merge(task_train_origin, objn):\n",
    "    for inout in task_train_origin:\n",
    "        inout['input'] = np.array(inout['input'])\n",
    "    task_train = copy.deepcopy(task_train_origin)\n",
    "    for i, inout in enumerate(task_train):    \n",
    "        arc = ARC_solver()\n",
    "        inp = inout['input']\n",
    "        inp = np.array(inp)\n",
    "        use_color = list(set(list(itertools.chain.from_iterable(inp))))\n",
    "        if len(use_color) != 2:\n",
    "            return task_train_origin\n",
    "        \n",
    "        try:\n",
    "            inp_o = copy.deepcopy(inp)\n",
    "            inp = np.where(inp_o==use_color[0], use_color[1], inp)\n",
    "            inp = np.where(inp_o==use_color[1], use_color[0], inp)\n",
    "            background = arc.get_background(inp)\n",
    "            arc.identify_object_by_isolation(inp, background)\n",
    "            if len(arc.identified_objects) == 4:\n",
    "                arc.identified_objects = arc.sort(arc.identified_objects, inp)\n",
    "                out = np.array(arc.identified_objects[objn])\n",
    "                out_o = copy.deepcopy(out)\n",
    "                out = np.where(out_o==use_color[0], use_color[1], out)\n",
    "                out = np.where(out_o==use_color[1], use_color[0], out)\n",
    "                task_train[i]['input'] = out\n",
    "        except:\n",
    "            return task_train_origin\n",
    "    return task_train\n",
    "\n",
    "def divide_block_and_merge1(task_train_origin):\n",
    "    return divide_block_and_merge(task_train_origin, 1)\n",
    "\n",
    "def divide_block_and_merge2(task_train_origin):\n",
    "    return divide_block_and_merge(task_train_origin, 2)\n",
    "\n",
    "def divide_block_and_merge3(task_train_origin):\n",
    "    return divide_block_and_merge(task_train_origin, 3)\n",
    "\n",
    "def add_block_merge(task_train):\n",
    "    arc = ARC_solver()\n",
    "    if len(task_train) > 2:\n",
    "        task_n = 2\n",
    "    else:\n",
    "        task_n = 0\n",
    "    inp = task_train[task_n]['input']\n",
    "    inp = np.array(inp)\n",
    "    use_color = list(set(list(itertools.chain.from_iterable(inp))))\n",
    "    if len(use_color) != 2:\n",
    "        return []\n",
    "    inp_o = copy.deepcopy(inp)\n",
    "    inp = np.where(inp_o==use_color[0], use_color[1], inp)\n",
    "    inp = np.where(inp_o==use_color[1], use_color[0], inp)\n",
    "    background = arc.get_background(inp)\n",
    "    arc.identify_object_by_isolation(inp, background)\n",
    "    if len(arc.identified_objects) == 4:\n",
    "        try:\n",
    "            arc.identified_objects = arc.sort(arc.identified_objects, inp)\n",
    "#             for i in arc.identified_objects:\n",
    "#                 p(i)\n",
    "            for i in range(4):\n",
    "                out = np.array(arc.identified_objects[i])\n",
    "                out_o = copy.deepcopy(out)\n",
    "                out = np.where(out_o==use_color[0], use_color[1], out)\n",
    "                out = np.where(out_o==use_color[1], use_color[0], out)    \n",
    "\n",
    "                if out.tolist() == task_train[task_n]['output']:\n",
    "                    return [divide_block_and_merge1, divide_block_and_merge2, divide_block_and_merge3]\n",
    "        except:\n",
    "            return []\n",
    "    return []\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add object detect2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_by_ele(objects, ele):\n",
    "    if ele == 'height':\n",
    "        max_height = 0\n",
    "        for obj in objects:\n",
    "            if len(obj) > max_height:\n",
    "                selected = obj\n",
    "                max_height = len(obj)\n",
    "    if ele == 'width':\n",
    "        max_width = 0\n",
    "        for obj in objects:\n",
    "            if len(obj[0]) > max_width:\n",
    "                selected = obj\n",
    "                max_width = len(obj[0])\n",
    "    if ele == 'area':\n",
    "        max_area = 0\n",
    "        for obj in objects:\n",
    "            if len(obj) * len(obj[0]) > max_area:\n",
    "                selected = obj\n",
    "                max_area = len(obj) * len(obj[0])\n",
    "        \n",
    "    return selected\n",
    "\n",
    "def add_object_detect2(task_train):\n",
    "    for select_ele in ['height', 'width', 'area']:\n",
    "        sucess = True\n",
    "        for inout in task_train:\n",
    "            arc = ARC_solver()\n",
    "            inp = copy.deepcopy(inout['input'])\n",
    "            inp = np.array(inp)\n",
    "            background = arc.get_background(inp)\n",
    "            arc.identify_object_by_isolation(inp, background)\n",
    "            obj = select_by_ele(arc.identified_objects, select_ele)\n",
    "            if (obj.shape != np.array(inout['output']).shape) & (obj.shape != np.array(inout['output']).T.shape):\n",
    "                sucess = False\n",
    "        if sucess:\n",
    "            def object_detect2(task_train_origin):\n",
    "                for inout in task_train_origin:\n",
    "                    inout['input'] = np.array(inout['input'])\n",
    "                task_train = copy.deepcopy(task_train_origin)\n",
    "                for i, inout in enumerate(task_train):   \n",
    "                    try:\n",
    "                        arc = ARC_solver()\n",
    "                        inp = copy.deepcopy(inout['input'])\n",
    "                        inp = np.array(inp)\n",
    "                        background = arc.get_background(inp)\n",
    "                        arc.identify_object_by_isolation(inp, background)\n",
    "                        obj = select_by_ele(arc.identified_objects, select_ele)\n",
    "                        task_train[i]['input'] = obj\n",
    "                    except:\n",
    "                        return task_train_origin\n",
    "                return task_train\n",
    "                    \n",
    "            return [object_detect2]\n",
    "    return []\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add crop_by_line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_crop_by_line(task_train):\n",
    "    success = True\n",
    "    for i, inout in enumerate(task_train):\n",
    "        inp = np.array(copy.deepcopy(inout['input']))\n",
    "        use_color = matrix_use_color(inp)\n",
    "        max_area = 0\n",
    "        max_enclosure_color = 0\n",
    "        include_line = False\n",
    "        uses = [0,0,0,0]\n",
    "        found = False\n",
    "        use_max_x = 0\n",
    "        use_max_y = 0\n",
    "        use_min_x = 0\n",
    "        use_min_y = 0\n",
    "        for color in use_color:\n",
    "            idx = [idx.tolist() for idx in np.array(np.where(inp==color)).T]\n",
    "\n",
    "            max_x = 0\n",
    "            max_y = 0\n",
    "            min_x = 100\n",
    "            min_y = 100\n",
    "            for i in idx:\n",
    "                if i[0] < min_x:\n",
    "                    min_x = i[0]\n",
    "                if i[1] < min_y:\n",
    "                    min_y = i[1]\n",
    "                if i[0] > max_x:\n",
    "                    max_x = i[0]\n",
    "                if i[1] > max_y:\n",
    "                    max_y = i[1]\n",
    "\n",
    "            enclosure_flag = True\n",
    "            for x in range(min_x, max_x+1):\n",
    "                if (inp[x][min_y] != color) or (inp[x][max_y] != color):\n",
    "                    enclosure_flag = False\n",
    "            for y in range(min_y, max_y+1):\n",
    "                if (inp[min_x][y] != color) or (inp[max_x][y] != color):\n",
    "                    enclosure_flag = False\n",
    "            for x in range(min_x+1, max_x):\n",
    "                for y in range(min_y+1, max_y):\n",
    "                    if inp[x][y] == color:\n",
    "                        enclosure_flag = False\n",
    "            if enclosure_flag & (max_x > 0) & (max_x - min_x > 1):\n",
    "                area = (max_x-min_x)*(max_y-min_y)\n",
    "                if max_area < area:\n",
    "                    max_area = area\n",
    "                    max_enclosure_color = color\n",
    "                    found = True\n",
    "                    use_max_x = max_x\n",
    "                    use_max_y = max_y\n",
    "                    use_min_x = min_x\n",
    "                    use_min_y = min_y\n",
    "        if not found:\n",
    "            return []\n",
    "        if i == 0:\n",
    "            if np.array(inout['output']).shape == (use_max_x-use_min_x-1, use_max_y-use_min_y-1):\n",
    "                include_line = False\n",
    "            elif np.array(inout['output']).shape == (use_max_x-use_min_x+1, use_max_y-use_min_y+1):\n",
    "                include_line = True\n",
    "            else:\n",
    "                success = False\n",
    "        else:\n",
    "            if (not include_line) & (np.array(inout['output']).shape == (use_max_x-use_min_x-1, use_max_y-use_min_y-1)) or ((include_line) & (np.array(inout['output']).shape == (use_max_x-use_min_x+1, use_max_y-use_min_y+1))):\n",
    "                pass\n",
    "            else:\n",
    "                success = False\n",
    "\n",
    "        if not success:\n",
    "            break\n",
    "\n",
    "    if success:\n",
    "        def crop_by_max_enclosure_color(task_train_origin):\n",
    "            for inout in task_train_origin:\n",
    "                inout['input'] = np.array(inout['input'])\n",
    "            task_train = copy.deepcopy(task_train_origin)\n",
    "            for task_n, inout in enumerate(task_train):\n",
    "                inp = np.array(copy.deepcopy(inout['input']))\n",
    "                use_color = matrix_use_color(inp)\n",
    "                max_area = 0\n",
    "                max_enclosure_color = 0\n",
    "                include_line = False\n",
    "                uses = [0,0,0,0]\n",
    "                found = False\n",
    "                use_max_x = 0\n",
    "                use_max_y = 0\n",
    "                use_min_x = 0\n",
    "                use_min_y = 0\n",
    "                for color in use_color:\n",
    "                    idx = [idx.tolist() for idx in np.array(np.where(inp==color)).T]\n",
    "\n",
    "                    max_x = 0\n",
    "                    max_y = 0\n",
    "                    min_x = 100\n",
    "                    min_y = 100\n",
    "                    for i in idx:\n",
    "                        if i[0] < min_x:\n",
    "                            min_x = i[0]\n",
    "                        if i[1] < min_y:\n",
    "                            min_y = i[1]\n",
    "                        if i[0] > max_x:\n",
    "                            max_x = i[0]\n",
    "                        if i[1] > max_y:\n",
    "                            max_y = i[1]\n",
    "                    enclosure_flag = True\n",
    "                    for x in range(min_x, max_x+1):\n",
    "                        if (inp[x][min_y] != color) or (inp[x][max_y] != color):\n",
    "                            enclosure_flag = False\n",
    "                    for y in range(min_y, max_y+1):\n",
    "                        if (inp[min_x][y] != color) or (inp[max_x][y] != color):\n",
    "                            enclosure_flag = False\n",
    "                    for x in range(min_x+1, max_x):\n",
    "                        for y in range(min_y+1, max_y):\n",
    "                            if inp[x][y] == color:\n",
    "                                enclosure_flag = False\n",
    "                    if enclosure_flag & (max_x > 0) & (max_x - min_x > 1):\n",
    "                        area = (max_x-min_x)*(max_y-min_y)\n",
    "                        if max_area < area:\n",
    "                            max_area = area\n",
    "                            max_enclosure_color = color\n",
    "                            found = True\n",
    "                            use_max_x = max_x\n",
    "                            use_max_y = max_y\n",
    "                            use_min_x = min_x\n",
    "                            use_min_y = min_y                \n",
    "                if include_line:\n",
    "                    out = inp[use_min_x:use_max_x+1, use_min_y:use_max_y+1]\n",
    "                else:\n",
    "                    out = inp[use_min_x+1:use_max_x, use_min_y+1:use_max_y]\n",
    "                task_train[task_n]['input'] = out\n",
    "            return task_train\n",
    "                \n",
    "        return [crop_by_max_enclosure_color]\n",
    "    return []"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# add back_to_black"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "metadata": {},
   "outputs": [],
   "source": [
    "def back_to_black(task_train_origin):\n",
    "    for inout in task_train_origin:\n",
    "        inout['input'] = np.array(inout['input'])\n",
    "    task_train = copy.deepcopy(task_train_origin)\n",
    "    for task_n, inout in enumerate(task_train):\n",
    "        inp = inout['input']\n",
    "        inp_o = copy.deepcopy(inp)\n",
    "        i = list(itertools.chain.from_iterable(inp))\n",
    "        most_use_color = collections.Counter(i).most_common()[0][0]\n",
    "        inp = np.where(inp_o==most_use_color, 0, inp)\n",
    "        inp = np.where(inp_o==0, most_use_color, inp)\n",
    "        task_train[task_n]['input'] = inp\n",
    "    return task_train\n",
    "\n",
    "def add_back_to_black_funcs(task_train):\n",
    "    change_back = True\n",
    "    for inout in task_train:\n",
    "        i = list(itertools.chain.from_iterable(inout['input']))\n",
    "        if collections.Counter(i).most_common()[0][0] == 0:\n",
    "            change_back = False\n",
    "    \n",
    "    if change_back:\n",
    "        return [back_to_black, inouts_array]\n",
    "    else:\n",
    "        return [inouts_array]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ARC_solver:\n",
    "    def __init__(self):\n",
    "        self.identified_objects = []\n",
    "        self.io_inx = [] # the original index of the identified objects (io)\n",
    "        self.io_height = [] # height of io\n",
    "        self.io_width = [] # width of io\n",
    "        self.io_pixel_count = [] # count of non-background pixels\n",
    "        self.io_size = [] # overall grid size\n",
    "        self.io_unique_colors = [] # number of unique colors\n",
    "        self.io_main_color = [] # the dominating color\n",
    "\n",
    "    def reset(self):\n",
    "        self.identified_objects = []\n",
    "        self.io_inx = []\n",
    "        self.io_height = []\n",
    "        self.io_width = []\n",
    "        self.io_pixel_count = []\n",
    "        self.io_size = []\n",
    "        self.io_unique_colors = []\n",
    "        self.io_main_color = []\n",
    "\n",
    "    def get_background(self, image):\n",
    "        # if image contains 0\n",
    "        if 0 in image:\n",
    "            background = 0\n",
    "        # else use the most frequent pixel color\n",
    "        else:\n",
    "            unique_colors, counts = np.unique(image, return_counts = True)\n",
    "            background = unique_colors[np.argmax(counts)]\n",
    "        return background\n",
    "\n",
    "    def check_pairs(self, inx_pairs, this_pair, return_inx = False):\n",
    "        # check if this_pair is in inx_pairs\n",
    "        match = []\n",
    "        for pair in inx_pairs:\n",
    "            if pair[0] == this_pair[0] and pair[1] == this_pair[1]:\n",
    "                match.append(True)\n",
    "            else:\n",
    "                match.append(False)\n",
    "        if return_inx:\n",
    "            return any(match), np.where(match)\n",
    "        else:\n",
    "            return any(match)\n",
    "\n",
    "    def check_neighbors(self, all_pairs, this_pair, objectness, this_object):\n",
    "        # all_pairs: an array of index pairs for all nonzero/colored pixels\n",
    "        # this_pair: the index pair whose neighbors will be checked\n",
    "        # objectness: an array with the shape of original image, storage for how much objectness has been identified\n",
    "        # this_object: the current object we are looking at\n",
    "        row_inx = this_pair[0]\n",
    "        col_inx = this_pair[1]\n",
    "        objectness[row_inx, col_inx] = this_object\n",
    "        # find if any neighboring pixels contain color\n",
    "        if self.check_pairs(all_pairs, [row_inx-1, col_inx-1]): # up-left\n",
    "            objectness[row_inx-1, col_inx-1] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx-1, col_inx]): # up\n",
    "            objectness[row_inx-1, col_inx] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx-1, col_inx+1]): # up-right\n",
    "            objectness[row_inx-1, col_inx+1] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx, col_inx-1]): # left\n",
    "            objectness[row_inx, col_inx-1] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx, col_inx+1]): # right\n",
    "            objectness[row_inx, col_inx+1] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx+1, col_inx-1]): # down-left\n",
    "            objectness[row_inx+1, col_inx-1] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx+1, col_inx]): # down\n",
    "            objectness[row_inx+1, col_inx] = this_object\n",
    "        if self.check_pairs(all_pairs, [row_inx+1, col_inx+1]): # down-right\n",
    "            objectness[row_inx+1, col_inx+1] = this_object\n",
    "        return objectness\n",
    "\n",
    "    def identify_object_by_color(self, true_image, background = 0):\n",
    "        # identify obeject by the color only\n",
    "        unique_colors = np.unique(true_image)\n",
    "        for i, color in enumerate(unique_colors):\n",
    "            image = np.copy(true_image) # make a copy from original first\n",
    "            if color == background:\n",
    "                continue\n",
    "            image[image != color] = background\n",
    "            inx = np.where(image == color)\n",
    "            obj = image[np.min(inx[0]):np.max(inx[0])+1, np.min(inx[1]):np.max(inx[1])+1]\n",
    "            # append the object attributes\n",
    "            self.identified_objects.append(obj)\n",
    "            self.io_inx.append(inx)\n",
    "            self.io_height.append(obj.shape[0])\n",
    "            self.io_width.append(obj.shape[1])\n",
    "            self.io_pixel_count.append(obj[obj != background].shape[0])\n",
    "            self.io_size.append(obj.size)\n",
    "            nc, c = np.unique(obj, return_counts = True)\n",
    "            self.io_unique_colors.append(nc)\n",
    "            self.io_main_color.append(nc[np.argmax(c)])\n",
    "\n",
    "    def identify_object_by_isolation(self, image, background = 0):\n",
    "        # identify all objects by physical isolation on the given image\n",
    "        all_pairs = np.array(np.where(image != background)).T\n",
    "        objectness = np.zeros(image.shape)\n",
    "        this_object = 1\n",
    "        while len(all_pairs) >= 1:\n",
    "            init_pair = all_pairs[0] # start with the first pair\n",
    "            objectness = self.check_neighbors(all_pairs, init_pair, objectness, this_object)\n",
    "            # get a list of index pairs whose neghbors haven't been checked\n",
    "            unchecked_pairs = np.array(np.where(objectness == this_object)).T\n",
    "            checked_pairs = np.zeros((0,2))\n",
    "            # check all the index pairs in the expanding unchecked_pairs untill all have been checked\n",
    "            while len(unchecked_pairs) != 0:\n",
    "                this_pair = unchecked_pairs[0]\n",
    "                objectness = self.check_neighbors(all_pairs, this_pair, objectness, this_object)\n",
    "                # append the checked_pairs\n",
    "                checked_pairs = np.vstack((checked_pairs, this_pair))\n",
    "                # get all index pairs for the currently identified object\n",
    "                current_object_pairs = np.array(np.where(objectness == this_object)).T\n",
    "                # delete the checked pairs from current object pairs\n",
    "                checked_inx = []\n",
    "                for pair in checked_pairs:\n",
    "                    _, inx = self.check_pairs(current_object_pairs, pair, return_inx = True)\n",
    "                    checked_inx.append(inx[0][0])\n",
    "                unchecked_pairs = np.delete(current_object_pairs, checked_inx, axis = 0)\n",
    "\n",
    "            # store this object to identified_objects\n",
    "            current_object_pairs = np.array(np.where(objectness == this_object)).T\n",
    "            cop = current_object_pairs.T\n",
    "            obj = image[np.min(cop[0]):np.max(cop[0])+1, np.min(cop[1]):np.max(cop[1])+1]\n",
    "            # delete the current object pairs from all_pairs\n",
    "            cop_inx = []\n",
    "            for pair in current_object_pairs:\n",
    "                _, this_cop_inx = self.check_pairs(all_pairs, pair, return_inx = True)\n",
    "                cop_inx.append(this_cop_inx[0][0])\n",
    "            all_pairs = np.delete(all_pairs, cop_inx, axis = 0)\n",
    "            # append the object attribute\n",
    "            # p(obj)\n",
    "            self.identified_objects.append(obj)\n",
    "            self.io_inx.append(inx)\n",
    "            self.io_height.append(obj.shape[0])\n",
    "            self.io_width.append(obj.shape[1])\n",
    "            self.io_pixel_count.append(obj[obj != background].shape[0])\n",
    "            self.io_size.append(obj.size)\n",
    "            nc, c = np.unique(obj, return_counts = True)\n",
    "            self.io_unique_colors.append(nc)\n",
    "            self.io_main_color.append(nc[np.argmax(c)])\n",
    "            # start identifying a new object\n",
    "            this_object += 1\n",
    "        return objectness\n",
    "\n",
    "    def identify_object_by_color_isolation(self, true_image, background = 0):\n",
    "        # identify objects first by color then by physical isolation\n",
    "        unique_colors = np.unique(true_image)\n",
    "        for i, color in enumerate(unique_colors):\n",
    "            image = np.copy(true_image) # make a copy from the original first\n",
    "            if color == background:\n",
    "                continue\n",
    "            # identify objects by isolation in this color only\n",
    "            image[image != color] = background\n",
    "            self.identify_object_by_isolation(image, background = background)\n",
    "            \n",
    "\n",
    "    def sort(self, objs, inp):\n",
    "        xs = []\n",
    "        ys = []\n",
    "        for i, o in enumerate(objs):\n",
    "            _, m, n = sliding_window_search(inp, o)\n",
    "            xs.append(m)\n",
    "            ys.append(n)\n",
    "\n",
    "        ans = [[[]],[[]],[[]],[[]]]\n",
    "        left = np.array(ys).argsort()[0:2] # 1,3\n",
    "        right = np.array(ys).argsort()[2:4] # 1,3\n",
    "        if xs[left[0]] <= xs[left[1]]:\n",
    "            ans[0] = objs[left[0]]\n",
    "            ans[2] = objs[left[1]]\n",
    "        else:\n",
    "            ans[2] = objs[left[0]]\n",
    "            ans[0] = objs[left[1]]    \n",
    "        if xs[right[0]] <= xs[right[1]]:\n",
    "            ans[1] = objs[right[0]]\n",
    "            ans[3] = objs[right[1]]\n",
    "        else:\n",
    "            ans[3] = objs[right[0]]\n",
    "            ans[1] = objs[right[1]]        \n",
    "        return ans        \n",
    "            \n",
    "    def merge(self, objects, belt, use_color):\n",
    "        ans=objects\n",
    "        ans=[[[]],[[]],[[]],[[]]]\n",
    "        for o in objects:\n",
    "            o = np.array(o)\n",
    "            max_total = 0\n",
    "            for x in [0,1]:\n",
    "                for y in [0,1]:\n",
    "                    if max_total < o[x:x+len(o)-1, y:y+len(o[0])-1].sum():\n",
    "                        max_total = o[x:x+len(o)-1, y:y+len(o[0])-1].sum()\n",
    "                        max_xy = (x, y)\n",
    "            if max_xy == (0,0):\n",
    "                ans[3] = o\n",
    "            elif max_xy == (0,1):\n",
    "                ans[2] = o\n",
    "            elif max_xy == (1,0):\n",
    "                ans[1] = o\n",
    "            else:\n",
    "                ans[0] = o\n",
    "\n",
    "        if belt == 0:\n",
    "            belt_list = [[use_color]]*len(ans[0])\n",
    "            u=np.hstack([ans[0], ans[1]])\n",
    "            u\n",
    "            s=np.hstack([ans[2], ans[3]])\n",
    "            return np.vstack([u,s])\n",
    "        else:\n",
    "            belt_list = [[use_color]*belt]*len(ans[0])\n",
    "\n",
    "            u=np.hstack([ans[0], belt_list, ans[1]])\n",
    "            s=np.hstack([ans[2], belt_list, ans[3]])\n",
    "            belt_list = [[use_color]*len(s[0])]*belt\n",
    "            return np.vstack([u,belt_list,s])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "metadata": {},
   "outputs": [],
   "source": [
    "def about_color_for_test(task_test):\n",
    "    in_colors = []\n",
    "    out_colors = []\n",
    "    color_changed = False\n",
    "    for inout in task_test:\n",
    "        in_vec = list(itertools.chain.from_iterable(inout['input']))\n",
    "        in_colors += list(set(in_vec))\n",
    "    return list(set(in_colors))\n",
    "\n",
    "def all_inputs_same_shape_and_all_outputs_same_shape(task_train):\n",
    "    m1, n1 = np.array(task_train[0]['input']).shape\n",
    "    m2, n2 = np.array(task_train[0]['output']).shape\n",
    "    all_inputs_same_shape = True\n",
    "    all_outputs_same_shape = True\n",
    "    for inout in task_train:\n",
    "        m1_, n1_ = np.array(inout['input']).shape\n",
    "        m2_, n2_ = np.array(inout['output']).shape\n",
    "        if (m1_ != m1) or (n1_ != n1):\n",
    "            all_inputs_same_shape = False\n",
    "        if (m2_ != m2) or (n2_ != n2):\n",
    "            all_outputs_same_shape = False\n",
    "    return all_inputs_same_shape, all_outputs_same_shape\n",
    "\n",
    "\n",
    "\n",
    "def change_color_for_p(inp_origin, in_use):\n",
    "    inp = copy.deepcopy(inp_origin)\n",
    "    color_map = {}\n",
    "    use_color_num = len(in_use)\n",
    "    out_colors = range(use_color_num)\n",
    "    for i,o in zip(sorted(in_use), sorted(out_colors)):\n",
    "        color_map[i] = o\n",
    "\n",
    "    for i, o in color_map.items():\n",
    "        inp = np.where(np.array(inp_origin) == i, o, inp)\n",
    "    return inp.tolist()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "metadata": {},
   "outputs": [],
   "source": [
    "def all_inouts_same_shape(task_train):\n",
    "    for inout in task_train:\n",
    "        if np.array(inout['input']).shape != np.array(inout['output']).shape:\n",
    "            return False\n",
    "    return True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_by_teacher_row_n(inp_o, back=False, n=0):\n",
    "    try:\n",
    "        one_color_rows = []\n",
    "        not_one_color_rows = []\n",
    "        inp = copy.deepcopy(inp_o)\n",
    "        for row in inp:\n",
    "            if len(set(row)) != 1:\n",
    "                not_one_color_rows.append(row)\n",
    "            else:\n",
    "                one_color_rows.append(row)\n",
    "        c = collections.Counter(np.array(inp).flatten().tolist())\n",
    "        back_color = c.most_common()[0][0]\n",
    "        for row_n, row in enumerate(inp):\n",
    "            if len(set(row)) == 1:\n",
    "                continue\n",
    "            tea = copy.deepcopy(not_one_color_rows[n])\n",
    "            success = True\n",
    "            for ele_n, ele in enumerate(row):\n",
    "                if (ele != back_color) & (tea[ele_n] != ele):\n",
    "                    success = False\n",
    "            if success:\n",
    "                inp[row_n] = tea\n",
    "            else:\n",
    "                inp[row_n] = [back_color] * len(row)\n",
    "\n",
    "        return np.array(inp).tolist()\n",
    "    except:\n",
    "        return [[0]]    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "metadata": {},
   "outputs": [],
   "source": [
    "def perfect_same(big, small):\n",
    "    for x in range(len(big[0])-len(small[0])+1):\n",
    "        for y in range(len(big)-len(small)+1):\n",
    "            if np.array(big)[y:y+len(small), x:x+len(small[0])].tolist() == small.tolist():\n",
    "                return True\n",
    "    small = np.flipud(small)\n",
    "    for x in range(len(big[0])-len(small[0])+1):\n",
    "        for y in range(len(big)-len(small)+1):\n",
    "            if np.array(big)[y:y+len(small), x:x+len(small[0])].tolist() == small.tolist():\n",
    "                return True    \n",
    "    small = np.flip(small)\n",
    "    for x in range(len(big[0])-len(small[0])+1):\n",
    "        for y in range(len(big)-len(small)+1):\n",
    "            if np.array(big)[y:y+len(small), x:x+len(small[0])].tolist() == small.tolist():\n",
    "                return True    \n",
    "    return False            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "def connect_two_point(inp_o, fill=False):\n",
    "    try:\n",
    "        counter = Counter(np.array(inp_o).flatten().tolist())\n",
    "        back = counter.most_common()[0][0]\n",
    "        inp = copy.deepcopy(np.array(inp_o))\n",
    "        for row_n, row in enumerate(inp_o):\n",
    "            start = -1\n",
    "            for ele_n, ele in enumerate(row):\n",
    "                if ele != back:\n",
    "                    if start == -1:\n",
    "                        start = ele_n\n",
    "                    else:\n",
    "                        end = ele_n\n",
    "                        back_pos = (start + end) // 2\n",
    "                        for i in range(back_pos - start - 1):\n",
    "                            inp[row_n, start+1+i] = row[start]\n",
    "                            inp[row_n, end-1-i] = row[end]\n",
    "                        if ((end - start) % 2 == 1) & fill:\n",
    "                            i += 1\n",
    "                            inp[row_n, start+1+i] = row[start]\n",
    "                            inp[row_n, end-1-i] = row[end]                          \n",
    "                        start = ele_n\n",
    "\n",
    "        for row_n, row in enumerate(np.transpose(inp_o)):\n",
    "            start = -1\n",
    "            for ele_n, ele in enumerate(row):\n",
    "                if ele != back:\n",
    "                    if start == -1:\n",
    "                        start = ele_n\n",
    "                    else:\n",
    "                        end = ele_n\n",
    "                        back_pos = (start + end) // 2\n",
    "                        for i in range(back_pos - start - 1):\n",
    "                            inp[start+1+i, row_n] = row[start]\n",
    "                            inp[end-1-i, row_n] = row[end]\n",
    "                        if ((end - start) % 2 == 1) & fill:\n",
    "                            i += 1\n",
    "                            inp[start+1+i, row_n] = row[start]\n",
    "                            inp[end-1-i, row_n] = row[end]                        \n",
    "                        start = ele_n\n",
    "        return inp.tolist()\n",
    "    except:\n",
    "        return [[0]]\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of input rows: 11 cols: 14\n",
      "Number of target rows: 3 cols: 9\n"
     ]
    }
   ],
   "source": [
    "from lightgbm import LGBMClassifier\n",
    "import pdb\n",
    "\n",
    "def preds_to_str(preds_list, idx=''):\n",
    "    pred_strs = []\n",
    "#     st()\n",
    "    for i in range(len(preds_list[0])):\n",
    "        pred_str = ''\n",
    "        for j, preds in enumerate(reversed(preds_list)):\n",
    "            if j == 3:\n",
    "                break\n",
    "            pred_str += flattener(np.array(preds[i]).tolist()) + ' '\n",
    "        pred_strs.append(pred_str)\n",
    "    return pred_strs\n",
    "\n",
    "data_path = Path('../input/abstraction-and-reasoning-challenge/')\n",
    "test_path = data_path / 'test'\n",
    "# https://www.kaggle.com/inversion/abstraction-and-reasoning-starter-notebook\n",
    "def flattener(pred):\n",
    "    str_pred = str([row for row in pred])\n",
    "    str_pred = str_pred.replace(', ', '')\n",
    "    str_pred = str_pred.replace('[[', '|')\n",
    "    str_pred = str_pred.replace('][', '|')\n",
    "    str_pred = str_pred.replace(']]', '|')\n",
    "    return str_pred\n",
    "\n",
    "def get_moore_neighbours(color, cur_row, cur_col, nrows, ncols):\n",
    "\n",
    "    if cur_row<=0: top = -1\n",
    "    else: top = color[cur_row-1][cur_col]\n",
    "        \n",
    "    if cur_row>=nrows-1: bottom = -1\n",
    "    else: bottom = color[cur_row+1][cur_col]\n",
    "        \n",
    "    if cur_col<=0: left = -1\n",
    "    else: left = color[cur_row][cur_col-1]\n",
    "        \n",
    "    if cur_col>=ncols-1: right = -1\n",
    "    else: right = color[cur_row][cur_col+1]\n",
    "        \n",
    "    return top, bottom, left, right\n",
    "\n",
    "def get_tl_tr(color, cur_row, cur_col, nrows, ncols):\n",
    "        \n",
    "    if cur_row==0:\n",
    "        top_left = -1\n",
    "        top_right = -1\n",
    "    else:\n",
    "        if cur_col==0: top_left=-1\n",
    "        else: top_left = color[cur_row-1][cur_col-1]\n",
    "        if cur_col==ncols-1: top_right=-1\n",
    "        else: top_right = color[cur_row-1][cur_col+1]   \n",
    "        \n",
    "    return top_left, top_right\n",
    "\n",
    "def features(task, mode='train'):\n",
    "    cur_idx = 0\n",
    "    num_train_pairs = len(task[mode])\n",
    "    total_inputs = sum([len(task[mode][i]['input'])*len(task[mode][i]['input'][0]) for i in range(num_train_pairs)])\n",
    "    feat = np.zeros((total_inputs,nfeat))\n",
    "    target = np.zeros((total_inputs,), dtype=np.int)\n",
    "    \n",
    "    global local_neighb\n",
    "    for task_num in range(num_train_pairs):\n",
    "        input_color = np.array(task[mode][task_num]['input'])\n",
    "        target_color = task[mode][task_num]['output']\n",
    "        nrows, ncols = len(task[mode][task_num]['input']), len(task[mode][task_num]['input'][0])\n",
    "\n",
    "        target_rows, target_cols = len(task[mode][task_num]['output']), len(task[mode][task_num]['output'][0])\n",
    "        \n",
    "        if (target_rows!=nrows) or (target_cols!=ncols):\n",
    "            print('Number of input rows:',nrows,'cols:',ncols)\n",
    "            print('Number of target rows:',target_rows,'cols:',target_cols)\n",
    "            not_valid=1\n",
    "            return None, None, 1\n",
    "\n",
    "        for i in range(nrows):\n",
    "            for j in range(ncols):\n",
    "                feat[cur_idx,0] = i\n",
    "                feat[cur_idx,1] = j\n",
    "                feat[cur_idx,2] = input_color[i][j]\n",
    "                feat[cur_idx,3:7] = get_moore_neighbours(input_color, i, j, nrows, ncols)\n",
    "                feat[cur_idx,7:9] = get_tl_tr(input_color, i, j, nrows, ncols)\n",
    "                feat[cur_idx,9] = len(np.unique(input_color[i,:]))\n",
    "                feat[cur_idx,10] = len(np.unique(input_color[:,j]))\n",
    "                feat[cur_idx,11] = (i+j)\n",
    "                feat[cur_idx,12] = len(np.unique(input_color[i-local_neighb:i+local_neighb,\n",
    "                                                             j-local_neighb:j+local_neighb]))\n",
    "        \n",
    "                target[cur_idx] = target_color[i][j]\n",
    "                cur_idx += 1\n",
    "            \n",
    "    return feat, target, 0\n",
    "\n",
    "all_task_ids = sorted(os.listdir(test_path))\n",
    "lgb_range = [24]\n",
    "\n",
    "nfeat = 13\n",
    "local_neighb = 5\n",
    "valid_scores = {}\n",
    "for task_n, task_id in enumerate(all_task_ids):\n",
    "    if task_n not in lgb_range:\n",
    "        continue\n",
    "    task_file = str(test_path / task_id)\n",
    "    with open(task_file, 'r') as f:\n",
    "        task = json.load(f)\n",
    "\n",
    "    feat, target, not_valid = features(task)\n",
    "    if not_valid:\n",
    "        not_valid = 0\n",
    "        continue\n",
    "\n",
    "    nrows, ncols = len(task['train'][-1]['input']\n",
    "                       ), len(task['train'][-1]['input'][0])\n",
    "    # use the last train sample for validation\n",
    "    val_idx = len(feat) - nrows*ncols\n",
    "\n",
    "    train_feat = feat[:val_idx]\n",
    "    val_feat = feat[val_idx:, :]\n",
    "\n",
    "    train_target = target[:val_idx]\n",
    "    val_target = target[val_idx:]\n",
    "\n",
    "    #     check if validation set has a new color\n",
    "    #     if so make the mapping color independant\n",
    "    if len(set(val_target) - set(train_target)):\n",
    "        continue\n",
    "\n",
    "    lgb = LGBMClassifier(n_estimators=50, n_jobs=-1)\n",
    "    lgb.fit(feat, target,\n",
    "            verbose=-1)\n",
    "\n",
    "#     training on input pairs is done.\n",
    "#     test predictions begins here\n",
    "\n",
    "    num_test_pairs = len(task['test'])\n",
    "    for task_num in range(num_test_pairs):\n",
    "        cur_idx = 0\n",
    "        input_color = np.array(task['test'][task_num]['input'])\n",
    "        nrows, ncols = len(task['test'][task_num]['input']), len(\n",
    "            task['test'][task_num]['input'][0])\n",
    "        feat = np.zeros((nrows*ncols, nfeat))\n",
    "        unique_col = {col: i for i, col in enumerate(\n",
    "            sorted(np.unique(input_color)))}\n",
    "\n",
    "        for i in range(nrows):\n",
    "            for j in range(ncols):\n",
    "                feat[cur_idx, 0] = i\n",
    "                feat[cur_idx, 1] = j\n",
    "                feat[cur_idx, 2] = input_color[i][j]\n",
    "                feat[cur_idx, 3:7] = get_moore_neighbours(\n",
    "                    input_color, i, j, nrows, ncols)\n",
    "                feat[cur_idx, 7:9] = get_tl_tr(\n",
    "                    input_color, i, j, nrows, ncols)\n",
    "                feat[cur_idx, 9] = len(np.unique(input_color[i, :]))\n",
    "                feat[cur_idx, 10] = len(np.unique(input_color[:, j]))\n",
    "                feat[cur_idx, 11] = (i+j)\n",
    "                feat[cur_idx, 12] = len(np.unique(input_color[i-local_neighb:i+local_neighb,\n",
    "                                                              j-local_neighb:j+local_neighb]))\n",
    "\n",
    "                cur_idx += 1\n",
    "\n",
    "        preds = lgb.predict(feat).reshape(nrows, ncols)\n",
    "        preds = preds.astype(int).tolist()\n",
    "        submission.loc[submission.index == f'{task_id[:-5]}_{task_num}', 'output'] = flattener(preds) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def about_color(task_train):\n",
    "    in_colors = []\n",
    "    out_colors = []\n",
    "    color_changed = False\n",
    "    for inout in task_train:\n",
    "        in_vec = list(itertools.chain.from_iterable(inout['input']))\n",
    "        in_colors += list(set(in_vec))\n",
    "        out_vec = list(itertools.chain.from_iterable(inout['output']))\n",
    "        out_colors += list(set(out_vec))\n",
    "        if set(in_vec) != set(out_vec):\n",
    "            color_changed = True\n",
    "    return list(set(in_colors)), list(set(out_colors)), color_changed\n",
    "def paint_rolling2(inp, back_color, pos):\n",
    "    i = 1\n",
    "    while True:\n",
    "        if i % 4 == 1:\n",
    "            inp[pos[0]:pos[0]+i, pos[1]] = back_color\n",
    "            pos = [pos[0]+i, pos[1]]\n",
    "        elif i % 4 == 2:\n",
    "            if pos[1]-i+1 < 0:\n",
    "                inp[pos[0], :pos[1]+1] = back_color\n",
    "            else:\n",
    "                inp[pos[0], pos[1]-i+1:pos[1]+1] = back_color\n",
    "            pos = [pos[0], pos[1]-i]\n",
    "        elif i % 4 == 3:\n",
    "            inp[pos[0]-i+1:pos[0]+1, pos[1]] = back_color\n",
    "            pos = [pos[0]-i, pos[1]]\n",
    "        elif i % 4 == 0:\n",
    "            inp[pos[0], pos[1]:pos[1]+i] = back_color\n",
    "            pos = [pos[0], pos[1]+i]\n",
    "        i += 1\n",
    "        if (pos[0]<0) or (pos[1] < 0) or (pos[0] >= inp.shape[0]) or (pos[1] >= inp.shape[1]) or (i > 100):\n",
    "#             inp[:, -2] = back_color\n",
    "            # inp[0, :] = back_color\n",
    "            return inp\n",
    "def paint_each_and_vstack2(inp):\n",
    "    try:\n",
    "        for i in range(len(inp)):\n",
    "            if len(set(inp[i])) != 1:\n",
    "                a = np.array(inp[:i])\n",
    "                b = np.array(inp[i:])\n",
    "        for i in range(len(inp)):\n",
    "            if len(set(inp[i])) == 1:\n",
    "                back_color = inp[i][0]\n",
    "                break\n",
    "        use_color = list(set(list(itertools.chain.from_iterable(a)))-set([back_color]))[0]\n",
    "        pos = tuple(nd[0] for nd in np.where(a == use_color))\n",
    "        pos = [pos[0]+1, pos[1]+1]\n",
    "        a = [[use_color]*a.shape[1]]*a.shape[0]\n",
    "\n",
    "        use_color = list(set(list(itertools.chain.from_iterable(b)))-set([back_color]))[0]\n",
    "        b = [[use_color]*b.shape[1]]*b.shape[0]\n",
    "\n",
    "        mat = np.vstack([a,b])\n",
    "        ud_flag = False\n",
    "        if np.array(a).shape[0] > np.array(b).shape[0]:\n",
    "            mat = np.flipud(mat)\n",
    "            ud_flag = True\n",
    "\n",
    "        mat = paint_rolling2(mat, back_color, pos)\n",
    "\n",
    "        if ud_flag:\n",
    "            mat = np.flipud(mat)\n",
    "\n",
    "        return mat\n",
    "    except:\n",
    "        return inp \n",
    "\n",
    "def add_stack4(task):\n",
    "    skip = False\n",
    "    success = False\n",
    "    if task['test'][0]['input'] != np.transpose(task['test'][0]['input']).tolist():\n",
    "        return False\n",
    "    for inout in task['train']:\n",
    "        inp_min_n_shape = np.array(inout['input']).shape\n",
    "        out_min_n_shape = np.array(inout['output']).shape\n",
    "        if (inp_min_n_shape[0] * 2 - 1 != out_min_n_shape[0]) or (inp_min_n_shape[1] * 2 - 1 != out_min_n_shape[1]):\n",
    "            return False\n",
    "        print(3)\n",
    "        inp = inout['input']\n",
    "        out = inout['output']\n",
    "        if (np.flip(stack4(np.flip(inp))).tolist() == out) or (np.array(stack4(inp)).tolist() == out):\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "\n",
    "def rebuild_by_identified_objects(objs, background, x, pattern):\n",
    "    try:\n",
    "        size_map = {}\n",
    "        for i, o in enumerate(objs):\n",
    "            size_map[i] = len(np.array(o).flatten())\n",
    "        size_map = sorted(size_map.items(), key=lambda x:x[1])\n",
    "        out = copy.deepcopy(objs[size_map[2][0]])\n",
    "        out_color = out[1][1]\n",
    "        ele = np.array(objs[size_map[pattern[0]][0]])\n",
    "        ele = np.where(ele==background, out_color, ele)\n",
    "        cood = objs[size_map[pattern[1]][0]]\n",
    "        for row_n, row in enumerate(cood):\n",
    "            for col_n, r in enumerate(row):\n",
    "                if r != background:\n",
    "                    out[row_n*len(ele):(row_n+1)*len(ele), col_n*len(ele[0]):(col_n+1)*len(ele[0])] = ele\n",
    "        for i in range((x-len(out[0]))//2):\n",
    "            out = np.insert(out, 0, background,axis=0)\n",
    "            out = np.insert(out, 0, background,axis=1)\n",
    "            out = np.insert(out, len(out[0]), background,axis=1)\n",
    "            out = np.insert(out, len(out), background,axis=0)\n",
    "        return out\n",
    "    except:\n",
    "        return [[0]]\n",
    "\n",
    "def recolor_by_origin_placement(inp_o, obj, background):\n",
    "    inp = np.array(copy.deepcopy(inp_o))\n",
    "    coods = []\n",
    "    obj_coods = []\n",
    "    x=0\n",
    "    for i in range(len(obj)):\n",
    "        y=0\n",
    "        x+=1\n",
    "        for j in range(len(obj[0])):\n",
    "            y+=1\n",
    "            if np.all(inp[x+i*len(obj):x+(i+1)*len(obj), y+j*len(obj[0]):y+(j+1)*len(obj[0])] == obj):\n",
    "                coods.append([x+i*len(obj),y+j*len(obj[0])])\n",
    "                obj_coods.append([i,j])\n",
    "    inp = np.where(inp_o == background, obj[0][0], inp)\n",
    "    inp = np.where(inp_o == obj[0][0], background, inp)\n",
    "    print(coods)\n",
    "    for c in obj_coods:\n",
    "        obj[c[0]][c[1]] = background\n",
    "    for c in coods:\n",
    "        inp[c[0]:c[0]+len(obj), c[1]:c[1]+len(obj[0])] = obj\n",
    "    return inp\n",
    "def paint_rolling4(inp, back_color, pos):\n",
    "    i = 1\n",
    "    while True:\n",
    "        if i % 4 == 1:\n",
    "            inp[pos[0]:pos[0]+i, pos[1]] = back_color\n",
    "            pos = [pos[0]+i, pos[1]]\n",
    "        elif i % 4 == 2:\n",
    "            if pos[1]-i+1 < 0:\n",
    "                inp[pos[0], :pos[1]+1] = back_color\n",
    "            else:\n",
    "                inp[pos[0], pos[1]-i+1:pos[1]+1] = back_color\n",
    "            pos = [pos[0], pos[1]-i]\n",
    "        elif i % 4 == 3:\n",
    "            inp[pos[0]-i+1:pos[0]+1, pos[1]] = back_color\n",
    "            pos = [pos[0]-i, pos[1]]\n",
    "        elif i % 4 == 0:\n",
    "            inp[pos[0], pos[1]:pos[1]+i] = back_color\n",
    "            pos = [pos[0], pos[1]+i]\n",
    "        i += 1\n",
    "        if (pos[0]<0) or (pos[1] < 0) or (pos[0] >= inp.shape[0]) or (pos[1] >= inp.shape[1]) or (i > 100):\n",
    "            inp[:, -2:] = back_color\n",
    "            # inp[0, :] = back_color\n",
    "            return inp\n",
    "\n",
    "def paint_each_and_vstack4(inp):\n",
    "    try:\n",
    "        for i in range(len(inp)):\n",
    "            if len(set(inp[i])) != 1:\n",
    "                a = np.array(inp[:i])\n",
    "                b = np.array(inp[i:])\n",
    "        for i in range(len(inp)):\n",
    "            if len(set(inp[i])) == 1:\n",
    "                back_color = inp[i][0]\n",
    "                break\n",
    "        use_color = list(set(list(itertools.chain.from_iterable(a)))-set([back_color]))[0]\n",
    "        pos = tuple(nd[0] for nd in np.where(a == use_color))\n",
    "        pos = [pos[0]+1, pos[1]+1]\n",
    "        a = [[use_color]*a.shape[1]]*a.shape[0]\n",
    "\n",
    "        use_color = list(set(list(itertools.chain.from_iterable(b)))-set([back_color]))[0]\n",
    "        b = [[use_color]*b.shape[1]]*b.shape[0]\n",
    "\n",
    "        mat = np.vstack([a,b])\n",
    "        ud_flag = False\n",
    "        if np.array(a).shape[0] > np.array(b).shape[0]:\n",
    "            mat = np.flipud(mat)\n",
    "            ud_flag = True\n",
    "\n",
    "        mat = paint_rolling4(mat, back_color, pos)\n",
    "\n",
    "        if ud_flag:\n",
    "            mat = np.flipud(mat)\n",
    "\n",
    "        return mat\n",
    "    except:\n",
    "        return inp\n",
    "\n",
    "def paint_rolling3(inp, back_color, pos):\n",
    "    i = 1\n",
    "    while True:\n",
    "        if i % 4 == 1:\n",
    "            inp[pos[0]:pos[0]+i, pos[1]] = back_color\n",
    "            pos = [pos[0]+i, pos[1]]\n",
    "        elif i % 4 == 2:\n",
    "            if pos[1]-i+1 < 0:\n",
    "                inp[pos[0], :pos[1]+1] = back_color\n",
    "            else:\n",
    "                inp[pos[0], pos[1]-i+1:pos[1]+1] = back_color\n",
    "            pos = [pos[0], pos[1]-i]\n",
    "        elif i % 4 == 3:\n",
    "            inp[pos[0]-i+1:pos[0]+1, pos[1]] = back_color\n",
    "            pos = [pos[0]-i, pos[1]]\n",
    "        elif i % 4 == 0:\n",
    "            inp[pos[0], pos[1]:pos[1]+i] = back_color\n",
    "            pos = [pos[0], pos[1]+i]\n",
    "        i += 1\n",
    "        if (pos[0]<0) or (pos[1] < 0) or (pos[0] >= inp.shape[0]) or (pos[1] >= inp.shape[1]) or (i > 100):\n",
    "            inp[:, -2] = back_color\n",
    "            # inp[0, :] = back_color\n",
    "            return inp\n",
    "\n",
    "def paint_each_and_vstack3(inp):\n",
    "    try:\n",
    "        for i in range(len(inp)):\n",
    "            if len(set(inp[i])) != 1:\n",
    "                a = np.array(inp[:i])\n",
    "                b = np.array(inp[i:])\n",
    "        for i in range(len(inp)):\n",
    "            if len(set(inp[i])) == 1:\n",
    "                back_color = inp[i][0]\n",
    "                break\n",
    "        use_color = list(set(list(itertools.chain.from_iterable(a)))-set([back_color]))[0]\n",
    "        pos = tuple(nd[0] for nd in np.where(a == use_color))\n",
    "        pos = [pos[0]+1, pos[1]+1]\n",
    "        a = [[use_color]*a.shape[1]]*a.shape[0]\n",
    "\n",
    "        use_color = list(set(list(itertools.chain.from_iterable(b)))-set([back_color]))[0]\n",
    "        b = [[use_color]*b.shape[1]]*b.shape[0]\n",
    "\n",
    "        mat = np.vstack([a,b])\n",
    "        ud_flag = False\n",
    "        if np.array(a).shape[0] > np.array(b).shape[0]:\n",
    "            mat = np.flipud(mat)\n",
    "            ud_flag = True\n",
    "\n",
    "        mat = paint_rolling3(mat, back_color, pos)\n",
    "\n",
    "        if ud_flag:\n",
    "            mat = np.flipud(mat)\n",
    "\n",
    "        return mat\n",
    "    except:\n",
    "        return inp\n",
    "\n",
    "def stack4(inp_o):\n",
    "    try:\n",
    "        inp = np.array(copy.deepcopy(inp_o))\n",
    "#         inp = np.where(inp==inp.T, inp, inp[-1][-1])\n",
    "        a = inp\n",
    "        b = np.fliplr(inp)\n",
    "        c = np.flipud(inp)\n",
    "        d = np.flip(inp)\n",
    "        e = np.hstack([a,b[:, 1:]])\n",
    "        f = np.hstack([c,d[:, 1:]])\n",
    "        return np.vstack([e, f[1:, :]])\n",
    "    except:\n",
    "        return inp_o                \n",
    "\n",
    "def copy_by_belt_and_change_color(inp_o, change, to_back, to_belt=False, reverse=False, mirror=True):\n",
    "    try:\n",
    "        inp = copy.deepcopy(inp_o)\n",
    "        belt = inp[0][0]\n",
    "        one_color_col_colors = []\n",
    "        for col in np.transpose(inp).tolist():\n",
    "            if len(set(col)) == 1:\n",
    "                test_has_one_color_col = True\n",
    "                one_color_col_colors.append(col[0])\n",
    "        one_color_col_colors = list(set(one_color_col_colors))\n",
    "        back = inp[0][0]\n",
    "        if len(set(np.array(inp)[:, 0])) == 1:\n",
    "            back = inp[0][0]\n",
    "        elif len(set(np.array(inp)[:, -1])) == 1:\n",
    "            back = inp[-1][-1]\n",
    "        if one_color_col_colors[0] == back:\n",
    "            belt = one_color_col_colors[1]\n",
    "        else:\n",
    "            belt = one_color_col_colors[0]\n",
    "\n",
    "        belt_xs = []\n",
    "        for ele_n, ele in enumerate(inp[0]):\n",
    "            if ele == belt:\n",
    "                belt_xs.append(ele_n)\n",
    "        change_left = False\n",
    "        change_right = True\n",
    "        if np.array(inp)[:, :belt_xs[0]].flatten().tolist().count(back) / len(np.array(inp)[:, :belt_xs[0]].flatten().tolist()) > np.array(inp)[:, belt_xs[-1]+1:].flatten().tolist().count(back) / len(np.array(inp)[:, belt_xs[1]+1:].flatten().tolist()):\n",
    "            change_left = True\n",
    "            change_right = False\n",
    "\n",
    "        range_x = np.min([belt_xs[0], len(inp[0])-belt_xs[-1]-1])\n",
    "        inp = np.array(inp)\n",
    "        use_colors = list(set(inp.flatten().tolist()) - set([back, belt]))\n",
    "        if len(use_colors) == 0:\n",
    "            use_color = belt\n",
    "        else:\n",
    "            use_color = use_colors[0]\n",
    "\n",
    "        for x in range(range_x):\n",
    "            for y in range(len(inp)):\n",
    "                a, b = inp[y, belt_xs[0]-x-1], inp[y, belt_xs[-1]+x+1]\n",
    "                if (a != back) & (b != back):\n",
    "                    if a == b:\n",
    "                        if to_back:\n",
    "                            inp[y, belt_xs[-1]+x+1] = back\n",
    "                            inp[y, belt_xs[0]-x-1] = back\n",
    "                        elif change:\n",
    "                            if a == belt:\n",
    "                                inp[y, belt_xs[-1]+x+1] = use_color\n",
    "                                inp[y, belt_xs[0]-x-1] = use_color\n",
    "                            else:\n",
    "                                inp[y, belt_xs[-1]+x+1] = belt\n",
    "                                inp[y, belt_xs[0]-x-1] = belt\n",
    "                    else:\n",
    "                        if to_belt:\n",
    "                            inp[y, belt_xs[-1]+x+1] = belt\n",
    "                            inp[y, belt_xs[0]-x-1] = belt\n",
    "                        elif reverse:\n",
    "                            inp[y, belt_xs[-1]+x+1] = a\n",
    "                            inp[y, belt_xs[0]-x-1] = b\n",
    "                        else:\n",
    "                            if a == belt:\n",
    "                                inp[y, belt_xs[-1]+x+1] = use_color\n",
    "                            else:\n",
    "                                inp[y, belt_xs[-1]+x+1] = belt\n",
    "                            if b == belt:\n",
    "                                inp[y, belt_xs[0]-x-1] = use_color\n",
    "                            else:\n",
    "                                inp[y, belt_xs[0]-x-1] = belt\n",
    "                elif (a != back):\n",
    "                    if a == belt:\n",
    "                        inp[y, belt_xs[-1]+x+1] = use_color\n",
    "                    else:\n",
    "                        inp[y, belt_xs[-1]+x+1] = belt\n",
    "                elif inp[y, belt_xs[-1]+x+1] != back:\n",
    "                    if b == belt:\n",
    "                        inp[y, belt_xs[0]-x-1] = use_color\n",
    "                    else:\n",
    "                        inp[y, belt_xs[0]-x-1] = belt\n",
    "        if not mirror:\n",
    "            if change_left:\n",
    "                inp[:, belt_xs[0]-range_x:belt_xs[0]] = np.fliplr(inp[:, belt_xs[0]-range_x:belt_xs[0]])\n",
    "            else:\n",
    "                inp[:, belt_xs[1]+1:belt_xs[1]+1+range_x] = np.fliplr(inp[:, belt_xs[1]+1:belt_xs[1]+1+range_x])\n",
    "\n",
    "        return inp\n",
    "    except:\n",
    "        return [[0]]\n",
    "def add_recolor_by_origin_placement(task):\n",
    "    inp = task['test'][0]['input']\n",
    "    inp_train = task['train'][0]['input']\n",
    "    if (len(inp) != len(inp[0])) or (len(inp_train) != len(inp_train[0])):\n",
    "        return False\n",
    "\n",
    "    use_color = list(set(list(itertools.chain.from_iterable(inp))))\n",
    "    if len(use_color) != 2:\n",
    "        return False\n",
    "    for inout in task['train']:\n",
    "        success = False\n",
    "\n",
    "        i_test = np.array(inout['input'])\n",
    "        arc = ARC_solver()\n",
    "        background = arc.get_background(i_test)\n",
    "        arc.identify_object_by_isolation(i_test, background)\n",
    "\n",
    "        a = recolor_by_origin_placement(i_test, arc.identified_objects[0], background)\n",
    "        if np.array(a).tolist() == inout['output']:\n",
    "            success = True\n",
    "            break\n",
    "    return success\n",
    "\n",
    "def add_rebuild_by_identified_objects(task):\n",
    "    use_colors = list(set(np.array(task['test'][0]['input']).flatten().tolist()))\n",
    "    if len(use_colors) != 4:\n",
    "        return False\n",
    "    inp = task['train'][-2]['input']\n",
    "    out = task['train'][-2]['output']\n",
    "    if (len(inp[0]) != len(out[0])):\n",
    "        return False\n",
    "\n",
    "    success = False\n",
    "    for test_n, inout in enumerate(task['train']):\n",
    "        i_test = np.array(inout['input'])\n",
    "        arc = ARC_solver()\n",
    "        background = arc.get_background(i_test)\n",
    "        arc.identify_object_by_isolation(i_test, background)\n",
    "\n",
    "        a = rebuild_by_identified_objects(arc.identified_objects, background, len(i_test[0]), [0,1])\n",
    "        b = rebuild_by_identified_objects(arc.identified_objects, background, len(i_test[0]), [1,0])\n",
    "        if (np.array(a).tolist() == inout['output']) or (np.array(b).tolist() == inout['output']):\n",
    "            success = True\n",
    "            break\n",
    "    return success\n",
    "\n",
    "def add_copy_by_belt_and_change_color(task):\n",
    "    skip = False\n",
    "    inp = task['test'][0]['input']\n",
    "    for n, row in enumerate(inp):\n",
    "        if len(set(row)) == 1:\n",
    "            skip = True\n",
    "    if skip:\n",
    "        return False\n",
    "    unique_one_color_col_ns = []\n",
    "    for n, col in enumerate(np.transpose(inp)):\n",
    "        if len(set(col)) == 1:\n",
    "            unique_one_color_col_ns.append(col[0])\n",
    "    if len(set(unique_one_color_col_ns)) != 2:\n",
    "        return False\n",
    "    success = False\n",
    "    for test_n, inout in enumerate(task['train']):\n",
    "        i_test = np.transpose(inout['input']).tolist()\n",
    "        a = np.transpose(copy_by_belt_and_change_color(i_test,True,False, mirror=False)).tolist()\n",
    "        b = np.transpose(copy_by_belt_and_change_color(i_test,True,False, reverse=True)).tolist()\n",
    "        c = np.transpose(copy_by_belt_and_change_color(i_test,True,False, to_belt=True)).tolist()\n",
    "        if (a == inout['output']) or (b == inout['output']) or (c == inout['output']):\n",
    "            success = True\n",
    "            break\n",
    "    if not success:\n",
    "        return False\n",
    "    return True\n",
    "\n",
    "def add_paint_each_and_vstack(task):\n",
    "    inp = copy.deepcopy(task['train'][-1]['input'])\n",
    "    in_use, out_use, color_changed = about_color(copy.deepcopy(task['train'][-1:]))\n",
    "    if len(in_use) != 3:\n",
    "        return False\n",
    "    v=list(itertools.chain.from_iterable(inp))\n",
    "    if (len(v) - v.count(in_use[0]) != 2) & (len(v) - v.count(in_use[1]) != 2) & (len(v) - v.count(in_use[2]) != 2):\n",
    "        return False\n",
    "    if np.array(paint_each_and_vstack2(inp)).tolist() != task['train'][-1]['output']:\n",
    "        return False\n",
    "    return True\n",
    "\n",
    "def connect_two_point2(inp_o, fill=False):\n",
    "    try:\n",
    "        counter = Counter(np.array(inp_o).flatten().tolist())\n",
    "        back = counter.most_common()[0][0]\n",
    "        inp = copy.deepcopy(np.array(inp_o))\n",
    "\n",
    "        for row_n, row in enumerate(np.transpose(inp_o)):\n",
    "            start = -1\n",
    "            for ele_n, ele in enumerate(row):\n",
    "                if ele != back:\n",
    "                    if start == -1:\n",
    "                        start = ele_n\n",
    "                    else:\n",
    "                        end = ele_n\n",
    "                        back_pos = (start + end) // 2\n",
    "                        for i in range(back_pos - start - 1):\n",
    "                            inp[start+1+i, row_n] = row[start]\n",
    "                            inp[end-1-i, row_n] = row[end]\n",
    "                        if ((end - start) % 2 == 1) & fill:\n",
    "                            i += 1\n",
    "                            inp[start+1+i, row_n] = row[start]\n",
    "                            inp[end-1-i, row_n] = row[end]                        \n",
    "                        start = ele_n\n",
    "        return inp.tolist()\n",
    "    except:\n",
    "        return [[0]]  \n",
    "def add_connect_two_point2(task):\n",
    "    success = False\n",
    "    for inout in task['train']:\n",
    "        if (np.array(connect_two_point2(inout['input'], fill=True)).tolist() == inout['output']) or (np.transpose(connect_two_point2(np.transpose(inout['input']), fill=True)).tolist() == inout['output']):\n",
    "            success = True\n",
    "    return success\n",
    "def stack4_2(inp_o):\n",
    "    try:\n",
    "        inp = np.array(copy.deepcopy(inp_o))\n",
    "        inp[-1][-2] = inp[-1][-1]\n",
    "        inp[-2][-1] = inp[-1][-1]\n",
    "        a = inp\n",
    "        b = np.fliplr(inp)\n",
    "        c = np.flipud(inp)\n",
    "        d = np.flip(inp)\n",
    "        e = np.hstack([a,b[:, 1:]])\n",
    "        f = np.hstack([c,d[:, 1:]])\n",
    "        return np.vstack([e, f[1:, :]])\n",
    "    except:\n",
    "        return inp_o\n",
    "\n",
    "def add_several_funcs(task):\n",
    "    try:\n",
    "        if (len(task['test'][0]['input']) % 16 == 0) & (len(task['test'][0]['input'][0]) % 16 == 0):\n",
    "            return None, False\n",
    "        use_flag = add_recolor_by_origin_placement(task)\n",
    "        if use_flag:\n",
    "            return recolor_by_origin_placement, True\n",
    "        use_flag = add_rebuild_by_identified_objects(task)\n",
    "        if use_flag:\n",
    "            return rebuild_by_identified_objects, True\n",
    "        use_flag = add_copy_by_belt_and_change_color(task)\n",
    "        if use_flag:\n",
    "            return copy_by_belt_and_change_color, True\n",
    "        use_flag = add_paint_each_and_vstack(task)\n",
    "        if use_flag:\n",
    "            return paint_each_and_vstack3, True\n",
    "        use_flag = add_stack4(task)\n",
    "        if use_flag:\n",
    "            return stack4, True\n",
    "        use_flag = add_connect_two_point2(task)\n",
    "        if use_flag:\n",
    "            return connect_two_point2, True\n",
    "        if add_block_merge(task['train']):\n",
    "                return divide_block_and_merge3, True\n",
    "        return None, False\n",
    "    except:\n",
    "        return None, False\n",
    "\n",
    "def apply_several_func(task, func):\n",
    "    try:\n",
    "        if func == recolor_by_origin_placement:\n",
    "            for test_n, inout in enumerate(task['test']):\n",
    "                i_test = np.array(inout['input'])\n",
    "                arc = ARC_solver()\n",
    "                background = arc.get_background(i_test)\n",
    "                arc.identify_object_by_isolation(i_test, background)\n",
    "\n",
    "                preds = []\n",
    "                a = recolor_by_origin_placement(i_test, arc.identified_objects[0], background)\n",
    "                b = [[0]]\n",
    "                c = [[0]]\n",
    "                return [a,b,c]\n",
    "        elif func == rebuild_by_identified_objects:\n",
    "            for test_n, inout in enumerate(task['test']):\n",
    "                i_test = np.array(inout['input'])\n",
    "        #         p(i_test)\n",
    "                arc = ARC_solver()\n",
    "                background = arc.get_background(i_test)\n",
    "                arc.identify_object_by_isolation(i_test, background)\n",
    "\n",
    "                preds = []\n",
    "                a = rebuild_by_identified_objects(arc.identified_objects, background, len(i_test[0]), [0,1])\n",
    "                b = rebuild_by_identified_objects(arc.identified_objects, background, len(i_test[0]), [1,0])\n",
    "                c = [[0]]\n",
    "                return [a,b,c]\n",
    "        elif func == copy_by_belt_and_change_color:\n",
    "            for test_n, inout in enumerate(task['test']):\n",
    "                i_test = inout['input']\n",
    "                preds = []\n",
    "                a = copy_by_belt_and_change_color(i_test,True,False, mirror=False)\n",
    "                b = copy_by_belt_and_change_color(i_test,True,False, reverse=True)\n",
    "                c = copy_by_belt_and_change_color(i_test,True,False, to_belt=True)\n",
    "                return [a,b,c]\n",
    "        elif func == paint_each_and_vstack3:\n",
    "            for test_n, inout in enumerate(task['test']):\n",
    "                i_test = inout['input']\n",
    "                a=paint_each_and_vstack3(np.flip(i_test))\n",
    "                b=paint_each_and_vstack3(i_test)\n",
    "                c=paint_each_and_vstack4(i_test)\n",
    "                return [a,b,c]\n",
    "        elif func == stack4:\n",
    "            for test_n, inout in enumerate(task['test']):\n",
    "                i_test = inout['input']\n",
    "                if i_test[0][0] == i_test[0][1]:\n",
    "                    a=stack4_2(np.flip(i_test))\n",
    "                    b=stack4(np.flip(i_test))\n",
    "                    c=stack4(i_test)\n",
    "                else:\n",
    "                    a=stack4_2(i_test)\n",
    "                    b=stack4(i_test)\n",
    "                    c=stack4(np.flip(i_test))\n",
    "                return [a,b,c]\n",
    "        elif func == connect_two_point2:\n",
    "            for test_n, inout in enumerate(task['test']):\n",
    "                i_test = inout['input']\n",
    "                preds = []\n",
    "                a = connect_two_point2(inout['input'], fill=False)\n",
    "                b = connect_two_point2(inout['input'], fill=True)\n",
    "                c = np.transpose(connect_two_point2(np.transpose(inout['input']), fill=True)).tolist()\n",
    "                return [a, b, c]\n",
    "        elif func == divide_block_and_merge3:\n",
    "            t1=divide_block_and_merge3(task['test'], 1)\n",
    "            t2=divide_block_and_merge3(task['test'], 2)\n",
    "            t3=divide_block_and_merge3(task['test'], 3)\n",
    "            return [t1[0]['input'], t2[0]['input'], t3[0]['input']]\n",
    "\n",
    "    except:\n",
    "        return []\n",
    "    \n",
    "def add_block_merge(task_train):\n",
    "    try:\n",
    "        arc = ARC_solver()\n",
    "        inout = task_train[-1]\n",
    "        inp = copy.deepcopy(inout['input'])\n",
    "        inp = np.array(inp)\n",
    "        use_color = list(set(list(itertools.chain.from_iterable(inp))))\n",
    "        if len(use_color) != 2:\n",
    "            return False\n",
    "        inp_o = copy.deepcopy(inp)\n",
    "        inp = np.where(inp_o==use_color[0], use_color[1], inp)\n",
    "        inp = np.where(inp_o==use_color[1], use_color[0], inp)\n",
    "        background = arc.get_background(inp)\n",
    "        arc.identify_object_by_isolation(inp, background)\n",
    "        if len(arc.identified_objects) != 4:\n",
    "            return False\n",
    "\n",
    "#             arc.identified_objects = arc.sort(arc.identified_objects, inp)\n",
    "#             for i in arc.identified_objects:\n",
    "#                 p(i)\n",
    "#         out = arc.merge(arc.identified_objects, 1, use_color[1])\n",
    "        for i in range(4):\n",
    "            out = arc.identified_objects[i]\n",
    "            out_o = copy.deepcopy(out)\n",
    "            out = np.where(out_o==use_color[0], use_color[1], out)\n",
    "            out = np.where(out_o==use_color[1], use_color[0], out)\n",
    "            if out.tolist() == inout['output']:\n",
    "                return True\n",
    "    except:\n",
    "        pass\n",
    "    return False    \n",
    "def divide_block_and_merge3(task_train_origin, obj_numb):\n",
    "    for inout in task_train_origin:\n",
    "        inout['input'] = np.array(inout['input'])\n",
    "    task_train = copy.deepcopy(task_train_origin)\n",
    "    for i, inout in enumerate(task_train):    \n",
    "        arc = ARC_solver()\n",
    "        inp = inout['input']\n",
    "        inp = np.array(inp)\n",
    "        use_color = list(set(list(itertools.chain.from_iterable(inp))))\n",
    "        if len(use_color) != 2:\n",
    "            return task_train_origin\n",
    "#         try:\n",
    "        inp_o = copy.deepcopy(inp)\n",
    "        inp = np.where(inp_o==use_color[0], use_color[1], inp)\n",
    "        inp = np.where(inp_o==use_color[1], use_color[0], inp)\n",
    "        background = arc.get_background(inp)\n",
    "        arc.identify_object_by_isolation(inp, background)\n",
    "        if len(arc.identified_objects) == 4:\n",
    "            arc.identified_objects = arc.sort(arc.identified_objects, inp)\n",
    "            out = arc.identified_objects[obj_numb]\n",
    "            out_o = copy.deepcopy(out)\n",
    "            out = np.where(out_o==use_color[0], use_color[1], out)\n",
    "            out = np.where(out_o==use_color[1], use_color[0], out)\n",
    "            task_train[i]['input'] = out\n",
    "#         except:\n",
    "#             return task_train_origin\n",
    "    return task_train\n",
    "def main(idx, task, env='dev'):\n",
    "    several_f = False\n",
    "    func_combi_map = defaultdict(list)\n",
    "    result = pd.Series()\n",
    "    preprocess_best_score_map = {}\n",
    "    best_aug_score_map = {}\n",
    "    success_list = []\n",
    "    final_score_map = {}\n",
    "    pre_final_score_map = {}\n",
    "    promising_map = defaultdict(bool)\n",
    "    time_map = {}\n",
    "    skip = False\n",
    "    back_to_black = False\n",
    "    origin_back_color = 1\n",
    "    preprocess_best_score = 0\n",
    "    best_func_combi = []\n",
    "    task_n = 0\n",
    "    correct_only_preprocess_flag = False\n",
    "    use_several_func = False\n",
    "    start = time()\n",
    "    print('--------------')\n",
    "    print(f'{task_n}：{idx}')\n",
    "    flip_funcs = [inouts_array, inouts_flip, inouts_flipud, inouts_fliplr]\n",
    "    back_to_black_funcs = add_back_to_black_funcs(copy.deepcopy(task['train']))\n",
    "    func, use_several_func_flag = add_several_funcs(task)\n",
    "    if (len(task['test'][0]['input']) % 16 != 0) & (len(task['test'][0]['input'][0]) % 16 != 0) & (not use_several_func_flag):\n",
    "        skip = True\n",
    "\n",
    "    for back_to_black_func in back_to_black_funcs:\n",
    "        if use_several_func_flag:\n",
    "            outputs = apply_several_func(task, func)\n",
    "            if len(outputs) != 0:\n",
    "                use_several_func = True\n",
    "                break\n",
    "            else:\n",
    "                use_several_func_flag = False\n",
    "        if correct_only_preprocess_flag or use_several_func or skip:\n",
    "            break\n",
    "        train_copy0 = back_to_black_func(copy.deepcopy(task['train']))\n",
    "        size_change_funcs = add_size_change_funcs(train_copy0, task_n)\n",
    "\n",
    "        for size_change_func in size_change_funcs:\n",
    "            if correct_only_preprocess_flag or use_several_func:\n",
    "                break\n",
    "            shaped_train = size_change_func(copy.deepcopy(train_copy0))\n",
    "            # print(type(shaped_train))\n",
    "            transpose_funcs = add_transpose(shaped_train)\n",
    "            for transpose_func in transpose_funcs:\n",
    "                if correct_only_preprocess_flag or use_several_func:\n",
    "                    break\n",
    "                shaped_train1 = transpose_func(copy.deepcopy(shaped_train))\n",
    "#                     if size_change_func == divide_block_and_merge1:\n",
    "#                         st()\n",
    "\n",
    "                shape_different_flag = False\n",
    "#                     print(size_change_funcs)\n",
    "                for shaped_inout in shaped_train1:\n",
    "                    if shaped_inout['input'].shape != np.array(shaped_inout['output']).shape:\n",
    "                        shape_different_flag = True\n",
    "                        break\n",
    "                if shape_different_flag:\n",
    "                    break\n",
    "\n",
    "                train4_funcs = add_train4_growth(shaped_train1)\n",
    "                for train4_func in train4_funcs:\n",
    "                    if correct_only_preprocess_flag:\n",
    "                        break\n",
    "                    shaped_train2 = train4_func(copy.deepcopy(shaped_train1))\n",
    "                    # print(type(shaped_train2))\n",
    "                    fill_closed_area_funcs = add_fill_closed_area(shaped_train2.copy())\n",
    "                    for fill_closed_area_func in fill_closed_area_funcs:\n",
    "                        if correct_only_preprocess_flag:\n",
    "                            break\n",
    "                        shaped_train3 = fill_closed_area_func(copy.deepcopy(shaped_train2))\n",
    "                        # print(type(shaped_train3))\n",
    "                        for flip_func_num, flip_func in enumerate(flip_funcs):\n",
    "                            if correct_only_preprocess_flag:\n",
    "                                break\n",
    "                            shaped_train4 = flip_func(copy.deepcopy(shaped_train3))\n",
    "                            patch_funcs = add_patch_funcs(shaped_train4, idx)\n",
    "                            for patch_func in patch_funcs:\n",
    "                                if correct_only_preprocess_flag:\n",
    "                                    break\n",
    "                                shaped_train5 = patch_func(copy.deepcopy(shaped_train4))\n",
    "                                task_train6_funcs = add_task_train6(shaped_train5)\n",
    "                                for train6_funcs in task_train6_funcs:\n",
    "                                    if correct_only_preprocess_flag:\n",
    "                                        break\n",
    "                                    shaped_train6 = train6_funcs(copy.deepcopy(shaped_train5))\n",
    "                                    move_object_funcs = add_move_object(shaped_train6)\n",
    "                                    for move_object_func in move_object_funcs:\n",
    "                                        if correct_only_preprocess_flag:\n",
    "                                            break\n",
    "                                        shaped_train7 = move_object_func(copy.deepcopy(shaped_train6))\n",
    "                                        recolor_funcs = add_recolor(shaped_train7, task_n)\n",
    "                                        for recolor_func in recolor_funcs:\n",
    "                                            if correct_only_preprocess_flag:\n",
    "                                                break\n",
    "                                            shaped_train8 = recolor_func(copy.deepcopy(shaped_train7))\n",
    "                                            kneighbor_funcs = add_kneighbors(shaped_train8)\n",
    "                                            for kneighbor_func in kneighbor_funcs:\n",
    "                                                if correct_only_preprocess_flag:\n",
    "                                                    break\n",
    "                                                shaped_train9 = kneighbor_func(copy.deepcopy(shaped_train8))\n",
    "\n",
    "                                                change_color_funcs = add_change_color_funcs(shaped_train9)\n",
    "                                                for change_color_func in change_color_funcs:\n",
    "                                                    if correct_only_preprocess_flag:\n",
    "                                                        break\n",
    "                                                    shaped_train10 = change_color_func(copy.deepcopy(shaped_train9))\n",
    "                                                    second_shape_different_flag = False\n",
    "                                                    shaped_train_copy = shaped_train10\n",
    "                                                    func_combi = [func for func in [back_to_black_func, size_change_func, patch_func, flip_func, transpose_func, train4_func, fill_closed_area_func, train6_funcs, move_object_func, recolor_func, kneighbor_func, change_color_func] if func != inouts_array]\n",
    "                                                    func_combi += [inouts_array] if len(func_combi) == 0 else []\n",
    "                                                    for train_num, in_out in enumerate(copy.deepcopy(shaped_train_copy)):\n",
    "                                                        if in_out['input'].shape != np.array(in_out['output']).shape:\n",
    "                                                            second_shape_different_flag = True\n",
    "                                                            break\n",
    "                                                        # st()\n",
    "                                                        if in_out['input'].tolist() == in_out['output']:\n",
    "            #                                                 print(func_combi)\n",
    "                                                            correct_only_preprocess_flag = True\n",
    "            #                                                 st()\n",
    "                                                            for idx_minus_num in [1, 2]:\n",
    "                                                                another_in_out = shaped_train_copy[train_num - idx_minus_num]\n",
    "                                                                if another_in_out['input'].tolist() != another_in_out['output']:\n",
    "                                                                    correct_only_preprocess_flag = False\n",
    "                                                            if correct_only_preprocess_flag:\n",
    "                                                                func_combi_map[idx].append(func_combi)\n",
    "                                                                preprocess_best_score = 0.999\n",
    "                                                    if second_shape_different_flag or correct_only_preprocess_flag:\n",
    "                                                        continue\n",
    "                #                                     st()\n",
    "                                                    similarity = get_similarity(shaped_train_copy, [], idx)\n",
    "                    #                                 print(func_combi)\n",
    "                    #                                 print(similarity)\n",
    "                                                    if similarity > preprocess_best_score:\n",
    "                                                        func_combi_map[idx].append(func_combi)\n",
    "                                                        preprocess_best_score = similarity\n",
    "                                                        best_func_combi = func_combi\n",
    "                                                        preprocess_best_score_map[idx] = preprocess_best_score\n",
    "\n",
    "    if use_several_func:\n",
    "        result[f'{idx}_0'] = ''\n",
    "        several_f = True\n",
    "        for out in outputs:\n",
    "            result[f'{idx}_0'] += flattener(np.array(out).tolist()) + ' '\n",
    "\n",
    "        success_list.append(task_n)\n",
    "    elif correct_only_preprocess_flag:\n",
    "        # TODO: 一回目はこれでやるとして、2回目以降を考える\n",
    "        print('↓correct_only_preprocess!↓')\n",
    "        print(f'idx: {idx}, func: {func_combi_map[idx]}')\n",
    "        success_list.append(task_n)\n",
    "\n",
    "        preds0, preds1, preds2 = [], [], []\n",
    "        if divide_block_and_merge1 in func_combi_map[idx][0]:\n",
    "            funcs0 = [divide_block_and_merge1]\n",
    "            funcs1 = [divide_block_and_merge2]\n",
    "            funcs2 = [divide_block_and_merge3]\n",
    "        else:\n",
    "            funcs0 = func_combi_map[idx][-1 % len(func_combi_map[idx])]\n",
    "            funcs1 = func_combi_map[idx][-2 % len(func_combi_map[idx])]\n",
    "            funcs2 = func_combi_map[idx][-3 % len(func_combi_map[idx])]\n",
    "#             task_test = copy.deepcopy(task['test'])\n",
    "#             for f in funcs0:\n",
    "#                 task_test = f(task_test)\n",
    "#             st()\n",
    "        success = False\n",
    "        final_score_map[idx] = 0\n",
    "        for i, _ in enumerate(task['test']):\n",
    "            result[f'{idx}_{i}'] = ''\n",
    "        for funcs in [funcs0, funcs1, funcs2]:\n",
    "            task_test = copy.deepcopy(task['test'])\n",
    "            for func in funcs:\n",
    "                task_test = func(task_test)\n",
    "            for i, sample in enumerate(task_test):\n",
    "                if 'output' in sample:\n",
    "                    if sample['input'].tolist() == sample['output']:\n",
    "                        preprocess_best_score_map[idx] = 1.0\n",
    "                        final_score_map[idx] = 1.0\n",
    "                pred = flattener(sample['input'].tolist())\n",
    "                result[f'{idx}_{i}'] += pred + ' '\n",
    "\n",
    "    elif (len(func_combi_map[idx]) > 0) or (input_output_shape_is_same(task)):\n",
    "        task_train = copy.deepcopy(task['train'])\n",
    "        task_test = copy.deepcopy(task['test'])\n",
    "        if len(func_combi_map[idx]) == 0:\n",
    "            func_combi_map[idx].append([inouts_array])\n",
    "        for func in func_combi_map[idx][-1]:\n",
    "            task_train = func(task_train)\n",
    "            task_test = func(task_test)\n",
    "\n",
    "        task_train2 = copy.deepcopy(task['train'])\n",
    "        task_test2 = copy.deepcopy(task['test'])\n",
    "        funcs2 = func_combi_map[idx][-2 % len(func_combi_map[idx])]\n",
    "        for func in funcs2:\n",
    "            task_train2 = func(task_train2)\n",
    "            task_test2 = func(task_test2)\n",
    "        task_train_aug = copy.deepcopy(task_train)\n",
    "        print(f'preprocess_best_score: {preprocess_best_score}, funcs: {func_combi_map[idx]}')\n",
    "        if preprocess_best_score > 0.99:\n",
    "            promising_map[idx] = True\n",
    "        if preprocess_best_score > 0.7:\n",
    "            if 'output' in task_test[0]:\n",
    "                pre_preds = final_train_and_predict(task_train, task_train2, task_train_aug, task_test, task_test2, idx=idx, success_map={}, final_score_map=pre_final_score_map, origin_task=task)\n",
    "            use_transpose_flag = apply_transpose_aug(task_train)\n",
    "            color_inouts, color_aug_func = apply_color_aug(task_train, preprocess_best_score, best_aug_score_map, idx, promising_map)\n",
    "            print(f'color_aug_func: {color_aug_func}')\n",
    "            mirror_aug_funcs = apply_mirror_aug(task_train, preprocess_best_score, idx, use_transpose_flag, promising_map)\n",
    "            print(f'mirror_aug_funcs: {mirror_aug_funcs}')\n",
    "            # こちらも一応試したい\n",
    "            # mirror_augs = [flipud_aug, fliplr_aug, flip_aug, transpose_aug]\n",
    "            task_train_aug = task_train + color_inouts\n",
    "            for mirror_aug_func in mirror_aug_funcs:\n",
    "                mirror_inouts = mirror_aug_func(task_train)\n",
    "                task_train_aug += mirror_inouts\n",
    "#             st()\n",
    "        print(f'final_train_length： {len(task_train_aug)}')\n",
    "        preds = final_train_and_predict(task_train, task_train2, task_train_aug, task_test, task_test2, idx=idx, success_map={}, final_score_map=final_score_map, final=True, promising=promising_map[idx], origin_task=task)\n",
    "        for i, pred in enumerate(preds):\n",
    "            result[f'{idx}_{i}'] = pred\n",
    "\n",
    "    else:\n",
    "        task_test = copy.deepcopy(task['test'])\n",
    "        inputs = [el['input'] for el in task_test]\n",
    "        for i, inp in enumerate(inputs):\n",
    "            result[f'{idx}_{i}'] = getDefaultPred(inp)\n",
    "    t = time() - start\n",
    "    print(f'{round(t)}秒')\n",
    "    time_map[idx] = t\n",
    "    if (task_n in success_list):\n",
    "        print(f'-------------------------------------------------------------------------------------success!! idx: {idx}')\n",
    "    if env == 'production':\n",
    "        os.system(f'echo {idx}: best_score_with_preprocess: {preprocess_best_score_map.get(idx, \"different shapes...\")}')\n",
    "    return result, func_combi_map, success_list, preprocess_best_score_map, final_score_map, best_aug_score_map, pre_final_score_map, time_map, several_f\n",
    "\n",
    "ca_skip = True\n",
    "move_skip = True\n",
    "patch_skip = True\n",
    "update_tasks = []\n",
    "\n",
    "private = '00576224' not in data.keys()\n",
    "if private:\n",
    "    for idx, task in data.items():\n",
    "        preds, func_combi_map, success_list, preprocess_best_score_map, final_score_map, best_aug_score_map, pre_final_score_map, time_map, several_f = main(idx, task, 'production')\n",
    "        if len(success_list) != 0:\n",
    "            update_tasks.append(i)\n",
    "            for idx, pred_str in preds.items():\n",
    "                submission.loc[submission.index == idx, 'output'] = pred_str\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "submission.to_csv('submission.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}