{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🏗️ FairRecovery++: Fair Long-Horizon Disaster Recovery RL\n",
    "\n",
    "**FairRecovery++** is a production-grade multi-agent OpenEnv environment designed to train LLMs to make fair, strategic, and resilient decisions in complex post-disaster systems. It combines adaptive world-modeling, multi-agent dynamics (citizens, NGOs, adversaries), and a curriculum-weighted reward engine to teach agents to balance speed with equity.\n",
    "\n",
    "### 🏆 Judge's Demo Guide\n",
    "This notebook provides a complete environment to:\n",
    "1. **Validate Integrity**: Run the full suite of 38 unit tests.\n",
    "2. **Fast Evaluation**: Generate all performance reports and reward/fairness plots using a heuristic baseline.\n",
    "3. **Real LLM Run**: (Optional) Execute a live training session if a GPU is available.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup Environment\n",
    "We install the core dependencies and clone the official repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install dependencies (approx. 1-2 mins)\n",
    "!pip install -U trl transformers accelerate peft bitsandbytes datasets matplotlib pytest pydantic uvicorn fastapi\n",
    "\n",
    "# Clone the repository\n",
    "!git clone https://github.com/joshua400/FairRecovery-PlusPlus.git\n",
    "\n",
    "# Install the project in editable mode\n",
    "%cd FairRecovery-PlusPlus\n",
    "!pip install -e .[dev]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Integrity Check (38 Tests)\n",
    "Verify that the environment core, multi-agent manager, reward engine, and safety shield are all functioning perfectly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pytest tests/ -v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Fast Performance Demo (No LLM Required)\n",
    "This cell runs a **heuristic-driven smoke test** that simulates 24 episodes using greedy and fair policies. This generates all the project's analytical plots instantaneously for judge review."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the fast evaluation (uses heuristic agents to simulate environment dynamics)\n",
    "!python train_sarvam_online.py --no-llm --episodes 24\n",
    "\n",
    "# Generate the final summary plots (Heatmap, Reward Curves, etc.)\n",
    "!python generate_summary_plots.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Visual Reports & Graphs\n",
    "The following plots demonstrate how the agent learns to balance **Utility** (recovery speed) with **Fairness** (equity across zones) over a long horizon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "import os\n",
    "\n",
    "plots = [\n",
    "    ('assets/training_results.png', 'Baseline vs Trained Reward Comparison'),\n",
    "    ('assets/score_heatmap.png', 'Per-Episode Reward Stability'),\n",
    "    ('assets/training_loss.png', 'Curriculum Reward Curve (Utility -> Fairness)'),\n",
    "    ('assets/utility_vs_fairness.png', 'Utility-Fairness Pareto Scatter')\n",
    "]\n",
    "\n",
    "for path, title in plots:\n",
    "    if os.path.exists(path):\n",
    "        print(f\"### {title}\")\n",
    "        display(Image(filename=path))\n",
    "        print(\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. (Optional) Real LLM Policy Evaluation\n",
    "If you have a GPU (T4/A10) and want to see the environment interact with a real language model, you can run this cell. It uses a lightweight `Qwen2.5-0.5B` model for fast inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: Requires GPU and ~5-10 mins to run\n",
    "# !python train_real_llm_colab.py --episodes 5 --difficulty medium"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "### 🔎 What to look for in the results:\n",
    "- **Fairness Delta**: Notice how the 'Trained' policy maintains higher service parity across zones.\n",
    "- **Safety Shield**: The logs will show the Shield blocking 'early-submit' exploits, forcing the agent to plan for the long horizon.\n",
    "- **Stage Discipline**: The agent correctly moves through *Analyze -> Allocate -> Execute -> Adapt* stages."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}