{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initial imports and data download if you need the KPF input files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "KP.20250208.17485.59.fits already exists. Skipping download.\n",
      "KP.20250208.17485.59_L1.fits already exists. Skipping download.\n"
     ]
    }
   ],
   "source": [
    "# Jupyter auto-reload source code\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from astropy.io import fits\n",
    "from rvdata.core.models.level2 import RV2\n",
    "\n",
    "\n",
    "#To download the input files for this example, set download=1 to run the code block below\n",
    "download=1\n",
    "\n",
    "if download == 1:\n",
    "    import requests\n",
    "\n",
    "    def download_file(url, filename):\n",
    "        response = requests.get(url)\n",
    "        response.raise_for_status()  # Check if the request was successful\n",
    "        with open(filename, \"wb\") as file:\n",
    "            file.write(response.content)\n",
    "\n",
    "    #Download the native, input FITS files for this example if not already on your computer\n",
    "    file_urls = {\n",
    "        \"KPF\": [\n",
    "            \"http://grinnell.as.arizona.edu/~rvdata/kpf/KP.20250208.17485.59.fits\",\n",
    "            \"http://grinnell.as.arizona.edu/~rvdata/kpf/KP.20250208.17485.59_L1.fits\",\n",
    "        ]\n",
    "    }\n",
    "\n",
    "    for file in file_urls[\"KPF\"]:\n",
    "        filename = file.split(\"/\")[-1].split(\"?\")[0]\n",
    "        if not os.path.exists(filename):\n",
    "            print(f\"Downloading {filename}...\")\n",
    "            download_file(file, filename)\n",
    "        else:\n",
    "            print(f\"{filename} already exists. Skipping download.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the L2 translator to go from native KPF data format to the EPRV Standard Format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Translating a KPF File [change file destination to where you have placed the L0 and L1 KPF files from Zenodo]\n",
    "l1file = 'KP.20250208.17485.59_L1.fits'\n",
    "l0file = 'KP.20250208.17485.59.fits'\n",
    "\n",
    "with fits.open(l1file) as hdul:\n",
    "        datetime = hdul[0].header['DATE-BEG']\n",
    "        timestamp, seconds = datetime.split('.')\n",
    "\n",
    "filetime = timestamp.replace(\":\",\"\").replace(\"-\",\"\")\n",
    "kpf_l2_filename = './docs/tutorials/KPFL2_'+filetime.replace(\"-:\",\"\")+'.fits'\n",
    "\n",
    "kpf_l2 = RV2.from_fits(l1file, l0file=l0file, instrument=\"KPF\")\n",
    "kpf_l2.to_fits(kpf_l2_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examining What's Inside the Community Standard L2 File\n",
    "\n",
    "Once you have created the community standard L2 fits file, you can read it into python either as a traditional FITS object using astropy's fits.open() function, or as a python object using the from_fits() function within RV2: \n",
    "\n",
    "`from astropy.io import fits` <br>\n",
    "`l2 = fits.open(kpf_l2_filename)`\n",
    "\n",
    "OR \n",
    "\n",
    "`from rvdata.core.models.level2 import RV2` <br>\n",
    "`l2_obj = RV2.from_fits(l2_standard)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#In this example we'll use the astropy fits.open() approach\n",
    "\n",
    "l2 = fits.open(kpf_l2_filename)\n",
    "\n",
    "#The primary header contains useful info about what was observed, what's in each trace, etc. The keywords will be the same regardless\n",
    "#of what facility was used to take the data, so you can always query the file in the same way.\n",
    "primary_header = l2[0].header\n",
    "print('What telescope produced this observation?  ',primary_header['TELESCOP'])\n",
    "print('What instrument was used?  ',primary_header['INSTRUME'])\n",
    "print('What object was observed?  ',primary_header['OBJECT'],)\n",
    "print('How many traces are produced by the spectrograph?  ',primary_header['NUMTRACE'] )\n",
    "print('What is in each of the traces? ')\n",
    "for i in range(1,primary_header['NUMTRACE']+1):\n",
    "    print('Trace',i,primary_header['TRACE'+str(i)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Extension Description table, located in the \"EXT_DESCRIPT\" extension, lists all of the FITS extensions in the L2 file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EXT_DESCRIPT = pd.DataFrame(l2['EXT_DESCRIPT'].data)\n",
    "EXT_DESCRIPT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### And the Order Table, located in the \"ORDER_TABLE\" extension, details the start and end wavelengths of each echelle order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ORDER_TABLE = pd.DataFrame(l2['ORDER_TABLE'].data)\n",
    "ORDER_TABLE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting Examples of L2 Spectra\n",
    "\n",
    "From the Extension Description table, we know that the stellar data is contained in Traces 2, 3, and 4 -- so we can take a look at those"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#specify what order you'd like to plot\n",
    "order = 20\n",
    "\n",
    "#standardize the y-axis range based on the order with this highest flux measurements so it's easier to compare\n",
    "flux_max = np.max([np.max(l2['TRACE2_FLUX'].data[order]), np.max(l2['TRACE3_FLUX'].data[order]), np.max(l2['TRACE4_FLUX'].data[order])])\n",
    "ymax = flux_max*1.05\n",
    "\n",
    "#Calculate scaling factors between the blaze extensions and the flux extensions for easier visual comparison\n",
    "blz_scale2 = np.nanmax(l2['TRACE2_FLUX'].data[order]/np.nanmax(l2['TRACE2_BLAZE'].data[order]))\n",
    "blz_scale3 = np.nanmax(l2['TRACE3_FLUX'].data[order]/np.nanmax(l2['TRACE3_BLAZE'].data[order]))\n",
    "blz_scale4 = np.nanmax(l2['TRACE4_FLUX'].data[order]/np.nanmax(l2['TRACE4_BLAZE'].data[order]))\n",
    "\n",
    "fig = plt.figure(1,[10,8])\n",
    "\n",
    "ax1 = fig.add_subplot(3,1,1)\n",
    "ax1.scatter(l2['TRACE2_WAVE'].data[order],l2['TRACE2_FLUX'].data[order],marker='.',s=5)\n",
    "ax1.plot(l2['TRACE2_WAVE'].data[order],l2['TRACE2_BLAZE'].data[order]*blz_scale2,color='orange')\n",
    "ax1.set_ylim(0,ymax)\n",
    "ax1.set_ylabel('Counts',fontsize=15,fontweight='bold')\n",
    "\n",
    "ax2 = fig.add_subplot(3,1,2)\n",
    "ax2.scatter(l2['TRACE3_WAVE'].data[order],l2['TRACE3_FLUX'].data[order],marker='.',s=5)\n",
    "ax2.plot(l2['TRACE3_WAVE'].data[order],l2['TRACE3_BLAZE'].data[order]*blz_scale3,color='orange')\n",
    "ax2.set_ylim(0,ymax)\n",
    "ax2.set_ylabel('Counts',fontsize=15,fontweight='bold')\n",
    "\n",
    "ax3 = fig.add_subplot(3,1,3)\n",
    "ax3.scatter(l2['TRACE4_WAVE'].data[order],l2['TRACE4_FLUX'].data[order],marker='.',s=5)\n",
    "ax3.plot(l2['TRACE4_WAVE'].data[order],l2['TRACE4_BLAZE'].data[order]*blz_scale4,color='orange')\n",
    "ax3.set_ylim(0,ymax)\n",
    "ax3.set_ylabel('Counts',fontsize=15,fontweight='bold')\n",
    "ax3.set_xlabel('Wavelength (A)',fontsize=15,fontweight='bold')\n",
    "\n",
    "filename_pieces = kpf_l2_filename.split('/')\n",
    "\n",
    "fig.suptitle('Science Traces for Order '+str(order)+' of file '+filename_pieces[-1],fontsize=15,fontweight='bold')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read in data from multiple instruments on the same star\n",
    "\n",
    "## Generate an EPRV Standard L2 file for NEID & KPF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rvdata.instruments.neid.level2 import RV2, NEIDRV2\n",
    "\n",
    "neid_L2_Filename = 'NEIDL2_20240921T052641.fits'\n",
    "\n",
    "neid_L2_obj = NEIDRV2.from_fits('neidL2_20240921T052641.fits', instrument=\"NEID\")\n",
    "neid_L2_obj.to_fits(neid_L2_Filename)\n",
    "\n",
    "kpf_l1file = 'KP.20240920.37616.85_L1.fits'\n",
    "kpf_l0file = 'KP.20240920.37616.85.fits'\n",
    "\n",
    "kpf_l2_filename = 'KPFL2_20240920T102656.fits'\n",
    "\n",
    "kpf_l2 = RV2.from_fits(l1file, l0file=l0file, instrument=\"KPF\")\n",
    "kpf_l2.to_fits(kpf_l2_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot an overlapping order, first with the blaze function and then after dividing it out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kpf_obs =  fits.open('./docs/tutorials/KPFL2_20240920T102656.fits')\n",
    "neid_obs = fits.open('./docs/tutorials/NEIDL2_20240921T052641.fits')\n",
    "\n",
    "kpf_order = 19\n",
    "neid_order1 = 55\n",
    "\n",
    "fig = plt.figure(1,[12,6])\n",
    "ax1 = fig.add_subplot(1,1,1)\n",
    "ax1.plot(kpf_obs['TRACE2_WAVE'].data[kpf_order],kpf_obs['TRACE2_FLUX'].data[kpf_order],color='blue')\n",
    "ax1.plot(neid_obs['TRACE1_WAVE'].data[neid_order1],neid_obs['TRACE1_FLUX'].data[neid_order1],color='orange')\n",
    "ax1.set_ylabel('Counts',fontsize=15,fontweight='bold')\n",
    "ax1.set_xlabel('Wavelength (A)',fontsize=15,fontweight='bold')\n",
    "\n",
    "ax1.text(5145,90000,'KPF obs of HD 4268',color='blue',fontsize=20,fontweight='bold')\n",
    "ax1.text(5145,80000,'NEID obs of HD 4268',color='darkorange',fontsize=20,fontweight='bold')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kpf_obs =  fits.open('./docs/tutorials/KPFL2_20240920T102656.fits')\n",
    "neid_obs = fits.open('./docs/tutorials/NEIDL2_20240921T052641.fits')\n",
    "\n",
    "kpf_order = 19\n",
    "neid_order = 55\n",
    "\n",
    "#Calculate scaling factors between the blaze extensions and the flux extensions for easier visual comparison\n",
    "blz_scale_kpf = np.nanmax(kpf_obs['TRACE2_FLUX'].data[kpf_order]/np.nanmax(kpf_obs['TRACE2_BLAZE'].data[kpf_order]))\n",
    "blz_scale_neid = np.nanmax(neid_obs['TRACE1_FLUX'].data[neid_order]/np.nanmax(neid_obs['TRACE1_BLAZE'].data[neid_order]))\n",
    "\n",
    "fig = plt.figure(1,[12,6])\n",
    "ax1 = fig.add_subplot(1,1,1)\n",
    "ax1.plot(neid_obs['TRACE1_WAVE'].data[neid_order],neid_obs['TRACE1_FLUX'].data[neid_order] / \n",
    "         (neid_obs['TRACE1_BLAZE'].data[neid_order] * blz_scale_neid) ,color='orange')\n",
    "ax1.plot(kpf_obs['TRACE2_WAVE'].data[kpf_order],kpf_obs['TRACE2_FLUX'].data[kpf_order] / \n",
    "         (kpf_obs['TRACE2_BLAZE'].data[kpf_order] * blz_scale_kpf) ,color='blue',alpha=.6)\n",
    "ax1.set_ylabel('Counts',fontsize=15,fontweight='bold')\n",
    "ax1.set_xlabel('Wavelength (A)',fontsize=15,fontweight='bold')\n",
    "ax1.set_ylim(0,1.3)\n",
    "\n",
    "ax1.text(5145,1.2,'KPF obs of HD 4268 - deblazed',color='blue',fontsize=16,fontweight='bold')\n",
    "ax1.text(5145,1.12,'NEID obs of HD 4268 - deblazed',color='darkorange',fontsize=16,fontweight='bold')\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}