# Licensed under a 3-clause BSD style license - see LICENSE.rst

import numpy as np

from ..utils import parse_input_data, parse_input_weights, parse_output_projection
from .background import determine_offset_matrix, solve_corrections_sgd
from .subset_array import ReprojectedArraySubset

__all__ = ["reproject_and_coadd"]


def reproject_and_coadd(
    input_data,
    output_projection,
    shape_out=None,
    input_weights=None,
    hdu_in=None,
    reproject_function=None,
    hdu_weights=None,
    combine_function="mean",
    match_background=False,
    background_reference=None,
    **kwargs,
):
    """
    Given a set of input images, reproject and co-add these to a single
    final image.

    This currently only works with 2-d images with celestial WCS.

    Parameters
    ----------
    input_data : iterable
        One or more input datasets to reproject and co-add. This should be an
        iterable containing one entry for each dataset, where a single dataset
        is one of:

            * The name of a FITS file
            * An `~astropy.io.fits.HDUList` object
            * An image HDU object such as a `~astropy.io.fits.PrimaryHDU`,
              `~astropy.io.fits.ImageHDU`, or `~astropy.io.fits.CompImageHDU`
              instance
            * A tuple where the first element is an `~numpy.ndarray` and the
              second element is either a `~astropy.wcs.WCS` or a
              `~astropy.io.fits.Header` object
            * An `~astropy.nddata.NDData` object from which the ``.data`` and
              ``.wcs`` attributes will be used as the input data.

    output_projection : `~astropy.wcs.WCS` or `~astropy.io.fits.Header`
        The output projection, which can be either a `~astropy.wcs.WCS`
        or a `~astropy.io.fits.Header` instance.
    shape_out : tuple, optional
        If ``output_projection`` is a `~astropy.wcs.WCS` instance, the
        shape of the output data should be specified separately.
    input_weights : iterable
        If specified, this should be an iterable with the same length as
        ``input_data``, where each item is one of:

            * The name of a FITS file
            * An `~astropy.io.fits.HDUList` object
            * An image HDU object such as a `~astropy.io.fits.PrimaryHDU`,
              `~astropy.io.fits.ImageHDU`, or `~astropy.io.fits.CompImageHDU`
              instance
            * An `~numpy.ndarray` array

    hdu_in : int or str, optional
        If one or more items in ``input_data`` is a FITS file or an
        `~astropy.io.fits.HDUList` instance, specifies the HDU to use.
    hdu_weights : int or str, optional
        If one or more items in ``input_weights`` is a FITS file or an
        `~astropy.io.fits.HDUList` instance, specifies the HDU to use.
    reproject_function : callable
        The function to use for the reprojection
    combine_function : { 'mean', 'sum', 'median' }
        The type of function to use for combining the values into the final
        image.
    match_background : bool
        Whether to match the backgrounds of the images.
    background_reference : `None` or `int`
        If `None`, the background matching will make it so that the average of
        the corrections for all images is zero. If an integer, this specifies
        the index of the image to use as a reference.
    kwargs
        Keyword arguments to be passed to the reprojection function.
    """

    # TODO: add support for saving intermediate files to disk to avoid blowing
    # up memory usage. We could probably still have references to array
    # objects, but we'd just make sure these were memory mapped

    # TODO: add support for specifying output array

    # Validate inputs

    if combine_function not in ("mean", "sum", "median"):
        raise ValueError("combine_function should be one of mean/sum/median")

    if reproject_function is None:
        raise ValueError(
            "reprojection function should be specified with the reproject_function argument"
        )

    # Parse the output projection to avoid having to do it for each

    wcs_out, shape_out = parse_output_projection(output_projection, shape_out=shape_out)

    # Start off by reprojecting individual images to the final projection

    arrays = []

    for idata in range(len(input_data)):

        # We need to pre-parse the data here since we need to figure out how to
        # optimize/minimize the size of each output tile (see below).
        array_in, wcs_in = parse_input_data(input_data[idata], hdu_in=hdu_in)

        # We also get the weights map, if specified
        if input_weights is None:
            weights_in = None
        else:
            weights_in = parse_input_weights(input_weights[idata], hdu_weights=hdu_weights)
            if np.any(np.isnan(weights_in)):
                weights_in = np.nan_to_num(weights_in)

        # Since we might be reprojecting small images into a large mosaic we
        # want to make sure that for each image we reproject to an array with
        # minimal footprint. We therefore find the pixel coordinates of corners
        # in the initial image and transform this to pixel coordinates in the
        # final image to figure out the final WCS and shape to reproject to for
        # each tile. Note that in future if we are worried about significant
        # distortions of the edges in the reprojection process we could simply
        # add arbitrary numbers of midpoints to this list.
        ny, nx = array_in.shape
        xc = np.array([-0.5, nx - 0.5, nx - 0.5, -0.5])
        yc = np.array([-0.5, -0.5, ny - 0.5, ny - 0.5])
        xc_out, yc_out = wcs_out.world_to_pixel(wcs_in.pixel_to_world(xc, yc))

        # Determine the cutout parameters

        # In some cases, images might not have valid coordinates in the corners,
        # such as all-sky images or full solar disk views. In this case we skip
        # this step and just use the full output WCS for reprojection.

        if np.any(np.isnan(xc_out)) or np.any(np.isnan(yc_out)):
            imin = 0
            imax = shape_out[1]
            jmin = 0
            jmax = shape_out[0]
        else:
            imin = max(0, int(np.floor(xc_out.min() + 0.5)))
            imax = min(shape_out[1], int(np.ceil(xc_out.max() + 0.5)))
            jmin = max(0, int(np.floor(yc_out.min() + 0.5)))
            jmax = min(shape_out[0], int(np.ceil(yc_out.max() + 0.5)))

        if imax < imin or jmax < jmin:
            continue

        wcs_out_indiv = wcs_out[jmin:jmax, imin:imax]
        shape_out_indiv = (jmax - jmin, imax - imin)

        # TODO: optimize handling of weights by making reprojection functions
        # able to handle weights, and make the footprint become the combined
        # footprint + weight map

        array, footprint = reproject_function(
            (array_in, wcs_in),
            output_projection=wcs_out_indiv,
            shape_out=shape_out_indiv,
            hdu_in=hdu_in,
            **kwargs,
        )

        if weights_in is not None:
            weights, _ = reproject_function(
                (weights_in, wcs_in),
                output_projection=wcs_out_indiv,
                shape_out=shape_out_indiv,
                hdu_in=hdu_in,
                **kwargs,
            )

        # For the purposes of mosaicking, we mask out NaN values from the array
        # and set the footprint to 0 at these locations.
        reset = np.isnan(array)
        array[reset] = 0.0
        footprint[reset] = 0.0

        # Combine weights and footprint
        if weights_in is not None:
            weights[reset] = 0.0
            footprint *= weights

        array = ReprojectedArraySubset(array, footprint, imin, imax, jmin, jmax)

        # TODO: make sure we gracefully handle the case where the
        # output image is empty (due e.g. to no overlap).

        arrays.append(array)

    # If requested, try and match the backgrounds.
    if match_background:
        offset_matrix = determine_offset_matrix(arrays)
        corrections = solve_corrections_sgd(offset_matrix)
        if background_reference:
            corrections -= corrections[background_reference]
        for array, correction in zip(arrays, corrections):
            array.array -= correction

    # At this point, the images are now ready to be co-added.

    # TODO: provide control over final dtype

    final_array = np.zeros(shape_out)
    final_footprint = np.zeros(shape_out)

    if combine_function in ("mean", "sum"):

        for array in arrays:

            # By default, values outside of the footprint are set to NaN
            # but we set these to 0 here to avoid getting NaNs in the
            # means/sums.
            array.array[array.footprint == 0] = 0

            final_array[array.view_in_original_array] += array.array * array.footprint
            final_footprint[array.view_in_original_array] += array.footprint

        if combine_function == "mean":
            with np.errstate(invalid="ignore"):
                final_array /= final_footprint

    elif combine_function == "median":

        # Here we need to operate in chunks since we could otherwise run
        # into memory issues

        raise NotImplementedError("combine_function='median' is not yet implemented")

    return final_array, final_footprint
