{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "00df1373-a903-4288-9250-ab4f659ff9f4",
   "metadata": {},
   "source": [
    "## Adding new flavors to a `RailProject` and using them\n",
    "\n",
    "This notebook will show you the basics of adding new flavors to `RailProject` and using them.  It's quite simple."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79e2ee1c-035d-4ce5-b3ac-48176a48a92c",
   "metadata": {},
   "source": [
    "#### Import the usual suspects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "115dc360-66e4-4c7a-aea6-6fab872e2083",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from rail.projects import RailProject, RailFlavor, library"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5477d9a-f25c-47bb-9226-30be5869c70d",
   "metadata": {},
   "source": [
    "#### Setup environment, and move to correct directory.  This will also download the needed data files to run this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62165b86-80c0-46a7-9103-530e69112cc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from rail.projects import library\n",
    "\n",
    "check_dir = os.path.basename(os.path.abspath(os.curdir))\n",
    "if check_dir == 'examples':\n",
    "    os.chdir('..')\n",
    "\n",
    "print(os.path.abspath(os.curdir))\n",
    "\n",
    "setup = library.setup_mininal_example_files()\n",
    "assert setup == 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae0ce175-a2f1-46b7-9d14-5277ac0e56b6",
   "metadata": {},
   "source": [
    "#### Build a very minimal project, with only the baseline flavor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52ed253e-b027-4533-b9d9-cd01ce48e182",
   "metadata": {},
   "outputs": [],
   "source": [
    "project = RailProject.load_config('tests/minimal.yaml')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81b24292-a507-4186-93ee-5d7725790449",
   "metadata": {},
   "source": [
    "#### Add a new flavor to the project\n",
    "\n",
    "Note that we are specifying:\n",
    "\n",
    "1. a name for the new flavor\n",
    "2. the pipelines that we can run under the new flavor\n",
    "3. overrides for the pipelines with respect to the baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9d03fa6-cb58-4e18-ad24-452e352f770b",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_flavor = project.add_flavor(\n",
    "    name='new_flavor', \n",
    "    pipelines=['pz', 'tomography'],\n",
    "    pipeline_overrides=dict(\n",
    "        default=dict(\n",
    "            kwargs=dict(\n",
    "              algorithms=['gpz'],\n",
    "            )\n",
    "        ),\n",
    "        pz=dict(\n",
    "            inform_gpz=dict(\n",
    "              gpz_method='VC',\n",
    "            )\n",
    "        )\n",
    "    )\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c878d53d-d910-4806-ad23-9df9190f06ce",
   "metadata": {},
   "source": [
    "#### Make sure the flavor got added, and inspect it's configuration. \n",
    "\n",
    "Note that it picked up a few parameters from the baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c857f68-47a4-4606-a344-bee668f88f1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "project.get_flavors()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c5c9441-e60d-4435-a2ef-c557929c4b3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_flavor.config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ef2528c-f160-4a6a-a80e-1068538f05f4",
   "metadata": {},
   "source": [
    "Let's see where the pipeline files will go"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbe7011f-9df0-484c-af75-23b480f68b79",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The template for pipeline files is: {project.get_path('pipeline_path')}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "062e3d0f-20fc-424e-a2e9-b2c4f54e8338",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"The path for the 'pz' pipeline of flavor 'baseline' is {project.get_path('pipeline_path', pipeline='pz', flavor='baseline')}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f50cb046-d70d-474d-96c6-2feea6a3de9c",
   "metadata": {},
   "source": [
    "#### Build the pipelines for the few flavor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7d2c46b-6cad-4414-bfc6-133d3e2c52e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "project.build_pipelines('new_flavor', force=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6ce9b07-e03d-4520-a817-91340d3de6d9",
   "metadata": {},
   "source": [
    "#### Let's see what the inputs for the new pz pipeline are"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7e5a9ec-788d-4dfc-a558-ed76b03943f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline_info = project.get_pipeline('pz')\n",
    "pipeline_instance = pipeline_info.make_instance(project, 'new_flavor', {})\n",
    "pipeline_instance.get_input_files(project, selection='gold')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "254f1718-631f-4053-9503-ebdf41931604",
   "metadata": {},
   "source": [
    "#### All set, run the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "453333af-c139-4618-ac17-b63dcb34c3c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "project.run_pipeline_single('pz', flavor='new_flavor', selection='gold')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2c0f7b8-266b-4437-aaf1-25008c7c2251",
   "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
}