"""Uniform interface for source extractoin using several methods."""

import logging

import numpy
from astropy.io import fits
from scipy import spatial
from numpy.lib import recfunctions

from astrowisp.utils.file_utilities import (
    prepare_file_output,
    get_unpacked_fits,
)
from autowisp import source_finder_util
from autowisp.evaluator import Evaluator


# This still makes sense as a class
# pylint: disable=too-few-public-methods
class SourceFinder:
    """Find sources in an image of the night sky and repor properties."""

    @staticmethod
    def _create_mock_source_list(fits_fname, configuration):
        """Create a source list with randomly generated properties."""

        with fits.open(fits_fname) as fits_file:
            # False positive
            # pylint: disable=no-member
            hdu_index = 0 if fits_file[0].header["NAXIS"] else 1
            xresolution = fits_file[hdu_index].header["NAXIS1"]
            yresolution = fits_file[hdu_index].header["NAXIS2"]
            # pylint: disable=no-member
            med_pixel = numpy.median(fits_file[hdu_index].data)
        nsources = 1000
        result = numpy.empty(
            nsources,
            dtype=[
                (
                    name,
                    (numpy.int32 if name in ["id", "npix"] else numpy.float64),
                )
                for name in source_finder_util.get_srcextract_columns("fistar")
            ],
        )

        result["id"] = numpy.arange(nsources)
        # False positive
        # pylint: disable=no-member
        result["x"] = numpy.random.random(nsources) * xresolution
        result["y"] = numpy.random.random(nsources) * yresolution

        result["bg"] = (1.0 + 0.1 * numpy.random.random(nsources)) * med_pixel
        result["flux"] = (
            numpy.random.random(nsources)
            * configuration["brightness_threshold"]
        )
        result["amp"] = 0.2 * result["flux"]
        result["s/n"] = result["flux"] / result["bg"]

        result["s"] = 2.3 + 0.2 * numpy.random.random(nsources)
        result["d"] = 0.3 + 0.1 * numpy.random.random(nsources)
        result["k"] = 0.3 + 0.1 * numpy.random.random(nsources)

        result["npix"] = 20 + (5 * numpy.random.random(nsources)).astype(int)
        # pylint: enable=no-member
        return result

    def _add_saturation_flags(self, fits_fname, source_list):
        """Add array of how many pixels near source are saturated."""

        with fits.open(fits_fname, "readonly") as fits_file:
            saturated_kdtree = spatial.KDTree(
                numpy.column_stack(numpy.nonzero(fits_file[2].data)) + 0.5
            )

        source_coords = numpy.column_stack((source_list["y"], source_list["x"]))
        print(
            "Match radii: " + repr(numpy.sqrt(source_list["npix"] / numpy.pi))
        )
        saturated_in_range = saturated_kdtree.query_ball_point(
            source_coords, numpy.sqrt(source_list["npix"] / numpy.pi)
        )
        print(f"Saturated in range shape: {saturated_in_range.shape}")
        print(f"Saturated in range: {saturated_in_range!r}")

        return recfunctions.append_fields(
            source_list,
            "nsatpix",
            [len(saturated) for saturated in saturated_in_range],
            usemask=False,
        )

    def __init__(
        self,
        *,
        tool="hatphot",
        brightness_threshold=10,
        filter_sources="True",
        max_sources=0,
        allow_overwrite=False,
        allow_dir_creation=False,
        always_return_sources=False,
    ):
        """Prepare to use the specified tool and define faint limit."""

        self.configuration = {
            "tool": tool,
            "brightness_threshold": brightness_threshold,
            "allow_overwrite": allow_overwrite,
            "allow_dir_creation": allow_dir_creation,
            "always_return_sources": always_return_sources,
            "filter_sources": filter_sources,
            "max_sources": max_sources,
        }

    def __call__(self, fits_fname, source_fname=None, **configuration):
        """
        Extract the sources from the given frame and save or return them.

        Args:
            fits_fname(str):    The filename of the fits file to extract
                sources from. Can be packed or unpacked.

            source_fname(str or None):    If None, the extract sources are
                returned as numpy. field array. Otherwise, this specifies a file
                to save the source extractor output.

        Returns:
            field array or None:
                The extracted source from the image, if source_fname was None.
        """

        logger = logging.getLogger(__name__)
        configuration = {**self.configuration, **configuration}
        if configuration["tool"] == "mock":
            return self._create_mock_source_list(fits_fname, configuration)

        start_extraction = getattr(
            source_finder_util, "start_" + configuration["tool"]
        )
        with get_unpacked_fits(fits_fname) as unpacked_fname:
            logger.debug(
                f"Extracting sources from {unpacked_fname!r} to {source_fname!r}"
            )
            extraction_args = (
                unpacked_fname,
                configuration["brightness_threshold"],
            )
            if source_fname:
                prepare_file_output(
                    source_fname,
                    configuration["allow_overwrite"],
                    configuration["allow_dir_creation"],
                )
                with open(source_fname, "wb") as destination:
                    start_extraction(*extraction_args, destination).wait()
                if not configuration["always_return_sources"]:
                    return None
                sources_file = source_fname
            else:
                print("Creating extraction process.")
                extraction_process = start_extraction(*extraction_args)
                print("Extraction process started.")
                sources_file = extraction_process.stdout

            print("Parsing source list from " + repr(unpacked_fname))
            logger.debug("Parsing source list from %s", repr(unpacked_fname))

            result = numpy.genfromtxt(
                sources_file,
                names=source_finder_util.get_srcextract_columns(
                    configuration["tool"]
                ),
                dtype=None,
                deletechars="",
            )
            if configuration["tool"] == "hatphot":
                result["x"] -= 0.5
                result["y"] -= 0.5
            if configuration["filter_sources"] != "True":
                result = result[
                    Evaluator(result)(configuration["filter_sources"])
                ]
            logger.debug(f"Sorting {unpacked_fname!r} sources")

            result.sort(order="flux")
            result = numpy.flip(result)
            if configuration["max_sources"] > 0:
                result = result[: configuration["max_sources"]]

            if not source_fname:
                logger.debug(
                    "Waiting for source extraction process for %s to finish",
                    repr(unpacked_fname),
                )
                extraction_process.communicate()
            logger.debug(
                f"Adding saturation flags for {unpacked_fname!r} sources"
            )

            return self._add_saturation_flags(unpacked_fname, result)


# pylint: enable=too-few-public-methods
