diff --git a/src/sIArena/terrain/generator/Generator.py b/src/sIArena/terrain/generator/Generator.py
index 2d2e3fe..a8ca6a3 100644
--- a/src/sIArena/terrain/generator/Generator.py
+++ b/src/sIArena/terrain/generator/Generator.py
@@ -19,7 +19,9 @@ def generate_random_terrain(
                 abruptness: float = 0.2,
                 seed: int = None,
                 origin: Coordinate = None,
-                destination: Coordinate = None
+                destination: Coordinate = None,
+                terrain_ctor: Terrain = Terrain,
+                cost_function: callable = None
             ) -> Terrain:
         # Max and min abruptness
         abruptness = min(1, max(0, abruptness))
@@ -45,7 +47,10 @@ def generate_random_terrain(
 
         final_m *= min_step
 
-        return Terrain(final_m, origin=origin, destination=destination)
+        if cost_function is not None:
+            return terrain_ctor(final_m, origin=origin, destination=destination, cost_function=cost_function)
+        else:
+            return terrain_ctor(final_m, origin=origin, destination=destination)
 
 
     @pure_virtual
diff --git a/src/sIArena/terrain/generator/test.ipynb b/src/sIArena/terrain/generator/test.ipynb
deleted file mode 100644
index a0eb3ff..0000000
--- a/src/sIArena/terrain/generator/test.ipynb
+++ /dev/null
@@ -1,326 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 27,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "+----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----+\n",
-      "|                                  ||                                                                                        |\n",
-      "|                                  ||                                                                                        |\n",
-      "+    ++----++    ++    ++    ++    ++    ++----++----++    ++----++    ++----++    ++    ++----++----++----++    ++----++    +\n",
-      "+    ++----++    ++    ++    ++    ++    ++----++----++    ++----++    ++----++    ++    ++----++----++----++    ++----++    +\n",
-      "|          ||    ||    ||    ||                      ||          ||          ||    ||                      ||          ||    |\n",
-      "|          ||    ||    ||    ||                      ||          ||          ||    ||                      ||          ||    |\n",
-      "+    ++----++----++    ++    ++----++----++----++----++----++----++----++----++    ++    ++    ++    ++----++----++----++----+\n",
-      "+    ++----++----++    ++    ++----++----++----++----++----++----++----++----++    ++    ++    ++    ++----++----++----++----+\n",
-      "|    ||          ||    ||          ||          ||          ||          ||          ||    ||    ||                ||          |\n",
-      "|    ||          ||    ||          ||          ||          ||          ||          ||    ||    ||                ||          |\n",
-      "+    ++    ++    ++----++    ++----++    ++    ++    ++----++    ++    ++----++    ++    ++----++    ++----++    ++    ++----+\n",
-      "+    ++    ++    ++----++    ++----++    ++    ++    ++----++    ++    ++----++    ++    ++----++    ++----++    ++    ++----+\n",
-      "|          ||          ||                ||                ||    ||    ||    ||    ||          ||          ||                |\n",
-      "|          ||          ||                ||                ||    ||    ||    ||    ||          ||          ||                |\n",
-      "+    ++----++    ++    ++    ++----++----++----++----++----++    ++----++    ++----++----++----++----++    ++----++    ++----+\n",
-      "+    ++----++    ++    ++    ++----++----++----++----++----++    ++----++    ++----++----++----++----++    ++----++    ++----+\n",
-      "|          ||    ||    ||    ||                                                                      ||          ||          |\n",
-      "|          ||    ||    ||    ||                                                                      ||          ||          |\n",
-      "+    ++----++----++    ++    ++    ++----++    ++----++    ++----++----++    ++----++----++----++----++----++----++----++    +\n",
-      "+    ++----++----++    ++    ++    ++----++    ++----++    ++----++----++    ++----++----++----++----++----++----++----++    +\n",
-      "|          ||    ||    ||                ||    ||          ||          ||                ||          ||          ||          |\n",
-      "|          ||    ||    ||                ||    ||          ||          ||                ||          ||          ||          |\n",
-      "+    ++    ++    ++----++----++    ++----++----++----++----++    ++----++----++    ++----++----++    ++    ++    ++----++    +\n",
-      "+    ++    ++    ++----++----++    ++----++----++----++----++    ++----++----++    ++----++----++    ++    ++    ++----++    +\n",
-      "|    ||                ||                                                    ||                            ||    ||          |\n",
-      "|    ||                ||                                                    ||                            ||    ||          |\n",
-      "+----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----++----+\n",
-      "\n"
-     ]
-    }
-   ],
-   "source": [
-    "import random\n",
-    "random.seed(5)\n",
-    "\n",
-    "\n",
-    "class Tile:\n",
-    "\n",
-    "    def __init__(self):\n",
-    "        self.up = False\n",
-    "        self.down = False\n",
-    "        self.left = False\n",
-    "        self.right = False\n",
-    "        self.visited = False\n",
-    "\n",
-    "    def add_up(self):\n",
-    "        self.up = True\n",
-    "\n",
-    "    def add_down(self):\n",
-    "        self.down = True\n",
-    "\n",
-    "    def add_left(self):\n",
-    "        self.left = True\n",
-    "\n",
-    "    def add_right(self):\n",
-    "        self.right = True\n",
-    "\n",
-    "    def visit(self):\n",
-    "        self.visited = True\n",
-    "\n",
-    "\n",
-    "\n",
-    "class Maze:\n",
-    "\n",
-    "    def __init__(self, n, m):\n",
-    "        self.n = n\n",
-    "        self.m = m\n",
-    "        self.matrix = [[Tile() for _ in range(m)] for _ in range(n)]\n",
-    "\n",
-    "        current = (0, 0)\n",
-    "        self.matrix[current[0]][current[1]].visit()\n",
-    "        to_visit = self.surrounding_coordinates(current)\n",
-    "\n",
-    "\n",
-    "        while to_visit:\n",
-    "\n",
-    "            # Get a random tile from to_visit list\n",
-    "            r = random.randint(0, len(to_visit) - 1)\n",
-    "            next_tile = to_visit.pop(r)\n",
-    "\n",
-    "            if self.tile(next_tile).visited:\n",
-    "                continue\n",
-    "\n",
-    "            self.matrix[next_tile[0]][next_tile[1]].visit()\n",
-    "\n",
-    "            # Get all surrounding tiles\n",
-    "            surrounding = self.surrounding_coordinates(next_tile)\n",
-    "\n",
-    "            # Separate visited from non visited\n",
-    "            visited = []\n",
-    "            non_visited = []\n",
-    "            for s in surrounding:\n",
-    "                if self.tile(s).visited:\n",
-    "                    visited.append(s)\n",
-    "                else:\n",
-    "                    non_visited.append(s)\n",
-    "\n",
-    "            # Join with a random visited one\n",
-    "            if visited:\n",
-    "                r = random.randint(0, len(visited) - 1)\n",
-    "                visited_tile = visited[r]\n",
-    "                self.join_tiles(next_tile, visited_tile)\n",
-    "                self.tile(visited_tile).visit()\n",
-    "            else:\n",
-    "                raise Exception(\"No visited tiles around\")\n",
-    "\n",
-    "            # Add non visited to to_visit\n",
-    "            for s in non_visited:\n",
-    "                to_visit.append(s)\n",
-    "\n",
-    "\n",
-    "    def tile(self, coor):\n",
-    "        return self.matrix[coor[0]][coor[1]]\n",
-    "\n",
-    "    def surrounding_coordinates(self, coor):\n",
-    "        n, m = coor\n",
-    "        coords = []\n",
-    "        if n > 0:\n",
-    "            coords.append((n - 1, m))\n",
-    "        if n < self.n - 1:\n",
-    "            coords.append((n + 1, m))\n",
-    "        if m > 0:\n",
-    "            coords.append((n, m - 1))\n",
-    "        if m < self.m - 1:\n",
-    "            coords.append((n, m + 1))\n",
-    "        return coords\n",
-    "\n",
-    "    def join_tiles(self, tile1, tile2):\n",
-    "        # Check if tile2 is up, down, left or right of tile1\n",
-    "        n1, m1 = tile1\n",
-    "        n2, m2 = tile2\n",
-    "\n",
-    "        if n1 == n2:\n",
-    "            if m1 < m2:\n",
-    "                self.matrix[n1][m1].add_right()\n",
-    "                self.matrix[n2][m2].add_left()\n",
-    "            else:\n",
-    "                self.matrix[n1][m1].add_left()\n",
-    "                self.matrix[n2][m2].add_right()\n",
-    "\n",
-    "        elif m1 == m2:\n",
-    "            if n1 < n2:\n",
-    "                self.matrix[n1][m1].add_down()\n",
-    "                self.matrix[n2][m2].add_up()\n",
-    "            else:\n",
-    "                self.matrix[n1][m1].add_up()\n",
-    "                self.matrix[n2][m2].add_down()\n",
-    "\n",
-    "        else:\n",
-    "            raise Exception(\"Tiles are not adjacent\")\n",
-    "\n",
-    "    def __str__(self):\n",
-    "        s = \"\"\n",
-    "        for i in range(self.n):\n",
-    "            for j in range(self.m):\n",
-    "                if self.matrix[i][j].up:\n",
-    "                    s += \"+    +\"\n",
-    "                else:\n",
-    "                    s += \"+----+\"\n",
-    "            s += \"\\n\"\n",
-    "            for _ in range(2):\n",
-    "                for j in range(self.m):\n",
-    "                    if self.matrix[i][j].left:\n",
-    "                        s += \" \"\n",
-    "                    else:\n",
-    "                        s += \"|\"\n",
-    "                    s += \"    \"\n",
-    "                    if self.matrix[i][j].right:\n",
-    "                        s += \" \"\n",
-    "                    else:\n",
-    "                        s += \"|\"\n",
-    "                s += \"\\n\"\n",
-    "            for j in range(self.m):\n",
-    "                if self.matrix[i][j].down:\n",
-    "                    s += \"+    +\"\n",
-    "                else:\n",
-    "                    s += \"+----+\"\n",
-    "            s += \"\\n\"\n",
-    "\n",
-    "        return s\n",
-    "\n",
-    "\n",
-    "m = Maze(7, 21)\n",
-    "print(m)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 46,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import numpy as np\n",
-    "\n",
-    "def generate_random_matrix_(\n",
-    "            n: int,\n",
-    "            m: int,\n",
-    "            abruptness: float = 0.5,\n",
-    "            seed: int = None,\n",
-    "            origin = None,\n",
-    "            destination = None\n",
-    "        ) -> np.matrix:\n",
-    "    \"\"\"\n",
-    "    TODO\n",
-    "    \"\"\"\n",
-    "\n",
-    "    if seed is not None:\n",
-    "        random.seed(seed)\n",
-    "\n",
-    "    maze_n = (1 + n) // 2\n",
-    "    maze_m = (1 + m) // 2\n",
-    "\n",
-    "    maze = Maze(maze_n, maze_m)\n",
-    "\n",
-    "    matrix = np.ones((n, m))\n",
-    "    max_height = n*m\n",
-    "\n",
-    "    print(f\"n: {n}, m: {m}, maze_n: {maze_n}, maze_m: {maze_m}, max_height: {max_height}\")\n",
-    "    print(maze)\n",
-    "\n",
-    "    for i in range(maze_n):\n",
-    "        for j in range(maze_m):\n",
-    "            up_left = (2*i, 2*j)        # always 0\n",
-    "            down_right = (2*i+1, 2*j+1) # always 1\n",
-    "            up_right = (2*i, 2*j+1)     # 1 when no right wall\n",
-    "            down_left = (2*i+1, 2*j)    # 1 when no down wall\n",
-    "\n",
-    "            # print(f\"i: {i}, j: {j}, up_left: {up_left}, down_right: {down_right}, up_right: {up_right}, down_left: {down_left}\")\n",
-    "\n",
-    "            if 2*j+1 < m and not maze.tile((i, j)).right:\n",
-    "                matrix[up_right] = max_height\n",
-    "            if 2*i+1 < n and not maze.tile((i, j)).down:\n",
-    "                matrix[down_left] = max_height\n",
-    "            if 2*i+1 < n and 2*j+1 < m:\n",
-    "                matrix[down_right] = max_height\n",
-    "\n",
-    "    return matrix"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 49,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "n: 7, m: 7, maze_n: 4, maze_m: 4, max_height: 49\n",
-      "+----++----++----++----+\n",
-      "|                      |\n",
-      "|                      |\n",
-      "+    ++----++    ++    +\n",
-      "+    ++----++    ++    +\n",
-      "|          ||    ||    |\n",
-      "|          ||    ||    |\n",
-      "+    ++    ++----++    +\n",
-      "+    ++    ++----++    +\n",
-      "|    ||          ||    |\n",
-      "|    ||          ||    |\n",
-      "+----++    ++    ++    +\n",
-      "+----++    ++    ++    +\n",
-      "|          ||    ||    |\n",
-      "|          ||    ||    |\n",
-      "+----++----++----++----+\n",
-      "\n",
-      "[[ 1.  1.  1.  1.  1.  1.  1.]\n",
-      " [ 1. 49. 49. 49.  1. 49.  1.]\n",
-      " [ 1.  1.  1. 49.  1. 49.  1.]\n",
-      " [ 1. 49.  1. 49. 49. 49.  1.]\n",
-      " [ 1. 49.  1.  1.  1. 49.  1.]\n",
-      " [49. 49.  1. 49.  1. 49.  1.]\n",
-      " [ 1.  1.  1. 49.  1. 49.  1.]]\n"
-     ]
-    }
-   ],
-   "source": [
-    "seed = 0\n",
-    "random.seed(seed)\n",
-    "\n",
-    "g = generate_random_matrix_(7, 7, seed=seed)\n",
-    "print(g)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": ".default_venv",
-   "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.12.3"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}