import sys
import os
import tempfile
import signal

import numpy as np

from trm import cline
from trm.cline import Cline

import hipercam as hcam
from hipercam import utils, spooler

__all__ = [
    "makeflat",
]

####################################################
#
# makeflat -- makes flat fields from a set of frames
#
####################################################


def makeflat(args=None):
    """``makeflat [source] (run first last [twait tmax] | flist)
    bias dark ngroup ccd lower upper [clobber] output``

    Averages a set of images to make a flat field.

    Typically flat-fields for HiPERCAM and ULTRA(CAM|SPEC) are taken with a
    strongly time-variable twilight sky as the Sun sets or rises. A typical
    flat field run may start out bright, or even saturated, but by the end be
    only a few thousand counts above bias. Moreover, there are very often
    stars visible in the images, so we usually take them while offsetting the
    telescope in a spiral pattern. The challenge is to combine these images
    while rejecting the stars and saturated frames and giving due weight to
    the better exposed images. This moreover has to be done for each CCD which
    vary significantly in sensitivity.

    'makeflat' does this as follows: given an input list of files (or
    optionally a single run), it reads them all in, debiases them
    (optionally), and calculates the mean count level in each CCD,
    normalises by the mean and writes out the results to temporary
    files. For each CCD it then sorts the files by their (original)
    mean level, and for those that lie between defined limits it takes
    the median of the mean-mormalised frames in groups of defined
    size. Thus, say one had 75 OK images, then these would be divided
    into 10 groups, the first 9 having 7 frames, the last having
    16. The median average of each of these would be taken. In each
    case the mean levels would be adjusted to be the same before
    taking the average to overcome the problem of taking a median of a
    time-variable sky. The assumption is that while the level may
    vary, the pattern of the image does not. It is up to the user to
    check that this is correct. Each of the medians is adjusted to
    have a mean equal to the sum of the means of the input
    frames. Finally the normal average of all of these median frames
    is taken and the mean level of the final output normalised to
    1. The first step, taking the median in groups is designed to
    remove the stars assuming that the telescope was spiralled. The
    size of the groups ('ngroup' below is a crucial parameter in
    whether this works). A good strategy is to run makeflat for a
    succession of ever larger 'ngroup' and then to divide the results
    into each other to see if stars are visible.

    The final step, the average of the medians with adjusted mean
    levels, is to ensure that the flats are combined in a way that
    reflects the level of signal that they have, i.e. to avoid giving
    equal weights to the median of a series of flats with 20,000 counts
    per pixel and another series with 1,000 counts per pixel. This
    somewhat complex procedure is implemented through a series of
    temporary files which are written and read as the script runs, but
    deleted at its end. This allows very large numbers to be combined
    as long as there is enough memory to load 'ngroup' CCDs
    simultaneously, which should usually be fine.

    Parameters:

        source : str [hidden]
           Data source, five options:

               | 'hs' : HiPERCAM server
               | 'hl' : local HiPERCAM FITS file
               | 'us' : ULTRACAM server
               | 'ul' : local ULTRACAM .xml/.dat files
               | 'hf' : list of HiPERCAM hcm FITS-format files

           'hf' is used to look at sets of frames generated by 'grab'
           or converted from foreign data formats. The standard
           start-off default for ``source`` can be set using the
           environment variable HIPERCAM_DEFAULT_SOURCE. e.g. in bash
           :code:`export HIPERCAM_DEFAULT_SOURCE="us"` would ensure it
           always started with the ULTRACAM server by default. If
           unspecified, it defaults to 'hl'.

        run : str [if source ends 's' or 'l']
           run number to access, e.g. 'run034'

        first : int [if source ends 's' or 'l']
           exposure number to start from. 1 = first frame ('0' is
           not supported).

        last : int [if source ends 's' or 'l']
           last exposure number must be >= first or 0 for the whole lot.

        twait : float [if source ends 's' or 'l'; hidden]
           time to wait between attempts to find a new exposure, seconds.

        tmax : float [if source ends 's' or 'l'; hidden]
           maximum time to wait between attempts to find a new exposure,
           seconds.

        flist : str [if source ends 'f']
           name of file list. Assumed that these are dias and dark corrected.

        bias : str
           Name of bias frame to subtract, 'none' to ignore.

        dark : str
           Name of dark frame to subtract, 'none' to ignore. Note that
           it is assumed all CCDs have the same exposure time when making
           a dark correction.

        ngroup : int
           the number of frames. Probably should be at least 5, preferably
           more. Experiment to see its effect.

        ccd : str
           CCD(s) to process, '0' for all, '1 3' for '1' and '3' only, etc.
           Would almost always expect this to be set = '0'.

        lower : list of floats
           Lower limits to the mean count level for a flat to be included. The
           count level is determined after bias subtraction.  Should be the
           same number as the selected CCDs, and will be assumed to be in the
           same order. Use this to elminate frames that are of so low a level
           that the accuracy of the bias subtraction could be a worry.
           Suggested hipercam values: 3000 for each CCD. Enter values separated
           by spaces.

        upper : list of floats
           Upper limits to the mean count level for a flat to be included. The
           count level is determined *after* bias subtraction.  Should be the
           same number as the selected CCDs, and will be assumed to be in the
           same order. Use this to eliminate saturated, peppered or non-linear
           frames. Suggested hipercam values: 58000, 58000, 58000, 40000 and
           40000 for CCDs 1, 2, 3, 4 and 5. Enter values separated by spaces.
           ULTRACAM values 49000, 29000, 27000 for CCDs 1, 2 and 3.

        clobber : bool [hidden]
           clobber any pre-existing output files

        output : str
           output file. will be set by default to match the input name.

      .. Note::

         This routine writes the files returned by 'grab' to
         automatically generated files, typically in .hipercam/tmp, to
         avoid polluting the working directory. These are removed at
         the end, but may not be if you ctrl-C. You should check
         .hipercam/tmp for redundant files every so often

    """

    command, args = cline.script_args(args)

    # get the inputs
    with Cline("HIPERCAM_ENV", ".hipercam", command, args) as cl:

        # register parameters
        cl.register("source", Cline.GLOBAL, Cline.HIDE)
        cl.register("run", Cline.GLOBAL, Cline.PROMPT)
        cl.register("first", Cline.LOCAL, Cline.PROMPT)
        cl.register("last", Cline.LOCAL, Cline.PROMPT)
        cl.register("twait", Cline.LOCAL, Cline.HIDE)
        cl.register("tmax", Cline.LOCAL, Cline.HIDE)
        cl.register("flist", Cline.LOCAL, Cline.PROMPT)
        cl.register("bias", Cline.LOCAL, Cline.PROMPT)
        cl.register("dark", Cline.LOCAL, Cline.PROMPT)
        cl.register("ngroup", Cline.LOCAL, Cline.PROMPT)
        cl.register("ccd", Cline.LOCAL, Cline.PROMPT)
        cl.register("lower", Cline.LOCAL, Cline.PROMPT)
        cl.register("upper", Cline.LOCAL, Cline.PROMPT)
        cl.register("clobber", Cline.LOCAL, Cline.HIDE)
        cl.register("output", Cline.LOCAL, Cline.PROMPT)

        # get inputs
        default_source = os.environ.get('HIPERCAM_DEFAULT_SOURCE','hl')
        source = cl.get_value(
            "source",
            "data source [hs, hl, us, ul, hf]",
            default_source,
            lvals=("hs", "hl", "us", "ul", "hf"),
        )

        # set a flag
        server_or_local = source.endswith("s") or source.endswith("l")

        if server_or_local:
            resource = cl.get_value("run", "run name", "run005")
            root = os.path.basename(resource)
            cl.set_default('output', cline.Fname(root, hcam.HCAM))
            first = cl.get_value("first", "first frame to average", 1, 1)
            last = cl.get_value("last", "last frame to average (0 for all)", first, 0)
            if last < first and last != 0:
                sys.stderr.write("last must be >= first or 0")
                sys.exit(1)
            twait = cl.get_value(
                "twait", "time to wait for a new frame [secs]", 1.0, 0.0
            )
            tmax = cl.get_value(
                "tmax", "maximum time to wait for a new frame [secs]", 10.0, 0.0
            )

        else:
            resource = cl.get_value(
                "flist", "file list", cline.Fname("files.lis", hcam.LIST)
            )
            first = 1

        # bias frame (if any)
        bias = cl.get_value(
            "bias",
            "bias frame ['none' to ignore]",
            cline.Fname("bias", hcam.HCAM),
            ignore="none",
        )

        # dark frame (if any)
        dark = cl.get_value(
            "dark",
            "dark frame ['none' to ignore]",
            cline.Fname("dark", hcam.HCAM),
            ignore="none",
        )

        ngroup = cl.get_value(
            "ngroup", "number of frames per median average group", 3, 1
        )

        ccdinf = spooler.get_ccd_pars(source, resource)

        if len(ccdinf) > 1:
            ccd = cl.get_value("ccd", "CCD(s) to process [0 for all]", "0")
            if ccd == "0":
                ccds = list(ccdinf.keys())
            else:
                ccds = ccd.split()
        else:
            ccds = list(ccdinf.keys())

        # need to check that the default has the right number of items, if not
        # overr-ride it
        lowers = cl.get_default("lower")
        if lowers is not None and len(lowers) != len(ccds):
            cl.set_default("lower", len(ccds) * (5000,))

        lowers = cl.get_value(
            "lower",
            "lower limits on mean count level for included flats, 1 per CCD",
            len(ccds) * (5000,)
        )

        uppers = cl.get_default("upper")
        if uppers is not None and len(uppers) != len(ccds):
            cl.set_default("upper", len(ccds) * (50000,))

        uppers = cl.get_value(
            "upper",
            "lower limits on mean count level for included flats, 1 per CCD",
            len(ccds) * (50000,)
        )

        clobber = cl.get_value(
            "clobber", "clobber any pre-existing files on output", False
        )

        output = cl.get_value(
            "output",
            "output average",
            cline.Fname(
                "hcam", hcam.HCAM, cline.Fname.NEW if clobber else cline.Fname.NOCLOBBER
            )
        )

    # inputs done with.

    if server_or_local or bias is not None or dark is not None:

        print("\nCalling 'grab' ...")

        args = [None, "prompt", source, "yes", resource]
        if server_or_local:
            args += [str(first), str(last),str(twait), str(tmax)]
        args += [
            "no",
            "none" if bias is None else bias,
            "none" if dark is None else dark,
            "none", "none", "f32",
        ]
        resource = hcam.scripts.grab(args)

    # at this point 'resource' is a list of files, no matter the input
    # method. 'fnames' below will be used to store still more temporaries

    fnames = []
    with CleanUp(
            resource, fnames,
            server_or_local or bias is not None or dark is not None
    ) as cleanup:

        # Read all the files to determine mean levels (after bias
        # subtraction) save the bias-subtracted, mean-level normalised
        # results to temporary files
        print("Reading all files in to determine their mean levels")
        bframe, dframe = None, None
        means = {}
        for cnam in ccds:
            means[cnam] = {}

        # We might have a load of temporaries from grab, but we are about to
        # make some more to save the normalised versions.
        tdir = utils.temp_dir()

        with spooler.HcamListSpool(resource) as spool:

            for mccd in spool:

                # here we determine the mean levels, store them
                # then normalise the CCDs by them and save the files
                # to disk

                # generate the name to save to automatically
                fd, fname = tempfile.mkstemp(suffix=hcam.HCAM, dir=tdir)

                for cnam in ccds:
                    # its unlikely that flats would be taken with skips, but
                    # you never know. Eliminate them from consideration now.
                    ccd = mccd[cnam]
                    if ccd.is_data():
                        cmean = mccd[cnam].mean()
                        means[cnam][fname] = cmean
                        mccd[cnam] /= cmean

                # write the disk, save the name, close the filehandle
                fnames.append(fname)
                mccd.write(fname)
                os.close(fd)

                # a bit of progress info
                print(f"Saved processed flat to {fname}")

        # now we go through CCD by CCD, using the first as a template
        # for the window names in which we will also store the results.
        template = hcam.MCCD.read(fnames[0])

        for cnam, lower, upper in zip(ccds, lowers, uppers):
            tccd = template[cnam]

            # get the keys (filenames) and corresponding mean values
            mkeys = np.array(list(means[cnam].keys()))
            mvals = np.array(list(means[cnam].values()))

            # chop down to acceptable ones
            ok = (mvals > lower) & (mvals < upper)

            mkeys = mkeys[ok]
            mvals = mvals[ok]

            # some more progress info
            print("Found {:d} frames for CCD {:s}".format(len(mkeys), cnam))
            if len(mkeys) == 0:
                print(
                    (".. cannot average 0 frames;" " will skip CCD {:s}").format(cnam)
                )
                continue

            elif len(mkeys) < ngroup:
                print(
                    (
                        "WARNING: fewer than ngroup = {:d} frames"
                        " found. Output for CCD {:s} could be poor"
                    ).format(ngroup, cnam)
                )

            nchunk = len(mkeys) // ngroup
            if nchunk == 0:
                nchunk = 1

            # sort by mean value
            isort = mvals.argsort()
            mvals = mvals[isort]
            mkeys = mkeys[isort]

            # wsum used to sum all the eight factors to allow overall
            # normalisation at the end of the loop
            wsum = 0.0

            for n in range(nchunk):
                # loop through in chunks of ngroup at a time with a
                # potentially larger group to sweep up the end ones.
                n1 = ngroup * n
                n2 = n1 + ngroup
                if n == nchunk:
                    n2 = len(mkeys)

                # load the CCDs of this group
                ccdgroup = []
                with spooler.HcamListSpool(list(mkeys[n1:n2]), cnam) as spool:
                    for ccd in spool:
                        ccdgroup.append(ccd)

                # take median of the group to get rid of jumping
                # stars. 'weight' used to weight the results when summing the
                # results together. this stage is like the 'n' option of
                # 'combine' except we have already cut out any junk frames and
                # we have normalised the remainder
                weight = mvals[n1:n2].sum()
                wsum += weight

                for wnam, wind in tccd.items():
                    # go through each window, building a list of all data
                    # arrays
                    arrs = [ccd[wnam].data for ccd in ccdgroup]
                    arr3d = np.stack(arrs)

                    # at this point, arr3d is a 3D array, with the first
                    # dimension (axis=0) running over the images. We take the
                    # median over this axis. The first time through we put
                    # this straight into the output Window.  afterwards we add
                    # it in (with the appropriate weight)
                    if n == 0:
                        wind.data = weight * np.median(arr3d, axis=0)
                    else:
                        wind.data += weight * np.median(arr3d, axis=0)

            # Normalise the final result to a mean = 1.
            tccd /= wsum

            # Add some history
            tccd.head.add_history(
                ("result of makeflat on {:d}" " frames, ngroup = {:d}").format(
                    len(mkeys), ngroup
                )
            )

        # Remove any CCDs not included to avoid impression of having done
        # something to them
        dcnams = []
        for cnam in template.keys():
            if cnam not in ccds:
                dcnams.append(cnam)
        for cnam in dcnams:
            del template[cnam]

        # write out
        template.write(output, clobber)
        print("\nFinal result written to {:s}".format(output))
        print('makeflat finished')

class CleanUp:
    """
    Context manager to handle temporary files
    """
    def __init__(self, flist, fnames, temp):
        self.flist = flist
        self.fnames = fnames
        self.temp = temp

    def _sigint_handler(self, signal_received, frame):
        print("\nmakeflat aborted")
        sys.exit(1)

    def __enter__(self):
        signal.signal(signal.SIGINT, self._sigint_handler)

    def __exit__(self, type, value, traceback):
        if self.temp:
            with open(self.flist) as fp:
                for line in fp:
                    os.remove(line.strip())
            os.remove(self.flist)

        for fname in self.fnames:
            if os.path.exists(fname):
                os.remove(fname)

        print('temporary files removed')

