{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "61d0f54d-c069-41d3-a139-f85492b6267b",
   "metadata": {},
   "source": [
    "# Nonogram Solver"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9c68d4a-6dd2-452f-b87c-5a1f0872ccde",
   "metadata": {},
   "source": [
    "## Setup and Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2cc41aff-bd5b-4c5b-9b91-84cff5d5d534",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "from tqdm.notebook import tqdm\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import os\n",
    "import glob\n",
    "\n",
    "# Configure device to use GPU if available, otherwise use CPU\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d35719ec-2f74-4d28-a29c-84d759c3dceb",
   "metadata": {},
   "source": [
    "## Helper Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "700617b0-cee2-4a35-8ea0-ba5d1a73f84c",
   "metadata": {},
   "source": [
    "### Nonogram Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0faad77e-0a1a-489b-b1e5-1840be042c17",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_unique_nonogram(grid_size, batch_size, existing_solutions=set()):\n",
    "    \"\"\"\n",
    "    Generate unique Nonogram puzzles with corresponding clues.\n",
    "    \n",
    "    Parameters:\n",
    "    - grid_size (int): Size of the Nonogram grid.\n",
    "    - batch_size (int): Number of Nonogram puzzles to generate.\n",
    "    - existing_solutions (set): Set of existing solutions to avoid duplicates.\n",
    "    \n",
    "    Returns:\n",
    "    - solutions (ndarray): Generated Nonogram solutions.\n",
    "    - row_clues (list): Row clues for the Nonogram puzzles.\n",
    "    - col_clues (list): Column clues for the Nonogram puzzles.\n",
    "    - existing_solutions (set): Updated set of existing solutions.\n",
    "    \"\"\"\n",
    "    solutions = []\n",
    "    while len(solutions) < batch_size:\n",
    "        new_solutions = np.random.randint(2, size=(batch_size, grid_size, grid_size))\n",
    "        for solution in new_solutions:\n",
    "            solution_tuple = tuple(map(tuple, solution))\n",
    "            if solution_tuple not in existing_solutions:\n",
    "                solutions.append(solution)\n",
    "                existing_solutions.add(solution_tuple)\n",
    "            if len(solutions) == batch_size:\n",
    "                break\n",
    "    solutions = np.array(solutions)\n",
    "    row_clues = [[list(map(len, ''.join(map(str, row)).split('0'))) for row in solution] for solution in solutions]\n",
    "    col_clues = [[list(map(len, ''.join(map(str, col)).split('0'))) for col in solution.T] for solution in solutions]\n",
    "    row_clues = [[[clue for clue in clues if clue > 0] or [0] for clues in row] for row in row_clues]\n",
    "    col_clues = [[[clue for clue in clues if clue > 0] or [0] for clues in col] for col in col_clues]\n",
    "    \n",
    "    return solutions, row_clues, col_clues, existing_solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2e6469e-82bf-4b20-b800-b839d89beff7",
   "metadata": {},
   "source": [
    "### Clue Padding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f9f0ac6d-2d41-4184-9c15-2a11b1571afc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_clues(clues, max_len):\n",
    "    \"\"\"\n",
    "    Pad clues to the maximum length.\n",
    "    \n",
    "    Parameters:\n",
    "    - clues (list): List of clues to pad.\n",
    "    - max_len (int): Maximum length to pad the clues to.\n",
    "    \n",
    "    Returns:\n",
    "    - padded_clues (list): Padded clues.\n",
    "    \"\"\"\n",
    "    return [clue + [0] * (max_len - len(clue)) for clue in clues]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46a57583-6cb8-4c56-b5aa-b019f6cc947d",
   "metadata": {},
   "source": [
    "### Correct Guess Calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4d534101-90af-4d96-b739-2f2f38b7f904",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_correct_guesses(states, solutions):\n",
    "    \"\"\"\n",
    "    Calculate the number of correct guesses.\n",
    "    \n",
    "    Parameters:\n",
    "    - states (ndarray): Current states of the Nonogram grids.\n",
    "    - solutions (ndarray): Solution grids of the Nonogram puzzles.\n",
    "    \n",
    "    Returns:\n",
    "    - correct_guesses (ndarray): Number of correct guesses.\n",
    "    \"\"\"\n",
    "    return np.sum(states == solutions, axis=(1, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02c6055a-dfc1-4ea4-9cb4-6253e54bfff5",
   "metadata": {},
   "source": [
    "### Checkpoint Saving"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "79470031-bd71-43f3-8e33-b20cee93ff72",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_checkpoint(agent, optimizer, episode, reward_list, correct_guess_percent_list, clue_max_len, clue_dim, directory='models'):\n",
    "    \"\"\"\n",
    "    Save the training checkpoint.\n",
    "    \n",
    "    Parameters:\n",
    "    - agent (NonogramAgent): The agent being trained.\n",
    "    - optimizer (torch.optim.Optimizer): Optimizer used for training.\n",
    "    - episode (int): Current episode number.\n",
    "    - reward_list (list): List of rewards.\n",
    "    - correct_guess_percent_list (list): List of correct guess percentages.\n",
    "    - clue_max_len (int): Maximum length of the clues.\n",
    "    - clue_dim (int): Dimensionality of the clues.\n",
    "    - directory (str): Directory to save the checkpoint.\n",
    "    \"\"\"\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)\n",
    "    \n",
    "    checkpoint_path = os.path.join(directory, f'checkpoint_{episode}.pth')\n",
    "    \n",
    "    torch.save({\n",
    "        'episode': episode,\n",
    "        'model_state_dict': agent.policy_net.state_dict(),\n",
    "        'optimizer_state_dict': optimizer.state_dict(),\n",
    "        'reward_list': reward_list,\n",
    "        'correct_guess_percent_list': correct_guess_percent_list,\n",
    "        'clue_max_len': clue_max_len,\n",
    "        'clue_dim': clue_dim\n",
    "    }, checkpoint_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "176cb831-8dcd-477c-914e-c743ecada642",
   "metadata": {},
   "source": [
    "### Checkpoint Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "75e185cc-029a-487d-9fd6-781ce13568d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_checkpoint(agent, optimizer, directory='models'):\n",
    "    \"\"\"\n",
    "    Load the latest training checkpoint.\n",
    "    \n",
    "    Parameters:\n",
    "    - agent (NonogramAgent): The agent being trained.\n",
    "    - optimizer (torch.optim.Optimizer): Optimizer used for training.\n",
    "    - directory (str): Directory to load the checkpoint from.\n",
    "    \n",
    "    Returns:\n",
    "    - episode (int): Episode number to resume from.\n",
    "    - reward_list (list): List of rewards.\n",
    "    - correct_guess_percent_list (list): List of correct guess percentages.\n",
    "    - clue_max_len (int): Maximum length of the clues.\n",
    "    - clue_dim (int): Dimensionality of the clues.\n",
    "    \"\"\"\n",
    "    checkpoints = sorted(glob.glob(os.path.join(directory, 'checkpoint_*.pth')), key=lambda x: int(x.split('_')[-1].split('.')[0]), reverse=True)\n",
    "    if checkpoints:\n",
    "        checkpoint = torch.load(checkpoints[0])\n",
    "        agent.policy_net.load_state_dict(checkpoint['model_state_dict'])\n",
    "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "        return checkpoint['episode'], checkpoint['reward_list'], checkpoint['correct_guess_percent_list'], checkpoint['clue_max_len'], checkpoint['clue_dim']\n",
    "    \n",
    "    return 0, [], [], None, None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "576367aa-ed0c-4327-8821-5fac05cfcbfd",
   "metadata": {},
   "source": [
    "### Reward Discounting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "76b68d5a-114a-4aa7-8ec2-eba43b2abb44",
   "metadata": {},
   "outputs": [],
   "source": [
    "def discount_rewards(rewards, gamma=0.995):\n",
    "    \"\"\"\n",
    "    Compute discounted rewards.\n",
    "    \n",
    "    Parameters:\n",
    "    - rewards (list): List of rewards.\n",
    "    - gamma (float): Discount factor.\n",
    "    \n",
    "    Returns:\n",
    "    - discounted_rewards (list): List of discounted rewards.\n",
    "    \"\"\"\n",
    "    discounted_rewards = []\n",
    "    for reward in rewards:\n",
    "        cumulative_rewards = 0\n",
    "        discounted = []\n",
    "        for r in reversed(reward):\n",
    "            cumulative_rewards = r + gamma * cumulative_rewards\n",
    "            discounted.insert(0, cumulative_rewards)\n",
    "        discounted_rewards.append(torch.tensor(discounted, dtype=torch.float32).to(device))\n",
    "    return discounted_rewards"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a37dc65d-c0c8-44fc-b504-ff237d8deb97",
   "metadata": {},
   "source": [
    "### Divide and Round Up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c060b09d-0e98-4252-ac57-65d16178efed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def divide_and_round_up(n):\n",
    "    \"\"\"\n",
    "    Divide and round up the number.\n",
    "    \n",
    "    Parameters:\n",
    "    - n (int): Number to divide and round up.\n",
    "    \n",
    "    Returns:\n",
    "    - result (int): Result of the division and rounding up.\n",
    "    \"\"\"\n",
    "    return (n + 1) // 2 if n % 2 != 0 else n // 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d107cc88-5177-45d8-8b21-82ffbe7c1d1b",
   "metadata": {},
   "source": [
    "### Visualize Nonogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9e05f4fd-e151-4f3f-9001-efc5af6606e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_nonogram(board):\n",
    "    \"\"\"\n",
    "    Visualize the Nonogram board.\n",
    "    \n",
    "    Parameters:\n",
    "    - board (ndarray): The current state of the Nonogram board.\n",
    "    \"\"\"\n",
    "    grid_size = len(board)\n",
    "    for row in range(grid_size):\n",
    "        print(\" \".join(str(cell) if cell != -1 else \"?\" for cell in board[row]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eccd28c-c0ed-41f7-b534-ebdf46b1d1a0",
   "metadata": {},
   "source": [
    "### Visualize Clues"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b8afd7a5-d2fd-4794-9f8f-14b2d9d6057f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_clues(clues):\n",
    "    \"\"\"\n",
    "    Visualize the clues for the Nonogram puzzle.\n",
    "    \n",
    "    Parameters:\n",
    "    - clues (list): List of clues for the puzzle.\n",
    "    \"\"\"\n",
    "    for clue in clues:\n",
    "        clue = [c for c in clue if c != 0]  # Remove padding zeros\n",
    "        if not clue:  # If no clues, add a single zero\n",
    "            clue = [0]\n",
    "        print(clue)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e786985c-28b6-470a-9bf4-c0614e7c6224",
   "metadata": {},
   "source": [
    "### Update Puzzle State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c7862751-7b5d-4ef6-8117-e6a277bae9c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_puzzle_state(agent, env, states, row_clues, col_clues, solutions):\n",
    "    \"\"\"\n",
    "    Update the puzzle state based on the agent's actions.\n",
    "    \n",
    "    Parameters:\n",
    "    - agent (NonogramAgent): The agent solving the puzzle.\n",
    "    - env (NonogramEnvironment): The environment of the Nonogram puzzle.\n",
    "    - states (ndarray): The current states of the puzzles.\n",
    "    - row_clues (list): The row clues for the puzzles.\n",
    "    - col_clues (list): The column clues for the puzzles.\n",
    "    - solutions (ndarray): The solutions for the puzzles.\n",
    "    \"\"\"\n",
    "    move_counter = 0\n",
    "    done = False\n",
    "    while not done:\n",
    "        actions, _ = agent.select_actions(states, row_clues, col_clues)\n",
    "        states, rewards, done = env.step(actions)\n",
    "        done = done[0]  # Since batch_size is 1\n",
    "\n",
    "        move = actions[0]\n",
    "        row, col, state_value = move\n",
    "        move_counter += 1\n",
    "        print(f\"\\nMove: {move_counter}, Guess: ({row + 1}, {col + 1}), State: {'1' if state_value == 1 else '0'}\")\n",
    "        \n",
    "        # Visualize current state\n",
    "        print(\"Current Puzzle State:\")\n",
    "        visualize_nonogram(states[0])\n",
    "\n",
    "        # Comparison with actual solution\n",
    "        current_state = states[0]\n",
    "        actual_state = solutions[0]\n",
    "        \n",
    "        # Calculate mismatches excluding cells with -1\n",
    "        mismatches = np.argwhere((current_state != actual_state) & (current_state != -1) & (actual_state != -1))\n",
    "        if mismatches.size > 0:\n",
    "            print(\"Mismatches:\")\n",
    "            for (i, j) in mismatches:\n",
    "                print(f\"  Mismatch at ({i}, {j}): Algorithm State = {current_state[i][j]}, Actual State = {actual_state[i][j]}\")\n",
    "\n",
    "    print(\"Puzzle Solved!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2228089-657a-4af1-b4ee-664850796383",
   "metadata": {},
   "source": [
    "## Nonogram Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8abc1188-e531-4feb-a336-92e4a2da152d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NonogramEnvironment:\n",
    "    def __init__(self, grid_size, batch_size, streak_cap=5):\n",
    "        \"\"\"\n",
    "        Initialize the Nonogram environment.\n",
    "        \n",
    "        Parameters:\n",
    "        - grid_size (int): Size of the Nonogram grid.\n",
    "        - batch_size (int): Number of puzzles in a batch.\n",
    "        - streak_cap (int): Maximum streak for unique guesses.\n",
    "        \"\"\"\n",
    "        self.grid_size = grid_size\n",
    "        self.batch_size = batch_size\n",
    "        self.streak_cap = streak_cap\n",
    "        self.solution, self.row_clues, self.col_clues = self.generate_initial_nonogram(grid_size, batch_size)\n",
    "        self.state = np.full((batch_size, grid_size, grid_size), -1, dtype=int)\n",
    "        self.steps = np.zeros(batch_size, dtype=int)\n",
    "        self.chosen_cells = [set() for _ in range(batch_size)]\n",
    "        self.correct_guesses = [set() for _ in range(batch_size)]\n",
    "        self.unique_guesses_streak = np.zeros(batch_size, dtype=int)\n",
    "\n",
    "    def generate_initial_nonogram(self, grid_size, batch_size):\n",
    "        \"\"\"\n",
    "        Generate initial Nonogram puzzles.\n",
    "        \n",
    "        Parameters:\n",
    "        - grid_size (int): Size of the Nonogram grid.\n",
    "        - batch_size (int): Number of puzzles in a batch.\n",
    "        \n",
    "        Returns:\n",
    "        - solution, row_clues, col_clues (tuple): Initial puzzles and their clues.\n",
    "        \"\"\"\n",
    "        return generate_unique_nonogram(grid_size, batch_size)[0:3]\n",
    "\n",
    "    def reset(self):\n",
    "        \"\"\"\n",
    "        Reset the Nonogram environment to its initial state.\n",
    "        \n",
    "        Returns:\n",
    "        - state, row_clues, col_clues (tuple): Reset state and clues.\n",
    "        \"\"\"\n",
    "        self.state = np.full((self.batch_size, self.grid_size, self.grid_size), -1, dtype=int)\n",
    "        self.steps = np.zeros(self.batch_size, dtype=int)\n",
    "        self.chosen_cells = [set() for _ in range(self.batch_size)]\n",
    "        self.correct_guesses = [set() for _ in range(self.batch_size)]\n",
    "        self.unique_guesses_streak = np.zeros(self.batch_size, dtype=int)\n",
    "        return self.state, self.row_clues, self.col_clues\n",
    "\n",
    "    def reset_with_solutions(self, solutions, row_clues, col_clues):\n",
    "        \"\"\"\n",
    "        Reset the environment with specified solutions and clues.\n",
    "        \n",
    "        Parameters:\n",
    "        - solutions (ndarray): Solution grids of the Nonogram puzzles.\n",
    "        - row_clues (list): Row clues for the Nonogram puzzles.\n",
    "        - col_clues (list): Column clues for the Nonogram puzzles.\n",
    "        \n",
    "        Returns:\n",
    "        - state, row_clues, col_clues (tuple): Reset state and clues.\n",
    "        \"\"\"\n",
    "        self.solution = solutions\n",
    "        self.row_clues = row_clues\n",
    "        self.col_clues = col_clues\n",
    "        return self.reset()\n",
    "\n",
    "    def step(self, actions):\n",
    "        \"\"\"\n",
    "        Take a step in the Nonogram environment.\n",
    "        \n",
    "        Parameters:\n",
    "        - actions (list): List of actions to take.\n",
    "        \n",
    "        Returns:\n",
    "        - state, rewards, done (tuple): Updated state, rewards, and done flags.\n",
    "        \"\"\"\n",
    "        rewards = np.zeros(self.batch_size, dtype=float)\n",
    "        done = np.zeros(self.batch_size, dtype=bool)\n",
    "\n",
    "        for i, action in enumerate(actions):\n",
    "            row, col, value = action\n",
    "            self.steps[i] += 1\n",
    "\n",
    "            if (row, col) in self.chosen_cells[i]:\n",
    "                rewards[i] = -5\n",
    "                self.unique_guesses_streak[i] = 0\n",
    "            else:\n",
    "                self.chosen_cells[i].add((row, col))\n",
    "                self.unique_guesses_streak[i] += 1\n",
    "                rewards[i] = min(self.unique_guesses_streak[i], self.streak_cap)\n",
    "\n",
    "                if self.solution[i, row, col] == value:\n",
    "                    rewards[i] += 2\n",
    "                    self.correct_guesses[i].add((row, col))\n",
    "                else:\n",
    "                    rewards[i] -= 2\n",
    "\n",
    "                self.state[i, row, col] = self.solution[i, row, col]\n",
    "                \n",
    "                if all(self.state[i, row, c] != -1 for c in range(self.grid_size)) and \\\n",
    "                   all(self.state[i, row, c] == self.solution[i, row, c] for c in range(self.grid_size)):\n",
    "                    rewards[i] += 10\n",
    "\n",
    "                if all(self.state[i, r, col] != -1 for r in range(self.grid_size)) and \\\n",
    "                   all(self.state[i, r, col] == self.solution[i, r, col] for r in range(self.grid_size)):\n",
    "                    rewards[i] += 10\n",
    "\n",
    "                if all(self.state[i, r, c] != -1 for r in range(self.grid_size) for c in range(self.grid_size)) and \\\n",
    "                   all(self.state[i, r, c] == self.solution[i, r, c] for r in range(self.grid_size) for c in range(self.grid_size)):\n",
    "                    rewards[i] += 100\n",
    "\n",
    "            done[i] = self._check_done(i)\n",
    "        return self.state, rewards, done\n",
    "\n",
    "    def _check_done(self, index):\n",
    "        \"\"\"\n",
    "        Check if the puzzle is solved or maximum steps reached.\n",
    "        \n",
    "        Parameters:\n",
    "        - index (int): Index of the puzzle.\n",
    "        \n",
    "        Returns:\n",
    "        - done (bool): Whether the puzzle is solved or maximum steps reached.\n",
    "        \"\"\"\n",
    "        return np.array_equal(self.state[index], self.solution[index]) or self.steps[index] >= self.grid_size ** 2\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04433ebf-4800-43c3-a35b-7754c7be09bb",
   "metadata": {},
   "source": [
    "## Model Definition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7105171c-e578-4975-9fa9-df33552855f1",
   "metadata": {},
   "source": [
    "### Clue Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6d3429ef-e33a-44f0-a48d-b6831f0a9997",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClueTransformer(nn.Module):\n",
    "    def __init__(self, grid_size, clue_max_len, clue_dim, num_heads, num_layers, model_dim):\n",
    "        \"\"\"\n",
    "        Initialize the Clue Transformer.\n",
    "        \n",
    "        Parameters:\n",
    "        - grid_size (int): Size of the Nonogram grid.\n",
    "        - clue_max_len (int): Maximum length of the clues.\n",
    "        - clue_dim (int): Dimensionality of the clues.\n",
    "        - num_heads (int): Number of attention heads.\n",
    "        - num_layers (int): Number of transformer layers.\n",
    "        - model_dim (int): Dimensionality of the model.\n",
    "        \"\"\"\n",
    "        super(ClueTransformer, self).__init__()\n",
    "        self.grid_size = grid_size\n",
    "        self.embedding = nn.Embedding(clue_dim + 1, model_dim)\n",
    "        self.model_dim = model_dim\n",
    "        self.positional_encoding = nn.Parameter(torch.randn(1, clue_max_len*grid_size, model_dim))\n",
    "        encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=num_heads, batch_first=True)\n",
    "        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)\n",
    "\n",
    "    def forward(self, clues):\n",
    "        \"\"\"\n",
    "        Forward pass of the Clue Transformer.\n",
    "        \n",
    "        Parameters:\n",
    "        - clues (Tensor): Clues for the Nonogram puzzles.\n",
    "        \n",
    "        Returns:\n",
    "        - transformer_output (Tensor): Output of the transformer.\n",
    "        \"\"\"\n",
    "        batch_size, num_clues, clue_len = clues.size()\n",
    "        clues = clues.view(batch_size, -1)\n",
    "\n",
    "        # Ensure clues are within the valid range\n",
    "        assert torch.max(clues) <= self.embedding.num_embeddings - 1, f\"Index {torch.max(clues)} out of bounds for embedding with size {self.embedding.num_embeddings}\"\n",
    "\n",
    "        embedded_clues = self.embedding(clues)\n",
    "        \n",
    "        max_len = embedded_clues.size(1)\n",
    "        if self.positional_encoding.size(1) < max_len:\n",
    "            self.positional_encoding = nn.Parameter(torch.randn(1, max_len, self.model_dim).to(embedded_clues.device))\n",
    "\n",
    "        embedded_clues = embedded_clues + self.positional_encoding[:, :embedded_clues.size(1), :]\n",
    "        transformer_output = self.transformer(embedded_clues)\n",
    "        return transformer_output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21bb870b-7598-4cfd-afce-9c46b3d1ad99",
   "metadata": {},
   "source": [
    "### Policy Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0f4e00e3-2cc9-458a-b80c-b6fbb906c555",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PolicyNetwork(nn.Module):\n",
    "    def __init__(self, grid_size, clue_max_len, clue_dim):\n",
    "        \"\"\"\n",
    "        Initialize the Policy Network.\n",
    "        \n",
    "        Parameters:\n",
    "        - grid_size (int): Size of the Nonogram grid.\n",
    "        - clue_max_len (int): Maximum length of the clues.\n",
    "        - clue_dim (int): Dimensionality of the clues.\n",
    "        \"\"\"\n",
    "        super(PolicyNetwork, self).__init__()\n",
    "        self.grid_size = grid_size\n",
    "        self.conv1 = nn.Conv2d(1, 4, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv2 = nn.Conv2d(4, 8, kernel_size=3, stride=1, padding=1)\n",
    "        self.fc1 = nn.Linear(8 * grid_size * grid_size, 16)\n",
    "        self.row_clue_transformer = ClueTransformer(grid_size, clue_max_len, clue_dim, num_heads=2, num_layers=1, model_dim=16)\n",
    "        self.col_clue_transformer = ClueTransformer(grid_size, clue_max_len, clue_dim, num_heads=2, num_layers=1, model_dim=16)\n",
    "        self.fc2 = nn.Linear(16 * 2 + 16, 32)\n",
    "        self.fc3 = nn.Linear(32, grid_size * grid_size * 2)\n",
    "\n",
    "    def forward(self, state, row_clues, col_clues):\n",
    "        \"\"\"\n",
    "        Forward pass of the Policy Network.\n",
    "        \n",
    "        Parameters:\n",
    "        - state (Tensor): Current state of the Nonogram puzzles.\n",
    "        - row_clues (Tensor): Row clues for the Nonogram puzzles.\n",
    "        - col_clues (Tensor): Column clues for the Nonogram puzzles.\n",
    "        \n",
    "        Returns:\n",
    "        - output (Tensor): Action logits for the Nonogram puzzles.\n",
    "        \"\"\"\n",
    "        state = state.to(device)\n",
    "        row_clues = row_clues.to(device)\n",
    "        col_clues = col_clues.to(device)\n",
    "    \n",
    "        x = state.unsqueeze(1).float()\n",
    "        x = F.relu(self.conv1(x))\n",
    "        x = F.relu(self.conv2(x))\n",
    "        x = x.view(-1, 8 * self.grid_size * self.grid_size)\n",
    "        x = F.relu(self.fc1(x))\n",
    "    \n",
    "        row_clues[row_clues >= clue_dim + 1] = clue_dim\n",
    "        col_clues[row_clues >= clue_dim + 1] = clue_dim\n",
    "    \n",
    "        row_clues = self.row_clue_transformer(row_clues).mean(dim=1)\n",
    "        col_clues = self.col_clue_transformer(col_clues).mean(dim=1)\n",
    "    \n",
    "        clues = torch.cat((row_clues, col_clues), dim=1)\n",
    "    \n",
    "        x = torch.cat((x, clues), dim=1)\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x.view(-1, self.grid_size * self.grid_size, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ef3db61-be52-4aa1-b76c-d16ce99d9008",
   "metadata": {},
   "source": [
    "## Nonogram Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "cf6434a0-dd6b-429d-a7d8-69d089d888bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NonogramAgent:\n",
    "    def __init__(self, grid_size, clue_max_len, clue_dim):\n",
    "        \"\"\"\n",
    "        Initialize the Nonogram Agent.\n",
    "        \n",
    "        Parameters:\n",
    "        - grid_size (int): Size of the Nonogram grid.\n",
    "        - clue_max_len (int): Maximum length of the clues.\n",
    "        - clue_dim (int): Dimensionality of the clues.\n",
    "        \"\"\"\n",
    "        self.policy_net = PolicyNetwork(grid_size, clue_max_len, clue_dim).to(device)\n",
    "        self.optimizer = optim.Adam(self.policy_net.parameters(), lr=0.001)\n",
    "        self.grid_size = grid_size\n",
    "\n",
    "    def select_actions(self, states, row_clues, col_clues):\n",
    "        \"\"\"\n",
    "        Select actions based on the current state and clues.\n",
    "        \n",
    "        Parameters:\n",
    "        - states (ndarray): Current states of the Nonogram puzzles.\n",
    "        - row_clues (ndarray): Row clues for the Nonogram puzzles.\n",
    "        - col_clues (ndarray): Column clues for the Nonogram puzzles.\n",
    "        \n",
    "        Returns:\n",
    "        - actions (list): List of selected actions.\n",
    "        - log_probs (Tensor): Log probabilities of the selected actions.\n",
    "        \"\"\"\n",
    "        states = torch.tensor(states, dtype=torch.float32).to(device)\n",
    "        row_clues = torch.tensor(row_clues, dtype=torch.long).to(device)\n",
    "        col_clues = torch.tensor(col_clues, dtype=torch.long).to(device)\n",
    "        logits = self.policy_net(states, row_clues, col_clues)\n",
    "                \n",
    "        action_probs = torch.softmax(logits.view(states.size(0), -1), dim=-1)\n",
    "                \n",
    "        action_dist = torch.distributions.Categorical(action_probs)\n",
    "        flat_actions = action_dist.sample()\n",
    "        log_probs = action_dist.log_prob(flat_actions)\n",
    "\n",
    "        actions = []\n",
    "        for flat_action in flat_actions:\n",
    "            flat_action_idx = flat_action.item()\n",
    "            position_idx = flat_action_idx // 2\n",
    "            value = flat_action_idx % 2\n",
    "            row = position_idx // self.grid_size\n",
    "            col = position_idx % self.grid_size\n",
    "            actions.append((row, col, value))\n",
    "\n",
    "        return actions, log_probs\n",
    "\n",
    "    def update_policy(self, log_probs, rewards):\n",
    "        \"\"\"\n",
    "        Update the policy network based on the collected log probabilities and rewards.\n",
    "        \n",
    "        Parameters:\n",
    "        - log_probs (list): Log probabilities of the selected actions.\n",
    "        - rewards (list): Rewards obtained from the actions.\n",
    "        \"\"\"\n",
    "        discounted_rewards = discount_rewards(rewards)\n",
    "        \n",
    "        discounted_rewards = torch.cat(discounted_rewards)\n",
    "        discounted_rewards = (discounted_rewards - discounted_rewards.mean()) / (discounted_rewards.std() + 1e-8)\n",
    "\n",
    "        log_probs = torch.cat([torch.stack(lps) for lps in log_probs])\n",
    "\n",
    "        loss = -torch.sum(log_probs * discounted_rewards.to(device))\n",
    "        self.optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        self.optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0789787e-7299-4e4c-a633-7fa00f23fe78",
   "metadata": {},
   "source": [
    "## Training Procedure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a7ef69b7-51bd-407b-8e49-03e129e388ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_agent(grid_size, clue_max_len, clue_dim, num_episodes, train_batch_size, val_batch_size, save_interval=100):\n",
    "    \"\"\"\n",
    "    Train the Nonogram Agent.\n",
    "    \n",
    "    Parameters:\n",
    "    - grid_size (int): Size of the Nonogram grid.\n",
    "    - clue_max_len (int): Maximum length of the clues.\n",
    "    - clue_dim (int): Dimensionality of the clues.\n",
    "    - num_episodes (int): Number of training episodes.\n",
    "    - train_batch_size (int): Batch size for training.\n",
    "    - val_batch_size (int): Batch size for validation.\n",
    "    - save_interval (int): Interval for saving checkpoints.\n",
    "    \"\"\"\n",
    "    validation_solutions, validation_row_clues, validation_col_clues, existing_solutions = generate_unique_nonogram(grid_size, val_batch_size)\n",
    "    validation_row_clues = [pad_clues(rc, clue_max_len) for rc in validation_row_clues]\n",
    "    validation_col_clues = [pad_clues(cc, clue_max_len) for cc in validation_col_clues]\n",
    "\n",
    "    env = NonogramEnvironment(grid_size, train_batch_size)\n",
    "    agent = NonogramAgent(grid_size, clue_max_len, clue_dim)\n",
    "    optimizer = agent.optimizer\n",
    "\n",
    "    writer = SummaryWriter('runs/nonogram_experiment')\n",
    "\n",
    "    total_cells = grid_size * grid_size\n",
    "\n",
    "    start_episode, reward_list, correct_guess_percent_list, saved_clue_max_len, saved_clue_dim = load_checkpoint(agent, optimizer)\n",
    "    \n",
    "    if saved_clue_max_len is not None:\n",
    "        clue_max_len = saved_clue_max_len\n",
    "    if saved_clue_dim is not None:\n",
    "        clue_dim = saved_clue_dim\n",
    "\n",
    "    progress_bar = tqdm(range(start_episode, num_episodes), desc=\"Training\")\n",
    "\n",
    "    for episode in progress_bar:\n",
    "        train_solutions, train_row_clues, train_col_clues, existing_solutions = generate_unique_nonogram(grid_size, train_batch_size, existing_solutions)\n",
    "        train_row_clues = [pad_clues(rc, clue_max_len) for rc in train_row_clues]\n",
    "        train_col_clues = [pad_clues(cc, clue_max_len) for cc in train_col_clues]\n",
    "\n",
    "        states, row_clues, col_clues = env.reset_with_solutions(train_solutions, train_row_clues, train_col_clues)\n",
    "\n",
    "        log_probs = [[] for _ in range(train_batch_size)]\n",
    "        rewards = [[] for _ in range(train_batch_size)]\n",
    "        done = np.zeros(train_batch_size, dtype=bool)\n",
    "\n",
    "        while not np.all(done):\n",
    "            actions, log_prob = agent.select_actions(states, row_clues, col_clues)\n",
    "            next_states, reward, done_step = env.step(actions)\n",
    "            for i in range(train_batch_size):\n",
    "                if not done[i]:\n",
    "                    log_probs[i].append(log_prob[i])\n",
    "                    rewards[i].append(reward[i])\n",
    "\n",
    "            states = next_states\n",
    "            done = np.logical_or(done, done_step)\n",
    "\n",
    "        agent.update_policy(log_probs, rewards)\n",
    "        correct_guesses = calculate_correct_guesses(states, env.solution)\n",
    "        correct_guess_percent = correct_guesses / total_cells\n",
    "        total_reward = np.mean([sum(r) for r in rewards])\n",
    "\n",
    "        reward_list.append(total_reward)\n",
    "        correct_guess_percent_list.append(correct_guess_percent.mean())\n",
    "\n",
    "        writer.add_scalar('Train_Reward', total_reward, episode)\n",
    "        writer.add_scalar('Train_Correct_Guess_Percent', correct_guess_percent.mean(), episode)\n",
    "\n",
    "        if episode % save_interval == 0:\n",
    "            save_checkpoint(agent, optimizer, episode, reward_list, correct_guess_percent_list, clue_max_len, clue_dim)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            validation_env = NonogramEnvironment(grid_size, val_batch_size)\n",
    "            validation_env.reset_with_solutions(validation_solutions, validation_row_clues, validation_col_clues)\n",
    "\n",
    "            validation_states, validation_row_clues, validation_col_clues = validation_env.reset()\n",
    "            validation_row_clues = [pad_clues(rc, clue_max_len) for rc in validation_row_clues]\n",
    "            validation_col_clues = [pad_clues(cc, clue_max_len) for cc in validation_col_clues]\n",
    "\n",
    "            validation_done = np.zeros(val_batch_size, dtype=bool)\n",
    "            validation_log_probs = [[] for _ in range(val_batch_size)]\n",
    "            validation_rewards = [[] for _ in range(val_batch_size)]\n",
    "\n",
    "            while not np.all(validation_done):\n",
    "                validation_actions, validation_log_prob = agent.select_actions(\n",
    "                    validation_states, validation_row_clues, validation_col_clues)\n",
    "                validation_next_states, validation_reward, validation_done_step = validation_env.step(validation_actions)\n",
    "                for i in range(val_batch_size):\n",
    "                    if not validation_done[i]:\n",
    "                        validation_log_probs[i].append(validation_log_prob[i])\n",
    "                        validation_rewards[i].append(validation_reward[i])\n",
    "                \n",
    "                validation_states = validation_next_states\n",
    "                validation_done = np.logical_or(validation_done, validation_done_step)\n",
    "\n",
    "            validation_correct_guesses = calculate_correct_guesses(validation_states, validation_solutions)\n",
    "            validation_correct_guess_percent = validation_correct_guesses / total_cells\n",
    "            validation_total_reward = np.mean([sum(r) for r in validation_rewards])\n",
    "\n",
    "            writer.add_scalar('Validation_Reward', validation_total_reward, episode)\n",
    "            writer.add_scalar('Validation_Correct_Guess_Percent', validation_correct_guess_percent.mean(), episode)\n",
    "\n",
    "        progress_bar.set_description(f\"Train Reward: {total_reward:.2f}, Train Correct: {correct_guess_percent.mean():.2%}, Val Reward: {validation_total_reward:.2f}, Val Correct: {validation_correct_guess_percent.mean():.2%}\")\n",
    "\n",
    "    writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f3469c2-c79a-4c07-a741-5c000ebd87c8",
   "metadata": {},
   "source": [
    "## Main Execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd8fb5be-ada0-4508-8719-ae5567df70ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    # Parameters for training\n",
    "    grid_size = 5\n",
    "    clue_max_len = divide_and_round_up(grid_size)\n",
    "    clue_dim = grid_size\n",
    "    num_episodes = 100000\n",
    "    train_batch_size = 512\n",
    "    val_batch_size = 128\n",
    "    save_interval = 1000\n",
    "\n",
    "    # Train the agent\n",
    "    train_agent(grid_size, clue_max_len, clue_dim, num_episodes, train_batch_size, val_batch_size, save_interval)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbabe8b2-1acd-41ce-a1a6-cf938d3cd603",
   "metadata": {},
   "source": [
    "## Testing Main Execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "10da98b7-63ab-43da-a181-cdbd69baa3fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Row Clues:\n",
      "[[1, 3, 0], [2, 1, 0], [2, 1, 0], [2, 2, 0], [2, 2, 0]]\n",
      "Column Clues:\n",
      "[[1, 2, 0], [4, 0, 0], [3, 0, 0], [1, 2, 0], [5, 0, 0]]\n",
      "Initial Puzzle State:\n",
      "? ? ? ? ?\n",
      "? ? ? ? ?\n",
      "? ? ? ? ?\n",
      "? ? ? ? ?\n",
      "? ? ? ? ?\n",
      "\n",
      "Move: 1, Guess: (4, 2), State: 0\n",
      "Current Puzzle State:\n",
      "? ? ? ? ?\n",
      "? ? ? ? ?\n",
      "? ? ? ? ?\n",
      "? 1 ? ? ?\n",
      "? ? ? ? ?\n",
      "\n",
      "Move: 2, Guess: (4, 4), State: 0\n",
      "Current Puzzle State:\n",
      "? ? ? ? ?\n",
      "? ? ? ? ?\n",
      "? ? ? ? ?\n",
      "? 1 ? 1 ?\n",
      "? ? ? ? ?\n",
      "\n",
      "Move: 3, Guess: (4, 3), State: 1\n",
      "Current Puzzle State:\n",
      "? ? ? ? ?\n",
      "? ? ? ? ?\n",
      "? ? ? ? ?\n",
      "? 1 0 1 ?\n",
      "? ? ? ? ?\n",
      "\n",
      "Move: 4, Guess: (2, 4), State: 1\n",
      "Current Puzzle State:\n",
      "? ? ? ? ?\n",
      "? ? ? 0 ?\n",
      "? ? ? ? ?\n",
      "? 1 0 1 ?\n",
      "? ? ? ? ?\n",
      "\n",
      "Move: 5, Guess: (3, 2), State: 0\n",
      "Current Puzzle State:\n",
      "? ? ? ? ?\n",
      "? ? ? 0 ?\n",
      "? 1 ? ? ?\n",
      "? 1 0 1 ?\n",
      "? ? ? ? ?\n",
      "\n",
      "Move: 6, Guess: (4, 5), State: 1\n",
      "Current Puzzle State:\n",
      "? ? ? ? ?\n",
      "? ? ? 0 ?\n",
      "? 1 ? ? ?\n",
      "? 1 0 1 1\n",
      "? ? ? ? ?\n",
      "\n",
      "Move: 7, Guess: (2, 2), State: 1\n",
      "Current Puzzle State:\n",
      "? ? ? ? ?\n",
      "? 1 ? 0 ?\n",
      "? 1 ? ? ?\n",
      "? 1 0 1 1\n",
      "? ? ? ? ?\n",
      "\n",
      "Move: 8, Guess: (1, 2), State: 0\n",
      "Current Puzzle State:\n",
      "? 0 ? ? ?\n",
      "? 1 ? 0 ?\n",
      "? 1 ? ? ?\n",
      "? 1 0 1 1\n",
      "? ? ? ? ?\n",
      "\n",
      "Move: 9, Guess: (2, 5), State: 0\n",
      "Current Puzzle State:\n",
      "? 0 ? ? ?\n",
      "? 1 ? 0 1\n",
      "? 1 ? ? ?\n",
      "? 1 0 1 1\n",
      "? ? ? ? ?\n",
      "\n",
      "Move: 10, Guess: (2, 3), State: 0\n",
      "Current Puzzle State:\n",
      "? 0 ? ? ?\n",
      "? 1 1 0 1\n",
      "? 1 ? ? ?\n",
      "? 1 0 1 1\n",
      "? ? ? ? ?\n",
      "\n",
      "Move: 11, Guess: (5, 3), State: 1\n",
      "Current Puzzle State:\n",
      "? 0 ? ? ?\n",
      "? 1 1 0 1\n",
      "? 1 ? ? ?\n",
      "? 1 0 1 1\n",
      "? ? 0 ? ?\n",
      "\n",
      "Move: 12, Guess: (1, 3), State: 1\n",
      "Current Puzzle State:\n",
      "? 0 1 ? ?\n",
      "? 1 1 0 1\n",
      "? 1 ? ? ?\n",
      "? 1 0 1 1\n",
      "? ? 0 ? ?\n",
      "\n",
      "Move: 13, Guess: (3, 3), State: 0\n",
      "Current Puzzle State:\n",
      "? 0 1 ? ?\n",
      "? 1 1 0 1\n",
      "? 1 1 ? ?\n",
      "? 1 0 1 1\n",
      "? ? 0 ? ?\n",
      "\n",
      "Move: 14, Guess: (3, 4), State: 0\n",
      "Current Puzzle State:\n",
      "? 0 1 ? ?\n",
      "? 1 1 0 1\n",
      "? 1 1 0 ?\n",
      "? 1 0 1 1\n",
      "? ? 0 ? ?\n",
      "\n",
      "Move: 15, Guess: (4, 1), State: 1\n",
      "Current Puzzle State:\n",
      "? 0 1 ? ?\n",
      "? 1 1 0 1\n",
      "? 1 1 0 ?\n",
      "1 1 0 1 1\n",
      "? ? 0 ? ?\n",
      "\n",
      "Move: 16, Guess: (1, 5), State: 1\n",
      "Current Puzzle State:\n",
      "? 0 1 ? 1\n",
      "? 1 1 0 1\n",
      "? 1 1 0 ?\n",
      "1 1 0 1 1\n",
      "? ? 0 ? ?\n",
      "\n",
      "Move: 17, Guess: (1, 4), State: 1\n",
      "Current Puzzle State:\n",
      "? 0 1 1 1\n",
      "? 1 1 0 1\n",
      "? 1 1 0 ?\n",
      "1 1 0 1 1\n",
      "? ? 0 ? ?\n",
      "\n",
      "Move: 18, Guess: (5, 2), State: 1\n",
      "Current Puzzle State:\n",
      "? 0 1 1 1\n",
      "? 1 1 0 1\n",
      "? 1 1 0 ?\n",
      "1 1 0 1 1\n",
      "? 1 0 ? ?\n",
      "\n",
      "Move: 19, Guess: (3, 1), State: 0\n",
      "Current Puzzle State:\n",
      "? 0 1 1 1\n",
      "? 1 1 0 1\n",
      "0 1 1 0 ?\n",
      "1 1 0 1 1\n",
      "? 1 0 ? ?\n",
      "\n",
      "Move: 20, Guess: (2, 1), State: 1\n",
      "Current Puzzle State:\n",
      "? 0 1 1 1\n",
      "0 1 1 0 1\n",
      "0 1 1 0 ?\n",
      "1 1 0 1 1\n",
      "? 1 0 ? ?\n",
      "\n",
      "Move: 21, Guess: (1, 1), State: 0\n",
      "Current Puzzle State:\n",
      "1 0 1 1 1\n",
      "0 1 1 0 1\n",
      "0 1 1 0 ?\n",
      "1 1 0 1 1\n",
      "? 1 0 ? ?\n",
      "\n",
      "Move: 22, Guess: (5, 4), State: 0\n",
      "Current Puzzle State:\n",
      "1 0 1 1 1\n",
      "0 1 1 0 1\n",
      "0 1 1 0 ?\n",
      "1 1 0 1 1\n",
      "? 1 0 1 ?\n",
      "\n",
      "Move: 23, Guess: (5, 5), State: 1\n",
      "Current Puzzle State:\n",
      "1 0 1 1 1\n",
      "0 1 1 0 1\n",
      "0 1 1 0 ?\n",
      "1 1 0 1 1\n",
      "? 1 0 1 1\n",
      "\n",
      "Move: 24, Guess: (3, 5), State: 0\n",
      "Current Puzzle State:\n",
      "1 0 1 1 1\n",
      "0 1 1 0 1\n",
      "0 1 1 0 1\n",
      "1 1 0 1 1\n",
      "? 1 0 1 1\n",
      "\n",
      "Move: 25, Guess: (5, 1), State: 0\n",
      "Current Puzzle State:\n",
      "1 0 1 1 1\n",
      "0 1 1 0 1\n",
      "0 1 1 0 1\n",
      "1 1 0 1 1\n",
      "1 1 0 1 1\n",
      "Puzzle Solved!\n"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    # Load the pretrained model\n",
    "    grid_size = 5\n",
    "    clue_max_len = 3\n",
    "    clue_dim = grid_size\n",
    "    agent = NonogramAgent(grid_size, clue_max_len, clue_dim)\n",
    "    optimizer = agent.optimizer\n",
    "    _, _, _, clue_max_len, clue_dim = load_checkpoint(agent, optimizer)\n",
    "\n",
    "    # Generate a random puzzle\n",
    "    solutions, row_clues, col_clues, _ = generate_unique_nonogram(grid_size, 1)\n",
    "    row_clues = [pad_clues(rc, clue_max_len) for rc in row_clues]\n",
    "    col_clues = [pad_clues(cc, clue_max_len) for cc in col_clues]\n",
    "\n",
    "    # Initialize the environment with the generated puzzle\n",
    "    env = NonogramEnvironment(grid_size, 1)\n",
    "    states, row_clues, col_clues = env.reset_with_solutions(solutions, row_clues, col_clues)\n",
    "\n",
    "    # Display clues\n",
    "    print(\"Row Clues:\")\n",
    "    visualize_clues(row_clues)\n",
    "    print(\"Column Clues:\")\n",
    "    visualize_clues(col_clues)\n",
    "\n",
    "    # Display initial puzzle state\n",
    "    print(\"Initial Puzzle State:\")\n",
    "    visualize_nonogram(states[0])\n",
    "\n",
    "    # Solve the puzzle\n",
    "    update_puzzle_state(agent, env, states, row_clues, col_clues, solutions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17b51eb8-a338-4aca-a01f-61a08c6c9278",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}