{
 "cells": [
  {
   "cell_type": "raw",
   "id": "0",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ".. _notebooks-pickling:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "# Pickling objects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Pickling an object means saving it in a binary format (usually a `PKL` file). Unpickling means converting the `PKL` back to a Python object. It can be really useful when some calculations take a lot of time and you do not want to recompute them every time."
   ]
  },
  {
   "cell_type": "raw",
   "id": "3",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ".. warning\n",
    "   Do not use this option for long terme storage. Every time the object source code changes, the unpickled object will have undetermined behavior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from importlib import resources\n",
    "import os\n",
    "from pprint import pprint\n",
    "from pathlib import Path\n",
    "\n",
    "from lightwin.config.config_manager import process_config\n",
    "from lightwin.core.accelerator.accelerator import Accelerator\n",
    "from lightwin.ui.workflow_setup import run_simulation\n",
    "from lightwin.util.pickling import MyCloudPickler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "Remove previous pickles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "pickle_folder = Path(\"pickles\")\n",
    "if not pickle_folder.is_dir():\n",
    "    os.mkdir(pickle_folder)\n",
    "else:\n",
    "    for file in Path(\"pickles\").iterdir():\n",
    "        os.remove(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Manual pickling/unpickling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "### Perform first calculation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "We start with an example configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "toml_filepath = resources.files(\"lightwin.data.ads\") / \"lightwin.toml\"\n",
    "\n",
    "toml_keys = {\n",
    "    \"files\": \"files\",\n",
    "    \"beam\": \"beam\",\n",
    "    \"beam_calculator\": \"generic_envelope1d\",\n",
    "    \"wtf\": \"generic_wtf\",\n",
    "    \"design_space\": \"fit_phi_s_design_space\",\n",
    "    \"plots\": \"plots_minimal\",\n",
    "}\n",
    "override = {\n",
    "    \"plots\": {\n",
    "        \"kwargs\": {\"lw\": 5},\n",
    "        \"emittance\": False,\n",
    "        \"energy\": False,\n",
    "        \"envelopes\": False,\n",
    "        \"phase\": False,\n",
    "        \"twiss\": False,\n",
    "    },\n",
    "}\n",
    "\n",
    "config = process_config(toml_filepath, toml_keys, override=override)\n",
    "fault_scenarios = run_simulation(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "### Pickle the Accelerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "pickler = MyCloudPickler()\n",
    "fault_scenario = fault_scenarios[0]\n",
    "\n",
    "reference = fault_scenario.ref_acc\n",
    "ref_pickle_path = reference.pickle(pickler, path=pickle_folder / \"reference-accelerator.pkl\")\n",
    "\n",
    "fixed = fault_scenario.fix_acc\n",
    "fix_pickle_path = fixed.pickle(pickler, path=pickle_folder / \"fixed-accelerator.pkl\")\n",
    "\n",
    "del fault_scenario, reference, fixed"
   ]
  },
  {
   "cell_type": "raw",
   "id": "13",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ".. warning::\n",
    "    SimulationOutput instances created by Cython do not seem to be pickable."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Unpickle the Accelerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "reference = Accelerator.from_pickle(pickler, ref_pickle_path)\n",
    "fixed = Accelerator.from_pickle(pickler, fix_pickle_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "We can quickly check that we recovered the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = None\n",
    "for accelerator in (reference, fixed):\n",
    "    ax = accelerator.plot(\"v_cav_mv\", marker=\"o\", ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "## Configure pickling/unpickling from the `TOML`"
   ]
  },
  {
   "cell_type": "raw",
   "id": "19",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "The `pickle_paths` entry in the `TOML` should associate every accelerator name with a `PKL` filepath."
   ]
  },
  {
   "cell_type": "raw",
   "id": "20",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ".. code-block::\n",
    "\n",
    "    [files]\n",
    "    dat_file = \"ads.dat\"\n",
    "    project_folder = \"lw_results/\"\n",
    "\n",
    "    [files.pickle_paths]\n",
    "    Reference = \"reference.pkl\"\n",
    "\n",
    "    # Scenario 1: pre-computed solution (skips optimization)\n",
    "    [files.pickle_paths.000001]\n",
    "    Solution = \"solution-000001.pkl\"\n",
    "\n",
    "    # Scenario 2: alternatives with custom names (optimization still runs, the pickled Accelerators will be appended)\n",
    "    [files.pickle_paths.000002]\n",
    "    \"Conservative approach\" = \"design-conservative.pkl\"\n",
    "    \"Aggressive tuning\" = \"design-aggressive.pkl\"\n",
    "\n",
    "    # Scenario 3: solution + alternatives\n",
    "    [files.pickle_paths.000003]\n",
    "    Solution = \"solution-000003.pkl\"\n",
    "    \"Tweaked design\" = \"tweaked.pkl\"\n",
    "    \"Experimental config\" = \"experimental.pkl\""
   ]
  },
  {
   "cell_type": "raw",
   "id": "21",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "- If the path does not exist, the simulation is performed normally and the associated :class:`.Accelerator` object is pickled.\n",
    "- If the path exists, the :class:`.Accelerator` is unpickled and associated calculations are skipped.\n",
    "\n",
    ".. warning::\n",
    "   After every update/source code editing, you should remove your pickles in order to avoid hard-to-track down bugs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "### First simple scenario"
   ]
  },
  {
   "cell_type": "raw",
   "id": "23",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Here, we define a pickle path for both the `Reference` :class:`.Accelerator` and the `Solution` of the first (and unique) :class:`.FaultScenario`.\n",
    "\n",
    ".. note::\n",
    "   `\"Reference\"` and `\"Solution\"` are reserved keywords for `Reference` and `Solution` :class:`.Accelerator`\\s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "pickle_paths = {\n",
    "    \"Reference\": pickle_folder / \"reference.pkl\",\n",
    "    \"000001\": {\n",
    "        \"Solution\": pickle_folder / \"solution-000001-3cavs.pkl\"\n",
    "    },\n",
    "}\n",
    "override[\"files\"] = {\"pickle_paths\": pickle_paths}\n",
    "config = process_config(toml_filepath, toml_keys, override=override)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "25",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Both files `\"reference.pkl\"` and `\"solution-000001-3cavs.pkl\"` do not exist. So a normal simulation will be run, and corresponding :class:`.Accelerator`\\s will be pickled to this paths. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "fault_scenarios = run_simulation(config)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "27",
   "metadata": {
    "editable": true,
    "raw_mimetype": "text/restructuredtext",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "Now that `\"reference.pkl\"` and `\"solution-000001-3cavs.pkl\"` do exist, :class:`.Accelerator`\\s are unpickled and simulations are skipped."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "fault_scenarios = run_simulation(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29",
   "metadata": {},
   "source": [
    "### Comparing several scenarios"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "pickle_paths = {\n",
    "    \"Reference\": pickle_folder / \"reference.pkl\",\n",
    "    \"000001\": {\n",
    "        \"Solution\": pickle_folder / \"solution-000001-4cavs.pkl\",\n",
    "        \"Alternative (3 compensating)\": pickle_folder / \"solution-000001-3cavs.pkl\"\n",
    "    }\n",
    "}\n",
    "\n",
    "override[\"files\"] = {\"pickle_paths\": pickle_paths}\n",
    "override[\"wtf\"] = {\"k\": 4}\n",
    "config = process_config(toml_filepath, toml_keys, override=override)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "fault_scenarios = run_simulation(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32",
   "metadata": {},
   "source": [
    "Clean pickles folder for the next executions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "for file in pickle_folder.iterdir():\n",
    "    os.remove(file)"
   ]
  }
 ],
 "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.14.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}