{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NeuPI: Self-Supervised Training of a Neural Surrogate\n",
    "\n",
    "This notebook demonstrates the core functionality of the `neupi` library: training a neural network to solve inference tasks on a Probabilistic Models (PMs) in a **self-supervised** manner. \n",
    "\n",
    "The key idea is to use the negative log likelihood score of the assignment from the PM as the loss function. The neural network generates candidate solutions (variable assignments), and the PGM evaluates their quality by computing their log-likelihood. The network's goal is to learn to produce assignments that maximize this likelihood (minimize the loss), effectively solving the Most Probable Explanation (MPE) inference task.\n",
    "\n",
    "We will cover:\n",
    "1. Loading a `MarkovNetwork` to act as the evaluator (the \"teacher\").\n",
    "2. Creating a synthetic dataset for training.\n",
    "3. Defining an `MLP` model (the \"neural solver\").\n",
    "4. Configuring the `SelfSupervisedTrainer` to manage the training process.\n",
    "5. Running the training loop and observing the decrease in loss."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We import the necessary components from PyTorch and `neupi`. This includes the PGM, the neural model, the trainer, and the loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n",
      "Markov Network path: networks/mn/Grids_17.uai\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from pathlib import Path\n",
    "import os\n",
    "\n",
    "# Import neupi components\n",
    "from neupi import (\n",
    "    MLP,\n",
    "    DiscreteEmbedder,\n",
    "    MarkovNetwork,\n",
    "    SelfSupervisedTrainer,\n",
    "    mpe_log_likelihood_loss,\n",
    ")\n",
    "\n",
    "# Define the device for computation\n",
    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using device: {DEVICE}\")\n",
    "\n",
    "# --- Path Setup ---\n",
    "ROOT_PATH = Path(os.getcwd()).parent\n",
    "# We use the 'Grids_17.uai' network as seen in the test suite\n",
    "UAI_PATH = Path(\"networks\") / \"mn\" / \"Grids_17.uai\"\n",
    "\n",
    "print(f\"Markov Network path: {UAI_PATH}\")\n",
    "assert UAI_PATH.exists(), f\"File not found: {UAI_PATH}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Load the PGM Evaluator\n",
    "\n",
    "First, we load the Markov Network that will provide the supervisory signal. Its evaluation function is the basis for our loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using 1d factors: False\n",
      "PGM is pairwise.\n",
      "Successfully loaded Markov Network with 400 variables.\n"
     ]
    }
   ],
   "source": [
    "mn_evaluator = MarkovNetwork(uai_file=str(UAI_PATH), device=DEVICE)\n",
    "num_vars = mn_evaluator.num_variables\n",
    "\n",
    "print(f\"Successfully loaded Markov Network with {num_vars} variables.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Create a Dummy DataLoader\n",
    "\n",
    "The trainer expects a `DataLoader` that yields batches of data. In our self-supervised setup, we don't need labeled data. The dataloader provides:\n",
    "- `inputs`: A source of randomness (e.g., noise) for the neural network to generate diverse solutions from.\n",
    "- `evidence_data`: A tensor representing partial assignments.\n",
    "- `evidence_mask`: A boolean mask indicating which variables in `evidence_data` are observed. We ignore the rest of the variables.\n",
    "\n",
    "For this example, we create placeholder tensors for a simple unsupervised MPE case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created a DataLoader with 64 samples and batch size 16.\n"
     ]
    }
   ],
   "source": [
    "num_samples = 64\n",
    "batch_size = 16\n",
    "\n",
    "# 1. Evidence Data: Placeholder, not strictly needed for this example but required by the API.\n",
    "evidence_data = torch.zeros(num_samples, num_vars, device=DEVICE, dtype=torch.float32)\n",
    "\n",
    "# 2. Evidence Mask: A mask of all False indicates no variables are observed (no observed variables in MPE).\n",
    "evidence_mask = torch.zeros(num_samples, num_vars, device=DEVICE, dtype=torch.bool)\n",
    "\n",
    "# 3. Query Mask: A mask of all True indicates all variables are query variables.\n",
    "query_mask = torch.ones(num_samples, num_vars, device=DEVICE, dtype=torch.bool)\n",
    "\n",
    "# 4. Unobserved Mask: A mask of all False indicates all variables are observed.\n",
    "unobs_mask = torch.zeros(num_samples, num_vars, device=DEVICE, dtype=torch.bool)\n",
    "\n",
    "\n",
    "dataset = TensorDataset(evidence_data, evidence_mask, query_mask, unobs_mask)\n",
    "dataloader = DataLoader(dataset, batch_size=batch_size)\n",
    "\n",
    "print(f\"Created a DataLoader with {len(dataset)} samples and batch size {batch_size}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Define the Neural Network Model\n",
    "\n",
    "We use a simple Multi-Layer Perceptron (MLP) as our surrogate model. It will take random noise as input and output a probability for each variable being in state 1. The `input_size` and `output_size` must match the number of variables in the PGM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MLP model initialized:\n",
      "MLP(\n",
      "  (hidden_layers): Sequential(\n",
      "    (0): Linear(in_features=800, out_features=64, bias=True)\n",
      "    (1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): ReLU()\n",
      "    (3): Linear(in_features=64, out_features=32, bias=True)\n",
      "    (4): BatchNorm1d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (5): ReLU()\n",
      "  )\n",
      "  (output_layer): Linear(in_features=32, out_features=400, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "embedding = DiscreteEmbedder(num_vars)\n",
    "model = MLP(hidden_sizes=[64, 32], output_size=num_vars, embedding=embedding).to(DEVICE)\n",
    "\n",
    "print(\"MLP model initialized:\")\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Configure the Trainer\n",
    "\n",
    "The `SelfSupervisedTrainer` orchestrates the training process. It brings together the model, the PGM evaluator, the loss function, and the optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trainer configured successfully.\n"
     ]
    }
   ],
   "source": [
    "# Initialize a standard PyTorch optimizer\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "\n",
    "# The loss function is the negative log-likelihood of the generated assignments.\n",
    "# Minimizing this loss is equivalent to maximizing the likelihood.\n",
    "loss_function = mpe_log_likelihood_loss\n",
    "\n",
    "# Initialize the trainer\n",
    "trainer = SelfSupervisedTrainer(\n",
    "    model=model,\n",
    "    pgm_evaluator=mn_evaluator,\n",
    "    loss_fn=loss_function,\n",
    "    optimizer=optimizer,\n",
    "    device=DEVICE,\n",
    ")\n",
    "\n",
    "print(\"Trainer configured successfully.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5: Run the Training\n",
    "\n",
    "Now, we can train the model using the `fit()` method. This will iterate through the provided data for a specified number of epochs. We'll also manually inspect the loss before and after to confirm that the model is learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial Loss on one batch: 0.0004\n",
      "Initial log-likelihood: -0.0004\n",
      "\n",
      "Starting training for 50 epochs...\n",
      "Epoch 1/50, Average Loss: -0.2512\n",
      "Epoch 2/50, Average Loss: -0.6543\n",
      "Epoch 3/50, Average Loss: -1.0587\n",
      "Epoch 4/50, Average Loss: -1.4651\n",
      "Epoch 5/50, Average Loss: -1.8745\n",
      "Epoch 6/50, Average Loss: -2.2881\n",
      "Epoch 7/50, Average Loss: -2.7073\n",
      "Epoch 8/50, Average Loss: -3.1334\n",
      "Epoch 9/50, Average Loss: -3.5679\n",
      "Epoch 10/50, Average Loss: -4.0122\n",
      "Epoch 11/50, Average Loss: -4.4678\n",
      "Epoch 12/50, Average Loss: -4.9363\n",
      "Epoch 13/50, Average Loss: -5.4191\n",
      "Epoch 14/50, Average Loss: -5.9177\n",
      "Epoch 15/50, Average Loss: -6.4336\n",
      "Epoch 16/50, Average Loss: -6.9681\n",
      "Epoch 17/50, Average Loss: -7.5227\n",
      "Epoch 18/50, Average Loss: -8.0987\n",
      "Epoch 19/50, Average Loss: -8.6974\n",
      "Epoch 20/50, Average Loss: -9.3200\n",
      "Epoch 21/50, Average Loss: -9.9679\n",
      "Epoch 22/50, Average Loss: -10.6423\n",
      "Epoch 23/50, Average Loss: -11.3443\n",
      "Epoch 24/50, Average Loss: -12.0751\n",
      "Epoch 25/50, Average Loss: -12.8358\n",
      "Epoch 26/50, Average Loss: -13.6274\n",
      "Epoch 27/50, Average Loss: -14.4510\n",
      "Epoch 28/50, Average Loss: -15.3075\n",
      "Epoch 29/50, Average Loss: -16.1978\n",
      "Epoch 30/50, Average Loss: -17.1228\n",
      "Epoch 31/50, Average Loss: -18.0832\n",
      "Epoch 32/50, Average Loss: -19.0798\n",
      "Epoch 33/50, Average Loss: -20.1134\n",
      "Epoch 34/50, Average Loss: -21.1847\n",
      "Epoch 35/50, Average Loss: -22.2944\n",
      "Epoch 36/50, Average Loss: -23.4431\n",
      "Epoch 37/50, Average Loss: -24.6315\n",
      "Epoch 38/50, Average Loss: -25.8602\n",
      "Epoch 39/50, Average Loss: -27.1298\n",
      "Epoch 40/50, Average Loss: -28.4407\n",
      "Epoch 41/50, Average Loss: -29.7937\n",
      "Epoch 42/50, Average Loss: -31.1891\n",
      "Epoch 43/50, Average Loss: -32.6276\n",
      "Epoch 44/50, Average Loss: -34.1095\n",
      "Epoch 45/50, Average Loss: -35.6352\n",
      "Epoch 46/50, Average Loss: -37.2052\n",
      "Epoch 47/50, Average Loss: -38.8199\n",
      "Epoch 48/50, Average Loss: -40.4795\n",
      "Epoch 49/50, Average Loss: -42.1844\n",
      "Epoch 50/50, Average Loss: -43.9348\n",
      "Training complete.\n",
      "Final Loss on the same batch: -45.0503\n",
      "Final log-likelihood: 45.0503\n"
     ]
    }
   ],
   "source": [
    "# Get a single batch to check the initial loss\n",
    "initial_batch = next(iter(dataloader))\n",
    "initial_loss = trainer.step(initial_batch)\n",
    "print(f\"Initial Loss on one batch: {initial_loss:.4f}\")\n",
    "print(f\"Initial log-likelihood: {-initial_loss:.4f}\")\n",
    "\n",
    "# Train the model for a few epochs\n",
    "num_epochs = 50\n",
    "print(f\"\\nStarting training for {num_epochs} epochs...\")\n",
    "trained_model = trainer.fit(dataloader, num_epochs=num_epochs)\n",
    "print(\"Training complete.\")\n",
    "\n",
    "# Check the loss again on the same initial batch to see the improvement\n",
    "final_loss = trainer.step(initial_batch)\n",
    "print(f\"Final Loss on the same batch: {final_loss:.4f}\")\n",
    "print(f\"Final log-likelihood: {-final_loss:.4f}\")\n",
    "\n",
    "\n",
    "# Verify that the loss has decreased\n",
    "assert final_loss < initial_loss, \"Loss did not decrease after training!\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this notebook, we successfully trained a neural network to generate high-likelihood solutions for a Markov Network's MPE problem. The entire process was self-supervised, requiring no pre-existing labeled data—only the structure of the PGM itself.\n",
    "\n",
    "The next step is to **use our newly trained model for inference**, which will be the focus of the next notebook."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "neupi",
   "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.9.23"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}