{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7f9536c9",
   "metadata": {},
   "source": [
    "# Dirac-3S Planted Benchmarking"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fb8e94e",
   "metadata": {},
   "source": [
    "### Import libraries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e44af7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "from eqc_direct.client import EqcClient\n",
    "from eqc_direct.utils import *\n",
    "\n",
    "\n",
    "ip_address=\"172.18.15.110\"\n",
    "client = EqcClient(ip_address=ip_address) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4c20d5a",
   "metadata": {},
   "source": [
    "## Set up Dirac-3 solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7da6eb82",
   "metadata": {},
   "outputs": [],
   "source": [
    "def direct_dirac3(indices, coefficients, relaxation_schedule, sum_constraint, num_samples, ip_address):\n",
    "    eqc_client = EqcClient(ip_address=ip_address) #S9 \"172.18.41.45\", #S3 \"172.18.41.173\"\n",
    "    lock_id, start_ts, end_ts=eqc_client.wait_for_lock()\n",
    "    try:\n",
    "        result_dict = eqc_client.solve_sum_constrained(\n",
    "            lock_id=lock_id,\n",
    "            poly_indices = indices,\n",
    "            poly_coefficients = coefficients,\n",
    "            relaxation_schedule = relaxation_schedule,\n",
    "            sum_constraint = sum_constraint,\n",
    "            num_samples = num_samples,\n",
    "        )\n",
    "        total_time = np.array(result_dict[\"preprocessing_time\"])+np.array(result_dict[\"postprocessing_time\"])+np.array(result_dict[\"runtime\"])\n",
    "        print(f\"Total execution time(s):{total_time}\")\n",
    "    finally:\n",
    "        # release lock when finished using the device\n",
    "        lock_release_out = eqc_client.release_lock(lock_id=lock_id)\n",
    "    return result_dict\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c79d6aa",
   "metadata": {},
   "source": [
    "### Load data file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "aa3e9737",
   "metadata": {},
   "outputs": [],
   "source": [
    "def loadQ(path, n):\n",
    "    tri = np.load(path, mmap_mode=\"r\")\n",
    "    Q = np.empty((n,n), dtype=tri.dtype)\n",
    "    iu = np.triu_indices(n)\n",
    "    Q[iu] = tri\n",
    "    Q.T[iu] = tri\n",
    "    return Q"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a69ffcc",
   "metadata": {},
   "source": [
    "### Solve with Dirac-3S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b1a40a1d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loaded file sucessfully.\n"
     ]
    }
   ],
   "source": [
    "# Instance details\n",
    "optimal_energy = 12000\n",
    "sum_constraint = 100\n",
    "num_var = 2000\n",
    "k = 44\n",
    "ub =10\n",
    "seed =100\n",
    "\n",
    "name = f\"STQP_n_{num_var}_k_{k}_R_{sum_constraint}_seed_{seed}_ub_{ub}\"\n",
    "instance_path = os.path.join(f\"Instances/{name}.npy\")\n",
    "c = np.zeros(num_var)\n",
    "try:\n",
    "    Q= loadQ(instance_path,num_var)\n",
    "    print(\"loaded file sucessfully.\")\n",
    "except FileNotFoundError:\n",
    "    print(f\"File {instance_path} does not exist.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e04a5baa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# solver parameters\n",
    "relaxation_schedule = 1\n",
    "num_samples = 10\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "71c8581e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitting job to Dirac-3S\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:Max precision for EQC device is float32 input type was dtype float64. Input matrix will be rounded\n",
      "/home/sutapa/Documents/eqc-direc-2.0.3/.venv/lib/python3.10/site-packages/eqc_direct/client.py:565: Warning: Max precision for EQC device is float32 input type was dtype float64. Input matrix will be rounded\n",
      "  warnings.warn(warn_dtype_msg, Warning)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total execution time(s):[29.90855955 28.52003086 26.55735099 18.47528472 23.55796547 25.75906217\n",
      " 18.2335756  15.7012367  18.30005711 26.98886002]\n",
      "Dirac-3 Run complete. Response received.\n"
     ]
    }
   ],
   "source": [
    "# solve using Dirac-3S\n",
    "poly_indices, poly_coefficients = convert_hamiltonian_to_poly_format(\n",
    "                linear_terms=c,\n",
    "                quadratic_terms=Q,\n",
    "            )\n",
    "print(\"Submitting job to Dirac-3S\")\n",
    "response = direct_dirac3(indices=poly_indices,\n",
    "                        coefficients=poly_coefficients,\n",
    "                        relaxation_schedule=relaxation_schedule,\n",
    "                        sum_constraint=sum_constraint,\n",
    "                        ip_address=ip_address,\n",
    "                        num_samples=num_samples)\n",
    "print(f\"Dirac-3 Run complete. Response received.\")\n",
    "energies = response['energy']\n",
    "solutions = response['solution']\n",
    "runtimes = response['runtime']\n",
    "preprocessing_time = response['preprocessing_time']\n",
    "postprocessing_time = response['postprocessing_time']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2b5d9a93",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dirac-3S all samples energies:[12000.0, 13021.03, 12969.42, 13030.29, 13038.24, 13035.86, 13048.64, 13118.84, 13120.48, 13079.18]\n",
      "Dirac-3 all samples runtimes:[27.94, 20.14, 18.12, 16.52, 21.61, 17.34, 16.28, 13.75, 16.34, 25.03]\n"
     ]
    }
   ],
   "source": [
    "print(f\"Dirac-3S all samples energies:{[round(e, 2) for e in energies]}\")\n",
    "print(f\"Dirac-3 all samples runtimes:{[round(t, 2) for t in runtimes]}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "71b5ec6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "best = min(zip(energies, solutions, runtimes), key=lambda x: x[0])\n",
    "best_energy, best_solution, best_runtime = best\n",
    "arr = np.array(best_solution)\n",
    "indices = np.where(arr> 1e-6)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c56f660d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "samples:10, Relaxation schedule:1\n",
      "best energy over 10 samples:11999.999023\n",
      "support size:44\n",
      "time taken by best smaple 27.940969467163086\n"
     ]
    }
   ],
   "source": [
    "print(f\"samples:{num_samples}, Relaxation schedule:{relaxation_schedule}\")\n",
    "print(f\"best energy over {num_samples} samples:{best_energy:.6f}\")\n",
    "print(f\"support size:{len(indices)}\")\n",
    "print(f\"time taken by best smaple {best_runtime}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bcdd646e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute relative gap\n",
    "tol = 1e-7\n",
    "abs_gap = best_energy-optimal_energy\n",
    "relative_gap = (round(abs_gap,6)*100)/round(optimal_energy,6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3dd2cfbb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimal solution found:11999.9990234.\n"
     ]
    }
   ],
   "source": [
    "# print result\n",
    "if abs_gap>tol:\n",
    "    print(f\"Dirac-3S solution is not optimal\")\n",
    "    print(f\"Absolute Gap:{abs_gap}\")\n",
    "    print(f\"Relative Gap(%):{abs(relative_gap)}\")\n",
    "else:\n",
    "    print(f\"Optimal solution found:{best_energy}.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95702667",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}