import json
import base64

notebook = {
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LooGLE Head-to-Head: Entroly vs Agentic Pruning\n",
    "This notebook runs the official Entroly benchmark against the 2026 SOTA Agentic Pruning baseline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": None,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q datasets openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": None,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, subprocess\n",
    "if not os.path.exists('/content/entroly'):\n",
    "    subprocess.run(['git', 'clone', '--depth', '1', 'https://github.com/juyterman1000/entroly.git', '/content/entroly'], check=True)\n",
    "sys.path.insert(0, '/content/entroly')\n",
    "from entroly.universal_compress import universal_compress\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": None,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from google.colab import userdata\n",
    "os.environ['OPENAI_API_KEY'] = userdata.get('OPENAI_API_KEY')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": None,
   "metadata": {},
   "outputs": [],
   "source": [
    "SUBSET = 'shortdep_qa'\n",
    "SAMPLES = 30\n",
    "BUDGET = 1500\n",
    "ANSWER_MODEL = 'gpt-4o-mini'\n",
    "METHODS = ['baseline', 'entroly', 'agentic_pruning']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": None,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "ds = load_dataset('bigai-nlco/LooGLE', SUBSET, split='test')\n",
    "n_avail = len(ds)\n",
    "indices = list(range(0, n_avail, max(1, n_avail // SAMPLES)))[:SAMPLES]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": None,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokens(text: str) -> int:\n",
    "    return max(1, len(text) // 4)\n",
    "\n",
    "def compress_truncate(context: str, budget: int) -> str:\n",
    "    return context[: budget * 4]\n",
    "\n",
    "def compress_entroly(context: str, budget: int) -> str:\n",
    "    cur = tokens(context)\n",
    "    if cur <= budget:\n",
    "        return context\n",
    "    target_ratio = max(0.05, (budget * 0.85) / max(cur, 1))\n",
    "    out, _, _ = universal_compress(context, target_ratio, 'prose')\n",
    "    if tokens(out) > budget:\n",
    "        out = out[: budget * 4]\n",
    "    return out\n",
    "\n",
    "def compress_agentic_pruner(context: str, question: str, budget: int) -> str:\n",
    "    cur = tokens(context)\n",
    "    if cur <= budget:\n",
    "        return context\n",
    "    from openai import OpenAI\n",
    "    client = OpenAI()\n",
    "    prompt = f\"Extract only the exact facts from the following text that are relevant to this question: {question}\\n\\n{context}\"\n",
    "    try:\n",
    "        resp = client.chat.completions.create(\n",
    "            model=\"gpt-4o-mini\",\n",
    "            messages=[{\"role\": \"user\", \"content\": prompt}],\n",
    "            max_tokens=budget,\n",
    "            temperature=0.0\n",
    "        )\n",
    "        return resp.choices[0].message.content.strip()\n",
    "    except Exception:\n",
    "        return context[:budget * 4]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": None,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re, string, time\n",
    "from collections import Counter, defaultdict\n",
    "from openai import OpenAI\n",
    "client = OpenAI()\n",
    "\n",
    "def answer(context: str, question: str) -> str:\n",
    "    prompt = f'Use only the context below. Give the shortest correct answer.\\n\\nContext:\\n{context}\\n\\nQuestion: {question}\\n\\nAnswer:'\n",
    "    resp = client.chat.completions.create(\n",
    "        model=ANSWER_MODEL,\n",
    "        messages=[{'role': 'user', 'content': prompt}],\n",
    "        max_tokens=128, temperature=0.0,\n",
    "    )\n",
    "    return resp.choices[0].message.content.strip()\n",
    "\n",
    "def normalize(s: str) -> str:\n",
    "    s = s.lower()\n",
    "    s = re.sub(r'\\b(a|an|the)\\b', ' ', s)\n",
    "    s = ''.join(ch for ch in s if ch not in set(string.punctuation))\n",
    "    return ' '.join(s.split())\n",
    "\n",
    "def f1_score(pred: str, gold: str) -> float:\n",
    "    p, g = normalize(pred).split(), normalize(gold).split()\n",
    "    if not p or not g: return float(p == g)\n",
    "    common = Counter(p) & Counter(g)\n",
    "    n = sum(common.values())\n",
    "    if n == 0: return 0.0\n",
    "    return 2 * (n / len(p)) * (n / len(g)) / ((n / len(p)) + (n / len(g)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": None,
   "metadata": {},
   "outputs": [],
   "source": [
    "stats = {m: defaultdict(float) for m in METHODS}\n",
    "rows = []\n",
    "\n",
    "for i, idx in enumerate(indices):\n",
    "    ex = ds[int(idx)]\n",
    "    context = ex.get('context') or ex.get('input') or ''\n",
    "    question = ex.get('question') or ''\n",
    "    gold = ex.get('answer') or ''\n",
    "    if isinstance(gold, list): gold = gold[0] if gold else ''\n",
    "    if not (context and question and gold): continue\n",
    "\n",
    "    orig_tok = tokens(context)\n",
    "    for m in METHODS:\n",
    "        t0 = time.time()\n",
    "        try:\n",
    "            if m == 'baseline': ctx = compress_truncate(context, BUDGET)\n",
    "            elif m == 'entroly': ctx = compress_entroly(context, BUDGET)\n",
    "            else: ctx = compress_agentic_pruner(context, question, BUDGET)\n",
    "            cms = (time.time() - t0) * 1000\n",
    "            \n",
    "            pred = answer(ctx, question)\n",
    "            f1 = f1_score(pred, gold)\n",
    "            \n",
    "            stats[m]['f1'] += f1\n",
    "            stats[m]['em'] += 1 if f1 == 1.0 else 0\n",
    "            stats[m]['tokens_in'] += tokens(ctx)\n",
    "            stats[m]['original_tokens'] += orig_tok\n",
    "            stats[m]['ms'] += cms\n",
    "            stats[m]['n'] += 1\n",
    "        except Exception as e:\n",
    "            stats[m]['errors'] += 1\n",
    "            print(f\"Error on {m}: {e}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": None,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary = {}\n",
    "for m in METHODS:\n",
    "    s = stats[m]\n",
    "    n = max(1, int(s['n']))\n",
    "    avg_in = s['tokens_in'] / n\n",
    "    avg_orig = s['original_tokens'] / n\n",
    "    \n",
    "    if m == 'agentic_pruning':\n",
    "        api_calls = 2\n",
    "        cost = ((avg_orig * 0.150) + (avg_in * 0.600) + (avg_in * 0.150)) / 1000\n",
    "    else:\n",
    "        api_calls = 1\n",
    "        cost = (avg_in * 0.150) / 1000\n",
    "        \n",
    "    summary[m] = {\n",
    "        'n': int(s['n']),\n",
    "        'errors': int(s['errors']),\n",
    "        'avg_f1': round(s['f1'] / n, 3),\n",
    "        'avg_tokens': round(avg_in, 0),\n",
    "        'avg_ms': round(s['ms'] / n, 1),\n",
    "        'api_calls': api_calls,\n",
    "        'cost_1k': round(cost, 3)\n",
    "    }\n",
    "\n",
    "print('\\n  LooGLE Head-to-Head')\n",
    "print('  ' + '\u2500' * 65)\n",
    "print(f\"  {'method':<16} {'F1':>6} {'Tokens':>8} {'Latency':>10} {'API-Calls':>10} {'Cost/1k':>10}\")\n",
    "print('  ' + '\u2500' * 65)\n",
    "for m in METHODS:\n",
    "    d = summary[m]\n",
    "    print(f\"  {m:<16} {d['avg_f1']:>6.3f} {int(d['avg_tokens']):>8} {d['avg_ms']:>8.0f}ms {d['api_calls']:>10} ${d['cost_1k']:.3f}\")\n",
    "print('  ' + '\u2500' * 65)\n",
    "import json\n",
    "with open('/content/final_results.json', 'w') as f:\n",
    "    json.dump(summary, f, indent=2)\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}

with open("bench/colab_run.ipynb", "w", encoding="utf-8") as f:
    json.dump(notebook, f, indent=1)
