Source code for ctao_datamodel._fits

"""Serialization of models to and from FITS headers."""

import re
import warnings
from pathlib import Path

from astropy import units as u
from astropy.io import fits
from astropy.io.fits import Card, Header
from astropy.table import Table
from pydantic import BaseModel

from ._core import QuantityFormat
from ._table import (
    ColumnValidationIssue,
    TableValidationError,
    model_validate_astropy_table,
)
from ._visitor import (
    extract_keyword_mapping,
    flatten_model_instance,
    get_field_metadata,
    get_field_unit,
    unflatten_model_instance,
)

__all__ = [
    "instance_to_fits_header",
    "fits_header_to_flat_dict",
    "fits_header_to_instance",
    "validate_fits_bintable_hdu",
]


def _get_field_from_instance(model_instance: BaseModel, flat_key: str, sep: str = "."):
    """Return FieldInfo for an instance at a flattened key."""
    path = flat_key.split(sep)
    field = path.pop()

    instance = model_instance
    for attr in path:
        # handle expanded list, where the key is an integer:
        if isinstance(instance, list) and re.match("^[0-9]+$", attr):
            list_index = int(attr)
            instance = instance[list_index]
            continue
        # normal case
        instance = getattr(instance, attr)

    # now instance is the sub-model, so get the field:
    return instance.__class__.model_fields[field]




def instance_to_fits_header(
    model_instance: BaseModel, use_short: bool = True, hierarch_namespace: str = "CTAO"
) -> Header:
    """
    Convert a model instance to a FITS header.

    The resulting header will have a Card for each keyword. Long keywords will use the
    FITS HIERARCH standard in the given namespace

    Parameters
    ----------
    model_instance: BaseModel
       model instance to serialize
    use_short: bool
       if True, replace any log keywords that have fits_keyword mapping with their short form.
    hierarch_namespace: str
       starting string for the HIERARCH keyword

    Returns
    -------
    Header:
       FITS header suitable for writing to a file.
    """
    flat = flatten_model_instance(
        model_instance,
        separator=" ",
        to_string=True,
        quantity_format=QuantityFormat.FLOAT,
    )

    cards = []
    hierarch_cards = []
    for k, v in flat.items():
        field = _get_field_from_instance(model_instance, k, sep=" ")
        unit = get_field_unit(field=field)
        desc = get_field_metadata(field=field, metadata_key="description")
        fits_keyword = get_field_metadata(field=field, metadata_key="fits_keyword")

        if unit:
            unit = u.Unit(unit)
            comment = f"[{unit:fits}] {desc}"
        else:
            comment = desc

        if fits_keyword and use_short:
            keyword = fits_keyword
            cards.append(Card(keyword=keyword, value=v, comment=comment))
        else:
            keyword = f"HIERARCH {hierarch_namespace} {k.upper()}"
            hierarch_cards.append(Card(keyword=keyword, value=v, comment=comment))

    return Header(cards + hierarch_cards)



def get_fits_to_ctao_mapping(model: type[BaseModel], sep: str = ".") -> dict[str, str]:
    """Return a dict mapping FITS keyword to flat CTAO keyword."""
    # get the mapping between FITS key and CTAO key
    ctao_to_fits = extract_keyword_mapping(model, metadata_key="fits_keyword", sep=sep)

    # make reverse mapping:
    fits_to_ctao = dict()
    for ctao_key, fits_key_list in ctao_to_fits.items():
        for fits_key in fits_key_list:
            fits_to_ctao[fits_key] = ctao_key

    return fits_to_ctao




def fits_header_to_flat_dict(
    header: Header, model: type[BaseModel], ignore_extra_keys: bool = True
) -> dict[str, str]:
    """
    Turn a FITS header back into a flat dict of CTAO keywords.

    Parameters
    ----------
    header: Header
        FITS header with keys and values to extract
    model: type[BaseModel]
        Model to use for schema and key mapping
    ignore_extra_keys: bool
        If False, issue warnings for keys in header that do not
        map to model.

    Returns
    -------
    dict[str,str]:
       mapping of CTAO keyword to string value.
    """
    # get the mapping

    fits_to_ctao = get_fits_to_ctao_mapping(model=model)

    # build the flat dict

    flattened = dict()

    for fits_key in header.keys():
        if fits_key == "COMMENT":
            # comments are handled specially by Astropy, and they span multiple
            # keys. However, converting them to strings joins them into one.
            value = str(header[fits_key])
        else:
            value = header[
                fits_key
            ]  #  use this to get value right for handles multi-card values

        if fits_key.startswith("CTAO "):
            # Hierarchical keys can just be converted:
            ctao_key = ".".join(fits_key.replace("CTAO ", "").lower().split(" "))
        elif fits_key in fits_to_ctao:
            ctao_key = fits_to_ctao[fits_key]
        else:
            if not ignore_extra_keys:
                warnings.warn(f"Key '{fits_key}' is not in model {model.__name__}")
            continue

        flattened[ctao_key] = value

    return flattened





def fits_header_to_instance(
    header: Header, model: type[BaseModel], ignore_extra_keys: bool = True
) -> BaseModel:
    """
    Turn a FITS header back into an instance of a model.

    Parameters
    ----------
    header: Header
        FITS header with keys and values to extract
    model: type[BaseModel]
        Model to use for schema and key mapping
    ignore_extra_keys: bool
        If False, issue warnings for keys in header that do not
        map to model.

    Returns
    -------
    BaseModel:
       instance of model provided
    """
    from ._versioning import FITSHeaderMigration, VersionedModel

    # first apply any migrations to the header
    if issubclass(model, VersionedModel):
        fits_migration = FITSHeaderMigration(header)

        for info in model._migrations:
            info.apply(fits_migration)

    # Then convert from FITS mapping and unflatten it:
    flattened = fits_header_to_flat_dict(
        header=header, model=model, ignore_extra_keys=ignore_extra_keys
    )
    return unflatten_model_instance(
        flattened, model=model, parent_key="", separator="."
    )





def validate_fits_bintable_hdu(
    fits_file: Path | str,
    hdu: int | str,
    column_model: type[BaseModel],
    header_model: type[BaseModel],
) -> tuple[Table, BaseModel]:
    """
    Validate a FITS Bintable HDU.

    Parameters
    ----------
    fits_file: Path | str
        the file to validate
    hdu: int | str
        name or number of the HDU to validate in the file
    column_model: type[BaseModel]
        The model defining the schema for the table columns. It must have fields
        that have the fits_column_dtype set.
    header_model: type[BaseModel]
        The model defining the schema for the header metadata to validate.

    Returns
    -------
    tuple[Table,BaseModel]:
        The table in Astropy Table format, and the metadata as a
        deserialized model.
    """
    issues = []

    # Check the table schema:
    table = Table.read(Path(fits_file), hdu=hdu)

    try:
        model_validate_astropy_table(table, model=column_model)
    except TableValidationError as err:
        issues += err.issues

    # check the metadata
    with fits.open(fits_file) as fitsfile:
        try:
            metadata = fits_header_to_instance(fitsfile["EVENTS"].header, header_model)
        except ValueError as err:
            issues.append(
                ColumnValidationIssue(
                    column="metadata", kind="keyword", message=str(err)
                )
            )

    if issues:
        raise TableValidationError(issues=issues)

    return table, metadata