Source code for ctao_datamodel._fits

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

import re
import warnings

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

from ._visitor import (
    extract_keyword_mapping,
    flatten_model_instance,
    get_field_metadata,
    unflatten_model_instance,
)

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


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, hieararch_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)

    cards = []
    for k, v in flat.items():
        field = _get_field_from_instance(model_instance, k, sep=" ")
        unit = get_field_metadata(field=field, metadata_key="unit")
        desc = get_field_metadata(field=field, metadata_key="description")
        fits = 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 and use_short:
            cards.append(Card(keyword=fits, value=v, comment=comment))
        else:
            cards.append(
                Card(
                    keyword=f"HIERARCH {hieararch_namespace} {k.upper()}",
                    value=v,
                    comment=desc,
                )
            )

    return Header(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 card in header.cards:
        fits_key = card.keyword
        value = card.value

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

    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
    """
    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="."
    )