"""
This module is designed to make romancal tweakreg and resample functions accessible without
installing the original package due to their complex dependencies. As such, it is only subtly modified from
the original to accommodate the less stringent install requirements.


romancal copyright notice:

Copyright (C) 2010 Association of Universities for Research in Astronomy (AURA)

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    1. Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.

    2. Redistributions in binary form must reproduce the above
      copyright notice, this list of conditions and the following
      disclaimer in the documentation and/or other materials provided
      with the distribution.

    3. The name of AURA and its representatives may not be used to
      endorse or promote products derived from this software without
      specific prior written permission.

THIS SOFTWARE IS PROVIDED BY AURA ``AS IS'' AND ANY EXPRESS OR IMPLIED
WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL AURA BE LIABLE FOR ANY DIRECT, INDIRECT,
INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
DAMAGE.

Original: https://github.com/spacetelescope/romancal/blob/main/romancal/tweakreg/tweakreg_step.py
"""


import os
from pathlib import Path

import numpy as np
from astropy.table import Table
from roman_datamodels import datamodels as rdm
import spike.stcal.tweakreg_tweakreg as tweakreg
from spike.stcal.tweakreg_tweakreg import _SINGLE_GROUP_REFCAT_STR, SINGLE_GROUP_REFCAT

from spike.romancal.assign_wcs_utils import add_s_region

# LOCAL
from spike.romancal.datamodels_library import ModelLibrary
from spike.romancal.stpipe_core import RomanStep

DEFAULT_ABS_REFCAT = SINGLE_GROUP_REFCAT[0]

__all__ = ["TweakRegStep"]


class TweakRegStep(RomanStep):
    """
    TweakRegStep: Image alignment based on catalogs of sources detected in
    input images.
    """

    class_alias = "tweakreg"

    spec = f"""
        use_custom_catalogs = boolean(default=False) # Use custom user-provided catalogs?
        catalog_format = string(default='ascii.ecsv') # Catalog output file format
        catfile = string(default='') # Name of the file with a list of custom user-provided catalogs
        catalog_path = string(default='') # Catalog output file path
        enforce_user_order = boolean(default=False) # Align images in user specified order?
        expand_refcat = boolean(default=False) # Expand reference catalog with new sources?
        minobj = integer(default=15) # Minimum number of objects acceptable for matching
        searchrad = float(default=2.0) # The search radius in arcsec for a match
        use2dhist = boolean(default=True) # Use 2d histogram to find initial offset?
        separation = float(default=1.0) # Minimum object separation in arcsec
        tolerance = float(default=0.7) # Matching tolerance for xyxymatch in arcsec
        fitgeometry = option('shift', 'rshift', 'rscale', 'general', default='rshift') # Fitting geometry
        nclip = integer(min=0, default=3) # Number of clipping iterations in fit
        sigma = float(min=0.0, default=3.0) # Clipping limit in sigma units
        abs_refcat = string(default='{DEFAULT_ABS_REFCAT}')  # Absolute reference
        # catalog. Options: {_SINGLE_GROUP_REFCAT_STR}
        save_abs_catalog = boolean(default=False)  # Write out used absolute astrometric reference catalog as a separate product
        abs_minobj = integer(default=15) # Minimum number of objects acceptable for matching when performing absolute astrometry
        abs_searchrad = float(default=6.0) # The search radius in arcsec for a match when performing absolute astrometry
        # We encourage setting this parameter to True. Otherwise, xoffset and yoffset will be set to zero.
        abs_use2dhist = boolean(default=True) # Use 2D histogram to find initial offset when performing absolute astrometry?
        abs_separation = float(default=1.0) # Minimum object separation in arcsec when performing absolute astrometry
        abs_tolerance = float(default=0.7) # Matching tolerance for xyxymatch in arcsec when performing absolute astrometry
        # Fitting geometry when performing absolute astrometry
        abs_fitgeometry = option('shift', 'rshift', 'rscale', 'general', default='rshift')
        abs_nclip = integer(min=0, default=3) # Number of clipping iterations in fit when performing absolute astrometry
        abs_sigma = float(min=0.0, default=3.0) # Clipping limit in sigma units when performing absolute astrometry
        output_use_model = boolean(default=True)  # When saving use `DataModel.meta.filename`
        update_source_catalog_coordinates = boolean(default=False) # Update source catalog file with tweaked coordinates?
    """  # noqa: E501

    reference_file_types = []

    def process(self, input):

        # properly handle input
        try:
            if isinstance(input, rdm.DataModel):
                images = ModelLibrary([input])
            elif str(input).endswith(".asdf"):
                images = ModelLibrary([rdm.open(input)])
            elif isinstance(input, ModelLibrary):
                images = input
            else:
                images = ModelLibrary(input)
        except TypeError as e:
            e.args = (
                "Input to tweakreg must be a list of DataModels, an "
                "association, or an already open ModelLibrary "
                "containing one or more DataModels.",
            ) + e.args[1:]
            raise e

        if not images:
            raise ValueError("Input must contain at least one image model.")

        self.log.info(
            f"Number of image groups to be aligned: {len(images.group_indices):d}."
        )
        self.log.info("Image groups:")
        for name in images.group_names:
            self.log.info(f"  {name}")
        # set the first image as reference
        with images:
            ref_image = images.borrow(0)
            images.shelve(ref_image, 0, modify=False)

        catdict = _parse_catfile(self.catfile)

        use_custom_catalogs = self.use_custom_catalogs
        # if user requested the use of custom catalogs and provided a
        # valid 'catfile' file name that has no custom catalogs,
        # turn off the use of custom catalogs:
        if catdict is not None and not catdict:
            self.log.warning(
                "'use_custom_catalogs' is set to True but 'catfile' "
                "contains no user catalogs."
            )
            use_custom_catalogs = False

        if use_custom_catalogs and catdict:
            with images:
                for i, member in enumerate(images.asn["products"][0]["members"]):
                    filename = member["expname"]
                    if filename in catdict:
                        # FIXME: I'm not sure if this captures all the possible combinations
                        # for example, meta.tweakreg_catalog is set by the container (when
                        # it's present in the association). However the code in this step
                        # checks meta.source_catalog.tweakreg_catalog. I think this means
                        # that setting a catalog via an association does not work. Is this
                        # intended? If so, the container can be updated to not support that.
                        model = images.borrow(i)
                        model.meta["source_catalog"] = {
                            "tweakreg_catalog_name": catdict[filename],
                        }
                        images.shelve(model, i)
                    else:
                        images.shelve(model, i, modify=False)

        # set path where the source catalog will be saved to
        if len(self.catalog_path) == 0:
            self.catalog_path = os.getcwd()
        self.catalog_path = Path(self.catalog_path).as_posix()
        self.log.info(f"All source catalogs will be saved to: {self.catalog_path}")

        # set reference catalog name
        if not self.abs_refcat:
            self.abs_refcat = DEFAULT_ABS_REFCAT.strip().upper()
        if self.abs_refcat != DEFAULT_ABS_REFCAT:
            self.expand_refcat = True

        # build the catalogs for input images
        imcats = []
        with images:
            for i, image_model in enumerate(images):
                exposure_type = image_model.meta.exposure.type
                if exposure_type != "WFI_IMAGE":
                    self.log.info("Skipping TweakReg for spectral exposure.")
                    image_model.meta.cal_step.tweakreg = "SKIPPED"
                else:
                    source_catalog = getattr(image_model.meta, "source_catalog", None)
                    if source_catalog is None:
                        images.shelve(image_model, i, modify=False)
                        raise AttributeError(
                            "Attribute 'meta.source_catalog' is missing. "
                            "Please either run SourceCatalogStep or provide a custom source catalog."
                        )

                    try:
                        catalog = self.get_tweakreg_catalog(source_catalog, image_model)
                    except AttributeError as e:
                        self.log.error(f"Failed to retrieve tweakreg_catalog: {e}")
                        images.shelve(image_model, i, modify=False)
                        raise e

                    try:
                        # validate catalog columns
                        _validate_catalog_columns(catalog)
                    except ValueError as e:
                        self.log.error(f"Failed to validate catalog columns: {e}")
                        images.shelve(image_model, i, modify=False)
                        raise e

                    catalog = tweakreg.filter_catalog_by_bounding_box(
                        catalog, image_model.meta.wcs.bounding_box
                    )

                    if self.save_abs_catalog:
                        output_name = os.path.join(
                            self.catalog_path, f"fit_{self.abs_refcat.lower()}_ref.ecsv"
                        )
                        catalog.write(
                            output_name, format=self.catalog_format, overwrite=True
                        )

                    image_model.meta["tweakreg_catalog"] = catalog.as_array()
                    nsources = len(catalog)
                    self.log.info(
                        f"Detected {nsources} sources in {image_model.meta.filename}."
                        if nsources
                        else f"No sources found in {image_model.meta.filename}."
                    )
                    # build image catalog
                    # catalog name
                    catalog_name = os.path.splitext(image_model.meta.filename)[0].strip(
                        "_- "
                    )
                    # catalog data
                    catalog_table = Table(image_model.meta.tweakreg_catalog)
                    catalog_table.meta["name"] = catalog_name

                    imcat = tweakreg.construct_wcs_corrector(
                        wcs=image_model.meta.wcs,
                        refang=image_model.meta.wcsinfo,
                        catalog=catalog_table,
                        group_id=image_model.meta.group_id,
                    )
                    imcat.meta["model_index"] = i
                    imcats.append(imcat)
                images.shelve(image_model, i)

        # run alignment only if it was possible to build image catalogs
        if len(imcats):
            # extract WCS correctors to use for image alignment
            if len(images.group_indices) > 1:
                self.do_relative_alignment(imcats)

            if self.abs_refcat in SINGLE_GROUP_REFCAT:
                self.do_absolute_alignment(ref_image, imcats)

            # finalize step
            with images:
                for imcat in imcats:
                    image_model = images.borrow(imcat.meta["model_index"])
                    image_model.meta.cal_step["tweakreg"] = "COMPLETE"
                    # remove source catalog
                    del image_model.meta["tweakreg_catalog"]

                    # retrieve fit status and update wcs if fit is successful:
                    if "SUCCESS" in imcat.meta.get("fit_info")["status"]:
                        # Update/create the WCS .name attribute with information
                        # on this astrometric fit as the only record that it was
                        # successful:

                        # NOTE: This .name attrib agreed upon by the JWST Cal
                        #       Working Group.
                        #       Current value is merely a place-holder based
                        #       on HST conventions. This value should also be
                        #       translated to the FITS WCSNAME keyword
                        #       IF that is what gets recorded in the archive
                        #       for end-user searches.
                        imcat.wcs.name = f"FIT-LVL2-{self.abs_refcat}"

                        # serialize object from tweakwcs
                        # (typecasting numpy objects to python types so that it doesn't cause an
                        # issue when saving datamodel to ASDF)
                        wcs_fit_results = {
                            k: (
                                v.tolist()
                                if isinstance(v, (np.ndarray, np.bool_))
                                else v
                            )
                            for k, v in imcat.meta["fit_info"].items()
                        }
                        # add fit results and new WCS to datamodel
                        image_model.meta["wcs_fit_results"] = wcs_fit_results
                        # remove unwanted keys from WCS fit results
                        for k in [
                            "eff_minobj",
                            "matched_ref_idx",
                            "matched_input_idx",
                            "fit_RA",
                            "fit_DEC",
                            "fitmask",
                        ]:
                            del image_model.meta["wcs_fit_results"][k]

                        # update WCS
                        image_model.meta.wcs = imcat.wcs
                        # update S_REGION
                        add_s_region(image_model)

                    images.shelve(image_model, imcat.meta["model_index"])

        return images

    def update_catalog_coordinates(self, tweakreg_catalog_name, tweaked_wcs):
        """
        Update the source catalog coordinates using the tweaked WCS.

        Parameters
        ----------
        tweakreg_catalog_name : str
            The name of the TweakReg catalog file produced by `SourceCatalog`.
        tweaked_wcs : `gwcs.wcs.WCS`
            The tweaked World Coordinate System (WCS) object.

        Returns
        -------
        None
        """
        # read in cat file
        with rdm.open(tweakreg_catalog_name) as source_catalog_model:
            # get catalog
            catalog = source_catalog_model.source_catalog

            # define mapping between pixel and world coordinates
            colname_mapping = {
                ("xcentroid", "ycentroid"): ("ra_centroid", "dec_centroid"),
                ("x_psf", "y_psf"): ("ra_psf", "dec_psf"),
            }

            for k, v in colname_mapping.items():
                # get column names
                x_colname, y_colname = k
                ra_colname, dec_colname = v

                # calculate new coordinates using tweaked WCS and update catalog coordinates
                catalog[ra_colname], catalog[dec_colname] = tweaked_wcs(
                    catalog[x_colname], catalog[y_colname]
                )

            # save updated catalog (overwrite cat file)
            self.save_model(
                source_catalog_model,
                output_file=source_catalog_model.meta.filename,
                suffix="cat",
                force=True,
            )

    def read_catalog(self, catalog_name):
        """
        Reads a source catalog from a specified file.

        This function determines the format of the catalog based on the file extension.
        If the file ends with "asdf", it uses a specific method to open and read the catalog;
        otherwise, it reads the catalog using a standard table format.

        Parameters
        ----------
        catalog_name : str
            The name of the catalog file to read.

        Returns
        -------
        Table
            The read catalog as a Table object.

        Raises
        ------
        ValueError
            If the catalog format is unsupported.
        """
        if catalog_name.endswith("asdf"):
            with rdm.open(catalog_name) as source_catalog_model:
                catalog = source_catalog_model.source_catalog
        else:
            catalog = Table.read(catalog_name, format=self.catalog_format)
        return catalog

    def get_tweakreg_catalog(self, source_catalog, image_model):
        """
        Retrieve the tweakreg catalog from source detection.

        This method checks the source detection metadata for the presence of a
        tweakreg catalog data or a string with its name. It returns the catalog
        as a Table object if either is found, or raises an error if neither is available.

        Parameters
        ----------
        source_catalog : object
            The source catalog metadata containing catalog information.
        image_model : DataModel
            The image model associated with the source detection.

        Returns
        -------
        Table
            The retrieved tweakreg catalog as a Table object.

        Raises
        ------
        AttributeError
            If the required catalog information is missing from the source detection.
        """
        if getattr(source_catalog, "tweakreg_catalog", None):
            tweakreg_catalog = Table(np.asarray(source_catalog.tweakreg_catalog))
            del image_model.meta.source_catalog["tweakreg_catalog"]
            return tweakreg_catalog

        if getattr(source_catalog, "tweakreg_catalog_name", None):
            return self.read_catalog(source_catalog.tweakreg_catalog_name)

        raise AttributeError(
            "Attribute 'meta.source_catalog.tweakreg_catalog' is missing. "
            "Please either run SourceCatalogStep or provide a custom source catalog."
        )

    def do_relative_alignment(self, imcats):
        """
        Perform relative alignment of images.

        This method performs relative alignment with the specified parameters,
        including search radius, separation, and fitting geometry.

        Parameters
        ----------
        imcats : list
            A list of image catalogs containing source information for alignment.

        Returns
        -------
        None
        """
        tweakreg.relative_align(
            imcats,
            searchrad=self.searchrad,
            separation=self.separation,
            use2dhist=self.use2dhist,
            tolerance=self.tolerance,
            xoffset=0,
            yoffset=0,
            enforce_user_order=self.enforce_user_order,
            expand_refcat=self.expand_refcat,
            minobj=self.minobj,
            fitgeometry=self.fitgeometry,
            nclip=self.nclip,
            sigma=self.sigma,
            clip_accum=True,
        )

    def do_absolute_alignment(self, ref_image, imcats):
        """
        Perform absolute alignment of images.

        This method retrieves a reference image and performs absolute alignment
        using the specified parameters, including reference WCS information and
        catalog details. It aligns the provided image catalogs to the absolute
        reference catalog.

        Parameters
        ----------
        ref_image : DataModel
            The reference image used for alignment, which contains WCS information.
        imcats : list
            A list of image catalogs containing source information for alignment.

        Returns
        -------
        None
        """
        tweakreg.absolute_align(
            imcats,
            self.abs_refcat,
            ref_wcs=ref_image.meta.wcs,
            ref_wcsinfo=ref_image.meta.wcsinfo,
            epoch=ref_image.meta.exposure.mid_time.decimalyear,
            abs_minobj=self.abs_minobj,
            abs_fitgeometry=self.abs_fitgeometry,
            abs_nclip=self.abs_nclip,
            abs_sigma=self.abs_sigma,
            abs_searchrad=self.abs_searchrad,
            abs_use2dhist=self.abs_use2dhist,
            abs_separation=self.abs_separation,
            abs_tolerance=self.abs_tolerance,
            save_abs_catalog=self.save_abs_catalog,
            abs_catalog_output_dir=self.output_dir,
            clip_accum=True,
        )


def _parse_catfile(catfile):
    """
    Parse a catalog file and return a dictionary mapping data models to catalog paths.

    This function reads a specified catalog file, extracting data model names and
    their associated catalog paths. It supports a format where each line contains
    a data model followed by an optional catalog path, and it ensures that the
    file adheres to the expected structure.

    Parameters
    ----------
    catfile : str
        The path to the catalog file to be parsed.

    Returns
    -------
    dict or None
        A dictionary mapping data model names to catalog paths, or None if the
        input file is empty or invalid.

    Raises
    ------
    ValueError
        If the catalog file contains more than two columns per line.
    """

    if catfile is None or not catfile.strip():
        return None

    catdict = {}

    with open(catfile) as f:
        catfile_dir = os.path.dirname(catfile)

        for line in f:
            sline = line.strip()
            if not sline or sline[0] == "#":
                continue

            data_model, *catalog = sline.split()
            catalog = list(map(str.strip, catalog))
            if len(catalog) == 1:
                catdict[data_model] = os.path.join(catfile_dir, catalog[0])
            elif not catalog:
                catdict[data_model] = None
            else:
                raise ValueError("'catfile' can contain at most two columns.")

    return catdict


def _validate_catalog_columns(catalog):
    """
    Validate the presence of required columns in the catalog.

    This method checks if the specified axis column exists in the catalog.
    If the axis is not found, it looks for a corresponding centroid column
    and renames it if present. If neither is found, it raises an error.

    Parameters
    ----------
    catalog : Table
        The catalog to validate, which should contain source information.
    axis : str
        The axis to check for in the catalog (e.g., 'x' or 'y').

    Returns
    -------
    None

    Raises
    ------
    ValueError
        If the required columns are missing from the catalog.
    """
    for axis in ["x", "y"]:
        if axis not in catalog.colnames:
            long_axis = f"{axis}centroid"
            if long_axis in catalog.colnames:
                catalog.rename_column(long_axis, axis)
            else:
                raise ValueError(
                    "'tweakreg' source catalogs must contain a header with "
                    "columns named either 'x' and 'y' or 'xcentroid' and 'ycentroid'."
                )
    return catalog
