{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d09ce366-f10d-4eb6-b8af-f6292c7ed926",
   "metadata": {},
   "source": [
    "# Embedding a 1D Flare Spectra in a 2D Image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "715d0eb7-fb85-42fc-bebb-839e60a797b9",
   "metadata": {},
   "source": [
    "We have a 1D spectra as a function of wavelength. To convert this to a 3D data cube (2 spatial dimensions, 1 wavelength dimension), we need to:\n",
    "\n",
    "1. Choose a spatial location for the flare\n",
    "2. Create an array to embed spectra in with dimensions $(N_\\lambda,2N+1,2N+1)$, where $N$ is the flare kernel size in pixel space\n",
    "3. Create a celestial WCS\n",
    "4. Interpolate spectra to constant wavelength bins\n",
    "5. Create the combined WCS\n",
    "6. Create spectral cube using blurred data and combined WCS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6dbf6cff-bae1-4b73-a994-3bc14d6a6b56",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib\n",
    "\n",
    "import asdf\n",
    "import ndcube\n",
    "import numpy as np\n",
    "import astropy.io.fits\n",
    "import astropy.units as u\n",
    "import astropy.wcs\n",
    "from fiasco.io import Parser\n",
    "from scipy.interpolate import PchipInterpolator\n",
    "import sunpy.map\n",
    "import sunpy.io._fits as sunpy_fits\n",
    "\n",
    "from sunpy.coordinates import get_earth, Helioprojective\n",
    "from astropy.coordinates import SkyCoord\n",
    "from astropy.wcs.utils import wcs_to_celestial_frame\n",
    "\n",
    "from synthesizAR.instruments.util import add_wave_keys_to_header\n",
    "\n",
    "from mocksipipeline.util import read_data_cube\n",
    "from mocksipipeline.instrument.configuration import moxsi_slot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bf472524-389e-4a17-ab91-6d1e3e6a051c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: UnitsWarning: The unit 'Angstrom' has been deprecated in the VOUnit standard. Suggested: 0.1nm. [astropy.units.format.utils]\n"
     ]
    }
   ],
   "source": [
    "goes_class = ['a1', 'b1', 'b7', 'm1', 'm5', 'x5']\n",
    "flare_spectra = []\n",
    "with asdf.open('../data/reference_spectra/caspi_flare_spectra.asdf', copy_arrays=True) as af:\n",
    "    for gc in goes_class:\n",
    "        flare_spectra.append((gc, ndcube.NDCube(af.tree[gc]['data'], wcs=af.tree[gc]['wcs'])))\n",
    "flare_spectra = ndcube.NDCollection(flare_spectra, aligned_axes=(0,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a3ee3a4-0490-43c0-9322-9b53712c2f42",
   "metadata": {},
   "source": [
    "Interpolate each of the spectra to a common wavelength axis because we need them to be sampled regularly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0b31bf04-8db9-45df-8dcd-7be69a0d3cb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "wave_min = 0.1 * u.AA\n",
    "wave_max = 80 * u.AA\n",
    "flare_spectra_interp = {}\n",
    "for k in flare_spectra:\n",
    "    energy_axis = flare_spectra[k].axis_world_coords(0)[0]\n",
    "    wavelength_axis = energy_axis.to('AA', equivalencies=u.spectral())\n",
    "    delta_wave = np.min(np.fabs(np.diff(wavelength_axis))) * 50  # Don't need the full resolution\n",
    "    wavelength_axis_norm = u.Quantity(np.arange(wave_min.to_value('AA'), wave_max.to_value('AA'), delta_wave.to_value('AA')), 'AA')\n",
    "    f_interp = PchipInterpolator(energy_axis.to_value('keV'), flare_spectra[k].data, extrapolate=True)\n",
    "    data_interp = f_interp(wavelength_axis_norm.to_value('keV', equivalencies=u.spectral())) * flare_spectra[k].unit\n",
    "    flare_spectra_interp[k] = (wavelength_axis_norm, data_interp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9d462fc2-3ad2-419d-8538-d62d55e0719e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: FITSFixedWarning: The WCS transformation has more axes (3) than the image it is associated with (2) [astropy.wcs.wcs]\n",
      "WARNING: FITSFixedWarning: The WCS transformation has more axes (3) than the image it is associated with (2) [astropy.wcs.wcs]\n",
      "WARNING: FITSFixedWarning: The WCS transformation has more axes (3) than the image it is associated with (2) [astropy.wcs.wcs]\n",
      "WARNING: FITSFixedWarning: The WCS transformation has more axes (3) than the image it is associated with (2) [astropy.wcs.wcs]\n",
      "WARNING: FITSFixedWarning: The WCS transformation has more axes (3) than the image it is associated with (2) [astropy.wcs.wcs]\n",
      "WARNING: FITSFixedWarning: The WCS transformation has more axes (3) than the image it is associated with (2) [astropy.wcs.wcs]\n"
     ]
    }
   ],
   "source": [
    "earth_observer = get_earth('2020-01-01')\n",
    "ref_coord = SkyCoord(0, 0, unit='arcsec', frame=Helioprojective(observer=earth_observer))\n",
    "array_size = (30, 130)  # this should always be even\n",
    "flare_spectra_interp_cubes = {}\n",
    "for k in flare_spectra_interp:\n",
    "    wavelength, data = flare_spectra_interp[k]\n",
    "    data_array = np.zeros(wavelength.shape+array_size) * data.unit\n",
    "    data_array[:, array_size[0]//2-1:array_size[0]//2+1, array_size[1]//2-1:array_size[1]//2+1] = data[:, None, None]\n",
    "    header = sunpy.map.make_fitswcs_header(\n",
    "        array_size,\n",
    "        ref_coord,\n",
    "        scale=moxsi_slot.optical_design.spatial_plate_scale,\n",
    "        unit=data.unit,\n",
    "    )\n",
    "    header = add_wave_keys_to_header(wavelength, header)\n",
    "    wcs = astropy.wcs.WCS(header=header)\n",
    "    flare_spectra_interp_cubes[k] = ndcube.NDCube(data_array, wcs=wcs, meta=header)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c1e700a8-cfc6-47ed-8068-f87a5394df5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_root = pathlib.Path('/Users/wtbarnes/Documents/codes/mocksipipeline/pipeline/results/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "004b3fad-b351-4bbd-a549-7f413a457c50",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: VerifyWarning: Verification reported errors: [astropy.io.fits.verify]\n",
      "WARNING: VerifyWarning: HDU 0: [astropy.io.fits.verify]\n",
      "WARNING: VerifyWarning:     'NAXIS3' card at the wrong place (card 35).  Fixed by moving it to the right place (card 5). [astropy.io.fits.verify]\n",
      "WARNING: VerifyWarning: Note: astropy.io.fits uses zero-based indexing.\n",
      " [astropy.io.fits.verify]\n"
     ]
    }
   ],
   "source": [
    "for k in flare_spectra_interp_cubes:\n",
    "    results_dir = results_root / f'caspi_spectra_{k}'\n",
    "    results_dir.mkdir(exist_ok=True)\n",
    "    sunpy_fits.write(results_dir / 'spectral_cube.fits',\n",
    "                     flare_spectra_interp_cubes[k].data,\n",
    "                     flare_spectra_interp_cubes[k].meta,\n",
    "                     overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d363f006-7950-4d4e-935a-38d56593589f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:mocksipipeline-dev]",
   "language": "python",
   "name": "conda-env-mocksipipeline-dev-py"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}