import sys
import os
import signal

import numpy as np

from trm import cline
from trm.cline import Cline

import hipercam as hcam
from hipercam import fringe

__all__ = [
    "averun",
]

#################################
#
# averun -- median average images
#
#################################


def averun(args=None):
    """``averun [source] (run first last twait tmax | flist) trim ([ncol
    nrow]) bias dark flat fmap (fpair [nhalf rmin rmax]) [method sigma
    adjust clobber] output``

    Averages images from a run using median combination, skipping the junk
    frames that result from NSKIP / NBLUE options in HiPERCAM and ULTRACAM
    data. `averun` is meant to be a simple tool to create median frames
    suitable prior to aperture selection with `setaper`. See `combine` if you
    want more fine-grained control over frame averaging. (`averun` uses a
    combination of `grab` [if needed] and `combine`).

    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'

        flist : str [if source ends 'f']
           name of file list

        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 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.

        trim : bool
           True to trim columns and/or rows off the edges of windows nearest
           the readout. Useful for ULTRACAM particularly.

        ncol : int [if trim, hidden]
           Number of columns to remove (on left of left-hand window, and right
           of right-hand windows)

        nrow : int [if trim, hidden]
           Number of rows to remove (bottom of windows)

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

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

        flat : str
           Name of flat field to divide by, 'none' to ignore.

        fmap : str
           Name of fringe map (see e.g. `makefringe`), 'none' to ignore.

        fpair : str [if fmap is not 'none']
           Name of fringe pair file (see e.g. `setfringe`). Required if
           a fringe map has been specified.

        nhalf : int [if fmap is not 'none', hidden]
           When calculating the differences for fringe measurement,
           a region extending +/-nhalf binned pixels will be used when
           measuring the amplitudes. Basically helps the stats.

        rmin : float [if fmap is not 'none', hidden]
           Minimum individual ratio to accept prior to calculating the overall
           median in order to reduce the effect of outliers. Although all ratios
           should be positive, you might want to set this a little below zero
           to allow for some statistical fluctuation.

        rmax : float [if fmap is not 'none', hidden]
           Maximum individual ratio to accept prior to calculating the overall
           median in order to reduce the effect of outliers. Probably typically
           < 1 if fringe map was created from longer exposure data.

        method : str [hidden, defaults to 'm']
           'm' for median, 'c' for clipped mean. See below for pros and cons.

        sigma : float [hidden; if method == 'c']
           With clipped mean combination, pixels that deviate by more than
           sigma RMS from the mean are kicked out. This is carried out in an
           iterative manner. sigma <= 0 implies no rejection, just a straight
           average. sigma=3 is typical.

        adjust : str [hidden; defaults to 'i']
           adjustments to make: 'i' = ignore; 'n' = normalise the mean of all
           frames to match the first; 'b' = add offsets so that the mean of
           all frames is the same as the first.  Option 'n' is useful for
           twilight flats; 'b' for combining biases.

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

        output : str
           output file. Set by default to match the last part of "run"
           (but it will have a different extension so they won't
           clash)

    """

    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("trim", Cline.GLOBAL, Cline.PROMPT)
        cl.register("ncol", Cline.GLOBAL, Cline.HIDE)
        cl.register("nrow", Cline.GLOBAL, Cline.HIDE)
        cl.register("flist", Cline.LOCAL, Cline.PROMPT)
        cl.register("bias", Cline.LOCAL, Cline.PROMPT)
        cl.register("dark", Cline.GLOBAL, Cline.PROMPT)
        cl.register("flat", Cline.GLOBAL, Cline.PROMPT)
        cl.register("fmap", Cline.GLOBAL, Cline.PROMPT)
        cl.register("fpair", Cline.GLOBAL, Cline.PROMPT)
        cl.register("nhalf", Cline.GLOBAL, Cline.HIDE)
        cl.register("rmin", Cline.GLOBAL, Cline.HIDE)
        cl.register("rmax", Cline.GLOBAL, Cline.HIDE)
        cl.register("method", Cline.LOCAL, Cline.HIDE)
        cl.register("sigma", Cline.LOCAL, Cline.HIDE)
        cl.register("adjust", Cline.LOCAL, Cline.HIDE)
        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", first, 0)

            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

        trim = cl.get_value("trim", "do you want to trim edges of windows?", True)
        if trim:
            ncol = cl.get_value("ncol", "number of columns to trim from windows", 0)
            nrow = cl.get_value("nrow", "number of rows to trim from windows", 0)

        # 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",
        )

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

        # fringe file (if any)
        fmap = cl.get_value(
            "fmap",
            "fringe map ['none' to ignore]",
            cline.Fname("fmap", hcam.HCAM),
            ignore="none",
        )

        if fmap is not None:
            fpair = cl.get_value(
                "fpair", "fringe pair file",
                cline.Fname("fringe", hcam.FRNG)
            )
            nhalf = cl.get_value(
                "nhalf", "half-size of fringe measurement regions",
                2, 0
            )
            rmin = cl.get_value(
                "rmin", "minimum fringe pair ratio", -0.2
            )
            rmax = cl.get_value(
                "rmax", "maximum fringe pair ratio", 1.0
            )

        cl.set_default("method", "m")
        method = cl.get_value(
            "method", "c(lipped mean), m(edian)", "c", lvals=("c", "m")
        )

        if method == "c":
            sigma = cl.get_value("sigma", "number of RMS deviations to clip", 3.0)

        cl.set_default("adjust", "i")
        adjust = cl.get_value(
            "adjust", "i(gnore), n(ormalise) b(ias offsets)", "i", lvals=("i", "n", "b")
        )

        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. Now do the work with 'grab' and 'combine'

    if server_or_local or bias is not None or dark is not None or \
       flat is not None or fmap is not None:

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

        # Build argument list
        args = [None,"prompt",source,"yes",resource]
        if server_or_local:
            args += [str(first), str(last), str(twait), str(tmax)]
        if trim:
            args += ["yes",str(ncol),str(nrow)]
        else:
            args += ["no"]
        args += [
            "none" if bias is None else bias,
            "none" if dark is None else dark,
            "none" if flat is None else flat,
        ]
        if fmap is None:
            args += ["none", "f32"]
        else:
            args += [fmap,fpair,str(nhalf),str(rmin),str(rmax),"false","f32"]

        resource = hcam.scripts.grab(args)

    # at this point 'resource' is a list of files, no matter the input
    # method.

    with CleanUp(
            resource, server_or_local or bias is not None or dark is not None
            or flat is not None or fmap is not None
    ) as cleanup:

        print("\nCalling 'combine' ...")
        args = [
            None, "prompt", resource,
            "none", "none", "none",
            method
        ]
        if method != "m":
            args += [str(sigma)]
        args += [
            adjust, "usemean=yes",
            "plot=no", "yes" if clobber else "no",
            output,
        ]
        hcam.scripts.combine(args)

        print("averun finished")

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

    def _sigint_handler(self, signal_received, frame):
        print("\naverun 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)
            print('temporary files removed')
