{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e11bcc56-0df9-4b13-a81c-cf6cce5438a8",
   "metadata": {},
   "source": [
    "```{eval-rst}\n",
    ".. currentmodule:: ctao_datamodel\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42329ace-c4fe-4b0f-8da9-388a5f566fa3",
   "metadata": {},
   "source": [
    "# Model Versions and Migrations\n",
    "\n",
    "Since models may evolve in time, it is important to be able to keep track of changes and to maintain backward compatibility when possible. ``ctao-datamodel`` provides an API for this purpose: the {py:class}`VersionedModel`.  Here is an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41495798-d1ba-4b20-aab5-e51fb3890e28",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy import units as u\n",
    "from astropy.io import fits\n",
    "from astropydantic import AstroPydanticTime\n",
    "from ctao_datamodel import (\n",
    "    AstroField,\n",
    "    Migration,\n",
    "    ModelBase,\n",
    "    Quantity,\n",
    "    VersionedModel,\n",
    "    fits_header_to_instance,\n",
    "    instance_to_fits_header,\n",
    "    migration,\n",
    "    models,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83b25db3-6203-49fa-990b-858d785430a4",
   "metadata": {},
   "source": [
    "## The original revision\n",
    "\n",
    "Imagine we had an old model that looked like this:\n",
    "\n",
    "```python3\n",
    "\n",
    "class TimeCov(ModelBase):\n",
    "    \"\"\"A sub-model example\"\"\"\n",
    "    reference: TimeReference | None = AstroField(\n",
    "        description=doc(TimeReference), default=None\n",
    "    )\n",
    "    t_start: float | AstroPydanticTime = AstroField(\n",
    "        \"Start of time range of the data as either an ISO string or a float in the\"\n",
    "        \" system defined by the time reference.\",\n",
    "        fits_keyword=\"TSTART\",\n",
    "    )\n",
    "    t_stop: float | AstroPydanticTime = AstroField(\n",
    "        \"End of time range of the data, as either an ISO string or a float in the\"\n",
    "        \" system defined by the time reference.\",\n",
    "        fits_keyword=\"TSTOP\",\n",
    "    )\n",
    "\n",
    "\n",
    "class Target(VersionedModel):\n",
    "    \"\"\"A simple example\"\"\"\n",
    "\n",
    "    object_name: str = AstroField(\"target name\", fits_keyword=\"OBJECT\")\n",
    "    ra: Quantity[\"deg\"] = AstroField(\n",
    "        \"Right ascension\", fits_keyword=\"RA2000\", fits_column_dtype=\"float64\", ucd=\"pos.eq.ra\"\n",
    "    )\n",
    "    dec: Quantity[\"deg\"] = AstroField(\n",
    "        \"Declination\", fits_column_dtype=\"float64\", fits_keyword=\"DE2000\", ucd=\"pos.eq.dec\"\n",
    "    )\n",
    "\n",
    "    coverage: TimeCov\n",
    "\n",
    "\n",
    "```\n",
    "\n",
    "But in the new version, we want to change some things: \n",
    "1. rename the field ``object_name`` to ``name``\n",
    "2. rename the sub-field ``coverage.t_start`` to ``coverage.t_min``, and similar for ``coverage.t_stop`` to ``coverage.t_max``\n",
    "3. change the FITS serialization keys of ``RA2000/DE2000`` to ``RA/DEC``\n",
    "4. drop support for ``t_min`` and ``t_max`` in an arbitrary time system, and allow only ISO strings\n",
    "5. make the time reference info deprecated\n",
    "\n",
    "And we want the new class to be backward compatible.  \n",
    "\n",
    "## The new version\n",
    "\n",
    "To achieve this, we created the  new model version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02342fea-99ab-4f39-8e5b-2eb409327710",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TimeCov(ModelBase):\n",
    "    \"\"\"A sub-model example\"\"\"\n",
    "\n",
    "    reference: models.common.TimeReference | None = AstroField(\n",
    "        \"DEPRECATED, do not use\", default=None\n",
    "    )\n",
    "\n",
    "    t_min: AstroPydanticTime = AstroField(\n",
    "        \"Start of time range of the data as either an ISO string or a float.\",\n",
    "        fits_keyword=\"TSTART\",\n",
    "    )\n",
    "    t_max: AstroPydanticTime = AstroField(\n",
    "        \"End of time range of the data, as either an ISO string or a float.\",\n",
    "        fits_keyword=\"TSTOP\",\n",
    "    )\n",
    "\n",
    "\n",
    "class Target(ModelBase):\n",
    "    \"\"\"A simple example\"\"\"\n",
    "\n",
    "    name: str = AstroField(\"target name\", fits_keyword=\"OBJECT\")\n",
    "    ra: Quantity[u.deg] = AstroField(\n",
    "        \"Right ascension\",\n",
    "        fits_keyword=\"RA\",\n",
    "        fits_column_dtype=\"float64\",\n",
    "        ucd=\"pos.eq.ra\",\n",
    "    )\n",
    "    dec: Quantity[u.deg] = AstroField(\n",
    "        \"Declination\", fits_column_dtype=\"float64\", fits_keyword=\"DEC\", ucd=\"pos.eq.dec\"\n",
    "    )\n",
    "\n",
    "    coverage: TimeCov"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62deadc2-f8b8-4d2d-ba09-aed4a98d365d",
   "metadata": {},
   "source": [
    "Let's set up a test case that was serialized with the old version into a FITS header:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9fc6902-38c7-46da-aea8-43f4cca8770b",
   "metadata": {},
   "outputs": [],
   "source": [
    "old_dict = dict(\n",
    "    object_name=\"Crab\",\n",
    "    ra=83.624 * u.deg,\n",
    "    dec=22.0174 * u.deg,\n",
    "    coverage=dict(\n",
    "        t_start=317788715.75,\n",
    "        t_stop=317788715.85,\n",
    "        reference={\n",
    "            \"position\": \"TOPOCENTER\",\n",
    "            \"time_mjd\": {\"value\": 58119.00042824074, \"unit\": \"d\"},\n",
    "            \"unit\": \"s\",\n",
    "            \"system\": \"TAI\",\n",
    "        },\n",
    "    ),\n",
    ")\n",
    "old_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "883709c2-8e73-44f7-b052-c957a5f78943",
   "metadata": {},
   "source": [
    "if we try to deserialize this, it fails:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb9d282a-cefa-4bb2-876b-7febe2774140",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    Target.model_validate(old_dict)\n",
    "except ValueError as err:\n",
    "    print(err)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d719e28-6225-48ec-a88d-0137340f2acc",
   "metadata": {},
   "source": [
    "## Adding Migration Information\n",
    "\n",
    "To fix this, we re-define our model and make it a subclass of {py:class}`VersionedModel`, and we need to add the steps that change the old version into the new version by using the {py:class}`Migration` class, which provides migrations for standard changes like renaming or dropping columns, changing metadata, and more complex user-defined transformations.\n",
    "\n",
    "To use this, we add a function for each version using the {py:func}`migration` decorator. \n",
    "\n",
    "Note however the step of disallowing times in MET with a time reference would break the conversion, so for that we need to define a more complex transformation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e216055a-bf51-4435-809f-6e900ab0947a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_met_time_to_iso(obj: dict, key: str, reference_key: str = \"reference\"):\n",
    "    \"\"\"Convert MET time assuming there is a TimeReference Object.\"\"\"\n",
    "\n",
    "    from astropy.time import Time\n",
    "\n",
    "    # check if it's already ISO:\n",
    "    if isinstance(obj.get(key, None), str):\n",
    "        return\n",
    "\n",
    "    reference = models.common.TimeReference.model_validate(obj.get(reference_key))\n",
    "    ref_time = Time(reference.time_mjd, format=\"mjd\", scale=reference.system.lower())\n",
    "    new_time = Time(ref_time + (obj[key] * reference.unit))\n",
    "    obj[key] = new_time.iso\n",
    "    return obj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dc77ade-c06a-43fb-ad19-cfb9c4b15efc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Target(VersionedModel):\n",
    "    \"\"\"A simple example\"\"\"\n",
    "\n",
    "    name: str = AstroField(\"target name\", fits_keyword=\"OBJECT\")\n",
    "    ra: Quantity[\"deg\"] = AstroField(\n",
    "        \"Right ascension\",\n",
    "        fits_keyword=\"RA\",\n",
    "        fits_column_dtype=\"float64\",\n",
    "        ucd=\"pos.eq.ra\",\n",
    "    )\n",
    "    dec: Quantity[\"deg\"] = AstroField(\n",
    "        \"Declination\", fits_column_dtype=\"float64\", fits_keyword=\"DEC\", ucd=\"pos.eq.dec\"\n",
    "    )\n",
    "\n",
    "    coverage: TimeCov\n",
    "\n",
    "    @migration(\"1.0.0\", \"2.0.0\")\n",
    "    def _(m: Migration):\n",
    "        m.rename(\"object_name\", \"name\")\n",
    "        m.rename(\"coverage.t_start\", \"coverage.t_min\")\n",
    "        m.rename(\"coverage.t_stop\", \"coverage.t_max\")\n",
    "        m.update_metadata(\"ra\", \"fits_keyword\", \"RA2000\", \"RA\")\n",
    "        m.update_metadata(\"dec\", \"fits_keyword\", \"DEC2000\", \"DEC\")\n",
    "        m.apply_transform(\n",
    "            \"coverage.t_min\", transform_met_time_to_iso, \"Convert MET to ISO\"\n",
    "        )\n",
    "        m.apply_transform(\n",
    "            \"coverage.t_max\", transform_met_time_to_iso, \"Convert MET to ISO\"\n",
    "        )\n",
    "        m.drop(\"coverage.reference\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bede1d87-46e2-4306-95c1-71f7f8ca5502",
   "metadata": {},
   "outputs": [],
   "source": [
    "Target.print_history()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbf043c2-bbe4-46a1-a0e5-1beedbf266aa",
   "metadata": {},
   "source": [
    "Now, we use ``VersionedModel.model_validate_versioned`` to do the migration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f63aa66-d2f5-4ef8-93fa-0052fa3631d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    validated = Target.model_validate_versioned(old_dict)\n",
    "except ValueError as err:\n",
    "    print(err)\n",
    "\n",
    "validated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93274196-00e6-4d48-8321-faaeecadb305",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(validated.model_dump_json(indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cb28cf5-2477-4511-89ea-254c98d70c95",
   "metadata": {},
   "source": [
    "It works! The old format was read and transformed to the the new one.\n",
    "\n",
    "## FITS Headers\n",
    "\n",
    "The migration should work for FITS header representation as well, since we used the {py:func}`Migration.update_metadata` migration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40442cf4-7d04-4083-b1e8-c46159a1d672",
   "metadata": {},
   "outputs": [],
   "source": [
    "old_header = fits.Header(\n",
    "    dict(\n",
    "        OBJECT=\"Crab\",\n",
    "        RA2000=83.624,\n",
    "        DEC2000=22.0174,\n",
    "        TREFPOS=\"TOPOCENTER\",\n",
    "        MJDREF=58119.00042824074,\n",
    "        TIMEUNIT=\"s       \",\n",
    "        TIMESYS=\"TAI     \",\n",
    "        TSTART=317788715.75,\n",
    "        TSTOP=317788715.85,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "199e3e82-96e5-4a53-95c6-24ac4abae0f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "validated_header = fits_header_to_instance(old_header, Target)\n",
    "validated_header"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb9bb049-ad74-4d72-82ee-18c4512ea0bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "instance_to_fits_header(validated_header)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee99abe3-6418-48f9-897f-9c64138aaeec",
   "metadata": {},
   "source": [
    "And you can see that even the FITS header was converted to the new format."
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}