{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NeuPI: Inference and Test-Time Refinement\n",
    "\n",
    "This notebook demonstrates the final step in the NeuPI pipeline: using a trained neural solver to perform Most Probable Explanation (MPE) (or MMAP) inference. We will explore two methods:\n",
    "\n",
    "1.  **`SinglePassInferenceEngine`**: A fast method that performs a single forward pass of the neural network to get the MPE assignments.\n",
    "2.  **`ITSELF_Engine`**: An advanced method that performs test-time refinement. It uses the PGM's feedback to fine-tune the model on each specific inference instance, often leading to significantly better results.\n",
    "\n",
    "We will cover:\n",
    "1.  Setting up a pre-trained model (recapping the updated training process from Notebook 2).\n",
    "2.  Creating a new dataset for inference.\n",
    "3.  Running the `SinglePassInferenceEngine` and evaluating its results.\n",
    "4.  Running the `ITSELF_Engine` to refine the predictions.\n",
    "5.  Comparing the log-likelihood scores to demonstrate the improvement from ITSELF."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We import all necessary components."
   ]
  },
  {
   "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",
    "    MarkovNetwork,\n",
    "    SelfSupervisedTrainer,\n",
    "    mpe_log_likelihood_loss,\n",
    "    DiscreteEmbedder,\n",
    "    ThresholdDiscretizer,\n",
    "    SinglePassInferenceEngine,\n",
    "    ITSELF_Engine,\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",
    "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: Recap - Get a Pre-Trained Model\n",
    "\n",
    "For this notebook to be self-contained, we'll quickly train a model, incorporating the library's recent updates (`DiscreteEmbedder` and the new data format). This will provide the `trained_model` we need for inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using 1d factors: False\n",
      "PGM is pairwise.\n",
      "Training a model for inference demonstration...\n",
      "Epoch 1/3, Average Loss: 6.0816\n",
      "Epoch 2/3, Average Loss: -6.7919\n",
      "Epoch 3/3, Average Loss: -14.1691\n",
      "Training complete. We now have a trained model.\n"
     ]
    }
   ],
   "source": [
    "# Load the PGM evaluator\n",
    "mn_evaluator = MarkovNetwork(uai_file=str(UAI_PATH), device=DEVICE)\n",
    "num_vars = mn_evaluator.num_variables\n",
    "\n",
    "# Create a dummy training dataloader\n",
    "num_samples_train = 64\n",
    "evidence_data_train = torch.randint(\n",
    "    0, 2, (num_samples_train, num_vars), device=DEVICE, dtype=torch.float32\n",
    ")\n",
    "evidence_mask_train = torch.rand(num_samples_train, num_vars, device=DEVICE) > 0.5\n",
    "query_mask_train = ~evidence_mask_train\n",
    "unobs_mask_train = torch.zeros_like(evidence_mask_train, dtype=torch.bool)\n",
    "train_dataset = TensorDataset(\n",
    "    evidence_data_train, evidence_mask_train, query_mask_train, unobs_mask_train\n",
    ")\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=16)\n",
    "\n",
    "# Setup model with the new DiscreteEmbedder\n",
    "embedding = DiscreteEmbedder(num_vars)\n",
    "model = MLP(hidden_sizes=[32, 16], output_size=num_vars, embedding=embedding).to(DEVICE)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "\n",
    "# Setup trainer\n",
    "trainer = SelfSupervisedTrainer(\n",
    "    model=model,\n",
    "    pgm_evaluator=mn_evaluator,\n",
    "    loss_fn=mpe_log_likelihood_loss,\n",
    "    optimizer=optimizer,\n",
    "    device=DEVICE,\n",
    ")\n",
    "\n",
    "# Train for a few epochs\n",
    "print(\"Training a model for inference demonstration...\")\n",
    "trained_model = trainer.fit(train_dataloader, num_epochs=3)\n",
    "print(\"Training complete. We now have a trained model.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Create an Inference DataLoader\n",
    "\n",
    "Now we create a new, unseen set of inference queries. The model has not seen this data during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created an inference DataLoader with 32 samples.\n"
     ]
    }
   ],
   "source": [
    "num_samples_inf = 32\n",
    "\n",
    "# The model takes evidence data and masks as input\n",
    "evidence_data_inf = torch.randint(\n",
    "    0, 2, (num_samples_inf, num_vars), device=DEVICE, dtype=torch.float32\n",
    ")\n",
    "evidence_mask_inf = torch.rand(num_samples_inf, num_vars, device=DEVICE) > 0.5\n",
    "query_mask_inf = ~evidence_mask_inf\n",
    "unobs_mask_inf = torch.zeros_like(evidence_mask_inf, dtype=torch.bool)\n",
    "\n",
    "inf_dataset = TensorDataset(evidence_data_inf, evidence_mask_inf, query_mask_inf, unobs_mask_inf)\n",
    "inf_dataloader = DataLoader(inf_dataset, batch_size=8)\n",
    "\n",
    "print(f\"Created an inference DataLoader with {len(inf_dataset)} samples.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Single-Pass Inference\n",
    "\n",
    "We first use the `SinglePassInferenceEngine`. It runs the model once, applies a discretizer to get binary assignments, and returns the result. This is the fastest method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running single-pass inference...\n",
      "Single-Pass Avg Log-Likelihood: 26.1420\n"
     ]
    }
   ],
   "source": [
    "# A discretizer is needed to convert the model's continuous outputs (probabilities) into binary assignments.\n",
    "discretizer = ThresholdDiscretizer(threshold=0.5)\n",
    "\n",
    "simple_inference_engine = SinglePassInferenceEngine(\n",
    "    model=trained_model, discretizer=discretizer, device=DEVICE\n",
    ")\n",
    "\n",
    "print(\"Running single-pass inference...\")\n",
    "initial_results = simple_inference_engine.run(inf_dataloader)\n",
    "initial_assignments = initial_results[\"final_assignments\"].to(DEVICE)\n",
    "\n",
    "# Evaluate the quality of these assignments using the PGM\n",
    "with torch.no_grad():\n",
    "    initial_ll = mn_evaluator(initial_assignments).mean()\n",
    "\n",
    "print(f\"Single-Pass Avg Log-Likelihood: {initial_ll.item():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: ITSELF Inference\n",
    "\n",
    "Now, we use the `ITSELF_Engine`. For each batch of data, it performs several optimization steps, fine-tuning the model's prediction specifically for that data. This test-time adaptation leverages the PGM evaluator to find better solutions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running ITSELF inference with test-time refinement...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ITSELF Refined Avg Log-Likelihood: 270.6185\n"
     ]
    }
   ],
   "source": [
    "itself_engine = ITSELF_Engine(\n",
    "    model=trained_model,\n",
    "    pgm_evaluator=mn_evaluator,\n",
    "    loss_fn=mpe_log_likelihood_loss,\n",
    "    optimizer_cls=torch.optim.Adam,  # The optimizer to use for refinement\n",
    "    discretizer=discretizer,\n",
    "    refinement_lr=1e-3,  # Learning rate for the refinement steps\n",
    "    refinement_steps=5,  # Number of refinement steps per instance\n",
    "    device=DEVICE,\n",
    ")\n",
    "\n",
    "print(\"Running ITSELF inference with test-time refinement...\")\n",
    "refined_results = itself_engine.run(inf_dataloader)\n",
    "refined_assignments = refined_results[\"final_assignments\"].to(DEVICE)\n",
    "\n",
    "# Evaluate the quality of the refined assignments\n",
    "with torch.no_grad():\n",
    "    refined_ll = mn_evaluator(refined_assignments).mean()\n",
    "\n",
    "print(f\"ITSELF Refined Avg Log-Likelihood: {refined_ll.item():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5: Comparison and Conclusion\n",
    "\n",
    "Finally, we compare the average log-likelihoods. A higher (less negative) log-likelihood indicates a better solution to the MPE problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial Avg Log-Likelihood (Single Pass): 26.1420\n",
      "Refined Avg Log-Likelihood (ITSELF):     270.6185\n",
      "\n",
      "Improvement from ITSELF: 244.4765\n",
      "\n",
      "Successfully demonstrated that ITSELF improves inference quality.\n"
     ]
    }
   ],
   "source": [
    "print(f\"Initial Avg Log-Likelihood (Single Pass): {initial_ll.item():.4f}\")\n",
    "print(f\"Refined Avg Log-Likelihood (ITSELF):     {refined_ll.item():.4f}\")\n",
    "\n",
    "improvement = refined_ll - initial_ll\n",
    "print(f\"\\nImprovement from ITSELF: {improvement.item():.4f}\")\n",
    "\n",
    "assert refined_ll > initial_ll, \"ITSELF failed to improve the log-likelihood!\"\n",
    "print(\"\\nSuccessfully demonstrated that ITSELF improves inference quality.\")"
   ]
  }
 ],
 "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
}