{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NeuPI: Advanced Discretization Methods\n",
    "\n",
    "This notebook demonstrates how to use the various discretization methods available in the **NeuPI** library. After a neural solver produces continuous probability outputs (typically from a sigmoid activation), a discretizer is needed to convert these probabilities into a final, binary MPE/MAP assignment.\n",
    "\n",
    "We will explore three methods, ranging from a simple baseline to sophisticated, search-based approaches:\n",
    "\n",
    "1.  **`ThresholdDiscretizer`**: The simplest method, which applies a fixed threshold (e.g., 0.5) to convert probabilities to binary values.\n",
    "2.  **`KNearestDiscretizer`**: A powerful method that uses a beam search (backed by a high-performance Cython module) to find the *k*-best binary assignments close to the continuous prediction.\n",
    "3.  **`HighUncertaintyDiscretizer`**: A smart heuristic that focuses its search on the *k* variables with the highest uncertainty (probabilities closest to 0.5), performing an exhaustive search over this reduced space.\n",
    "\n",
    "We will compare the quality of the assignments produced by each method by evaluating their average log-likelihood using a PGM evaluator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First, we import the necessary components from `neupi` and other libraries. We will also set up a PGM evaluator and generate some dummy continuous outputs from a hypothetical neural network to serve as our inference data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n",
      "Using 1d factors: False\n",
      "PGM is pairwise.\n",
      "Loaded Markov Network with 400 variables.\n",
      "Generated dummy probability outputs of shape: torch.Size([16, 400])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from pathlib import Path\n",
    "\n",
    "# Import neupi components\n",
    "from neupi.training.pm_ssl.pgm.mn import MarkovNetwork\n",
    "from neupi.discretize.threshold import ThresholdDiscretizer\n",
    "from neupi.discretize.kn import KNearestDiscretizer\n",
    "from neupi.discretize.oauai import OAUAI\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",
    "# Assuming this notebook is in 'examples/', and networks are in 'examples/networks/'\n",
    "UAI_PATH = Path(\"networks\") / \"mn\" / \"Grids_17.uai\"\n",
    "assert UAI_PATH.exists(), f\"File not found: {UAI_PATH}. Please run from the 'examples' directory.\"\n",
    "\n",
    "# --- Load PGM Evaluator ---\n",
    "mn_evaluator = MarkovNetwork(uai_file=str(UAI_PATH), device=DEVICE)\n",
    "num_vars = mn_evaluator.num_variables\n",
    "print(f\"Loaded Markov Network with {num_vars} variables.\")\n",
    "\n",
    "# --- Create Dummy Inference Data ---\n",
    "# This simulates the continuous probability outputs from a neural network's sigmoid layer.\n",
    "num_samples = 16\n",
    "prob_outputs = torch.rand(num_samples, num_vars, device=DEVICE, dtype=torch.float32)\n",
    "\n",
    "# Define a query mask (where we want to find the best assignment)\n",
    "# For this example, all variables are query variables.\n",
    "query_mask = torch.ones_like(prob_outputs, dtype=torch.bool)\n",
    "evidence_mask = torch.zeros_like(prob_outputs, dtype=torch.bool)\n",
    "unobs_mask = torch.zeros_like(prob_outputs, dtype=torch.bool)\n",
    "\n",
    "print(f\"Generated dummy probability outputs of shape: {prob_outputs.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method 1: Baseline with `ThresholdDiscretizer`\n",
    "\n",
    "This is our baseline. It's fast and simple, providing a good reference point for the more advanced methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ThresholdDiscretizer Avg. Log-Likelihood: 57.4206\n"
     ]
    }
   ],
   "source": [
    "# 1. Initialize the discretizer\n",
    "threshold_discretizer = ThresholdDiscretizer(threshold=0.5)\n",
    "\n",
    "# 2. Get the discrete assignments\n",
    "threshold_assignments = threshold_discretizer(prob_outputs)\n",
    "\n",
    "# 3. Evaluate the assignments to get their log-likelihood scores\n",
    "with torch.no_grad():\n",
    "    threshold_scores = mn_evaluator(threshold_assignments)\n",
    "\n",
    "avg_threshold_score = threshold_scores.mean().item()\n",
    "print(f\"ThresholdDiscretizer Avg. Log-Likelihood: {avg_threshold_score:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method 2: Beam Search with `KNearestDiscretizer`\n",
    "\n",
    "This method performs a beam search to find the *k* assignments that are closest (in terms of L1 distance) to the continuous predictions. It then scores each of these *k* candidates with the PGM evaluator and returns the best one. This is a powerful, model-aware search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                          "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "KNearestDiscretizer Avg. Log-Likelihood: 73.0331\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    }
   ],
   "source": [
    "# 1. Initialize the discretizer. It needs the PGM evaluator as its scoring function.\n",
    "knn_discretizer = KNearestDiscretizer(pgm_evaluator=mn_evaluator, k=10)  # Beam width\n",
    "\n",
    "# 2. Get the discrete assignments\n",
    "# The other masks are passed to maintain a consistent API\n",
    "knn_assignments = knn_discretizer(prob_outputs, evidence_mask, query_mask, unobs_mask)\n",
    "\n",
    "# 3. Evaluate the assignments\n",
    "with torch.no_grad():\n",
    "    knn_scores = mn_evaluator(knn_assignments)\n",
    "\n",
    "avg_knn_score = knn_scores.mean().item()\n",
    "print(f\"KNearestDiscretizer Avg. Log-Likelihood: {avg_knn_score:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method 3: `OAUAI`\n",
    "\n",
    "This method uses a smart heuristic: instead of searching over all variables, it identifies the *k* variables whose predicted probabilities are closest to 0.5 (i.e., the ones the network is least certain about). It then performs an exhaustive search over all 2^k possibilities for this small subset and picks the best one. Other orcale can be used to answer the query over the variables with the highest uncertainty (probabilities closest to 0.5)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                  "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HighUncertaintyDiscretizer Avg. Log-Likelihood: 99.4000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    }
   ],
   "source": [
    "# 1. Initialize the discretizer\n",
    "uncertainty_discretizer = OAUAI(\n",
    "    pgm_evaluator=mn_evaluator, k=5  # Number of uncertain variables to search over\n",
    ")\n",
    "\n",
    "# 2. Get the discrete assignments\n",
    "uncertainty_assignments = uncertainty_discretizer(\n",
    "    prob_outputs, evidence_mask, query_mask, unobs_mask\n",
    ")\n",
    "\n",
    "# 3. Evaluate the assignments\n",
    "with torch.no_grad():\n",
    "    uncertainty_scores = mn_evaluator(uncertainty_assignments)\n",
    "\n",
    "avg_uncertainty_score = uncertainty_scores.mean().item()\n",
    "print(f\"HighUncertaintyDiscretizer Avg. Log-Likelihood: {avg_uncertainty_score:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparison and Conclusion\n",
    "\n",
    "Let's compare the average log-likelihood scores from all three methods. A higher (less negative) score indicates a better set of MPE solutions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Discretization Performance Summary ---\n",
      "Baseline (Threshold):\t\t57.4206\n",
      "K-Nearest (Beam Search):\t\t73.0331\n",
      "OAUAI (Oracle Based):\t99.4000\n",
      "\n",
      "Improvement over baseline (KNN): +15.6125\n",
      "Improvement over baseline (OAUAI): +41.9794\n",
      "\n",
      "Successfully verified that the advanced methods provide scores greater than or equal to the baseline.\n"
     ]
    }
   ],
   "source": [
    "print(\"--- Discretization Performance Summary ---\")\n",
    "print(f\"Baseline (Threshold):\\t\\t{avg_threshold_score:.4f}\")\n",
    "print(f\"K-Nearest (Beam Search):\\t\\t{avg_knn_score:.4f}\")\n",
    "print(f\"OAUAI (Oracle Based):\\t{avg_uncertainty_score:.4f}\")\n",
    "\n",
    "improvement_knn = avg_knn_score - avg_threshold_score\n",
    "improvement_uncertainty = avg_uncertainty_score - avg_threshold_score\n",
    "\n",
    "print(f\"\\nImprovement over baseline (KNN): {improvement_knn:+.4f}\")\n",
    "print(f\"Improvement over baseline (OAUAI): {improvement_uncertainty:+.4f}\")\n",
    "\n",
    "assert avg_knn_score >= avg_threshold_score\n",
    "assert avg_uncertainty_score >= avg_threshold_score\n",
    "print(\n",
    "    \"\\nSuccessfully verified that the advanced methods provide scores greater than or equal to the baseline.\"\n",
    ")"
   ]
  }
 ],
 "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
}