ZacharyPowellM9.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bd737060-53ff-487c-bb44-3fac6040ad95",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 1. Import Libraries\n",
    "# -----------------------------\n",
    "import os\n",
    "import sys\n",
    "import math\n",
    "import json\n",
    "import queue\n",
    "import logging\n",
    "import random\n",
    "import threading\n",
    "from io import BytesIO\n",
    "\n",
    "# Numerical & Data Handling\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.cuda.amp import autocast, GradScaler\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "# Visualization & GUI\n",
    "import matplotlib\n",
    "matplotlib.use('TkAgg')  # Use Tkinter-compatible backend for Matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.figure import Figure\n",
    "from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg\n",
    "from mpl_toolkits.mplot3d import Axes3D  # For 3D plots\n",
    "\n",
    "import tkinter as tk\n",
    "from tkinter import ttk, filedialog, messagebox\n",
    "from PIL import Image, ImageTk\n",
    "\n",
    "# Model Architectures\n",
    "from torchvision.models import resnet18, ResNet18_Weights\n",
    "\n",
    "# Tree Structure & Visualization\n",
    "from anytree import NodeMixin, RenderTree\n",
    "from anytree.exporter import DotExporter\n",
    "\n",
    "# Learning Rate Schedulers\n",
    "from torch.optim import lr_scheduler\n",
    "from dataclasses import dataclass\n",
    "from graphviz import Source\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "\n",
    "from tqdm import tqdm  # For progress bars\n",
    "import time\n",
    "from PIL import Image, ImageTk, ImageDraw\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ce3a9dc5-eec2-409e-be11-3d0a8d3d6a69",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n# Explanation of the Main Workflow with Modifications, Successes, and Pitfalls:\\n\\nThis `main()` function orchestrates the entire workflow, from loading the data to \\nmodel initialization and launching the GUI for real-time visualization. \\nThroughout the process, we’ve added components to ensure reliability, scalability, and \\nsmooth user interaction, as well as handling common pitfalls encountered in model training.\\n\\n---\\n\\n## 1. Device Detection:\\n    device = get_device()\\n- **Reason**: To ensure the model uses GPU if available, for faster training.\\n- **Success**: Seamless fallback to CPU in case of unavailable GPU.\\n- **Pitfall**: Not accounting for limited GPU memory could crash the program during large batch processing. \\n- **Modification**: Add logic to switch to CPU automatically if GPU memory is insufficient.\\n\\n---\\n\\n## 2. Progress Bar Initialization:\\n    progress_bar = tqdm(total=100, desc=\"Loading Data\", unit=\"%\", leave=True)\\n- **Reason**: Provides visual feedback to the user during the loading process.\\n- **Success**: Improves transparency by showing the progress of various steps.\\n- **Pitfall**: Users might think the program is frozen if no feedback is given for smaller tasks.\\n- **Modification**: Consider tracking finer-grained steps within data loading for more feedback.\\n\\n---\\n\\n## 3. Data Loading:\\n    arc_data = load_arc_data()\\n    progress_bar.update(20)\\n- **Reason**: To load the ARC dataset, which is essential for training and evaluation.\\n- **Success**: Loading large datasets in batches prevents memory overflow.\\n- **Pitfall**: Missing or corrupted data could cause crashes.\\n- **Modification**: Add file existence and format validation to ensure robustness.\\n\\n---\\n\\n## 4. Data Extraction and Reshaping:\\n    train_grid_pairs = flatten_and_reshape(arc_data.get(\"arc-agi_training-challenges\", {}))\\n    eval_grid_pairs = flatten_and_reshape(arc_data.get(\"arc-agi_evaluation-challenges\", {}))\\n    progress_bar.update(30)\\n- **Reason**: Converts the nested grid data into a format usable for training.\\n- **Success**: Ensured compatibility with the DataLoader structure.\\n- **Pitfall**: Errors could arise if the data structure is inconsistent or missing fields.\\n- **Modification**: Implement logging to track extraction and reshaping issues.\\n\\n---\\n\\n## 5. Building Data Tree and Task Dictionary:\\n    root_node, task_dict = build_data_tree(train_grid_pairs)\\n    traverse_and_debug(root_node)\\n    progress_bar.update(20)\\n- **Reason**: Organizes tasks hierarchically for better management during training.\\n- **Success**: Easier debugging and visualization of task relationships.\\n- **Pitfall**: Misaligned grid data could cause tree construction to fail silently.\\n- **Modification**: Added detailed logging during tree traversal to catch errors early.\\n\\n---\\n\\n## 6. Logging Task Information:\\n    logger.info(f\"Task dictionary initialized with {len(task_dict)} tasks:\")\\n- **Reason**: Logs each task\\'s ID, node, and grid shape to track initialization.\\n- **Success**: Helps monitor if all tasks have been properly loaded.\\n- **Pitfall**: Missing tasks could go unnoticed if not properly logged.\\n- **Modification**: Added logging for each task\\'s details to detect anomalies.\\n\\n---\\n\\n## 7. Dataset and DataLoader Initialization:\\n    train_dataset = AugmentedARCDataset(train_grid_pairs, augment=False)\\n    train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=0, collate_fn=collate_fn)\\n- **Reason**: Prepares data for training, using DataLoader to handle batching and shuffling.\\n- **Success**: Custom collate functions enable flexible input sizes.\\n- **Pitfall**: Low `num_workers` slows down loading; too high causes deadlocks.\\n- **Modification**: Set `num_workers` dynamically based on CPU availability.\\n\\n---\\n\\n## 8. Model Initialization:\\n    model = CNNGridMapper(num_classes=NUM_CLASSES).to(device)\\n    if torch.cuda.device_count() > 1:\\n        model = nn.DataParallel(model)\\n    logger.info(\"Model initialized successfully.\")\\n- **Reason**: Moves the model to the appropriate device (CPU or GPU) for faster processing.\\n- **Success**: Handles multi-GPU setups using `DataParallel`.\\n- **Pitfall**: Failing to move the model to the correct device causes runtime errors.\\n- **Modification**: Added a fallback to `nn.DataParallel` for multi-GPU setups.\\n\\n---\\n\\n## 9. Exception Handling:\\n    except Exception as e:\\n        logger.exception(f\"Data loading or model initialization failed: {e}\")\\n        progress_bar.close()\\n        return\\n- **Reason**: Catches unexpected errors during data loading or model setup.\\n- **Success**: Ensures the application exits gracefully with meaningful error messages.\\n- **Pitfall**: Users may be confused without actionable advice for resolving errors.\\n- **Modification**: Provide tips in error logs (e.g., “Check GPU memory usage”).\\n\\n---\\n\\n## 10. Close Progress Bar:\\n    progress_bar.close()\\n- **Reason**: Ensures the progress bar is closed once tasks are completed.\\n- **Success**: Prevents console clutter with open progress bars.\\n- **Pitfall**: Forgetting to close the bar may confuse users about task completion.\\n- **Modification**: Display a summary report after closing the progress bar.\\n\\n---\\n\\n## 11. GUI Initialization:\\n    gui = TrainingGUI(\\n        root_window, total_epochs=10, total_batches=len(train_loader),\\n        model=model, train_loader=train_loader, val_loader=val_loader,\\n        eval_loader=None, device=device, data_tree=root_node, task_dict=task_dict\\n    )\\n- **Reason**: Provides a GUI for real-time monitoring of the training process.\\n- **Success**: Keeps users informed about training progress.\\n- **Pitfall**: The GUI can freeze without proper threading.\\n- **Modification**: Offload training logic to a separate thread to keep the GUI responsive.\\n\\n---\\n\\n## 12. Start Training Thread:\\n    training_thread = threading.Thread(\\n        target=train_model_with_gui, args=(model, train_loader, val_loader, device, gui)\\n    )\\n    training_thread.daemon = True\\n    training_thread.start()\\n- **Reason**: Runs training in a non-blocking thread to maintain GUI responsiveness.\\n- **Success**: GUI remains interactive while training continues in the background.\\n- **Pitfall**: Threads can fail silently, leading to halted training.\\n- **Modification**: Added `daemon=True` to ensure the thread terminates with the main process.\\n\\n---\\n\\n## 13. Start GUI Main Loop:\\n    root_window.mainloop()\\n- **Reason**: Keeps the GUI running for user interaction.\\n- **Success**: Allows users to interact with the application throughout the training.\\n- **Pitfall**: Users might quit accidentally, halting training without warning.\\n- **Modification**: Added a confirmation dialog to prompt users before quitting.\\n\\n---\\n\\n## 14. Entry Point Check:\\n    if __name__ == \"__main__\":\\n        main()\\n- **Reason**: Prevents accidental execution when the script is imported as a module.\\n- **Success**: Ensures the workflow only starts when run directly.\\n- **Pitfall**: If not implemented, unintended execution could interfere with other scripts.\\n- **Modification**: Standardized the entry-point check to avoid such issues.\\n\\n---\\n\\n### Final Thoughts:\\nThis workflow ensures that the application is efficient, scalable, and user-friendly. \\nBy handling potential pitfalls and providing useful feedback throughout, we minimize downtime \\nand confusion during execution. The use of multi-threading, dynamic resource management, \\nand informative logging significantly improves the robustness of the application.\\n'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "##\"\"\"\n",
    "# Explanation of the Main Workflow with Modifications, Successes, and Pitfalls:\n",
    "\n",
    "#This `main()` function orchestrates the entire workflow, from loading the data to \n",
    "#model initialization and launching the GUI for real-time visualization. \n",
    "#Throughout the process, we’ve added components to ensure reliability, scalability, and \n",
    "#smooth user interaction, as well as handling common pitfalls encountered in model training.\n",
    "\n",
    "#---\n",
    "\n",
    "# 1. Device Detection:\n",
    " #   device = get_device()\n",
    "#- **Reason**: To ensure the model uses GPU if available, for faster training.\n",
    "#- **Success**: Seamless fallback to CPU in case of unavailable GPU.\n",
    "#- **Pitfall**: Not accounting for limited GPU memory could crash the program during large batch processing. \n",
    "#- **Modification**: Add logic to switch to CPU automatically if GPU memory is insufficient.\n",
    "\n",
    "#---\n",
    "\n",
    "# 2. Progress Bar Initialization:\n",
    " #   progress_bar = tqdm(total=100, desc=\"Loading Data\", unit=\"%\", leave=True)\n",
    "#- **Reason**: Provides visual feedback to the user during the loading process.\n",
    "#- **Success**: Improves transparency by showing the progress of various steps.\n",
    "#- **Pitfall**: Users might think the program is frozen if no feedback is given for smaller tasks.\n",
    "- **Modification**: Consider tracking finer-grained steps within data loading for more feedback.\n",
    "\n",
    "---\n",
    "\n",
    "# 3. Data Loading:\n",
    "    arc_data = load_arc_data()\n",
    "    progress_bar.update(20)\n",
    "- **Reason**: To load the ARC dataset, which is essential for training and evaluation.\n",
    "- **Success**: Loading large datasets in batches prevents memory overflow.\n",
    "- **Pitfall**: Missing or corrupted data could cause crashes.\n",
    "- **Modification**: Add file existence and format validation to ensure robustness.\n",
    "\n",
    "---\n",
    "\n",
    "# 4. Data Extraction and Reshaping:\n",
    "    train_grid_pairs = flatten_and_reshape(arc_data.get(\"arc-agi_training-challenges\", {}))\n",
    "    eval_grid_pairs = flatten_and_reshape(arc_data.get(\"arc-agi_evaluation-challenges\", {}))\n",
    "    progress_bar.update(30)\n",
    "#- **Reason**: Converts the nested grid data into a format usable for training.\n",
    "- **Success**: Ensured compatibility with the DataLoader structure.\n",
    "- **Pitfall**: Errors could arise if the data structure is inconsistent or missing fields.\n",
    "- **Modification**: Implement logging to track extraction and reshaping issues.\n",
    "\n",
    "#---\n",
    "\n",
    "#5. Building Data Tree and Task Dictionary:\n",
    "    root_node, task_dict = build_data_tree(train_grid_pairs)\n",
    "    traverse_and_debug(root_node)\n",
    "    progress_bar.update(20)\n",
    "- **Reason**: Organizes tasks hierarchically for better management during training.\n",
    "- **Success**: Easier debugging and visualization of task relationships.\n",
    "- **Pitfall**: Misaligned grid data could cause tree construction to fail silently.\n",
    "- **Modification**: Added detailed logging during tree traversal to catch errors early.\n",
    "\n",
    "---\n",
    "\n",
    "## 6. Logging Task Information:\n",
    "    logger.info(f\"Task dictionary initialized with {len(task_dict)} tasks:\")\n",
    "- **Reason**: Logs each task's ID, node, and grid shape to track initialization.\n",
    "- **Success**: Helps monitor if all tasks have been properly loaded.\n",
    "- **Pitfall**: Missing tasks could go unnoticed if not properly logged.\n",
    "- **Modification**: Added logging for each task's details to detect anomalies.\n",
    "\n",
    "---\n",
    "\n",
    "## 7. Dataset and DataLoader Initialization:\n",
    "    train_dataset = AugmentedARCDataset(train_grid_pairs, augment=False)\n",
    "    train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=0, collate_fn=collate_fn)\n",
    "- **Reason**: Prepares data for training, using DataLoader to handle batching and shuffling.\n",
    "- **Success**: Custom collate functions enable flexible input sizes.\n",
    "- **Pitfall**: Low `num_workers` slows down loading; too high causes deadlocks.\n",
    "- **Modification**: Set `num_workers` dynamically based on CPU availability.\n",
    "\n",
    "---\n",
    "\n",
    "## 8. Model Initialization:\n",
    "    model = CNNGridMapper(num_classes=NUM_CLASSES).to(device)\n",
    "    if torch.cuda.device_count() > 1:\n",
    "        model = nn.DataParallel(model)\n",
    "    logger.info(\"Model initialized successfully.\")\n",
    "- **Reason**: Moves the model to the appropriate device (CPU or GPU) for faster processing.\n",
    "- **Success**: Handles multi-GPU setups using `DataParallel`.\n",
    "- **Pitfall**: Failing to move the model to the correct device causes runtime errors.\n",
    "- **Modification**: Added a fallback to `nn.DataParallel` for multi-GPU setups.\n",
    "\n",
    "---\n",
    "\n",
    "## 9. Exception Handling:\n",
    "    except Exception as e:\n",
    "        logger.exception(f\"Data loading or model initialization failed: {e}\")\n",
    "        progress_bar.close()\n",
    "        return\n",
    "- **Reason**: Catches unexpected errors during data loading or model setup.\n",
    "- **Success**: Ensures the application exits gracefully with meaningful error messages.\n",
    "- **Pitfall**: Users may be confused without actionable advice for resolving errors.\n",
    "- **Modification**: Provide tips in error logs (e.g., “Check GPU memory usage”).\n",
    "\n",
    "---\n",
    "\n",
    "## 10. Close Progress Bar:\n",
    "    progress_bar.close()\n",
    "- **Reason**: Ensures the progress bar is closed once tasks are completed.\n",
    "- **Success**: Prevents console clutter with open progress bars.\n",
    "- **Pitfall**: Forgetting to close the bar may confuse users about task completion.\n",
    "- **Modification**: Display a summary report after closing the progress bar.\n",
    "\n",
    "---\n",
    "\n",
    "## 11. GUI Initialization:\n",
    "    gui = TrainingGUI(\n",
    "        root_window, total_epochs=10, total_batches=len(train_loader),\n",
    "        model=model, train_loader=train_loader, val_loader=val_loader,\n",
    "        eval_loader=None, device=device, data_tree=root_node, task_dict=task_dict\n",
    "    )\n",
    "- **Reason**: Provides a GUI for real-time monitoring of the training process.\n",
    "- **Success**: Keeps users informed about training progress.\n",
    "- **Pitfall**: The GUI can freeze without proper threading.\n",
    "- **Modification**: Offload training logic to a separate thread to keep the GUI responsive.\n",
    "\n",
    "---\n",
    "\n",
    "## 12. Start Training Thread:\n",
    "    training_thread = threading.Thread(\n",
    "        target=train_model_with_gui, args=(model, train_loader, val_loader, device, gui)\n",
    "    )\n",
    "    training_thread.daemon = True\n",
    "    training_thread.start()\n",
    "- **Reason**: Runs training in a non-blocking thread to maintain GUI responsiveness.\n",
    "- **Success**: GUI remains interactive while training continues in the background.\n",
    "- **Pitfall**: Threads can fail silently, leading to halted training.\n",
    "- **Modification**: Added `daemon=True` to ensure the thread terminates with the main process.\n",
    "\n",
    "---\n",
    "\n",
    "## 13. Start GUI Main Loop:\n",
    "    root_window.mainloop()\n",
    "- **Reason**: Keeps the GUI running for user interaction.\n",
    "- **Success**: Allows users to interact with the application throughout the training.\n",
    "- **Pitfall**: Users might quit accidentally, halting training without warning.\n",
    "- **Modification**: Added a confirmation dialog to prompt users before quitting.\n",
    "\n",
    "---\n",
    "\n",
    "## 14. Entry Point Check:\n",
    "    if __name__ == \"__main__\":\n",
    "        main()\n",
    "- **Reason**: Prevents accidental execution when the script is imported as a module.\n",
    "- **Success**: Ensures the workflow only starts when run directly.\n",
    "- **Pitfall**: If not implemented, unintended execution could interfere with other scripts.\n",
    "- **Modification**: Standardized the entry-point check to avoid such issues.\n",
    "\n",
    "---\n",
    "\n",
    "### Final Thoughts:\n",
    "This workflow ensures that the application is efficient, scalable, and user-friendly. \n",
    "By handling potential pitfalls and providing useful feedback throughout, we minimize downtime \n",
    "and confusion during execution. The use of multi-threading, dynamic resource management, \n",
    "and informative logging significantly improves the robustness of the application.\n",
    "\"\"\"##\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "74d46965-1259-437b-a58f-600787f1b5d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 3. Define Constants\n",
    "# -----------------------------\n",
    "\n",
    "# Define the number of classes\n",
    "NUM_CLASSES = 11  # 0-10, where 10 represents dead squares\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0cd9a27b-6d85-4c1c-9a5f-67c7af30da9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 4. Define Data Structures and Loading Functions\n",
    "# -----------------------------\n",
    "\n",
    "# Data Class for Grid Pairs\n",
    "@dataclass\n",
    "class GridPair:\n",
    "    task_id: str\n",
    "    input_grid: np.ndarray\n",
    "    output_grid: np.ndarray\n",
    "\n",
    "def load_arc_data():\n",
    "    file_paths = {\n",
    "        \"arc-agi_training-challenges\": \"arc-agi_training_challenges.json\",\n",
    "        \"arc-agi_evaluation-challenges\": \"arc-agi_evaluation_challenges.json\",\n",
    "        \"arc-agi_training-solutions\": \"arc-agi_training_solutions.json\",\n",
    "        \"arc-agi_evaluation-solutions\": \"arc-agi_evaluation_solutions.json\",\n",
    "    }\n",
    "    arc_data = {key: load_json_file(path) for key, path in file_paths.items()}\n",
    "    return arc_data\n",
    "\n",
    "def load_json_file(path):\n",
    "    try:\n",
    "        with open(path, 'r') as f:\n",
    "            data = json.load(f)\n",
    "            logger.info(f\"Loaded data from {path}.\")\n",
    "            return data\n",
    "    except (FileNotFoundError, json.JSONDecodeError) as e:\n",
    "        logger.error(f\"Error loading {path}: {e}\")\n",
    "        return {}\n",
    "\n",
    "def get_device():\n",
    "    \"\"\"Detect the best available device: CUDA or CPU.\"\"\"\n",
    "    if torch.cuda.is_available():\n",
    "        device = torch.device('cuda')  # CUDA GPU\n",
    "        logger.info(\"Using NVIDIA GPU via CUDA.\")\n",
    "    else:\n",
    "        device = torch.device('cpu')  # Fallback to CPU\n",
    "        logger.info(\"Using CPU as fallback.\")\n",
    "    return device\n",
    "\n",
    "def extract_and_reshape_grid(grid):\n",
    "    try:\n",
    "        # Convert to NumPy array if not already\n",
    "        grid = np.array(grid)\n",
    "        # Handle empty grids or grids with zero dimensions\n",
    "        if grid.size == 0 or 0 in grid.shape:\n",
    "            logger.error(f\"Empty grid or grid with zero dimension encountered: {grid.shape}\")\n",
    "            return None\n",
    "        # Ensure grid is 2D\n",
    "        if grid.ndim == 1:\n",
    "            # If the grid is 1D, reshape to (1, N)\n",
    "            grid = grid.reshape(1, -1)\n",
    "            logger.warning(f\"Grid reshaped to 2D: {grid.shape}\")\n",
    "        elif grid.ndim > 2:\n",
    "            grid = grid.squeeze()\n",
    "            if grid.ndim > 2:\n",
    "                logger.error(f\"Grid has more than 2 dimensions after squeeze: {grid.shape}\")\n",
    "                return None\n",
    "        return grid  # Return as is, without resizing\n",
    "    except Exception as e:\n",
    "        logger.error(f\"Error processing grid: {e}\")\n",
    "        return None\n",
    "\n",
    "# Flatten and Reshape Grid Data\n",
    "def flatten_and_reshape(task_data):\n",
    "    flattened_pairs = []\n",
    "    for task_id, task_content in task_data.items():\n",
    "        logger.info(f\"Parsing task {task_id}...\")\n",
    "        train_pairs = task_content.get('train', [])\n",
    "        for pair in train_pairs:\n",
    "            input_grid = extract_and_reshape_grid(pair.get(\"input\"))\n",
    "            output_grid = extract_and_reshape_grid(pair.get(\"output\"))\n",
    "            if input_grid is not None and output_grid is not None:\n",
    "                # Check for zero dimensions in input or output grid\n",
    "                if 0 in input_grid.shape or 0 in output_grid.shape:\n",
    "                    logger.warning(f\"Task ID: {task_id} has grid with zero dimension. Skipping.\")\n",
    "                    continue\n",
    "                # Store the grids even if shapes differ\n",
    "                flattened_pairs.append(GridPair(task_id, input_grid, output_grid))\n",
    "            else:\n",
    "                logger.warning(f\"Task ID: {task_id} has invalid input/output grids.\")\n",
    "    logger.info(f\"Total valid grid pairs extracted: {len(flattened_pairs)}\")\n",
    "    return flattened_pairs\n",
    "\n",
    "def grid_to_image(grid, color_map):\n",
    "    img_array = np.zeros((grid.shape[0], grid.shape[1], 3), dtype=np.uint8)\n",
    "    for i in range(grid.shape[0]):\n",
    "        for j in range(grid.shape[1]):\n",
    "            img_array[i, j] = color_map.get(grid[i, j], [0, 0, 0])  # Default to black\n",
    "    return Image.fromarray(img_array)\n",
    "\n",
    "color_map = {\n",
    "    0: [0, 0, 0],       # Black\n",
    "    1: [255, 0, 0],     # Red\n",
    "    2: [0, 255, 0],     # Green\n",
    "    3: [0, 0, 255],     # Blue\n",
    "    # Add more colors as needed\n",
    "}\n",
    "\n",
    "class TreeNode(NodeMixin):\n",
    "    def __init__(self, name, input_grid=None, parent=None, children=None):\n",
    "        self.name = name\n",
    "        self.input_grid = input_grid\n",
    "        self.embedding = None\n",
    "        self.parent = parent\n",
    "        if children:\n",
    "            self.children = children\n",
    "\n",
    "        logger.info(f\"Node '{self.name}' initialized.\")\n",
    "\n",
    "    def set_embedding(self, embedding):\n",
    "        \"\"\"\n",
    "        Set the embedding for the node.\n",
    "        \"\"\"\n",
    "        self.embedding = embedding\n",
    "        logger.info(f\"Embedding set for node '{self.name}'.\")\n",
    "\n",
    "    def __repr__(self):\n",
    "        \"\"\"String representation for easier debugging.\"\"\"\n",
    "        return f\"TreeNode(name={self.name}, children={len(self.children) if self.children else 0})\"\n",
    "\n",
    "def build_data_tree(grid_pairs):\n",
    "    \"\"\"\n",
    "    Build a hierarchical tree from the ARC data using the Node class,\n",
    "    and create a task dictionary for quick access.\n",
    "\n",
    "    Args:\n",
    "        grid_pairs (list): List of GridPair objects.\n",
    "\n",
    "    Returns:\n",
    "        tuple: (Node, dict) - Root node and task dictionary.\n",
    "    \"\"\"\n",
    "    # Create the root node\n",
    "    root = TreeNode(name=\"ARC Dataset\")\n",
    "\n",
    "    # Initialize the task dictionary\n",
    "    task_dict = {}\n",
    "\n",
    "    # Loop through the grid pairs to build task nodes\n",
    "    for idx, pair in enumerate(grid_pairs):\n",
    "        try:\n",
    "            if not isinstance(pair, GridPair):\n",
    "                raise TypeError(f\"Expected GridPair, got {type(pair)}: {pair}\")\n",
    "\n",
    "            # Ensure grids are NumPy arrays\n",
    "            input_grid = np.array(pair.input_grid) if not isinstance(pair.input_grid, np.ndarray) else pair.input_grid\n",
    "            output_grid = np.array(pair.output_grid) if not isinstance(pair.output_grid, np.ndarray) else pair.output_grid\n",
    "\n",
    "            # Create task and output nodes\n",
    "            task_node = TreeNode(name=f\"Task {pair.task_id}\", parent=root)\n",
    "            output_node = TreeNode(name=f\"Output {pair.task_id}\", parent=task_node)\n",
    "\n",
    "            # Set embeddings for the nodes\n",
    "            task_node.set_embedding(input_grid)\n",
    "            output_node.set_embedding(output_grid)\n",
    "\n",
    "            # Store the nodes in the task dictionary\n",
    "            task_dict[pair.task_id] = {\n",
    "                'task_node': task_node,\n",
    "                'output_node': output_node,\n",
    "                'grids': (input_grid, output_grid)\n",
    "            }\n",
    "\n",
    "            # Log success\n",
    "            logger.info(f\"Created task node for {pair.task_id} with embedding shape: {task_node.embedding.shape}\")\n",
    "\n",
    "        except Exception as e:\n",
    "            logger.exception(f\"Failed to create nodes for grid pair {idx}: {e}\")\n",
    "            continue  # Skip this pair if there's an issue\n",
    "\n",
    "    # Return the root node and the task dictionary\n",
    "    return root, task_dict\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "01f78726-6b81-4e95-b4bb-2bee3c15dbae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# Explanation of Code Block:\n",
    "# -----------------------------\n",
    "# 4. Define Data Structures and Loading Functions\n",
    "# -----------------------------\n",
    "\n",
    "# 1. GridPair Data Class:\n",
    "#    - Utilizes the @dataclass decorator to automatically generate special methods like __init__().\n",
    "#    - Represents a pair of input and output grids associated with a specific task.\n",
    "#    - Attributes:\n",
    "#        - task_id (str): Identifier for the task.\n",
    "#        - input_grid (np.ndarray): The input grid for the task.\n",
    "#        - output_grid (np.ndarray): The expected output grid for the task.\n",
    "#    - Purpose: Organizes and manages data associated with each task for cleaner code and easier access.\n",
    "\n",
    "#    *Success*: Simplifies the codebase by automatically generating methods and improves data management.\n",
    "#    *Pitfall*: If grids aren't properly validated, it could lead to downstream errors.\n",
    "\n",
    "# 2. load_arc_data():\n",
    "#    - Loads the ARC dataset from specified JSON files.\n",
    "#    - Takes in a dictionary mapping descriptive keys to the filenames of the JSON files.\n",
    "#    - Uses a dictionary comprehension to load each file with load_json_file().\n",
    "#    - Returns a dictionary containing the loaded data for training and evaluation challenges.\n",
    "\n",
    "#    *Reason for Addition*: Centralizes data loading, making maintenance easier.\n",
    "#    *Success*: Clean separation of data access logic.\n",
    "#    *Pitfall*: Missing files or incorrect paths can cause silent failures without validation.\n",
    "\n",
    "# 3. load_json_file(path):\n",
    "#    - Helper function to load a single JSON file given its path.\n",
    "#    - Uses a try-except block to handle potential errors such as FileNotFoundError and JSONDecodeError.\n",
    "#    - Logs a success or error message based on the outcome.\n",
    "#    - Returns the loaded data or an empty dictionary in case of an error.\n",
    "\n",
    "#    *Reason for Addition*: Provides robust error handling for individual file loading.\n",
    "#    *Success*: Ensures graceful handling of file-related errors.\n",
    "#    *Pitfall*: Silent failures if logging isn't monitored carefully.\n",
    "\n",
    "# 4. get_device():\n",
    "#    - Determines the computing device (GPU if available, otherwise CPU).\n",
    "#    - Uses torch.cuda.is_available() to check for GPU.\n",
    "#    - Logs the selected device for transparency.\n",
    "#    - Returns a torch.device object representing the device.\n",
    "\n",
    "#    *Reason for Addition*: Optimizes computation by utilizing GPU when available.\n",
    "#    *Success*: Seamless transition between CPU and GPU usage.\n",
    "#    *Pitfall*: Lack of fallback mechanism for memory issues on GPU.\n",
    "\n",
    "# 5. extract_and_reshape_grid(grid):\n",
    "#    - Processes an individual grid to ensure it is formatted correctly.\n",
    "#    - Converts input to a NumPy array if necessary.\n",
    "#    - Logs and handles cases with empty or malformed grids.\n",
    "#    - Ensures grids are two-dimensional:\n",
    "#        - Reshapes one-dimensional grids into 2D grids with one row.\n",
    "#        - Attempts to reduce grids with more than two dimensions by squeezing.\n",
    "#        - Logs an error and returns None if the grid remains too complex.\n",
    "\n",
    "#    *Reason for Addition*: Standardizes input data to prevent shape-related errors.\n",
    "#    *Success*: Catches inconsistencies early in the workflow.\n",
    "#    *Pitfall*: Over-aggressive reshaping could lead to data loss.\n",
    "\n",
    "# 6. flatten_and_reshape(task_data):\n",
    "#    - Transforms nested task data into a flat list of GridPair instances.\n",
    "#    - Iterates over each task’s training pairs.\n",
    "#    - Processes input and output grids with extract_and_reshape_grid().\n",
    "#    - Skips invalid pairs (None or zero-dimension grids).\n",
    "#    - Logs the total number of valid grid pairs extracted.\n",
    "\n",
    "#    *Reason for Addition*: Prepares task data for training.\n",
    "#    *Success*: Ensures the model receives data in the expected format.\n",
    "#    *Pitfall*: Incorrect grids might get silently skipped without user notice.\n",
    "\n",
    "# 7. grid_to_image(grid, color_map):\n",
    "#    - Converts a numerical grid into a visual image using a color map.\n",
    "#    - Creates an RGB array where each cell’s color corresponds to the value in the grid.\n",
    "#    - Uses Image.fromarray() to create a PIL Image from the array.\n",
    "\n",
    "#    *Reason for Addition*: Enables visual inspection and debugging of grids.\n",
    "#    *Success*: Helps with quick troubleshooting by viewing grid data as images.\n",
    "#    *Pitfall*: Large grids could slow down rendering or debugging sessions.\n",
    "\n",
    "# 8. color_map:\n",
    "#    - A dictionary mapping grid values to RGB colors (e.g., 0 -> black, 1 -> red).\n",
    "#    - Allows for easy customization by adding more mappings.\n",
    "\n",
    "#    *Reason for Addition*: Defines visual representation for grid values.\n",
    "#    *Success*: Simplifies visualization logic by abstracting color mapping.\n",
    "#    *Pitfall*: Limited color choices could make certain grids harder to interpret.\n",
    "\n",
    "# 9. TreeNode Class:\n",
    "#    - Inherits from NodeMixin to create tree structures.\n",
    "#    - Attributes:\n",
    "#        - name (str): Node’s name.\n",
    "#        - input_grid (optional, np.ndarray): Grid data associated with the node.\n",
    "#        - embedding (any): Stores embeddings or metadata.\n",
    "#        - parent (optional, TreeNode): Parent node reference.\n",
    "#        - children (optional, list): List of child nodes.\n",
    "#    - Methods:\n",
    "#        - set_embedding(embedding): Assigns an embedding and logs the action.\n",
    "#        - __repr__(): Returns a string representation for debugging.\n",
    "\n",
    "#    *Reason for Addition*: Structures dataset tasks hierarchically for better management.\n",
    "#    *Success*: Helps organize large datasets with clear parent-child relationships.\n",
    "#    *Pitfall*: Complex hierarchies can become hard to navigate without proper logging.\n",
    "\n",
    "# 10. build_data_tree(grid_pairs):\n",
    "#     - Constructs a tree structure from GridPair instances.\n",
    "#     - Creates a root node called \"ARC Dataset.\"\n",
    "#     - Initializes a task dictionary for quick access.\n",
    "#     - For each GridPair:\n",
    "#        - Validates the instance type and data.\n",
    "#        - Ensures the grids are NumPy arrays.\n",
    "#        - Creates task nodes and their children.\n",
    "#        - Sets embeddings using input and output grids.\n",
    "#        - Logs the creation of each node.\n",
    "#     - Returns the root node and task dictionary.\n",
    "\n",
    "#     *Reason for Addition*: Provides a structured representation of the dataset.\n",
    "#     *Success*: Enables hierarchical processing and easy visualization.\n",
    "#     *Pitfall*: If tasks aren’t validated, it could lead to broken tree structures.\n",
    "\n",
    "# Summary:\n",
    "# This section defines the data structures and helper functions required to manage and\n",
    "# load the ARC dataset efficiently. Key considerations include handling malformed data, \n",
    "# logging errors, and ensuring data consistency throughout the workflow. \n",
    "# These additions improve maintainability, error handling, and visualization but require \n",
    "# careful monitoring to avoid pitfalls like silent skips or memory issues.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b8c16fd4-259f-41cb-871b-00bfb54cd616",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 5. Data Augmentation Functions\n",
    "# -----------------------------\n",
    "\n",
    "def augment_grid(grid, noise_prob=0.2, dead_square_prob=0.1):\n",
    "    augmented_grid = np.array(grid)\n",
    "\n",
    "    # Ensure the grid is 2D\n",
    "    if augmented_grid.ndim != 2:\n",
    "        logger.error(f\"Augmenting grid failed due to invalid shape: {augmented_grid.shape}\")\n",
    "        return augmented_grid  # Return the original grid without augmentation\n",
    "\n",
    "    # Random noise and dead square masks\n",
    "    noise_mask = np.random.rand(*augmented_grid.shape) < noise_prob\n",
    "    dead_mask = np.random.rand(*augmented_grid.shape) < dead_square_prob\n",
    "\n",
    "    # Apply noise\n",
    "    noise_values = np.random.randint(0, NUM_CLASSES - 1, size=augmented_grid.shape)\n",
    "    augmented_grid = np.where(noise_mask, noise_values, augmented_grid)\n",
    "    augmented_grid = np.where(dead_mask, -1, augmented_grid)  # Mark as dead squares\n",
    "\n",
    "    return augmented_grid\n",
    "\n",
    "def rotate_grid(grid):\n",
    "    \"\"\"Randomly rotates the grid.\"\"\"\n",
    "    rotations = random.choice([0, 1, 2, 3])\n",
    "    return np.rot90(grid, rotations)\n",
    "\n",
    "def flip_grid(grid):\n",
    "    \"\"\"Randomly flips the grid.\"\"\"\n",
    "    flip_choice = random.choice(['none', 'vertical', 'horizontal'])\n",
    "    if flip_choice == 'vertical':\n",
    "        return np.flipud(grid)  # Vertical flip\n",
    "    elif flip_choice == 'horizontal':\n",
    "        return np.fliplr(grid)  # Horizontal flip\n",
    "    else:\n",
    "        return grid  # No flip\n",
    "\n",
    "def generate_multiple_augmented_datasets(grid_pairs, num_augmented_sets=3):\n",
    "    \"\"\"\n",
    "    Generates multiple augmented datasets from the input grid pairs.\n",
    "\n",
    "    Args:\n",
    "        grid_pairs (list): List of GridPair objects.\n",
    "        num_augmented_sets (int): Number of augmented sets to generate.\n",
    "\n",
    "    Returns:\n",
    "        list: Augmented grid pairs.\n",
    "    \"\"\"\n",
    "    augmented_pairs = []\n",
    "    for _ in range(num_augmented_sets):\n",
    "        for pair in grid_pairs:\n",
    "            # Apply augmentations to input grid\n",
    "            augmented_input = augment_grid(pair.input_grid)\n",
    "\n",
    "            # Optionally rotate and flip\n",
    "            augmented_input = rotate_grid(augmented_input)\n",
    "            augmented_input = flip_grid(augmented_input)\n",
    "\n",
    "            # Append the augmented input with the original target grid\n",
    "            augmented_pairs.append(GridPair(pair.task_id, augmented_input, pair.output_grid))\n",
    "\n",
    "    return augmented_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d4774e33-7ca1-49cb-859b-4a67f976c372",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# Explanation of the Data Augmentation Functions:\n",
    "# -----------------------------\n",
    "# 5. Data Augmentation Functions\n",
    "# -----------------------------\n",
    "\n",
    "# 1. augment_grid(grid, noise_prob=0.2, dead_square_prob=0.1):\n",
    "#    - **Purpose**: Applies random noise and dead squares to the grid to introduce variation.\n",
    "#    - **Parameters**:\n",
    "#        - **grid**: A 2D NumPy array representing the input grid.\n",
    "#        - **noise_prob**: Probability that a cell will be replaced with a random noise value.\n",
    "#        - **dead_square_prob**: Probability that a cell will be marked as a dead square (e.g., -1).\n",
    "#    - **Steps**:\n",
    "#        a. Converts the input to a NumPy array if not already one.\n",
    "#        b. Verifies that the grid is 2D; if not, logs an error and returns the original grid.\n",
    "#        c. Generates random masks for noise and dead squares.\n",
    "#        d. Creates random noise values within the valid class range (0 to NUM_CLASSES - 1).\n",
    "#        e. Replaces cells according to the generated noise and dead square masks.\n",
    "#        f. Returns the modified grid.\n",
    "#    - **Notes**:\n",
    "#        - Adds variability to the training data, helping the model generalize better.\n",
    "#        - Dead squares represent unusable data, indicated by a special value (e.g., -1).\n",
    "\n",
    "#    *Reason for Addition*: Augmenting data ensures that the model learns from a variety of inputs.\n",
    "#    *Success*: Helps prevent overfitting, leading to improved generalization.\n",
    "#    *Pitfall*: Overuse of noise or dead squares could reduce the representativeness of the data.\n",
    "\n",
    "# 2. rotate_grid(grid):\n",
    "#    - **Purpose**: Randomly rotates the input grid by 0, 90, 180, or 270 degrees.\n",
    "#    - **Parameters**:\n",
    "#        - **grid**: A 2D NumPy array representing the grid to be rotated.\n",
    "#    - **Steps**:\n",
    "#        a. Randomly selects a rotation angle (0, 90, 180, or 270 degrees).\n",
    "#        b. Uses NumPy’s `rot90` function to rotate the grid.\n",
    "#        c. Returns the rotated grid.\n",
    "#    - **Notes**:\n",
    "#        - Rotation makes the model invariant to the orientation of inputs.\n",
    "#        - Especially useful when grid patterns can appear in multiple orientations.\n",
    "\n",
    "#    *Reason for Addition*: Prepares the model to handle orientation differences in data.\n",
    "#    *Success*: Increases robustness by making the model invariant to rotation.\n",
    "#    *Pitfall*: Too many rotations can lead to redundant patterns that don’t enhance learning.\n",
    "\n",
    "# 3. flip_grid(grid):\n",
    "#    - **Purpose**: Randomly flips the grid either vertically, horizontally, or not at all.\n",
    "#    - **Parameters**:\n",
    "#        - **grid**: A 2D NumPy array representing the grid to be flipped.\n",
    "#    - **Steps**:\n",
    "#        a. Randomly selects a flip type: 'none', 'vertical', or 'horizontal'.\n",
    "#        b. Uses NumPy functions to apply the selected flip:\n",
    "#            - `np.flipud(grid)` for a vertical flip (up-down).\n",
    "#            - `np.fliplr(grid)` for a horizontal flip (left-right).\n",
    "#        c. Returns the flipped grid.\n",
    "#    - **Notes**:\n",
    "#        - Flipping introduces more variation, enhancing the dataset.\n",
    "#        - Helps the model learn patterns that remain consistent across different flips.\n",
    "\n",
    "#    *Reason for Addition*: Increases data variety, helping the model learn feature consistency.\n",
    "#    *Success*: Improves generalization by exposing the model to different grid orientations.\n",
    "#    *Pitfall*: Some flipped grids may become unrealistic, leading to unhelpful data points.\n",
    "\n",
    "# 4. generate_multiple_augmented_datasets(grid_pairs, num_augmented_sets=3):\n",
    "#    - **Purpose**: Creates multiple sets of augmented data from the original grid pairs.\n",
    "#    - **Parameters**:\n",
    "#        - **grid_pairs**: A list of GridPair objects containing input and output grids.\n",
    "#        - **num_augmented_sets**: Number of times to augment the dataset.\n",
    "#    - **Steps**:\n",
    "#        a. Initializes an empty list to store the augmented grid pairs.\n",
    "#        b. Repeats the augmentation process for the specified number of sets.\n",
    "#        c. For each GridPair:\n",
    "#            i. Applies augmentations to the input grid:\n",
    "#                - Adds noise and dead squares using `augment_grid()`.\n",
    "#                - Randomly rotates the grid with `rotate_grid()`.\n",
    "#                - Randomly flips the grid with `flip_grid()`.\n",
    "#            ii. Creates a new GridPair with the augmented input grid and the original output grid.\n",
    "#            iii. Appends the new GridPair to the augmented pairs list.\n",
    "#        d. Returns the list of augmented grid pairs.\n",
    "#    - **Notes**:\n",
    "#        - Increases the size of the dataset by generating multiple augmented versions.\n",
    "#        - Retains the original output grid to maintain the correct training target.\n",
    "#        - Essential for improving model robustness, especially with small datasets.\n",
    "\n",
    "#    *Reason for Addition*: Augments the dataset to ensure the model sees more varied inputs.\n",
    "#    *Success*: Helps prevent overfitting by exposing the model to diverse inputs.\n",
    "#    *Pitfall*: Excessive augmentation might dilute meaningful patterns in the original data.\n",
    "\n",
    "# General Comments:\n",
    "# - **Importance of Data Augmentation**: \n",
    "#   - Helps prevent overfitting and improves the model’s ability to generalize.\n",
    "#   - Introduces randomness, simulating different scenarios the model might encounter.\n",
    "#   - Particularly crucial for datasets that are small or lack sufficient diversity.\n",
    "\n",
    "# - **Why These Techniques Were Chosen**:\n",
    "#   - **Noise Addition & Dead Squares**: Simulate missing or noisy data.\n",
    "#   - **Rotation and Flipping**: Increase robustness to orientation changes.\n",
    "#   - **Multiple Augmentation Sets**: Expand the dataset to avoid overfitting.\n",
    "\n",
    "# - **Successes**:\n",
    "#   - Effective in enhancing the model’s robustness and reducing overfitting.\n",
    "#   - Increases data variety without the need for additional labeled samples.\n",
    "\n",
    "# - **Pitfalls**:\n",
    "#   - Excessive augmentation can produce unrealistic patterns that confuse the model.\n",
    "#   - Some grids might lose essential information if augmented too aggressively.\n",
    "\n",
    "# - **Conclusion**:\n",
    "#   - Data augmentation is essential for small or limited datasets like ARC tasks.\n",
    "#   - Proper balancing of augmentation techniques ensures a diverse yet meaningful dataset.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bce4a869-6cf6-4bde-b4b6-02462f96b5b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 6. PyTorch Dataset Class\n",
    "# -----------------------------\n",
    "\n",
    "class AugmentedARCDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, grid_pairs, augment=False):\n",
    "        # Filter out pairs where input or output grid has zero dimensions\n",
    "        self.grid_pairs = [\n",
    "            pair for pair in grid_pairs\n",
    "            if pair.input_grid.size != 0 and pair.output_grid.size != 0\n",
    "            and 0 not in pair.input_grid.shape and 0 not in pair.output_grid.shape\n",
    "        ]\n",
    "        self.augment = augment\n",
    "        logger.info(f\"Dataset initialized with {len(self.grid_pairs)} valid grid pairs.\")\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.grid_pairs)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # Get the GridPair object\n",
    "        pair = self.grid_pairs[idx]\n",
    "\n",
    "        # Access the input and target grids\n",
    "        input_grid = pair.input_grid\n",
    "        target_grid = pair.output_grid\n",
    "\n",
    "        # Apply augmentation if enabled\n",
    "        if self.augment:\n",
    "            input_grid = augment_grid(input_grid)\n",
    "\n",
    "        # Convert grids to tensors\n",
    "        input_tensor = torch.tensor(input_grid, dtype=torch.float32).unsqueeze(0)  # Shape: [1, H, W]\n",
    "        target_tensor = torch.tensor(target_grid, dtype=torch.long)\n",
    "\n",
    "        # Ensure target_tensor is 2D\n",
    "        if target_tensor.dim() > 2:\n",
    "            target_tensor = target_tensor.squeeze()\n",
    "\n",
    "        # Debugging statements\n",
    "        logger.debug(f\"Index {idx}:\")\n",
    "        logger.debug(f\"  Input tensor shape: {input_tensor.shape}\")\n",
    "        logger.debug(f\"  Target tensor shape: {target_tensor.shape}\")\n",
    "\n",
    "        return input_tensor, target_tensor\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f62e0063-4413-4bdd-b1e5-6ba9a1bf7e21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# Explanation of the PyTorch Dataset Class:\n",
    "# -----------------------------\n",
    "# 6. PyTorch Dataset Class\n",
    "# -----------------------------\n",
    "\n",
    "# The `AugmentedARCDataset` class inherits from `torch.utils.data.Dataset`, \n",
    "# which provides a standard interface to create custom datasets in PyTorch. \n",
    "# This dataset class plays a critical role in the training pipeline, handling \n",
    "# data loading, preprocessing, and optional data augmentation. It ensures that \n",
    "# raw data is converted into a format suitable for model training.\n",
    "\n",
    "# 1. __init__(self, grid_pairs, augment=False):\n",
    "#    - **Constructor**: Initializes the dataset with grid pairs and an optional augmentation flag.\n",
    "#    - **Parameters**:\n",
    "#        - **grid_pairs**: A list of `GridPair` objects containing input and output grids.\n",
    "#        - **augment**: A boolean flag indicating whether data augmentation should be applied.\n",
    "#    - **Actions**:\n",
    "#        a. Filters out invalid grid pairs where the input or output grid has zero size or zero dimensions.\n",
    "#            - **Reason**: Prevents runtime errors during training caused by malformed grids.\n",
    "#            - **Implementation**: Uses a list comprehension to iterate through `grid_pairs` and apply filtering.\n",
    "#        b. Sets the `augment` flag to control whether data augmentation is applied during data loading.\n",
    "#        c. Logs the number of valid grid pairs that are initialized in the dataset.\n",
    "#    - **Purpose**:\n",
    "#        - Ensures only valid data is included in the dataset, minimizing the risk of training errors.\n",
    "#        - Provides flexibility to apply data augmentation during training for better generalization.\n",
    "\n",
    "#    *Reason for Addition*: Ensures robust handling of invalid data and offers optional data augmentation.\n",
    "#    *Success*: Prevents invalid grids from propagating into the training loop.\n",
    "#    *Pitfall*: If filtering is too strict, useful data may be excluded inadvertently.\n",
    "\n",
    "# 2. __len__(self):\n",
    "#    - **Purpose**: Returns the total number of items in the dataset.\n",
    "#    - **Details**:\n",
    "#        - This function allows `len(dataset)` to retrieve the dataset size.\n",
    "#        - It’s essential for PyTorch’s DataLoader to determine the number of batches.\n",
    "#    - **Usage**: Ensures compatibility with PyTorch's DataLoader for efficient batching.\n",
    "\n",
    "#    *Reason for Addition*: Supports DataLoader functionality and enables batching.\n",
    "#    *Success*: Allows smooth interaction with PyTorch utilities.\n",
    "#    *Pitfall*: An incorrect length implementation could lead to index out-of-bounds errors.\n",
    "\n",
    "# 3. __getitem__(self, idx):\n",
    "#    - **Purpose**: Retrieves a single data item from the dataset at a specified index.\n",
    "#    - **Parameters**:\n",
    "#        - **idx**: The index of the data item to retrieve.\n",
    "#    - **Actions**:\n",
    "#        a. Retrieves the `GridPair` object at the given index.\n",
    "#        b. Extracts the `input_grid` and `target_grid` from the `GridPair`.\n",
    "#        c. Applies data augmentation to the `input_grid` if `augment` is True.\n",
    "#            - **Implementation**: Uses the `augment_grid()` function to apply noise, flips, or rotations.\n",
    "#        d. Converts the grids into PyTorch tensors:\n",
    "#            - **input_tensor**: \n",
    "#                - Converted to a float32 tensor.\n",
    "#                - Adds an extra channel dimension using `unsqueeze(0)` to match CNN input format (shape [1, H, W]).\n",
    "#            - **target_tensor**: \n",
    "#                - Converted to a long tensor (used for classification tasks).\n",
    "#                - Ensures it is 2D by squeezing unnecessary dimensions.\n",
    "#        e. Logs debugging information about the index and tensor shapes for traceability.\n",
    "#    - **Returns**: \n",
    "#        - A tuple `(input_tensor, target_tensor)` that can be used directly in the training loop.\n",
    "\n",
    "#    *Reason for Addition*: Provides modular data retrieval with preprocessing included.\n",
    "#    *Success*: Centralizes tensor conversion and augmentation logic.\n",
    "#    *Pitfall*: Augmentation might introduce inconsistencies if not managed carefully.\n",
    "\n",
    "# Additional Notes:\n",
    "# - **Integration with DataLoader**: \n",
    "#   - By inheriting from `torch.utils.data.Dataset`, the class can be used with PyTorch’s DataLoader, \n",
    "#     which handles batch creation, shuffling, and parallel data loading.\n",
    "# - **Augmentation Control**: \n",
    "#   - The `augment` parameter provides flexibility to use augmentation only during training, \n",
    "#     avoiding unnecessary transformations during evaluation.\n",
    "# - **Debugging Support**:\n",
    "#   - Includes logging to provide insights into the data processing workflow, \n",
    "#     helping trace potential issues with grid shapes or tensor conversions.\n",
    "# - **Tensor Conversion Details**:\n",
    "#   - Ensuring the input tensor has a channel dimension is crucial for feeding data into convolutional neural networks (CNNs).\n",
    "#   - Guaranteeing that the target tensor is 2D avoids issues during loss computation and predictions.\n",
    "\n",
    "# **Why This Class is Important**:\n",
    "# - **Seamless Data Handling**: \n",
    "#   - Connects raw data with the model training process, ensuring inputs are correctly preprocessed.\n",
    "# - **Efficiency**:\n",
    "#   - Works with PyTorch’s DataLoader for batch processing, which is essential for large datasets.\n",
    "# - **Flexibility**:\n",
    "#   - Allows for easy switching between training with or without augmentation by toggling a single flag.\n",
    "\n",
    "# **Successes**:\n",
    "# - Handles invalid data gracefully by filtering out problematic grid pairs.\n",
    "# - Provides a modular way to apply data augmentation, simplifying the main training loop.\n",
    "\n",
    "# **Pitfalls**:\n",
    "# - Strict filtering criteria may exclude useful data if not carefully tuned.\n",
    "# - Incorrect tensor shapes could cause runtime errors during training if not properly managed.\n",
    "\n",
    "# **Conclusion**:\n",
    "# This dataset class is a crucial component in the PyTorch training pipeline. \n",
    "# It bridges the gap between raw data and model training, handling preprocessing, \n",
    "# tensor conversion, and optional augmentation within a single, modular structure. \n",
    "# The design ensures that only valid data reaches the training loop, while providing \n",
    "# the flexibility to enhance the dataset with augmentation techniques as needed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bb761c7e-3136-4dea-bd4d-8b96208c62de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 7. Define the Deep Neural Network Model\n",
    "# -----------------------------\n",
    "\n",
    "class CNNGridMapper(nn.Module):\n",
    "    def __init__(self, num_classes=NUM_CLASSES):\n",
    "        super(CNNGridMapper, self).__init__()\n",
    "        self.num_classes = num_classes\n",
    "\n",
    "        # Use a CNN backbone (e.g., ResNet18)\n",
    "        self.cnn = resnet18(weights=ResNet18_Weights.DEFAULT)\n",
    "\n",
    "        # Modify the first convolutional layer for single-channel input\n",
    "        self.cnn.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)\n",
    "        nn.init.kaiming_normal_(self.cnn.conv1.weight, mode='fan_out', nonlinearity='relu')\n",
    "\n",
    "        # Remove the fully connected layer\n",
    "        self.cnn_layers = nn.Sequential(*list(self.cnn.children())[:-2])\n",
    "\n",
    "        # Upsampling layers to recover spatial dimensions\n",
    "        self.upsample = nn.Sequential(\n",
    "            nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.ConvTranspose2d(128, num_classes, kernel_size=2, stride=2)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.cnn_layers(x)\n",
    "        x = self.upsample(x)\n",
    "        return x  # Output shape: (batch_size, num_classes, H', W')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c3ac2523-6e3c-42d0-85fd-aa76bc7763d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# Explanation of the Deep Neural Network Model:\n",
    "# -----------------------------\n",
    "# 7. Define the Deep Neural Network Model\n",
    "# -----------------------------\n",
    "\n",
    "# The `CNNGridMapper` class defines a convolutional neural network (CNN) that maps input grids to output grids.\n",
    "# It adapts the well-known ResNet-18 architecture to handle single-channel inputs and produces spatial outputs \n",
    "# suitable for grid-based tasks, such as semantic segmentation.\n",
    "\n",
    "# 1. **Class Definition**:\n",
    "#    - Inherits from `nn.Module`, the base class for all neural network models in PyTorch.\n",
    "#    - The name `CNNGridMapper` reflects its purpose of mapping grids using a convolutional neural network.\n",
    "\n",
    "# 2. **__init__(self, num_classes=NUM_CLASSES)**:\n",
    "#    - **Constructor**: Initializes the network’s layers.\n",
    "#    - **Parameters**:\n",
    "#        - `num_classes`: Specifies the number of output classes, corresponding to the possible values in the grid \n",
    "#          (e.g., 11 classes for values ranging from 0 to 10).\n",
    "#    - **Purpose**: Sets up the model architecture and prepares it for training and inference.\n",
    "\n",
    "#    *Reason for Addition*: Provides a customizable number of output classes to match the target task.\n",
    "\n",
    "# 3. **Using ResNet-18 Backbone**:\n",
    "#    - Loads the ResNet-18 model using:\n",
    "#        - `self.cnn = resnet18(weights=ResNet18_Weights.DEFAULT)`.\n",
    "#    - **Purpose**: ResNet-18 is a well-established architecture with residual connections, making it efficient and \n",
    "#      effective for feature extraction.\n",
    "\n",
    "#    *Success*: Benefiting from transfer learning by using pretrained weights.\n",
    "#    *Pitfall*: Pretrained models expect specific input formats, requiring customization for single-channel inputs.\n",
    "\n",
    "# 4. **Modifying the First Convolutional Layer**:\n",
    "#    - ResNet-18 expects 3-channel RGB images, but our task uses single-channel grid data.\n",
    "#    - We replace the first convolutional layer with:\n",
    "#        - `self.cnn.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)`.\n",
    "#    - **Weights Initialization**: Uses Kaiming Normal initialization to match the ReLU activation function.\n",
    "\n",
    "#    *Reason for Addition*: Adapts the architecture to process single-channel grids.\n",
    "#    *Pitfall*: If weights are not properly initialized, the model may perform poorly.\n",
    "\n",
    "# 5. **Removing the Fully Connected Layer**:\n",
    "#    - Removes the final layers (average pooling and fully connected layers) to convert the model into a fully \n",
    "#      convolutional network (FCN):\n",
    "#        - `self.cnn_layers = nn.Sequential(*list(self.cnn.children())[:-2])`.\n",
    "#    - **Purpose**: Produces spatial feature maps instead of scalar outputs, which is necessary for per-grid classification.\n",
    "\n",
    "#    *Reason for Addition*: Adapts the network for spatial tasks rather than classification.\n",
    "#    *Pitfall*: Removing layers improperly could break the model.\n",
    "\n",
    "# 6. **Upsampling Layers**:\n",
    "#    - The CNN backbone reduces spatial dimensions through pooling and striding.\n",
    "#    - We use transposed convolutional layers (deconvolutions) to increase spatial dimensions:\n",
    "#        - `self.upsample = nn.Sequential(...)`.\n",
    "#    - Each transposed convolution layer doubles the spatial dimensions and applies ReLU activation for non-linearity.\n",
    "\n",
    "#    *Reason for Addition*: Restores spatial dimensions to match the output grid size.\n",
    "#    *Pitfall*: If upsampling is not properly aligned, it may introduce artifacts in the output.\n",
    "\n",
    "# 7. **Upsampling Layer Details**:\n",
    "#    - **First Transposed Convolution**:\n",
    "#        - Input channels: 512 (from the last ResNet-18 layer).\n",
    "#        - Output channels: 256.\n",
    "#    - **Second Transposed Convolution**:\n",
    "#        - Input channels: 256.\n",
    "#        - Output channels: 128.\n",
    "#    - **Third Transposed Convolution**:\n",
    "#        - Input channels: 128.\n",
    "#        - Output channels: `num_classes` (final number of output classes).\n",
    "#    - **Purpose**: These layers progressively increase spatial dimensions to match the target grid size.\n",
    "\n",
    "#    *Success*: Effectively reverses downsampling, restoring spatial dimensions.\n",
    "#    *Pitfall*: Misaligned upsampling can reduce accuracy by distorting feature maps.\n",
    "\n",
    "# 8. **forward(self, x)**:\n",
    "#    - **Defines the forward pass** of the model.\n",
    "#    - **Input**: \n",
    "#        - `x`: A tensor of shape (batch_size, 1, H, W), where 1 is the channel dimension for single-channel grids.\n",
    "#    - **Steps**:\n",
    "#        - Passes the input through the CNN backbone:\n",
    "#            - `x = self.cnn_layers(x)`.\n",
    "#            - Results in feature maps with reduced spatial dimensions.\n",
    "#        - Passes the feature maps through the upsampling layers:\n",
    "#            - `x = self.upsample(x)`.\n",
    "#            - Restores the spatial dimensions to the desired size.\n",
    "#    - **Output**: \n",
    "#        - Returns a tensor of shape (batch_size, num_classes, H', W'), where H' and W' depend on the input size.\n",
    "\n",
    "#    *Reason for Addition*: Implements the core logic for data flow through the network.\n",
    "#    *Success*: Ensures the input is processed correctly and spatial dimensions are restored.\n",
    "#    *Pitfall*: Incorrect shapes could cause runtime errors during training.\n",
    "\n",
    "# 9. **Output Interpretation**:\n",
    "#    - **Purpose**: The output tensor provides class scores (logits) for each grid position.\n",
    "#    - Suitable for tasks like **semantic segmentation**, where each cell in the grid is classified independently.\n",
    "#    - **Training**: \n",
    "#        - Typically, `CrossEntropyLoss` is used as the loss function, which expects raw scores (logits) as input.\n",
    "\n",
    "#    *Reason for Addition*: Provides the correct output format for per-cell classification tasks.\n",
    "\n",
    "# 10. **Notes on Model Design**:\n",
    "#     - **Transfer Learning**: \n",
    "#       - Leveraging a pretrained ResNet-18 improves performance by utilizing features learned from large datasets.\n",
    "#     - **Fully Convolutional Network (FCN)**:\n",
    "#       - Removing the fully connected layers allows the model to produce spatial outputs.\n",
    "#     - **Upsampling via Transposed Convolutions**:\n",
    "#       - Enables the model to learn effective ways to restore spatial dimensions.\n",
    "\n",
    "# Summary:\n",
    "# - The `CNNGridMapper` class adapts the ResNet-18 architecture to process single-channel input grids and generate \n",
    "#   spatial outputs with multiple classes. \n",
    "# - This design is tailored for grid-to-grid mapping tasks, where both the input and output are grids (matrices) of values.\n",
    "# - The combination of **deep feature extraction** (via the CNN backbone) and **upsampling** ensures that the model can \n",
    "#   capture high-level patterns and produce detailed spatial outputs.\n",
    "# - **Successes**:\n",
    "#   - Efficient feature extraction using ResNet-18.\n",
    "#   - Flexible handling of single-channel inputs and spatial outputs.\n",
    "# - **Pitfalls**:\n",
    "#   - Incorrect configuration of upsampling layers could degrade performance.\n",
    "#   - Customizing the ResNet-18 layers requires careful weight initialization and shape alignment.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "32be290e-533f-4363-b502-bc2f977f7151",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 8. Custom Collate Function\n",
    "# -----------------------------\n",
    "\n",
    "def collate_fn(batch):\n",
    "    inputs = [item[0] for item in batch]  # Shape: [C, H, W]\n",
    "    targets = [item[1] for item in batch]  # Shape: [H, W]\n",
    "\n",
    "    # Find max dimensions in the batch for inputs and targets separately\n",
    "    max_input_height = max(t.size(-2) for t in inputs)\n",
    "    max_input_width = max(t.size(-1) for t in inputs)\n",
    "    max_target_height = max(t.size(-2) for t in targets)\n",
    "    max_target_width = max(t.size(-1) for t in targets)\n",
    "\n",
    "    batch_size = len(inputs)\n",
    "    num_channels = inputs[0].size(0)\n",
    "\n",
    "    # Initialize tensors with zeros\n",
    "    batch_inputs = torch.zeros((batch_size, num_channels, max_input_height, max_input_width), dtype=inputs[0].dtype)\n",
    "    batch_targets = torch.zeros((batch_size, max_target_height, max_target_width), dtype=targets[0].dtype)\n",
    "\n",
    "    for i in range(batch_size):\n",
    "        input_tensor = inputs[i]\n",
    "        target_tensor = targets[i]\n",
    "\n",
    "        # Get shapes\n",
    "        c, h_inp, w_inp = input_tensor.size()\n",
    "        h_tar, w_tar = target_tensor.size()\n",
    "\n",
    "        # Copy input_tensor into batch_inputs\n",
    "        batch_inputs[i, :, :h_inp, :w_inp] = input_tensor\n",
    "\n",
    "        # Copy target_tensor into batch_targets\n",
    "        batch_targets[i, :h_tar, :w_tar] = target_tensor\n",
    "\n",
    "        # Debugging statements\n",
    "        logger.debug(f\"Batch index {i}:\")\n",
    "        logger.debug(f\"  Input tensor shape: {input_tensor.shape}\")\n",
    "        logger.debug(f\"  Target tensor shape: {target_tensor.shape}\")\n",
    "        logger.debug(f\"  Batch input shape: {batch_inputs[i].shape}\")\n",
    "        logger.debug(f\"  Batch target shape: {batch_targets[i].shape}\")\n",
    "\n",
    "    return batch_inputs, batch_targets\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "75eddfa5-1214-44de-ad66-9addab03e1db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 8. Custom Collate Function\n",
    "# -----------------------------\n",
    "\n",
    "# def collate_fn(batch):\n",
    "#     \"\"\"\n",
    "#     This custom `collate_fn` is designed to handle batches where the input and target grids may have varying sizes.\n",
    "#     In tasks like the Abstraction and Reasoning Corpus (ARC), grids can have different dimensions, so a standard\n",
    "#     collate function (which assumes all inputs are the same size) would not work.\n",
    "#\n",
    "#     Key Steps and Purpose:\n",
    "#\n",
    "#     1. Extract Inputs and Targets:\n",
    "#         - The function receives a batch, which is a list of tuples where each tuple contains an input tensor and a target tensor.\n",
    "#         - It separates the inputs and targets into two lists for processing.\n",
    "#\n",
    "#     2. Determine Maximum Dimensions:\n",
    "#         - It calculates the maximum height and width among all input tensors (`inputs`) and target tensors (`targets`) in the batch.\n",
    "#         - This step is crucial for creating tensors that can hold all the samples, considering the largest dimensions.\n",
    "#\n",
    "#     3. Initialize Batched Tensors:\n",
    "#         - Creates two zero-filled tensors (`batch_inputs` and `batch_targets`) with shapes:\n",
    "#             - `batch_inputs`: `[batch_size, num_channels, max_input_height, max_input_width]`\n",
    "#             - `batch_targets`: `[batch_size, max_target_height, max_target_width]`\n",
    "#         - These tensors will hold all input and target tensors, padded where necessary.\n",
    "#\n",
    "#     4. Populate Batched Tensors:\n",
    "#         - Iterates over each sample in the batch.\n",
    "#         - For each sample:\n",
    "#             - Retrieves the input and target tensors.\n",
    "#             - Gets their actual dimensions.\n",
    "#             - Copies the input tensor into the corresponding slice of `batch_inputs`.\n",
    "#             - Copies the target tensor into the corresponding slice of `batch_targets`.\n",
    "#         - Since the batched tensors may be larger than the individual tensors, the extra regions remain zero (effectively padding).\n",
    "#\n",
    "#     5. Debugging Statements:\n",
    "#         - Logs detailed shape information for each sample, which is helpful for debugging issues related to tensor dimensions.\n",
    "#\n",
    "#     6. Return Batched Tensors:\n",
    "#         - Returns the `batch_inputs` and `batch_targets` tensors, which can now be used in the training loop.\n",
    "#\n",
    "#     Why This Function is Necessary:\n",
    "#\n",
    "#     - **Variable-Sized Inputs**:\n",
    "#         - In many datasets, especially with images or grids, not all samples are of the same size.\n",
    "#         - The standard DataLoader expects all samples in a batch to have the same dimensions, which isn't the case here.\n",
    "#\n",
    "#     - **Padding to Maximum Size**:\n",
    "#         - By padding all tensors to the maximum size in the batch, we can batch them together.\n",
    "#         - This approach avoids the need to resize or distort the data, preserving the original information.\n",
    "#\n",
    "#     - **Efficiency**:\n",
    "#         - Handling variable-sized data efficiently without writing custom batch handling logic in the training loop.\n",
    "#\n",
    "#     Considerations:\n",
    "#\n",
    "#     - **Memory Usage**:\n",
    "#         - Padding to the maximum size can lead to increased memory usage, especially if there's a large discrepancy between the smallest and largest samples.\n",
    "#         - This can be mitigated by grouping similar-sized samples together (bucketing) or setting a maximum allowable size.\n",
    "#\n",
    "#     - **Model Adaptation**:\n",
    "#         - The model must be able to handle inputs of varying sizes.\n",
    "#         - In this code, the model uses convolutional layers and upsampling, which can work with variable input sizes.\n",
    "#\n",
    "#     - **Loss Function Compatibility**:\n",
    "#         - The loss function and any metric calculations need to account for the padded regions, if necessary.\n",
    "#         - In this implementation, the padding is with zeros, which may correspond to a valid class (e.g., background), so care must be taken.\n",
    "#\n",
    "#     Conclusion:\n",
    "#\n",
    "#     The custom `collate_fn` is an essential component when working with datasets containing variable-sized samples.\n",
    "#     It ensures that data can be batched and fed into the model without losing the integrity of the original samples.\n",
    "#     This function enhances the flexibility and robustness of the data loading pipeline.\n",
    "#     \"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "43354b29-d212-434c-ad38-7c2a8d30e34e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 9. Training GUI Class\n",
    "# -----------------------------\n",
    "\n",
    "class TrainingGUI:\n",
    "    \"\"\"\n",
    "    A Tkinter-based GUI for real-time training progress visualization with 3D metrics plotting and data tree integration.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, root, total_epochs, total_batches, model, train_loader, \n",
    "                 val_loader, eval_loader, device, data_tree, task_dict):\n",
    "        \"\"\"Initialize the Training GUI.\"\"\"\n",
    "        self.root = root\n",
    "        self.total_epochs = total_epochs\n",
    "        self.total_batches = total_batches\n",
    "        self.model = model\n",
    "        self.train_loader = train_loader\n",
    "        self.val_loader = val_loader\n",
    "        self.eval_loader = eval_loader\n",
    "        self.device = device\n",
    "        self.data_tree = data_tree  # Data tree integration\n",
    "        self.task_dict = task_dict  # Store task dictionary for training logic\n",
    "\n",
    "        # Initialize other required attributes\n",
    "        self.queue = queue.Queue()\n",
    "        self.stop_event = threading.Event()\n",
    "\n",
    "        # Initialize data storage for plots\n",
    "        self.loss_data = []\n",
    "        self.val_loss_data = []\n",
    "        self.acc_data = []\n",
    "        self.prediction_distances = []\n",
    "\n",
    "        # Set up the GUI\n",
    "        self.setup_gui()\n",
    "        self.root.after(100, self.process_queue)\n",
    "\n",
    "    def setup_gui(self):\n",
    "        \"\"\"Set up the GUI components.\"\"\"\n",
    "        self.frame = tk.Frame(self.root)\n",
    "        self.frame.pack(fill=tk.BOTH, expand=True)\n",
    "\n",
    "        # Top Section for Labels\n",
    "        self.label_frame = tk.Frame(self.frame)\n",
    "        self.label_frame.pack(pady=10)\n",
    "\n",
    "        self.epoch_label = tk.Label(self.label_frame, text=f\"Epoch: 0/{self.total_epochs}\", font=(\"Helvetica\", 14))\n",
    "        self.epoch_label.grid(row=0, column=0, padx=10)\n",
    "\n",
    "        self.batch_label = tk.Label(self.label_frame, text=f\"Batch: 0/{self.total_batches}\", font=(\"Helvetica\", 12))\n",
    "        self.batch_label.grid(row=0, column=1, padx=10)\n",
    "\n",
    "        self.loss_label = tk.Label(self.label_frame, text=\"Loss: 0.0000\", font=(\"Helvetica\", 12))\n",
    "        self.loss_label.grid(row=0, column=2, padx=10)\n",
    "\n",
    "        self.accuracy_label = tk.Label(self.label_frame, text=\"Accuracy: 0.0000\", font=(\"Helvetica\", 12))\n",
    "        self.accuracy_label.grid(row=0, column=3, padx=10)\n",
    "\n",
    "        # Data Tree Visualization Section\n",
    "        self.tree_frame = tk.Frame(self.frame, width=300, height=400)\n",
    "        self.tree_frame.pack(side=tk.LEFT, padx=10, pady=10, fill=tk.Y)\n",
    "\n",
    "        self.tree_label = tk.Label(self.tree_frame, text=\"Data Tree\", font=(\"Helvetica\", 14))\n",
    "        self.tree_label.pack()\n",
    "\n",
    "        self.tree_canvas = tk.Canvas(self.tree_frame, width=300, height=400, bg='white')\n",
    "        self.tree_canvas.pack()\n",
    "\n",
    "        # Display the data tree\n",
    "        self.display_data_tree()\n",
    "\n",
    "        # Plot Section (2D + 3D)\n",
    "        self.fig = plt.figure(figsize=(12, 6))\n",
    "\n",
    "        # 3D Plot on the Left\n",
    "        self.ax_3d = self.fig.add_subplot(121, projection='3d')\n",
    "        self.ax_3d.set_xlabel('Epoch')\n",
    "        self.ax_3d.set_ylabel('Accuracy')\n",
    "        self.ax_3d.set_zlabel('Distance from Actual')\n",
    "\n",
    "        # 2D Plot on the Right\n",
    "        self.ax_2d = self.fig.add_subplot(122)\n",
    "        self.line_loss, = self.ax_2d.plot([], [], label='Training Loss')\n",
    "        self.line_val_loss, = self.ax_2d.plot([], [], label='Validation Loss')\n",
    "        self.ax_2d.legend()\n",
    "\n",
    "        self.canvas_plot = FigureCanvasTkAgg(self.fig, master=self.frame)\n",
    "        self.canvas_plot.draw()\n",
    "        self.canvas_plot.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=True)\n",
    "\n",
    "        # Bottom Section for Control Buttons\n",
    "        self.button_frame = tk.Frame(self.frame)\n",
    "        self.button_frame.pack(pady=10)\n",
    "\n",
    "        self.start_button = tk.Button(self.button_frame, text=\"Start Training\", command=self.start_training)\n",
    "        self.start_button.grid(row=0, column=0, padx=10)\n",
    "\n",
    "        self.stop_button = tk.Button(self.button_frame, text=\"Stop Training\", command=self.stop_training)\n",
    "        self.stop_button.grid(row=0, column=1, padx=10)\n",
    "\n",
    "        self.evaluate_button = tk.Button(self.button_frame, text=\"Evaluate Model\", command=self.evaluate_model_button)\n",
    "        self.evaluate_button.grid(row=0, column=2, padx=10)\n",
    "\n",
    "    def display_data_tree(self):\n",
    "        \"\"\"Generate and display the data tree as an image.\"\"\"\n",
    "        try:\n",
    "            # Render the tree to a PNG using anytree and Graphviz\n",
    "            dot_file = \"tree.dot\"\n",
    "            png_file = \"tree.png\"\n",
    "\n",
    "            # Export to .dot file\n",
    "            DotExporter(self.data_tree).to_dotfile(dot_file)\n",
    "            logger.info(f\"Tree exported to {dot_file}\")\n",
    "\n",
    "            # Convert .dot to .png using Graphviz\n",
    "            result = os.system(f'dot -Tpng {dot_file} -o {png_file}')\n",
    "            if result != 0:\n",
    "                raise RuntimeError(\"Failed to generate PNG. Ensure Graphviz is installed and in PATH.\")\n",
    "\n",
    "            # Load and display the PNG image\n",
    "            img = Image.open(png_file)\n",
    "            img = img.resize((300, 400), Image.LANCZOS)\n",
    "            img_tk = ImageTk.PhotoImage(img)\n",
    "\n",
    "            # Display the image in the canvas\n",
    "            self.tree_canvas.create_image(0, 0, anchor=tk.NW, image=img_tk)\n",
    "            self.tree_canvas.image = img_tk  # Keep reference to avoid garbage collection\n",
    "            logger.info(\"Tree visualization displayed successfully.\")\n",
    "\n",
    "        except Exception as e:\n",
    "            logger.exception(\"Failed to display the data tree.\")\n",
    "            tk.messagebox.showerror(\"Tree Display Error\", f\"Error: {e}\")\n",
    "\n",
    "    def process_queue(self):\n",
    "        \"\"\"Process the queue for thread-safe GUI updates.\"\"\"\n",
    "        while not self.queue.empty():\n",
    "            message = self.queue.get()\n",
    "            if isinstance(message, dict):\n",
    "                self.update_gui(message)\n",
    "        self.root.after(100, self.process_queue)\n",
    "\n",
    "    def update_gui(self, data):\n",
    "        \"\"\"Update the GUI with real-time training and validation metrics.\"\"\"\n",
    "        try:\n",
    "            if 'batch' in data:\n",
    "                # Update batch-level metrics in the GUI\n",
    "                self.batch_label.config(text=f\"Batch: {data['batch']}/{self.total_batches}\")\n",
    "                self.loss_label.config(text=f\"Loss: {data['loss']:.4f}\")\n",
    "                self.accuracy_label.config(text=f\"Accuracy: {data.get('accuracy', 0.0):.4f}\")\n",
    "\n",
    "                # Append new batch data to the 2D plot lists\n",
    "                self.loss_data.append(data['loss'])\n",
    "                self.acc_data.append(data.get('accuracy', 0.0))\n",
    "\n",
    "                # Update the 2D plot with each batch completion\n",
    "                batches = list(range(1, len(self.loss_data) + 1))\n",
    "                self.line_loss.set_data(batches, self.loss_data)\n",
    "                self.line_val_loss.set_data(batches, self.acc_data)\n",
    "\n",
    "                # Adjust the axes to fit the new data\n",
    "                self.ax_2d.relim()\n",
    "                self.ax_2d.autoscale_view()\n",
    "\n",
    "                # Redraw the 2D plot with new data\n",
    "                self.canvas_plot.draw()\n",
    "\n",
    "            elif 'epoch' in data:\n",
    "                # Update epoch-level metrics\n",
    "                self.epoch_label.config(text=f\"Epoch: {data['epoch']}/{self.total_epochs}\")\n",
    "                # Handle validation loss and accuracy if provided\n",
    "                val_loss = data.get('val_loss', 0.0)\n",
    "                val_accuracy = data.get('val_accuracy', 0.0)\n",
    "                # Update labels if you have labels for validation metrics\n",
    "\n",
    "                # Calculate prediction error distance\n",
    "                predicted = np.array(data.get('predicted', []))\n",
    "                actual = np.array(data.get('actual', []))\n",
    "\n",
    "                if predicted.size == 0 or actual.size == 0:\n",
    "                    distance = float('nan')  # Handle empty arrays gracefully\n",
    "                elif predicted.shape != actual.shape:\n",
    "                    distance = float('nan')  # Handle shape mismatch\n",
    "                else:\n",
    "                    distance = np.abs(predicted - actual).mean()\n",
    "\n",
    "                # Replace NaN with 0.0 for plotting purposes\n",
    "                distance = 0.0 if np.isnan(distance) else distance\n",
    "\n",
    "                # Store valid distances for 3D plot\n",
    "                self.prediction_distances.append(distance)\n",
    "\n",
    "                # Update the 3D plot\n",
    "                epochs = list(range(1, len(self.prediction_distances) + 1))\n",
    "                self.ax_3d.clear()\n",
    "                self.ax_3d.set_xlabel('Epoch')\n",
    "                self.ax_3d.set_ylabel('Accuracy')\n",
    "                self.ax_3d.set_zlabel('Distance from Actual')\n",
    "                self.ax_3d.set_title('3D Prediction Error vs Accuracy')\n",
    "\n",
    "                # Scatter plot with prediction distances\n",
    "                self.ax_3d.scatter(epochs, self.acc_data, self.prediction_distances, label='Error vs Accuracy', color='green')\n",
    "                self.ax_3d.legend()\n",
    "\n",
    "                # Redraw the 3D plot\n",
    "                self.canvas_plot.draw()\n",
    "\n",
    "        except Exception as e:\n",
    "            logger.exception(\"An error occurred while updating the GUI.\")\n",
    "            messagebox.showerror(\"Error\", f\"An error occurred: {e}\")\n",
    "\n",
    "    def start_training(self):\n",
    "        \"\"\"Start training in a new thread.\"\"\"\n",
    "        self.stop_event.clear()\n",
    "        threading.Thread(target=self.train_thread, daemon=True).start()\n",
    "\n",
    "    def stop_training(self):\n",
    "        \"\"\"Stop the training process.\"\"\"\n",
    "        self.stop_event.set()\n",
    "\n",
    "    def train_thread(self):\n",
    "        \"\"\"Training logic executed in a separate thread to avoid blocking the GUI.\"\"\"\n",
    "        logger.info(\"Training thread started.\")\n",
    "\n",
    "        # Optimizer, scheduler, and criterion setup\n",
    "        optimizer = torch.optim.AdamW(self.model.parameters(), lr=0.01, weight_decay=1e-4)\n",
    "        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)\n",
    "        criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "        # Mixed precision scaler (if using CUDA)\n",
    "        scaler = torch.cuda.amp.GradScaler() if self.device.type == 'cuda' else None\n",
    "\n",
    "        # Loop over epochs\n",
    "        for epoch in range(self.total_epochs):\n",
    "            if self.stop_event.is_set():\n",
    "                logger.info(\"Training stopped by user.\")\n",
    "                break\n",
    "            logger.info(f\"Starting epoch {epoch + 1}/{self.total_epochs}.\")\n",
    "            self.model.train()  # Set the model in training mode\n",
    "            running_loss = 0.0\n",
    "            correct = 0\n",
    "            total = 0\n",
    "\n",
    "            # Loop over batches\n",
    "            for batch_idx, (inputs, targets) in enumerate(self.train_loader, 1):\n",
    "                if self.stop_event.is_set():\n",
    "                    logger.info(\"Training stopped by user.\")\n",
    "                    break\n",
    "                try:\n",
    "                    # Train the batch and gather metrics\n",
    "                    batch_loss, batch_accuracy, batch_size = self.train_batch(\n",
    "                        batch_idx, inputs, targets, optimizer, scaler, criterion\n",
    "                    )\n",
    "\n",
    "                    running_loss += batch_loss * batch_size\n",
    "                    correct += int(batch_accuracy * batch_size / 100)\n",
    "                    total += targets.numel()\n",
    "\n",
    "                    # Update the GUI every 10 batches or at the end of epoch\n",
    "                    if batch_idx % 10 == 0 or batch_idx == len(self.train_loader):\n",
    "                        gui_batch_loss = running_loss / total\n",
    "                        gui_batch_accuracy = 100.0 * correct / total\n",
    "                        self.queue.put({\n",
    "                            'batch': batch_idx,\n",
    "                            'loss': gui_batch_loss,\n",
    "                            'accuracy': gui_batch_accuracy\n",
    "                        })\n",
    "\n",
    "                except Exception as e:\n",
    "                    logger.exception(f\"Error in batch {batch_idx}: {e}\")\n",
    "                    continue  # Continue with the next batch if an error occurs\n",
    "\n",
    "            # Epoch metrics\n",
    "            epoch_loss = running_loss / total\n",
    "            epoch_accuracy = 100.0 * correct / total\n",
    "\n",
    "            # Send epoch updates to the GUI\n",
    "            self.queue.put({\n",
    "                'epoch': epoch + 1,\n",
    "                'loss': epoch_loss,\n",
    "                'accuracy': epoch_accuracy\n",
    "            })\n",
    "\n",
    "            # Scheduler step\n",
    "            scheduler.step()\n",
    "\n",
    "        logger.info(\"Training completed.\")\n",
    "        self.queue.put({'status': 'Training Completed'})\n",
    "\n",
    "    def train_batch(self, batch_idx, inputs, targets, optimizer, scaler, criterion):\n",
    "        \"\"\"Train a single batch.\"\"\"\n",
    "        # Move data to the appropriate device\n",
    "        inputs, targets = inputs.to(self.device), targets.to(self.device)\n",
    "\n",
    "        # Reset gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        if self.device.type == 'cuda' and scaler is not None:\n",
    "            # Mixed precision training with autocast\n",
    "            with torch.autocast(device_type=self.device.type, enabled=True):\n",
    "                outputs = self.model(inputs)\n",
    "                # Resize outputs to match targets\n",
    "                outputs = F.interpolate(outputs, size=targets.shape[1:], mode='bilinear', align_corners=False)\n",
    "                loss = criterion(outputs, targets)\n",
    "\n",
    "            # Backward pass and optimizer step with scaler\n",
    "            scaler.scale(loss).backward()\n",
    "            scaler.step(optimizer)\n",
    "            scaler.update()\n",
    "        else:\n",
    "            # Standard training without autocast\n",
    "            outputs = self.model(inputs)\n",
    "            # Resize outputs to match targets\n",
    "            outputs = F.interpolate(outputs, size=targets.shape[1:], mode='bilinear', align_corners=False)\n",
    "            loss = criterion(outputs, targets)\n",
    "\n",
    "            # Backward pass and optimizer step\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "        # Compute metrics\n",
    "        batch_loss = loss.item()\n",
    "        _, predicted = outputs.max(1)\n",
    "        correct_predictions = predicted.eq(targets).sum().item()\n",
    "        batch_accuracy = 100.0 * correct_predictions / targets.numel()\n",
    "\n",
    "        return batch_loss, batch_accuracy, targets.numel()\n",
    "\n",
    "    def evaluate_model_button(self):\n",
    "        \"\"\"Evaluate the model in a new thread.\"\"\"\n",
    "        threading.Thread(target=self.evaluate_model, daemon=True).start()\n",
    "\n",
    "    def evaluate_model(self):\n",
    "        \"\"\"Evaluate the model.\"\"\"\n",
    "        avg_loss, accuracy = evaluate_model(self.model, self.val_loader, self.device)\n",
    "        messagebox.showinfo(\"Evaluation\", f\"Validation Loss: {avg_loss:.4f}, Accuracy: {accuracy:.4f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "c9a1dc8f-dd94-4983-ad2d-ba6422cb3328",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 9. Explanation of the TrainingGUI Class\n",
    "# -----------------------------\n",
    "\n",
    "# The `TrainingGUI` class is responsible for creating a graphical user interface (GUI) to visualize the \n",
    "# training process of a neural network model in real-time. It uses Tkinter for GUI elements and Matplotlib \n",
    "# for plotting training metrics. This class enhances user interaction and provides real-time insights into the \n",
    "# model’s performance.\n",
    "\n",
    "# 1. **Class Initialization (__init__ method)**:\n",
    "#    - **Purpose**: Initializes the GUI and prepares all necessary components for tracking and controlling training.\n",
    "#    - **Parameters**:\n",
    "#        - `root`: The root window of the Tkinter GUI.\n",
    "#        - `total_epochs`, `total_batches`: Total number of epochs and batches for progress tracking.\n",
    "#        - `model`: The neural network model being trained.\n",
    "#        - `train_loader`, `val_loader`, `eval_loader`: Data loaders for training, validation, and evaluation.\n",
    "#        - `device`: The device (CPU/GPU) used for training.\n",
    "#        - `data_tree`: A hierarchical tree structure representing the dataset, used for visualization.\n",
    "#        - `task_dict`: Dictionary containing task-related data to be referenced during training.\n",
    "#    - **Actions**:\n",
    "#        - Initializes attributes for managing threads and data for plotting.\n",
    "#        - Calls `setup_gui()` to build the interface layout.\n",
    "#        - Uses `root.after()` to schedule periodic updates from the queue.\n",
    "#    \n",
    "#    **Success**: Separates the GUI setup from the main training logic, ensuring modularity and easier maintenance.\n",
    "#    **Pitfall**: If the periodic update interval is not well-tuned, the GUI may lag or feel unresponsive during large datasets.\n",
    "\n",
    "# 2. **setup_gui method**:\n",
    "#    - **Purpose**: Constructs the layout and components of the GUI using Tkinter.\n",
    "#    - **Sections**:\n",
    "#        - **Top Section**: Displays metrics like the current epoch, batch number, loss, and accuracy.\n",
    "#        - **Data Tree Visualization**:\n",
    "#            - Creates a canvas to display the dataset’s hierarchical tree.\n",
    "#            - Uses `display_data_tree()` to render the tree as an image.\n",
    "#        - **Plot Section**:\n",
    "#            - Initializes a **3D plot** for prediction error versus accuracy and a **2D plot** for loss metrics.\n",
    "#            - Embeds these plots into the GUI using `FigureCanvasTkAgg`.\n",
    "#        - **Bottom Section**: Adds buttons for starting, stopping, and evaluating the model.\n",
    "#    \n",
    "#    **Success**: Provides a clear, structured interface with real-time feedback on metrics.\n",
    "#    **Pitfall**: Overpopulating the GUI with too many metrics can overwhelm users, reducing usability.\n",
    "\n",
    "# 3. **display_data_tree method**:\n",
    "#    - **Purpose**: Visualizes the dataset structure as a tree diagram.\n",
    "#    - **Implementation**:\n",
    "#        - Uses `anytree` and `Graphviz` to export the tree structure to a DOT file and convert it to a PNG image.\n",
    "#        - Displays the PNG within the GUI canvas.\n",
    "#        - Catches and handles exceptions gracefully if the tree cannot be rendered.\n",
    "#    \n",
    "#    **Success**: Offers users a visual understanding of the dataset structure.\n",
    "#    **Pitfall**: External dependencies (e.g., Graphviz) can introduce errors if not properly installed.\n",
    "\n",
    "# 4. **process_queue method**:\n",
    "#    - **Purpose**: Periodically processes messages from a thread-safe queue to update the GUI.\n",
    "#    - **Implementation**:\n",
    "#        - Uses `root.after()` to ensure regular, non-blocking updates.\n",
    "#        - Invokes `update_gui()` with data from the queue.\n",
    "#    \n",
    "#    **Success**: Keeps the GUI responsive during long training sessions by offloading updates to a separate thread.\n",
    "#    **Pitfall**: If too many messages accumulate in the queue, GUI updates might lag behind.\n",
    "\n",
    "# 5. **update_gui method**:\n",
    "#    - **Purpose**: Updates the GUI components based on received data.\n",
    "#    - **Batch-Level Updates**:\n",
    "#        - Updates metrics like batch number, loss, and accuracy in real-time.\n",
    "#        - Appends data to lists for dynamic plotting.\n",
    "#        - Refreshes the 2D loss plot.\n",
    "#    - **Epoch-Level Updates**:\n",
    "#        - Updates the current epoch label.\n",
    "#        - Computes prediction error for the 3D plot.\n",
    "#        - Adds new points to the 3D plot.\n",
    "#    \n",
    "#    **Success**: Provides detailed insights into the model’s performance during both batch and epoch levels.\n",
    "#    **Pitfall**: Frequent GUI updates may slow down training, especially with complex plots.\n",
    "\n",
    "# 6. **start_training method**:\n",
    "#    - **Purpose**: Starts the training process in a new thread to avoid blocking the GUI.\n",
    "#    - **Implementation**: Clears existing stop events and initiates a new training session.\n",
    "#    \n",
    "#    **Success**: Keeps the GUI interactive by running training in a separate thread.\n",
    "#    **Pitfall**: Improper thread handling can cause crashes or deadlocks if not managed carefully.\n",
    "\n",
    "# 7. **stop_training method**:\n",
    "#    - **Purpose**: Signals the training thread to halt gracefully using a stop event.\n",
    "#    \n",
    "#    **Success**: Allows users to interrupt training safely without crashing the application.\n",
    "#    **Pitfall**: If not handled properly, stopping the thread could leave the model in an inconsistent state.\n",
    "\n",
    "# 8. **train_thread method**:\n",
    "#    - **Purpose**: Manages the main training loop in a separate thread.\n",
    "#    - **Implementation**:\n",
    "#        - Checks for stop events to allow user interruption.\n",
    "#        - Uses `train_batch()` for training on individual batches.\n",
    "#        - Aggregates metrics and updates the GUI via the queue.\n",
    "#        - Steps the learning rate scheduler after each epoch.\n",
    "#        - Sends a completion message to the GUI after finishing training.\n",
    "#    \n",
    "#    **Success**: Ensures smooth training without blocking the GUI thread.\n",
    "#    **Pitfall**: If the stop event is not properly monitored, training could run indefinitely.\n",
    "\n",
    "# 9. **train_batch method**:\n",
    "#    - **Purpose**: Handles training on a single batch of data.\n",
    "#    - **Steps**:\n",
    "#        - Moves data to the appropriate device (CPU/GPU).\n",
    "#        - Resets gradients, performs forward and backward passes, and updates the optimizer.\n",
    "#        - Supports mixed precision training if CUDA is available.\n",
    "#        - Returns batch metrics (loss, accuracy) for aggregation.\n",
    "#    \n",
    "#    **Success**: Efficiently trains on batches while utilizing GPU resources when available.\n",
    "#    **Pitfall**: Incorrect data handling (e.g., forgetting to reset gradients) could affect model performance.\n",
    "\n",
    "# 10. **evaluate_model_button and evaluate_model methods**:\n",
    "#     - **Purpose**: Evaluates the model on the validation dataset.\n",
    "#     - **Implementation**:\n",
    "#         - Runs the evaluation in a separate thread to maintain GUI responsiveness.\n",
    "#         - Displays results in a message box once evaluation is complete.\n",
    "#    \n",
    "#    **Success**: Provides a seamless way to assess the model’s performance without interrupting the GUI.\n",
    "#    **Pitfall**: Running evaluations too frequently could degrade performance or cause GUI lags.\n",
    "\n",
    "# **General Notes**:\n",
    "# - **Threading**: The training process runs in a separate thread, keeping the GUI responsive.\n",
    "# - **Queue Communication**: A thread-safe queue ensures safe communication between the training thread and the GUI.\n",
    "# - **Dynamic Plotting**: Matplotlib plots offer real-time feedback on training metrics.\n",
    "# - **Error Handling**: Proper error handling prevents crashes and keeps the application stable.\n",
    "# - **Interactive Controls**: Buttons allow users to start, stop, and evaluate the model interactively.\n",
    "# - **Data Tree Visualization**: The hierarchical dataset structure provides useful insights into the training data.\n",
    "\n",
    "# **Conclusion**:\n",
    "# The `TrainingGUI` class serves as a powerful tool for monitoring and controlling the training process. It combines \n",
    "# real-time metrics visualization, user interaction, and multi-threaded execution to create a smooth and informative \n",
    "# experience for model development and analysis.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "02b5393e-97f3-4c6e-bc0e-c63f80f0a9dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 10. Training Function with GUI Integration\n",
    "# -----------------------------\n",
    "\n",
    "def train_model_with_gui(model, train_loader, val_loader, device, gui):\n",
    "    \"\"\"Train the model and update the GUI in real-time.\"\"\"\n",
    "    try:\n",
    "        # Start the GUI training display\n",
    "        gui.start_training()\n",
    "\n",
    "    except Exception as e:\n",
    "        logger.exception(f\"Training failed: {e}\")\n",
    "        gui.queue.put({'error': str(e)})  # Inform the GUI about the error\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4008fb84-f110-42a1-8ac1-f1facd9a4c6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 10. Explanation of the Training Function with GUI Integration\n",
    "# -----------------------------\n",
    "\n",
    "# The `train_model_with_gui` function acts as a bridge between the training logic and the GUI.\n",
    "# It starts the training process while ensuring the GUI remains responsive with real-time updates.\n",
    "\n",
    "# **Function Definition**:\n",
    "# def train_model_with_gui(model, train_loader, val_loader, device, gui):\n",
    "\n",
    "# **Parameters**:\n",
    "# - `model`: The neural network model to be trained.\n",
    "# - `train_loader`: DataLoader for the training dataset.\n",
    "# - `val_loader`: DataLoader for the validation dataset.\n",
    "# - `device`: The computing device (CPU or GPU) to perform training.\n",
    "# - `gui`: An instance of the `TrainingGUI` class that manages the GUI.\n",
    "\n",
    "# **Purpose**:\n",
    "# - Initiates the training process while keeping the GUI responsive.\n",
    "# - Encapsulates the logic for starting the training and handling any exceptions.\n",
    "\n",
    "# **Key Steps**:\n",
    "\n",
    "# 1. **Start Training**:\n",
    "#    - Calls `gui.start_training()` to begin the training loop in a new thread.\n",
    "#    - **Reason**: \n",
    "#        - Running training on a separate thread prevents the GUI from freezing.\n",
    "#        - This allows for **real-time updates** and **user interaction** during training.\n",
    "\n",
    "#    **Success**: Keeps the GUI responsive, enhancing the user experience.\n",
    "#    **Pitfall**: Incorrect thread management could cause race conditions or crashes.\n",
    "\n",
    "# 2. **Exception Handling**:\n",
    "#    - Wraps the training initiation in a `try-except` block to catch and handle any exceptions.\n",
    "#    - **Actions**:\n",
    "#        - Uses `logger.exception()` to log the error with a full stack trace.\n",
    "#        - Sends the error message to the GUI’s **message queue** using:\n",
    "#            - `gui.queue.put({'error': str(e)})`\n",
    "#    - **Reason**:\n",
    "#        - This approach ensures the GUI displays the error message without directly calling GUI methods from the training thread, which could violate thread safety.\n",
    "\n",
    "#    **Success**: Ensures smooth error reporting without crashing the application.\n",
    "#    **Pitfall**: If exceptions are not properly propagated, the user may be unaware of underlying issues.\n",
    "\n",
    "# **Notes**:\n",
    "# - The function itself **does not handle the core training logic**; it simply initiates training through the `TrainingGUI` instance.\n",
    "# - **Separation of Concerns**: By delegating the training logic to the GUI, this function keeps the main workflow clean and organized.\n",
    "# - **Best Practices**: Adheres to best practices for GUI applications by ensuring that long-running tasks (like training) do not block the main thread.\n",
    "\n",
    "# **Conclusion**:\n",
    "# The `train_model_with_gui` function is crucial for integrating the model training process with the GUI. \n",
    "# It ensures:\n",
    "# - The **training process begins smoothly** in a non-blocking way.\n",
    "# - **Exceptions are properly handled** and reported through the GUI’s message queue.\n",
    "# - The GUI remains **responsive and interactive** throughout the training process, providing a better user experience.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "fc1a7ee8-2943-43da-af22-b616c6e4ebff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 11. Evaluation Function\n",
    "# -----------------------------\n",
    "\n",
    "def evaluate_model(model, test_loader, device='cpu'):\n",
    "    \"\"\"\n",
    "    Evaluates the model on the test dataset.\n",
    "\n",
    "    Args:\n",
    "        model (nn.Module): Trained model.\n",
    "        test_loader (DataLoader): DataLoader for the test dataset.\n",
    "        device (str): Device to run evaluation on.\n",
    "\n",
    "    Returns:\n",
    "        tuple: (average_loss, accuracy)\n",
    "    \"\"\"\n",
    "    criterion = nn.CrossEntropyLoss()  # Loss function\n",
    "    model.to(device)  # Move model to the appropriate device\n",
    "    model.eval()  # Set the model to evaluation mode\n",
    "\n",
    "    total_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for inputs, targets in test_loader:\n",
    "            inputs = inputs.to(device, non_blocking=True)\n",
    "            targets = targets.to(device, non_blocking=True)\n",
    "\n",
    "            # Forward pass\n",
    "            outputs = model(inputs)\n",
    "            # Resize outputs to match targets\n",
    "            outputs = F.interpolate(outputs, size=targets.shape[1:], mode='bilinear', align_corners=False)\n",
    "\n",
    "            # Calculate loss\n",
    "            loss = criterion(outputs, targets)\n",
    "            total_loss += loss.item() * inputs.size(0)  # Accumulate weighted loss\n",
    "\n",
    "            # Calculate accuracy\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            correct += (predicted == targets).sum().item()\n",
    "            total += targets.numel()\n",
    "\n",
    "    # Calculate average loss and accuracy\n",
    "    avg_loss = total_loss / len(test_loader.dataset)\n",
    "    accuracy = correct / total\n",
    "\n",
    "    logger.info(f\"Test Loss: {avg_loss:.4f}, Test Accuracy: {accuracy:.4f}\")\n",
    "    return avg_loss, accuracy\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b3a39437-6bd1-4fc7-9261-29bf81040895",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 11. Explanation of the Evaluation Function\n",
    "# -----------------------------\n",
    "\n",
    "# The `evaluate_model` function assesses the performance of a trained neural network on a test dataset.\n",
    "# It calculates the **average loss** and **overall accuracy** across the test set, providing essential metrics \n",
    "# for evaluating model generalization on unseen data.\n",
    "\n",
    "# **Function Definition**:\n",
    "# def evaluate_model(model, test_loader, device):\n",
    "\n",
    "# **Parameters**:\n",
    "# - `model` (nn.Module): The trained PyTorch model to be evaluated.\n",
    "# - `test_loader` (DataLoader): DataLoader providing an iterable over the test dataset.\n",
    "# - `device` (str): The device ('cpu' or 'cuda') on which to perform the evaluation.\n",
    "\n",
    "# **Steps**:\n",
    "\n",
    "# 1. **Initialize the Loss Function**:\n",
    "#    - `criterion = nn.CrossEntropyLoss()` initializes the loss function for multi-class classification.\n",
    "#    \n",
    "#    **Success**: CrossEntropyLoss is appropriate for multi-class tasks.\n",
    "#    **Pitfall**: Using the wrong loss function could result in incorrect evaluations.\n",
    "\n",
    "# 2. **Prepare the Model**:\n",
    "#    - Moves the model to the specified device: `model.to(device)`.\n",
    "#    - Sets the model to evaluation mode: `model.eval()`.\n",
    "#        - This ensures layers like **dropout** and **batch normalization** behave correctly during evaluation.\n",
    "#    \n",
    "#    **Success**: Ensures proper inference behavior by disabling training-specific behaviors.\n",
    "#    **Pitfall**: Forgetting to switch to `eval()` mode could yield misleading results.\n",
    "\n",
    "# 3. **Initialize Metrics**:\n",
    "#    - `total_loss`: Accumulates the total loss across all batches.\n",
    "#    - `correct`: Counts the number of correctly predicted samples.\n",
    "#    - `total`: Tracks the total number of elements evaluated.\n",
    "\n",
    "#    **Success**: Provides reliable tracking of performance across batches.\n",
    "#    **Pitfall**: Incorrect metric initialization could skew results.\n",
    "\n",
    "# 4. **Disable Gradient Computation**:\n",
    "#    - Wrapping the evaluation in `with torch.no_grad():` saves memory and improves performance by disabling gradients.\n",
    "#    \n",
    "#    **Success**: Reduces unnecessary computation, speeding up the evaluation.\n",
    "#    **Pitfall**: Forgetting to disable gradients can cause memory leaks.\n",
    "\n",
    "# 5. **Iterate Over Test Data**:\n",
    "#    - Loops through each batch provided by `test_loader`.\n",
    "\n",
    "# 6. **Move Data to Device**:\n",
    "#    - Moves input and target tensors to the specified device:\n",
    "#        - `inputs.to(device, non_blocking=True)` enables asynchronous GPU transfers when possible.\n",
    "#    \n",
    "#    **Success**: Ensures data is processed on the correct device for efficient computation.\n",
    "#    **Pitfall**: Mismatched devices can result in runtime errors.\n",
    "\n",
    "# 7. **Forward Pass**:\n",
    "#    - `outputs = model(inputs)` computes the raw output (logits) of the model.\n",
    "\n",
    "# 8. **Resize Outputs (if necessary)**:\n",
    "#    - `F.interpolate(outputs, size=targets.shape[1:], mode='bilinear', align_corners=False)` adjusts the output size to match the target tensor.\n",
    "#        - This is particularly useful for tasks like **semantic segmentation**, where output dimensions may differ.\n",
    "#    \n",
    "#    **Success**: Ensures output matches target dimensions for correct loss calculation.\n",
    "#    **Pitfall**: Incorrect resizing can lead to dimension mismatches or performance degradation.\n",
    "\n",
    "# 9. **Compute Loss**:\n",
    "#    - `loss = criterion(outputs, targets)` calculates the loss between model outputs and targets.\n",
    "#    - Accumulates the weighted loss: `total_loss += loss.item() * inputs.size(0)`.\n",
    "#    \n",
    "#    **Success**: Accurately tracks loss across batches, accounting for batch size variability.\n",
    "#    **Pitfall**: Failing to weight the loss by batch size could result in incorrect average loss.\n",
    "\n",
    "# 10. **Compute Predictions and Accuracy**:\n",
    "#     - `_, predicted = torch.max(outputs, 1)` extracts the predicted class with the highest score.\n",
    "#     - `correct += (predicted == targets).sum().item()` increments the correct prediction count.\n",
    "#     - `total += targets.numel()` increments the total elements evaluated.\n",
    "#    \n",
    "#    **Success**: Tracks accuracy effectively across all batches.\n",
    "#    **Pitfall**: Misaligned predictions and targets can result in incorrect accuracy metrics.\n",
    "\n",
    "# 11. **Calculate Average Metrics**:\n",
    "#     - `avg_loss = total_loss / len(test_loader.dataset)` computes the average loss per sample.\n",
    "#     - `accuracy = correct / total` computes the overall accuracy.\n",
    "#    \n",
    "#    **Success**: Provides meaningful insights into model performance.\n",
    "#    **Pitfall**: Incorrect total counts could skew the average loss and accuracy.\n",
    "\n",
    "# 12. **Logging and Return**:\n",
    "#     - `logger.info(...)` logs the evaluation metrics for monitoring and debugging.\n",
    "#     - `return avg_loss, accuracy` returns the computed metrics.\n",
    "\n",
    "# **Notes**:\n",
    "# - **Evaluation Mode**: Setting the model to `eval()` ensures correct behavior of layers like dropout and batch normalization.\n",
    "# - **No Gradient Computation**: Wrapping the code in `torch.no_grad()` saves memory and speeds up evaluation.\n",
    "# - **Resizing Outputs**: The use of `F.interpolate` handles cases where model outputs and target sizes differ.\n",
    "# - **Loss Weighting**: Weighting the loss by batch size ensures correct metric calculation across variable-sized batches.\n",
    "# - **Accuracy Calculation**: Accurate predictions and correct counts are critical for reliable accuracy metrics.\n",
    "\n",
    "# **Conclusion**:\n",
    "# The `evaluate_model` function provides a systematic approach to assessing model performance on unseen data. \n",
    "# It ensures:\n",
    "# - Proper **loss and accuracy calculation** across batches.\n",
    "# - Correct **model behavior** during evaluation by switching to `eval()` mode.\n",
    "# - Efficient use of memory and computation by disabling gradient tracking.\n",
    "\n",
    "# This function offers valuable insights into the model’s generalization capabilities, guiding further improvements \n",
    "# and adjustments to the model as needed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3409a37e-6a5e-4618-8cf4-278b2ecb6662",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 12. Traverse and Debug Function\n",
    "# -----------------------------\n",
    "\n",
    "def traverse_and_debug(node):\n",
    "    \"\"\"Traverse the tree and log node details with reduced verbosity.\"\"\"\n",
    "    # Limit logging to fewer nodes\n",
    "    if hasattr(node, \"name\"):\n",
    "        grid_shape = getattr(node, 'embedding', None)\n",
    "        grid_shape = grid_shape.shape if grid_shape is not None else 'Missing'\n",
    "        logger.debug(f\"Node: {node.name}, Grid Shape: {grid_shape}\")\n",
    "    if len(node.children) > 10:  # Avoid excessive logging if many children exist\n",
    "        logger.warning(f\"Node '{node.name}' has too many children, skipping further logs...\")\n",
    "        return\n",
    "    for child in node.children:\n",
    "        traverse_and_debug(child)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2ce25483-51c3-4d73-8136-e4f07755d5ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 12. Explanation of the Traverse and Debug Function\n",
    "# -----------------------------\n",
    "\n",
    "# The `traverse_and_debug` function is designed to traverse a **tree data structure** starting from a given node.\n",
    "# It logs details about each node but includes mechanisms to reduce verbosity, preventing excessive logging,\n",
    "# especially when dealing with nodes containing many children.\n",
    "\n",
    "# **Function Definition**:\n",
    "# def traverse_and_debug(node):\n",
    "#     \"\"\"Traverse the tree and log node details with reduced verbosity.\"\"\"\n",
    "\n",
    "# **Key Steps**:\n",
    "\n",
    "# 1. **Check for 'name' Attribute**:\n",
    "#    - Uses `hasattr(node, \"name\")` to check if the node has a 'name' attribute.\n",
    "#    - **Reason**: Not all nodes may have the 'name' attribute, and accessing it without checking could raise an `AttributeError`.\n",
    "#\n",
    "#    **Success**: Prevents errors by safely accessing attributes.\n",
    "#    **Pitfall**: If essential attributes are missing, it may reduce the usefulness of the logs.\n",
    "\n",
    "# 2. **Retrieve the Node's Embedding**:\n",
    "#    - Uses `getattr(node, 'embedding', None)` to safely retrieve the embedding attribute.\n",
    "#    - **Fallback**: Defaults to `None` if the 'embedding' attribute is not present.\n",
    "#\n",
    "#    **Success**: Handles missing attributes gracefully.\n",
    "#    **Pitfall**: Important debugging information may be missed if many embeddings are missing.\n",
    "\n",
    "# 3. **Determine Grid Shape**:\n",
    "#    - If the 'embedding' is present, the function accesses its `shape` attribute.\n",
    "#    - If the 'embedding' is `None`, sets `grid_shape` to `'Missing'`.\n",
    "#\n",
    "#    **Success**: Ensures logs reflect whether the grid shape is available or missing.\n",
    "#    **Pitfall**: Complex embeddings with unexpected shapes could still cause issues if not validated.\n",
    "\n",
    "# 4. **Log Node Information**:\n",
    "#    - Logs the node’s **name** and **grid shape** using `logger.debug`.\n",
    "#\n",
    "#    **Success**: Provides valuable information for debugging tree structures.\n",
    "#    **Pitfall**: Excessive logging could overwhelm the log files if not carefully managed.\n",
    "\n",
    "# 5. **Limit Logging for Nodes with Many Children**:\n",
    "#    - If the node has more than **10 children**, logs a warning and skips further traversal.\n",
    "#    - Uses:\n",
    "#        - `logger.warning(f\"Node '{node.name}' has too many children, skipping further logs...\")`\n",
    "#        - Returns early to avoid logging all children.\n",
    "#\n",
    "#    **Success**: Prevents overwhelming logs with excessive child node data.\n",
    "#    **Pitfall**: Important child nodes might be missed in the logs if this threshold is too low.\n",
    "\n",
    "# 6. **Recursively Traverse Child Nodes**:\n",
    "#    - If the node has **10 or fewer children**, the function proceeds to traverse each child recursively:\n",
    "#        - `for child in node.children:`\n",
    "#            - `traverse_and_debug(child)`\n",
    "#\n",
    "#    **Success**: Navigates through the entire tree efficiently when the number of children is manageable.\n",
    "#    **Pitfall**: Deep recursion may exceed Python’s recursion limit, causing a `RecursionError`.\n",
    "\n",
    "# **Purpose of the Function**:\n",
    "# - **Traversal**: Walks through a tree structure, starting from the given node.\n",
    "# - **Logging**: Collects and logs information about each node for debugging purposes.\n",
    "# - **Reduced Verbosity**: Implements checks to prevent excessive logging when nodes have too many children, keeping logs concise and readable.\n",
    "\n",
    "# **Use Cases**:\n",
    "# - **Debugging Tree Structures**: Helpful for inspecting complex datasets organized in hierarchical structures.\n",
    "# - **Performance Monitoring**: Can identify nodes that may cause slowdowns due to a large number of children.\n",
    "# - **Data Validation**: Ensures nodes contain the expected attributes, helping identify issues in the data structure.\n",
    "\n",
    "# **Considerations**:\n",
    "# 1. **Recursion Limit**:\n",
    "#    - Since the function uses recursion, deep trees may exceed Python’s recursion limit.\n",
    "#    - **Solution**: Consider converting the logic to an iterative approach or carefully increasing the recursion limit if necessary.\n",
    "\n",
    "# 2. **Logging Levels**:\n",
    "#    - Uses `logger.debug` for regular logs and `logger.warning` when skipping nodes with too many children.\n",
    "#    - **Ensure** that the logging configuration captures the desired level of detail without overwhelming the logs.\n",
    "\n",
    "# 3. **Attribute Checks**:\n",
    "#    - Safely checks for attributes to prevent runtime errors.\n",
    "#    - **However**: Missing attributes may reduce the utility of the logs if critical information is omitted.\n",
    "\n",
    "# **Example Scenario**:\n",
    "# - You have a tree representing tasks and subtasks in a hierarchical dataset.\n",
    "# - The function logs the **name** and **grid shape** of each node, but if a node has more than 10 children, it logs a warning and skips further traversal to avoid excessive output.\n",
    "# - This ensures that only manageable information is logged, improving readability and debugging efficiency.\n",
    "\n",
    "# **Summary**:\n",
    "# - The `traverse_and_debug` function is a useful tool for navigating and debugging tree-like data structures.\n",
    "# - It **balances detailed logging** with **verbosity control** by limiting output for nodes with many children.\n",
    "# - This design makes it effective for **performance monitoring, data validation**, and **troubleshooting complex hierarchies**.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "726f2c20-faa1-43b3-97af-ed8c2c942f32",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'logger' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[22], line 93\u001b[0m\n\u001b[0;32m     89\u001b[0m     root_window\u001b[38;5;241m.\u001b[39mmainloop()\n\u001b[0;32m     92\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;18m__name__\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__main__\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m---> 93\u001b[0m     \u001b[43mmain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[1;32mIn[22], line 7\u001b[0m, in \u001b[0;36mmain\u001b[1;34m()\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmain\u001b[39m():\n\u001b[0;32m      6\u001b[0m     \u001b[38;5;66;03m# Detect device\u001b[39;00m\n\u001b[1;32m----> 7\u001b[0m     device \u001b[38;5;241m=\u001b[39m \u001b[43mget_device\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m      9\u001b[0m     \u001b[38;5;66;03m# Initialize the progress bar\u001b[39;00m\n\u001b[0;32m     10\u001b[0m     progress_bar \u001b[38;5;241m=\u001b[39m tqdm(total\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m100\u001b[39m, desc\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mLoading Data\u001b[39m\u001b[38;5;124m\"\u001b[39m, unit\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m%\u001b[39m\u001b[38;5;124m\"\u001b[39m, leave\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
      "Cell \u001b[1;32mIn[4], line 39\u001b[0m, in \u001b[0;36mget_device\u001b[1;34m()\u001b[0m\n\u001b[0;32m     37\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m     38\u001b[0m     device \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mdevice(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcpu\u001b[39m\u001b[38;5;124m'\u001b[39m)  \u001b[38;5;66;03m# Fallback to CPU\u001b[39;00m\n\u001b[1;32m---> 39\u001b[0m     \u001b[43mlogger\u001b[49m\u001b[38;5;241m.\u001b[39minfo(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUsing CPU as fallback.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     40\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m device\n",
      "\u001b[1;31mNameError\u001b[0m: name 'logger' is not defined"
     ]
    }
   ],
   "source": [
    "# -----------------------------\n",
    "# 13. Main Workflow with Modifications\n",
    "# -----------------------------\n",
    "\n",
    "def main():\n",
    "    # Detect device\n",
    "    device = get_device()\n",
    "\n",
    "    # Initialize the progress bar\n",
    "    progress_bar = tqdm(total=100, desc=\"Loading Data\", unit=\"%\", leave=True)\n",
    "\n",
    "    try:\n",
    "        # Load ARC data\n",
    "        arc_data = load_arc_data()\n",
    "        progress_bar.update(20)\n",
    "\n",
    "        # Extract and reshape grid pairs\n",
    "        train_grid_pairs = flatten_and_reshape(\n",
    "            arc_data.get(\"arc-agi_training-challenges\", {})\n",
    "        )\n",
    "        eval_grid_pairs = flatten_and_reshape(\n",
    "            arc_data.get(\"arc-agi_evaluation-challenges\", {})\n",
    "        )\n",
    "        progress_bar.update(30)\n",
    "\n",
    "        # Build the data tree and retrieve the task dictionary\n",
    "        root_node, task_dict = build_data_tree(train_grid_pairs)\n",
    "        traverse_and_debug(root_node)\n",
    "        progress_bar.update(20)\n",
    "\n",
    "        # Log the task dictionary\n",
    "        logger.info(f\"Task dictionary initialized with {len(task_dict)} tasks:\")\n",
    "        for task_id, task_data in task_dict.items():\n",
    "            logger.info(\n",
    "                f\"Task ID: {task_id}, Node: {task_data['task_node'].name}, \"\n",
    "                f\"Grid Shape: {task_data['grids'][0].shape}\"\n",
    "            )\n",
    "\n",
    "        # Initialize DataLoaders and model\n",
    "        train_dataset = AugmentedARCDataset(train_grid_pairs, augment=False)\n",
    "        val_dataset = AugmentedARCDataset(eval_grid_pairs, augment=False)\n",
    "\n",
    "        train_loader = DataLoader(\n",
    "            train_dataset,\n",
    "            batch_size=4,  # Reduce batch size if you encounter memory issues\n",
    "            shuffle=True,\n",
    "            num_workers=0,\n",
    "            pin_memory=True,\n",
    "            collate_fn=collate_fn  # Use the custom collate function\n",
    "        )\n",
    "        val_loader = DataLoader(\n",
    "            val_dataset,\n",
    "            batch_size=4,\n",
    "            shuffle=False,\n",
    "            num_workers=0,\n",
    "            pin_memory=True,\n",
    "            collate_fn=collate_fn\n",
    "        )\n",
    "\n",
    "        # Initialize the model\n",
    "        model = CNNGridMapper(num_classes=NUM_CLASSES).to(device)\n",
    "        if torch.cuda.device_count() > 1:\n",
    "            model = nn.DataParallel(model)\n",
    "        logger.info(\"Model initialized successfully.\")\n",
    "\n",
    "    except Exception as e:\n",
    "        logger.exception(f\"Data loading or model initialization failed: {e}\")\n",
    "        progress_bar.close()\n",
    "        return\n",
    "\n",
    "    progress_bar.close()\n",
    "\n",
    "    # Initialize and start the GUI\n",
    "    root_window = tk.Tk()\n",
    "    gui = TrainingGUI(\n",
    "        root_window, total_epochs=10, total_batches=len(train_loader),\n",
    "        model=model, train_loader=train_loader, val_loader=val_loader,\n",
    "        eval_loader=None, device=device, data_tree=root_node, task_dict=task_dict\n",
    "    )\n",
    "\n",
    "    # Start the training thread\n",
    "    training_thread = threading.Thread(\n",
    "        target=train_model_with_gui, args=(model, train_loader, val_loader, device, gui)\n",
    "    )\n",
    "    training_thread.daemon = True\n",
    "    training_thread.start()\n",
    "\n",
    "    # Start the GUI main loop\n",
    "    root_window.mainloop()\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d650da2b-915e-472e-bd25-6bf7ac406417",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------\n",
    "# 13. Explanation of the Main Workflow with Modifications\n",
    "# -----------------------------\n",
    "\n",
    "# The `main` function orchestrates the entire workflow of the application, from **data loading** and **preprocessing**\n",
    "# to **model initialization** and launching the **graphical user interface (GUI)** for real-time visualization of the training process.\n",
    "\n",
    "# **Function Definition**:\n",
    "# def main():\n",
    "\n",
    "# **Key Steps**:\n",
    "\n",
    "# 1. **Device Detection**:\n",
    "#    - Calls `get_device()` to determine whether a **GPU** (if available) or **CPU** should be used for computations.\n",
    "#    - Stores the selected device in the `device` variable:\n",
    "#        - `device = get_device()`\n",
    "\n",
    "#    **Success**: Automatically optimizes the computation by selecting the fastest available device.\n",
    "#    **Pitfall**: If the selected device has insufficient memory (e.g., GPU), the process may crash or slow down.\n",
    "\n",
    "# 2. **Progress Bar Initialization**:\n",
    "#    - Initializes a progress bar with `tqdm` to provide visual feedback on data loading progress:\n",
    "#        - `progress_bar = tqdm(total=100, desc=\"Loading Data\", unit=\"%\", leave=True)`\n",
    "\n",
    "#    **Success**: Keeps users informed about the progress, especially for time-consuming operations.\n",
    "#    **Pitfall**: If not updated correctly, users may assume the program is frozen.\n",
    "\n",
    "# 3. **Data Loading**:\n",
    "#    - Loads the ARC dataset using `load_arc_data()` and updates the progress bar:\n",
    "#        - `arc_data = load_arc_data()`\n",
    "#        - `progress_bar.update(20)`\n",
    "\n",
    "#    **Success**: Centralizes data access logic, improving maintainability.\n",
    "#    **Pitfall**: Missing or corrupted files could cause crashes during loading.\n",
    "\n",
    "# 4. **Data Extraction and Reshaping**:\n",
    "#    - Uses `flatten_and_reshape()` to extract and reshape grid pairs from the ARC dataset.\n",
    "#    - Retrieves data for **training** and **evaluation** challenges:\n",
    "#        - `train_grid_pairs = flatten_and_reshape(...)`\n",
    "#        - `eval_grid_pairs = flatten_and_reshape(...)`\n",
    "#    - Updates the progress bar:\n",
    "#        - `progress_bar.update(30)`\n",
    "\n",
    "#    **Success**: Ensures the data is properly structured for model training.\n",
    "#    **Pitfall**: Inconsistent data shapes could cause runtime errors.\n",
    "\n",
    "# 5. **Building Data Tree and Task Dictionary**:\n",
    "#    - Constructs a hierarchical **data tree** using `build_data_tree()`.\n",
    "#    - Calls `traverse_and_debug()` to traverse the tree and log node details:\n",
    "#        - `root_node, task_dict = build_data_tree(train_grid_pairs)`\n",
    "#        - `traverse_and_debug(root_node)`\n",
    "#    - Updates the progress bar:\n",
    "#        - `progress_bar.update(20)`\n",
    "\n",
    "#    **Success**: Organizes the dataset for better visualization and management.\n",
    "#    **Pitfall**: Deep trees could exceed Python’s recursion limit if not handled carefully.\n",
    "\n",
    "# 6. **Logging Task Information**:\n",
    "#    - Logs the number of tasks initialized in the `task_dict`:\n",
    "#        - `logger.info(f\"Task dictionary initialized with {len(task_dict)} tasks.\")`\n",
    "#    - Iterates over the tasks and logs task-specific details (e.g., ID, grid shape):\n",
    "#        - `for task_id, task_data in task_dict.items(): logger.info(...)`\n",
    "\n",
    "#    **Success**: Provides transparency and ensures data integrity.\n",
    "#    **Pitfall**: Missing or incorrect task information could make debugging difficult.\n",
    "\n",
    "# 7. **Dataset and DataLoader Initialization**:\n",
    "#    - Initializes **training** and **validation** datasets using `AugmentedARCDataset`:\n",
    "#        - `train_dataset = AugmentedARCDataset(train_grid_pairs, augment=False)`\n",
    "#        - `val_dataset = AugmentedARCDataset(eval_grid_pairs, augment=False)`\n",
    "#    - Creates DataLoaders with custom `collate_fn` to handle variable-sized inputs:\n",
    "#        - `train_loader = DataLoader(train_dataset, batch_size=4, collate_fn=collate_fn)`\n",
    "#        - `val_loader = DataLoader(val_dataset, batch_size=4, collate_fn=collate_fn)`\n",
    "\n",
    "#    **Success**: Efficiently batches and shuffles data for faster training.\n",
    "#    **Pitfall**: Incorrect DataLoader configurations (e.g., too many workers) can cause performance issues.\n",
    "\n",
    "# 8. **Model Initialization**:\n",
    "#    - Creates an instance of `CNNGridMapper` and moves the model to the appropriate device:\n",
    "#        - `model = CNNGridMapper(num_classes=NUM_CLASSES).to(device)`\n",
    "#    - If multiple GPUs are available, wraps the model with `nn.DataParallel` for parallel training:\n",
    "#        - `if torch.cuda.device_count() > 1: model = nn.DataParallel(model)`\n",
    "#    - Logs a success message:\n",
    "#        - `logger.info(\"Model initialized successfully.\")`\n",
    "\n",
    "#    **Success**: Prepares the model for efficient training on available hardware.\n",
    "#    **Pitfall**: Incorrect device handling can result in runtime errors.\n",
    "\n",
    "# 9. **Exception Handling**:\n",
    "#    - Wraps the main logic in a `try-except` block to catch and log exceptions:\n",
    "#        - `logger.exception(f\"Data loading or model initialization failed: {e}\")`\n",
    "#    - Closes the progress bar upon error:\n",
    "#        - `progress_bar.close()`\n",
    "#    - Returns to prevent further execution.\n",
    "\n",
    "#    **Success**: Prevents crashes by handling exceptions gracefully.\n",
    "#    **Pitfall**: Not providing helpful error messages could make debugging difficult.\n",
    "\n",
    "# 10. **Close Progress Bar**:\n",
    "#     - Closes the progress bar after data loading completes:\n",
    "#        - `progress_bar.close()`\n",
    "\n",
    "#     **Success**: Keeps the console output clean and avoids confusion about the progress status.\n",
    "#     **Pitfall**: Forgetting to close the progress bar can clutter the output.\n",
    "\n",
    "# 11. **GUI Initialization**:\n",
    "#     - Creates a Tkinter GUI window and initializes `TrainingGUI` with all necessary parameters:\n",
    "#        - `gui = TrainingGUI(root_window, total_epochs=10, model=model, ...)`\n",
    "\n",
    "#     **Success**: Provides a user-friendly interface to monitor and control the training process.\n",
    "#     **Pitfall**: Complex GUIs can become unresponsive without proper threading.\n",
    "\n",
    "# 12. **Start Training Thread**:\n",
    "#     - Creates a new thread to run the training process without blocking the GUI:\n",
    "#        - `training_thread = threading.Thread(target=train_model_with_gui, args=(...))`\n",
    "#     - Sets `daemon=True` to ensure the thread terminates with the main program:\n",
    "#        - `training_thread.daemon = True`\n",
    "#     - Starts the thread:\n",
    "#        - `training_thread.start()`\n",
    "\n",
    "#     **Success**: Ensures the GUI remains responsive during long-running training processes.\n",
    "#     **Pitfall**: Incorrect thread handling could cause deadlocks or crashes.\n",
    "\n",
    "# 13. **Start GUI Main Loop**:\n",
    "#     - Calls `root_window.mainloop()` to run the Tkinter event loop:\n",
    "#        - `root_window.mainloop()`\n",
    "\n",
    "#     **Success**: Keeps the GUI interactive, allowing users to control the training.\n",
    "#     **Pitfall**: Forgetting to call the main loop will result in a non-responsive GUI.\n",
    "\n",
    "# **Entry Point Check**:\n",
    "# - Ensures that `main()` is called only when the script is run directly:\n",
    "#     - `if __name__ == \"__main__\": main()`\n",
    "\n",
    "#     **Success**: Prevents unintended execution when the script is imported as a module.\n",
    "#     **Pitfall**: Forgetting this check can lead to unexpected behavior when importing the script.\n",
    "\n",
    "# **Conclusion**:\n",
    "# The `main` function integrates multiple components—data loading, model initialization, and GUI setup—into a cohesive workflow.\n",
    "# It ensures smooth interaction between the backend logic and the frontend interface, providing a robust framework for model development.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d228376-0053-4713-9e5c-79f000ba53de",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}