Source code for autowisp.light_curves.collect_light_curves

"""Functions for creating light curves from DR files."""

from string import Formatter
import os.path
import os
import logging

import numpy

from autowisp.hat.file_parsers import parse_fname_keywords
from autowisp import DataReductionFile
from autowisp.catalog import ensure_catalog
from autowisp.processing_steps.manual_util import get_catalog_config
from autowisp import LightCurveFile
from .lc_data_io import LCDataIO




class DecodingStringFormatter(Formatter):
    """Add one more conversion type: ``'d'`` that calls decode on the arg."""



    def convert_field(self, value, conversion):
        """If conversion is ``'d'`` -> ``value.decode()`` else pass to parent"""

        if conversion == "d":
            return value.decode()
        return super().convert_field(value, conversion)






def get_combined_sources(dr_filenames, **dr_path_substitutions):
    """Return all sources which are in at lest one of the DR files."""

    dr_sources = set()
    for dr_fname in dr_filenames:
        with DataReductionFile(dr_fname, "r") as dr_file:
            dr_sources.update(
                dr_file.get_source_ids(
                    string_source_ids=False, **dr_path_substitutions
                )
            )
    return dr_sources



# This is simple enough
# pylint: disable=too-many-locals


def collect_light_curves(
    dr_filenames,
    configuration,
    mark_start,
    mark_end,
    *,
    dr_fname_parser=parse_fname_keywords,
    optional_header=None,
    observatory=None,
    **path_substitutions,
):
    """
    Add the data from a collection of DR files to LCs, creating LCs if needed.

    Args:
        dr_filenames([str]):    The filenames of the data reduction files to add
            to LCs.

        configuration:    Object with attributes configuring the LC collection
            procedure.

        path_substitutions:    Any substitutions to resolve paths within DR and
            LC files to data to read/write (e.g. versions of various
            componenents).

        dr_fname_parser:    See same name argument to LCDataIO::create().

        optional_header:    See same name argument to LCDataIO::create().

    Returns:
        [(src ID part 1, src ID part 2, ...)];
            The sources for which new lightcurves were created.
    """

    logger = logging.getLogger(__name__)
    srcid_formatter = DecodingStringFormatter()
    with DataReductionFile(dr_filenames[0], "r") as first_dr:
        configuration["srcextract_psf_params"] = [
            param.decode()
            for param in first_dr.get_attribute(
                "srcextract.psf_map.cfg.psf_params", **path_substitutions
            )
        ]
        if not configuration["num_magfit_iterations"]:
            path_substitutions["magfit_iteration"] = (
                first_dr.get_num_magfit_iterations(**path_substitutions) - 1
            )
        with DataReductionFile(
            configuration["single_photref_dr_fname"], "r"
        ) as sphotref_dr:
            (catalog_sources, catalog_header), outliers = ensure_catalog(
                header=sphotref_dr.get_frame_header(),
                dr_files=dr_filenames,
                configuration=get_catalog_config(configuration, "lc"),
                return_metadata=True,
                skytoframe_version=configuration["skytoframe_version"],
            )[:2]
        for outlier_ind in reversed(outliers):
            outlier_dr = dr_filenames.pop(outlier_ind)
            logger.warning(
                "Data reduction file %s has outlier pointing. Discarding!",
                outlier_dr,
            )
            mark_end(outlier_dr, -2, final=True)
        data_io = LCDataIO.create(
            catalog_sources=catalog_sources,
            config=configuration,
            source_list=list(
                get_combined_sources(dr_filenames, **path_substitutions)
            ),
            source_id_parser=first_dr.parse_hat_source_id,
            dr_fname_parser=dr_fname_parser,
            optional_header=optional_header,
            observatory=observatory,
            **path_substitutions,
        )
    frame_chunk = data_io.max_dimension_size["frame"]
    logger.debug(
        "Generating LC filenames per: %s", repr(configuration["lc_fname"])
    )
    sources_lc_fnames = [
        (
            source_id,
            srcid_formatter.format(
                configuration["lc_fname"], *numpy.atleast_1d(source_id)
            ),
        )
        for source_id in data_io.source_destinations.keys()
    ]

    for dirname in {
        os.path.abspath(os.path.dirname(lc_fname))
        for _, lc_fname in sources_lc_fnames
    }:
        if not os.path.exists(dirname):
            os.makedirs(dirname)

    num_processed = 0
    while num_processed < len(dr_filenames):
        stop_processing = min(len(dr_filenames), num_processed + frame_chunk)
        data_io.prepare_for_reading()
        for dr_fname in dr_filenames[num_processed:stop_processing]:
            mark_start(dr_fname)
        config_skipped = list(
            map(
                data_io.read,
                enumerate(dr_filenames[num_processed:stop_processing]),
            )
        )

        data_io.prepare_for_writing([entry[0] for entry in config_skipped])
        # data_io.print_organized_configurations()

        for write_arg in sources_lc_fnames:
            data_io.write(write_arg)

        for dr_fname in dr_filenames[num_processed:stop_processing]:
            mark_end(dr_fname, status=1, final=False)

        for _, lc_fname in sources_lc_fnames:
            if os.path.exists(lc_fname):
                with LightCurveFile(lc_fname, "a") as lc_file:
                    lc_file.confirm_lc_length()

        for dr_fname in dr_filenames[num_processed:stop_processing]:
            mark_end(dr_fname, status=1, final=True)

        num_processed = stop_processing

    return catalog_sources, catalog_header



# pylint: enable=too-many-locals