{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "16d7dc30",
   "metadata": {},
   "source": [
    "# Create ROOT File from ASTERIA Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aff0bea2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.utils.exceptions import AstropyDeprecationWarning\n",
    "from asteria import set_rcparams\n",
    "from asteria.simulation import Simulation\n",
    "from snewpy.neutrino import Flavor\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import astropy.units as u\n",
    "\n",
    "import numpy as np\n",
    "import uproot\n",
    "import os\n",
    "import warnings\n",
    "warnings.simplefilter(action='ignore', category=(FutureWarning, AstropyDeprecationWarning))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70b7e598-5254-4cd6-b314-db64a24e7f7d",
   "metadata": {},
   "source": [
    "## Set up Plotting Defaults"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4c345e7-fda7-44cd-980e-e79bf94e7497",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_rcparams()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0e18359-fb1c-4eb4-9ebc-3f6e9b799c50",
   "metadata": {},
   "source": [
    "## Set up and Run a Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8be54655",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = {\n",
    "    'name': 'Sukhbold_2015',\n",
    "    'param':{\n",
    "        'progenitor_mass': 27*u.Msun,\n",
    "        'eos': 'LS220'}\n",
    "}\n",
    "\n",
    "Emin=0*u.MeV; Emax=100*u.MeV; dE=1*u.MeV\n",
    "tmin=-1*u.s;  tmax=10*u.s;    dt=2*u.ms\n",
    "         \n",
    "params = {\n",
    "    'model': model,\n",
    "    'distance':1*u.kpc, \n",
    "    'Emin':0*u.MeV, 'Emax':100*u.MeV, 'dE':1*u.MeV,\n",
    "    'tmin':-1*u.s, 'tmax':10*u.s, 'dt':2*u.ms,\n",
    "    'mixing_scheme':'AdiabaticMSW',\n",
    "    'hierarchy':'normal'\n",
    "    \n",
    "}\n",
    "\n",
    "sim = Simulation(**params)\n",
    "sim.run()\n",
    "\n",
    "time = np.append(sim.time, sim.time[-1] + sim._sim_dt)\n",
    "\n",
    "fig, ax = plt.subplots(1, figsize = (6,7))\n",
    "for flavor in sim.flavors:\n",
    "    ax.plot(sim.time, sim.avg_dom_signal(flavor=flavor), label=flavor.to_tex())\n",
    "ax.legend()\n",
    "ax.set(xlabel=r'$t-t_\\mathrm{bounce}$ [s]', ylabel='signal per DOM', xlim=(-0.025, 0.65));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21d51aa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, figsize=(9,5.5))\n",
    "\n",
    "for flavor in Flavor:\n",
    "    ax.plot(sim.time, sim.E_per_V[flavor], label=flavor.to_tex())\n",
    "ax.set(xlabel=r'$t-t_\\mathrm{bounce}$ [s]',\n",
    "       xlim=(sim.time[0].to_value('s'), sim.time[-1].to_value('s')),\n",
    "       ylabel=r'energy/volume [MeV m$^{-3}$]')\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e786673d",
   "metadata": {},
   "source": [
    "### Create ROOT file\n",
    "\n",
    "Create `np.histogram`, using the simulation time binning and total signal per DOM. Create ROOT file using histogram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e90f5f0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "outfile = 'Sukhbold_s27_LS220.root'\n",
    "with uproot.recreate(outfile) as file:\n",
    "    file['total_photonic_energy_distance_1kpc'] = np.histogram(sim.time.value, bins=time.value, weights=sim.total_E_per_V.value)\n",
    "    file['version'] = '1.2'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6f2954e",
   "metadata": {},
   "outputs": [],
   "source": [
    "variable = np.histogram(sim.time.value, bins=time.value, weights=sim.total_E_per_V.value)\n",
    "\n",
    "fig, ax = plt.subplots(1, figsize=(9,5.5)) \n",
    "ax.step(variable[1][:-1], variable[0])\n",
    "ax.set(xlabel=r'$t-t_\\mathrm{bounce}$ [s]',\n",
    "       xlim=(sim.time[0].to_value('s'), sim.time[-1].to_value('s')),\n",
    "       ylabel=r'energy/volume [MeV m$^{-3}$]');"
   ]
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}