{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b9638861-048f-4a85-aa3a-59040a549bae",
   "metadata": {},
   "source": [
    "# Astropy Table and FITS BINTABLE Schemas\n",
    "\n",
    "In this example, we will show how to use model classes to validate a FITS (or any Astropy table)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8b1468d-695a-4bd0-952d-b8a189a157c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import ClassVar\n",
    "\n",
    "import ctao_datamodel as dm\n",
    "import numpy as np\n",
    "from astropy import units as u\n",
    "from astropy.io import fits\n",
    "from astropy.table import Table\n",
    "from ctao_datamodel import AstroField, ModelBase, PlantUMLDiagram, Quantity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dac9f72c-50e6-4e1e-a83a-1af56d255786",
   "metadata": {},
   "source": [
    "## Define a table schema\n",
    "\n",
    "Here, we will implement a simple \"event-list\" table schema:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac243de0-6575-4b02-ac1f-4842a0a588ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EventListTable(ModelBase):\n",
    "    \"\"\"The column schema for an EventList\"\"\"\n",
    "\n",
    "    _namespace: ClassVar[str] = \"Example.DL3\"\n",
    "\n",
    "    event_id: int = AstroField(\n",
    "        \"Unique event identifier\", fits_column_dtype=\"int64\", fits_keyword=\"EVENT_ID\"\n",
    "    )\n",
    "    energy: Quantity[\"TeV\"] = AstroField(\n",
    "        \"Reconstructed energy\", fits_column_dtype=\"float32\", ucd=\"em.energy\"\n",
    "    )\n",
    "    ra: Quantity[\"deg\"] = AstroField(\n",
    "        \"Right ascension\", fits_column_dtype=\"float64\", ucd=\"pos.eq.ra\"\n",
    "    )\n",
    "    dec: Quantity[\"deg\"] = AstroField(\n",
    "        \"Declination\", fits_column_dtype=\"float64\", ucd=\"pos.eq.dec\"\n",
    "    )\n",
    "    gammaness: float | None = AstroField(\n",
    "        \"Gammaness\", default=None, fits_column_dtype=\"float64\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74b85dcc-4ae9-478b-841a-0949b080ae5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "PlantUMLDiagram(EventListTable)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4416cb89-6bc9-4a9e-ad40-068f89a2af7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "dm.model_to_table(EventListTable)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9121df03-d233-4b77-9082-16c12b57d82f",
   "metadata": {},
   "source": [
    "## Create an Empty Table\n",
    "\n",
    "First, we can create an empty table from the schema:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d77eda75-a0a2-4e29-b164-d94db876533f",
   "metadata": {},
   "outputs": [],
   "source": [
    "table = dm.model_to_astropy_table(EventListTable, n_rows=2)\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a62949e-3cd1-4371-9830-7bf142c47c8b",
   "metadata": {},
   "source": [
    "And convert it to a FITS file, to see the resulting headers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b53677b-b7d8-4717-95fb-eb4648c524df",
   "metadata": {},
   "outputs": [],
   "source": [
    "table.write(\"example_table.fits\", overwrite=True)\n",
    "fits.getheader(\"example_table.fits\", ext=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42a460ea-a660-488d-a651-ed099156d8bb",
   "metadata": {},
   "source": [
    "## Validate Existing Table\n",
    "\n",
    "Here, we take an existing table. Note that we don't have a GAMMANESS column, which shouldn't be a problem, since it is optional. We also use different (but compatible) units to what is in the schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40f0e566-3196-4328-990a-f61252d3c7ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "existing_table = Table(\n",
    "    dict(\n",
    "        EVENT_ID=[12345, 12346],\n",
    "        ENERGY=([1.0, 3.0] * u.TeV).astype(np.float32),\n",
    "        RA=[15.0, 17.4] * u.hourangle,\n",
    "        DEC=[-2.3, -2.4] * u.deg,\n",
    "    )\n",
    ")\n",
    "existing_table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b82791c-f709-4dfe-8050-6178d3a8e471",
   "metadata": {},
   "outputs": [],
   "source": [
    "dm.model_validate_astropy_table(existing_table, EventListTable)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4bc4a3d-dbf2-486c-91fc-34e37cc539b2",
   "metadata": {},
   "source": [
    "and, indeed this validates! (no errors)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a222627b-a540-4c03-b138-c7c0392fe73d",
   "metadata": {},
   "source": [
    "But if we have a wrong column unit for example, it should fail:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af5bca9f-d1b6-4f08-a0b7-3495ee15a09b",
   "metadata": {},
   "outputs": [],
   "source": [
    "bad_table = Table(\n",
    "    dict(\n",
    "        EVENT_ID=[12345, 12346],\n",
    "        ENERGY=([1.0, 3.0] * u.TeV).astype(np.float32),\n",
    "        # missing RA column\n",
    "        DEC=[-2.3, -2.4] * u.m,  # wrong unit\n",
    "    )\n",
    ")\n",
    "bad_table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8836833-c50b-4668-af6a-4fc5ea9f2faf",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    dm.model_validate_astropy_table(bad_table, EventListTable)\n",
    "except dm.TableValidationError as err:\n",
    "    print(err)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22b3edff-7984-4129-9619-5969fac2b5e3",
   "metadata": {},
   "source": [
    "## Validating A full FITS BINTABLE with Header metadata\n",
    "\n",
    "Table metadata can be validated in different ways dependending on its format.  \n",
    "* If the metadata is stored in a hierarchical  JSON representation, the usual pydantic validation can be used directly, e.g. just using ``SomeModel.model_validate(table.meta)``\n",
    "* If the metadata have been flattened, the same still works but they must be unflattened first, e.g. ``SomeModel.model_validate(dm.unflatten_model_instance(table.meta))``\n",
    "* if the metadata is in FITS keyword format,  `~ctao_datamodel.fits_header_to_instance` method can be used to validate and transform back into the hierarchical representation.\n",
    "\n",
    "Here is an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "030f4f3f-5eb1-4d02-8b84-b3939a0d450b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EventListMetadata(ModelBase):\n",
    "    \"\"\"Schema for some metadata keywords, can be hierarchical.\"\"\"\n",
    "\n",
    "    object: str = AstroField(\"object name\", fits_keyword=\"OBJECT\")\n",
    "    contact: dm.models.dataproducts.Contact = AstroField(\"contact info.\")\n",
    "\n",
    "\n",
    "dm.PlantUMLDiagram(EventListMetadata, details=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5e0a187-1257-4881-b914-8d9124776c58",
   "metadata": {},
   "outputs": [],
   "source": [
    "meta = EventListMetadata(\n",
    "    object=\"Crab Nebula\",\n",
    "    contact=dm.models.dataproducts.Contact(\n",
    "        name=\"Jane User\", organization=\"none\", email=\"jane@none.org\"\n",
    "    ),\n",
    ")\n",
    "header = dm.instance_to_fits_header(meta)\n",
    "header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5c78a2f-a012-4528-a0d8-40cc3abf9432",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_table = Table(\n",
    "    dict(\n",
    "        EVENT_ID=[12345, 12346],\n",
    "        ENERGY=([1.0, 3.0] * u.TeV).astype(np.float32),\n",
    "        RA=[15.0, 17.4] * u.deg,\n",
    "        DEC=[-2.3, -2.4] * u.deg,\n",
    "    )\n",
    ")\n",
    "\n",
    "hdu = fits.BinTableHDU(data=new_table, header=header, name=\"EVENTS\")\n",
    "hdu.writeto(\"test_with_header.fits\", overwrite=True)\n",
    "fits.getheader(\"test_with_header.fits\", \"EVENTS\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7014c48f-c864-446c-bbc6-33448f918088",
   "metadata": {},
   "source": [
    "Now, let's try validating the FITS file we wrote using a helper function that validates both the column schema and the metadata at once, and returns deserialized metadata:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32502eb2-4e7a-41ee-9c80-b6b0135c9130",
   "metadata": {},
   "outputs": [],
   "source": [
    "table, metadata = dm.validate_fits_bintable_hdu(\n",
    "    fits_file=\"test_with_header.fits\",\n",
    "    hdu=\"EVENTS\",\n",
    "    column_model=EventListTable,\n",
    "    header_model=EventListMetadata,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09897b7b-c14a-4dc2-a8b6-bbe55cea6e9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb9b881e-8881-428d-b906-6e09398e250d",
   "metadata": {},
   "outputs": [],
   "source": [
    "table.meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5cdf67d-549e-48b0-b080-cfff4e0c5bbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffcbebb8-b050-4e56-a2d2-cfbf960ab3e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata.contact.email"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dbabe60-97bd-4725-856d-7ffd91de63cd",
   "metadata": {},
   "source": [
    "And then you can see the table is valid (no errors) and the metadata has been checked and converted to a model instance.  If there had been a problem, you would get an exception with the details."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b960bf4c-45fc-46bd-b301-eaf8189b6212",
   "metadata": {},
   "source": [
    "## Other table formats\n",
    "\n",
    "Since the underlying implementation uses  *astropy* `~astropy.table.Table` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73b90477-ffa1-4168-802c-558c2775c80d",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_2 = dm.model_to_astropy_table(EventListTable)\n",
    "table_2.add_row([12345, 0.2, 10, -20.2, 0.2])\n",
    "table_2.add_row([12346, 100.2, 10.2, -20.1, 0.9])\n",
    "table_2.meta = dm.flatten_model_instance(metadata) \n",
    "table_2.meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "373b21e0-c4bf-4974-a66d-ace2aca91ff3",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_2[\"ENERGY\"].meta"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5445e67b-ac43-490c-9808-8b1f37a12214",
   "metadata": {},
   "source": [
    "For example, one can write the table to a VOTable as follows.  However, the user should be careful here, as currently astropy **does not** serialize table metadata to this format, so information is lost.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abcb3075-d0db-4088-ad45-e1f192cdffb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_2.write(\"test.votable\", format=\"votable\", overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5111b4d-a974-48ae-aca6-de7ba9678281",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat test.votable"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95ea3f31-35b0-4abf-a41e-aaccfb18a7ba",
   "metadata": {},
   "source": [
    "And here it is in ECSV format, which does serialize metadata, and has full round-tripping?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1048ac54-0951-4bda-b639-1e800945b973",
   "metadata": {},
   "outputs": [],
   "source": [
    "table_2.write(\"test.ecsv\", overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a805d38-c341-4474-98f8-b6686def1eaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat test.ecsv"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}