from functools import partial
import os
import sys
import multiprocessing
import numpy as np

from trm import cline
from trm.cline import Cline

import hipercam as hcam
from hipercam import spooler

from hipercam.reduction import (
    Rfile,
    initial_checks,
    update_plots,
    ProcessCCDs,
    setup_plots,
    setup_plot_buffers,
    LogWriter,
    ccdproc,
    extractFlux,
)


# get hipercam version to write into the reduce log file
from pkg_resources import get_distribution, DistributionNotFound

try:
    hipercam_version = get_distribution("hipercam").version
except DistributionNotFound:
    hipercam_version = "not found"

has_psf_option = True
try:
    from hipercam.psf_reduction import extractFluxPSF
except ModuleNotFoundError:
    has_psf_option = False

__all__ = [
    "reduce",
]

NaN = float("NaN")

################################################
#
# reduce -- reduces multi-CCD imaging photometry
#
################################################
def reduce(args=None):
    """``reduce [source] rfile (run first last [twait tmax]| flist) trim
    ([ncol nrow]) log lplot implot (ccd nx msub xlo xhi ylo yhi iset
    (ilo ihi | plo phi))``

    Reduces a sequence of multi-CCD images, plotting lightcurves as
    images come in. It can extract with either simple aperture
    photometry or Tim Naylor's optimal photometry, on specific targets
    defined in an aperture file using |setaper|.

    reduce can source data from both the ULTRACAM and HiPERCAM
    servers, from local 'raw' ULTRACAM and HiPERCAM files (i.e. .xml +
    .dat for ULTRACAM, 3D FITS files for HiPERCAM) and from lists of
    HiPERCAM '.hcm' files. If you have data from a different
    instrument you should convert into the FITS-based hcm format.

    reduce is primarily configured from a file with extension
    ".red". This contains a series of directives, e.g. to say how to
    re-position and re-size the apertures. An initial reduce file is
    best generated with the script |genred| after you have created an
    aperture file. This contains lots of help on what to do. |genred|
    tries to get something usable fast for observing purposes, but the
    parameters can nearly always be tweaked to produce better
    outcomes.

    A reduce run can be terminated at any point with ctrl-C without doing
    any harm. You may often want to do this at the start in order to adjust
    parameters of the reduce file.

    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' or a file list. If a run,
           then reduce and log below will be set to have the same name by
           default.

        first : int [if source ends 's' or 'l']
           first frame to reduce. 1 = first frame; set = 0 to always try to
           get the most recent frame (if it has changed).

        last : int [if source ends 's' or 'l', hidden]
           last frame to reduce. 0 to just continue until the end.  This is
           not prompted for by default and must be set explicitly.  It
           defaults to 0 if not set. Its purpose is to allow accurate
           profiling tests.

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

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

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

        trim : bool
           True to trim columns and/or rows off the edges of windows nearest
           the readout. Particularly useful with ULTRACAM windowed data where
           the first few rows and columns can contain bad data.

        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)

        rfile : str
           the "reduce" file, i.e. ASCII text file suitable for reading by
           ConfigParser. Best seen by example as it has many parts. If you
           are reducing a run, this will be set to have the same root name by
           default (but a different extension to avoid name clashes).

        log : str
           log file for the results. If you are reducing a run, this
           will be set to have the same root name by default (but a
           different extension to avoid name clashes)

        tkeep : float
           maximum number of minutes of data to store in internal buffers, 0
           for the lot. When large numbers of frames are stored, performance
           can be slowed (although I am not entirely clear why) in which case
           it makes sense to lose the earlier points (without affecting the
           saving to disk). This parameter also gives operation similar to that
           of "max_xrange" parameter in the ULTRACAM pipeline whereby just
           the last few minutes are shown.

        lplot : bool
           flag to indicate you want to plot the light curve. Saves time not
           to especially in high-speed runs.

        implot : bool
           flag to indicate you want to plot images.

        ccd : string [if implot]
           CCD(s) to plot, '0' for all, '1 3' to plot '1' and '3' only, etc.

        nx : int [if implot]
           number of panels across to display.

        msub : bool [if implot]
           subtract the median from each window before scaling for the
           image display or not. This happens after any bias subtraction.

        xlo : float [if implot]
           left-hand X-limit for plot

        xhi : float [if implot]
           right-hand X-limit for plot (can actually be < xlo)

        ylo : float [if implot]
           lower Y-limit for plot

        yhi : float [if implot]
           upper Y-limit for plot (can be < ylo)

        iset : string [if implot]
           determines how the intensities are determined. There are three
           options: 'a' for automatic simply scales from the minimum to the
           maximum value found on a per CCD basis. 'd' for direct just takes
           two numbers from the user. 'p' for percentile dtermines levels
           based upon percentiles determined from the entire CCD on a per CCD
           basis.

        ilo : float [if implot and iset='d']
           lower intensity level

        ihi : float [if implot and iset='d']
           upper intensity level

        plo : float [if implot and iset='p']
           lower percentile level

        phi : float [if implot and iset='p']
           upper percentile level

    .. Warning::

       The transmission plot generated with reduce is not reliable in the
       case of optimal photometry since it is highly correlated with the
       seeing. If you are worried about the transmission during observing,
       you should always use normal aperture photometry.

    """

    command, args = cline.script_args(args)

    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.HIDE)
        cl.register("twait", Cline.LOCAL, Cline.HIDE)
        cl.register("tmax", Cline.LOCAL, Cline.HIDE)
        cl.register("flist", Cline.LOCAL, Cline.PROMPT)
        cl.register("trim", Cline.GLOBAL, Cline.PROMPT)
        cl.register("ncol", Cline.GLOBAL, Cline.HIDE)
        cl.register("nrow", Cline.GLOBAL, Cline.HIDE)
        cl.register("rfile", Cline.GLOBAL, Cline.PROMPT)
        cl.register("log", Cline.GLOBAL, Cline.PROMPT)
        cl.register("tkeep", Cline.GLOBAL, Cline.PROMPT)
        cl.register("lplot", Cline.LOCAL, Cline.PROMPT)
        cl.register("implot", Cline.LOCAL, Cline.PROMPT)
        cl.register("ccd", Cline.LOCAL, Cline.PROMPT)
        cl.register("nx", Cline.LOCAL, Cline.PROMPT)
        cl.register("msub", Cline.GLOBAL, Cline.PROMPT)
        cl.register("iset", Cline.GLOBAL, Cline.PROMPT)
        cl.register("ilo", Cline.GLOBAL, Cline.PROMPT)
        cl.register("ihi", Cline.GLOBAL, Cline.PROMPT)
        cl.register("plo", Cline.GLOBAL, Cline.PROMPT)
        cl.register("phi", Cline.LOCAL, Cline.PROMPT)
        cl.register("xlo", Cline.GLOBAL, Cline.PROMPT)
        cl.register("xhi", Cline.GLOBAL, Cline.PROMPT)
        cl.register("ylo", Cline.GLOBAL, Cline.PROMPT)
        cl.register("yhi", Cline.GLOBAL, 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 some flags
        server_or_local = source.endswith("s") or source.endswith("l")

        if server_or_local:
            resource = cl.get_value("run", "run name", "run005")
            first = cl.get_value("first", "first frame to reduce", 1, 0)
            cl.set_default("last", 0)
            last = cl.get_value("last", "last frame to reduce", 0, 0)
            if last and last < first:
                print("Cannot set last < first unless last == 0")
                print("*** reduce aborted")
                exit(1)

            twait = cl.get_value(
                "twait", "time to wait for a new frame [secs]", 1.0, 0.0
            )
            tmx = cl.get_value(
                "tmax", "maximum time to wait for a new frame [secs]", 10.0, 0.0
            )

            # keep the reduce name and log in sync to save errors
            root = os.path.basename(resource)
            cl.set_default("rfile", cline.Fname(root, hcam.RED))
            cl.set_default("log", cline.Fname(root, hcam.LOG, cline.Fname.NEW))

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

        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)

        # the reduce file
        rfilen = cl.get_value(
            "rfile", "reduce file", cline.Fname("reduce.red", hcam.RED)
        )
        try:
            rfile = Rfile.read(rfilen)
        except hcam.HipercamError as err:
            # abort on failure to read as there are many ways to get reduce
            # files wrong
            print(err, file=sys.stderr)
            print("*** reduce aborted")
            exit(1)

        log = cl.get_value(
            "log",
            "name of log file to store results",
            cline.Fname("reduce.log", hcam.LOG, cline.Fname.NEW),
        )

        tkeep = cl.get_value(
            "tkeep",
            "number of minute of data to" " keep in internal buffers (0 for all)",
            0.0,
            0.0,
        )

        lplot = cl.get_value("lplot", "do you want to plot light curves?", True)

        implot = cl.get_value("implot", "do you want to plot images?", True)

        if implot:

            # define the panel grid. first get the labels and maximum
            # dimensions
            ccdinf = spooler.get_ccd_pars(source, resource)

            nxdef = cl.get_default("nx", 3)

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

                if len(ccds) > 1:
                    nxdef = min(len(ccds), nxdef)
                    cl.set_default("nx", nxdef)
                    nx = cl.get_value("nx", "number of panels in X", 3, 1)
                else:
                    nx = 1
            else:
                nx = 1
                ccds = list(ccdinf.keys())

            # define the display intensities
            msub = cl.get_value("msub", "subtract median from each window?", True)

            iset = cl.get_value(
                "iset",
                "set intensity a(utomatically)," " d(irectly) or with p(ercentiles)?",
                "a",
                lvals=["a", "d", "p"],
            )

            plo, phi = 5, 95
            ilo, ihi = 0, 1000
            if iset == "d":
                ilo = cl.get_value("ilo", "lower intensity limit", 0.0)
                ihi = cl.get_value("ihi", "upper intensity limit", 1000.0)
            elif iset == "p":
                plo = cl.get_value(
                    "plo", "lower intensity limit percentile", 5.0, 0.0, 100.0
                )
                phi = cl.get_value(
                    "phi", "upper intensity limit percentile", 95.0, 0.0, 100.0
                )

            # arguments that set region to plot
            for i, cnam in enumerate(ccds):
                nxtot, nytot, nxpad, nypad = ccdinf[cnam]
                if i == 0:
                    xmin, xmax = float(-nxpad), float(nxtot + nxpad + 1)
                    ymin, ymax = float(-nypad), float(nytot + nypad + 1)
                else:
                    xmin = min(xmin, float(-nxpad))
                    xmax = max(xmax, float(nxtot + nxpad + 1))
                    ymin = min(ymin, float(-nypad))
                    ymax = max(ymax, float(nytot + nypad + 1))

            xlo = cl.get_value(
                "xlo", "left-hand X value", xmin, xmin, xmax, enforce=False
            )
            xhi = cl.get_value(
                "xhi", "right-hand X value", xmax, xmin, xmax, enforce=False
            )
            ylo = cl.get_value("ylo", "lower Y value", ymin, ymin, ymax, enforce=False)
            yhi = cl.get_value("yhi", "upper Y value", ymax, ymin, ymax, enforce=False)

        else:
            ccds, nx, msub, iset = None, None, None, None
            ilo, ihi, plo, phi = None, None, None, None
            xlo, xhi, ylo, yhi = None, None, None, None

        # save list of parameter values for writing to the reduction file
        plist = cl.list()

    ################################################################
    #
    # all the inputs have now been obtained. Get on with doing stuff
    if implot:
        plot_lims = (xlo, xhi, ylo, yhi)
    else:
        plot_lims = None

    imdev, lcdev, spanel, tpanel, xpanel, ypanel, lpanel = setup_plots(
        rfile, ccds, nx, plot_lims, implot, lplot
    )

    # a couple of initialisations
    total_time = 0  # time waiting for new frame

    if lplot:
        lbuffer, xbuffer, ybuffer, tbuffer, sbuffer = setup_plot_buffers(rfile)
    else:
        lbuffer, xbuffer, ybuffer, tbuffer, sbuffer = None, None, None, None, None

    ############################################
    #
    # open the log file and write headers
    #
    with LogWriter(log, rfile, hipercam_version, plist) as logfile:

        ncpu = rfile["general"]["ncpu"]
        if ncpu > 1:
            pool = multiprocessing.Pool(processes=ncpu)
        else:
            pool = None

        # whether some parameters have been initialised
        initialised = False

        # containers for the processed, debiassed and raw MCCD groups
        # and their frame numbers
        pccds, bccds, mccds, nframes = [], [], [], []

        ##############################################
        #
        # Finally, start winding through the frames
        #

        with spooler.data_source(source, resource, first, full=False) as spool:

            # 'spool' is an iterable source of MCCDs
            for nf, mccd in enumerate(spool):

                if server_or_local:

                    # Handle the waiting game ...
                    give_up, try_again, total_time = spooler.hang_about(
                        mccd, twait, tmx, total_time
                    )

                    if give_up:
                        # Giving up, but need to handle any partially filled
                        # frame group

                        if len(mccds):
                            # finish processing remaining frames. This step
                            # will only occur if we have at least once passed
                            # to later stages during which read and gain will
                            # be set up
                            results = processor(pccds, mccds, nframes)

                            # write out results to the log file
                            alerts = logfile.write_results(results)

                            # print out any accumulated alert messages
                            if len(alerts):
                                print("\n".join(alerts))

                            update_plots(
                                results,
                                rfile,
                                implot,
                                lplot,
                                imdev,
                                lcdev,
                                pccd,
                                ccds,
                                msub,
                                nx,
                                iset,
                                plo,
                                phi,
                                ilo,
                                ihi,
                                xlo,
                                xhi,
                                ylo,
                                yhi,
                                lpanel,
                                xpanel,
                                ypanel,
                                tpanel,
                                spanel,
                                tkeep,
                                lbuffer,
                                xbuffer,
                                ybuffer,
                                tbuffer,
                                sbuffer,
                            )
                            mccds = []

                        print("reduce finished")
                        break

                    elif try_again:
                        continue

                # Trim the frames: ULTRACAM windowed data has bad
                # columns and rows on the sides of windows closest to
                # the readout which can badly affect reduction. This
                # option strips them.
                if trim:
                    hcam.ccd.trim_ultracam(mccd, ncol, nrow)

                # indicate progress
                if "NFRAME" in mccd.head:
                    nframe = mccd.head["NFRAME"]
                else:
                    nframe = nf + 1

                if source != "hf" and last and nframe > last:
                    # finite last frame number

                    if len(mccds):
                        # finish processing remaining frames
                        results = processor(pccds, bccds, mccds, nframes)

                        # write out results to the log file
                        alerts = logfile.write_results(results)

                        # print out any accumulated alert messages
                        if len(alerts):
                            print("\n".join(alerts))

                        update_plots(
                            results,
                            rfile,
                            implot,
                            lplot,
                            imdev,
                            lcdev,
                            pccd,
                            ccds,
                            msub,
                            nx,
                            iset,
                            plo,
                            phi,
                            ilo,
                            ihi,
                            xlo,
                            xhi,
                            ylo,
                            yhi,
                            lpanel,
                            xpanel,
                            ypanel,
                            tpanel,
                            spanel,
                            tkeep,
                            lbuffer,
                            xbuffer,
                            ybuffer,
                            tbuffer,
                            sbuffer,
                        )
                        mccds = []

                    print("\nHave reduced up to the last frame set.")
                    print("reduce finished")
                    break

                print(
                    "Frame {:d}: {:s} ({:s})".format(
                        nframe,
                        mccd.head["TIMSTAMP"],
                        "ok" if mccd.head.get("GOODTIME", True) else "nok",
                    ),
                    end="" if implot else "\n",
                )

                if not initialised:
                    # This is the first frame  which allows us to make
                    # some checks and initialisations.
                    read, gain, ok = initial_checks(mccd, rfile)

                    # choose between PSF and normal photometry
                    use_psf = rfile["psf_photom"].get("use_psf", "no")
                    if use_psf.lower().strip() != "no":
                        if not has_psf_option:
                            raise ValueError(
                                "'psf_photom' is set to 'yes', but photutils not installed"
                            )
                        # create a process function by wrapping ccdproc
                        # in a partial whose first argument is the PSF extraction
                        # function
                        process_func = partial(ccdproc, extractFluxPSF)
                    else:
                        # use a partial to encapsulate the normal extraction
                        process_func = partial(ccdproc, extractFlux)

                    # Define the CCD processor function object
                    processor = ProcessCCDs(rfile, read, gain, process_func, pool)

                    # set flag to show we are set
                    if not ok:
                        break
                    initialised = True

                # De-bias the data. Retain a copy of the raw data as 'mccd'
                # in order to judge saturation. Plain debiassed data called
                # bccd. Data for full processing called 'pccd'
                if rfile.bias is not None:
                    # subtract bias
                    bccd = mccd - rfile.bias
                    pccd = bccd.copy()
                    bexpose = rfile.bias.head.get("EXPTIME", 0.0)
                else:
                    # no bias subtraction
                    bccd = mccd.copy()
                    pccd = mccd.copy()
                    bexpose = 0.0

                if rfile.dark is not None:
                    # subtract dark, CCD by CCD
                    dexpose = rfile.dark.head["EXPTIME"]
                    for cnam in pccd:
                        ccd = pccd[cnam]
                        cexpose = ccd.head["EXPTIME"]
                        scale = (cexpose - bexpose) / dexpose
                        ccd -= scale * rfile.dark[cnam]

                if rfile.flat is not None:
                    # apply flat field to processed frame
                    pccd /= rfile.flat

                if rfile["focal_mask"]["demask"]:
                    # attempt to correct for poorly placed frame
                    # transfer mask causing a step illumination in the
                    # y-direction. Loop through all windows of all
                    # CCDs. Also include a stage where we average in
                    # the Y direction to try to eliminate high pixels.
                    dthresh = rfile["focal_mask"]["dthresh"]

                    for cnam, ccd in pccd.items():
                        for wnam, wind in ccd.items():

                            # form mean in Y direction, then try to
                            # mask out high pixels
                            ymean = np.mean(wind.data, 0)
                            xmask = ymean == ymean
                            while 1:
                                # rejection cycle, rejecting
                                # overly positive pixels
                                ave = ymean[xmask].mean()
                                rms = ymean[xmask].std()
                                diff = ymean - ave
                                diff[~xmask] = 0
                                imax = np.argmax(diff)
                                if diff[imax] > dthresh * rms:
                                    xmask[imax] = False
                                else:
                                    break

                            # form median in X direction
                            xmedian = np.median(wind.data[:, xmask], 1)

                            # subtract it's median to avoid removing
                            # general background
                            xmedian -= np.median(xmedian)

                            # now subtract from 2D image using
                            # broadcasting rules
                            wind.data -= xmedian.reshape((len(xmedian), 1))

                if rfile.fmap is not None:
                    # fringe correction
                    calsec = rfile["calibration"]

                    # apply only to CCDs present in both fmap and fpair
                    for cnam in rfile.fmap:
                        if cnam in rfile.fpair:
                            ccd = pccd[cnam]
                            fscale = rfile.fpair[cnam].scale(
                                ccd,
                                rfile.fmap[cnam],
                                calsec["nhalf"],
                                calsec["rmin"],
                                calsec["rmax"],
                            )
                            ccd -= fscale * rfile.fmap[cnam]

                # Accumulate frames into processing groups for faster
                # parallelisation
                pccds.append(pccd)
                mccds.append(mccd)
                bccds.append(bccd)
                nframes.append(nframe)

                if len(pccds) == rfile["general"]["ngroup"]:
                    # parallel processing. This should usually be the first
                    # points at which it takes place
                    results = processor(pccds, bccds, mccds, nframes)

                    # write out results to the log file
                    alerts = logfile.write_results(results)

                    # print out any accumulated alert messages
                    if len(alerts):
                        print("\n".join(alerts))

                    update_plots(
                        results,
                        rfile,
                        implot,
                        lplot,
                        imdev,
                        lcdev,
                        pccds[-1],
                        ccds,
                        msub,
                        nx,
                        iset,
                        plo,
                        phi,
                        ilo,
                        ihi,
                        xlo,
                        xhi,
                        ylo,
                        yhi,
                        lpanel,
                        xpanel,
                        ypanel,
                        tpanel,
                        spanel,
                        tkeep,
                        lbuffer,
                        xbuffer,
                        ybuffer,
                        tbuffer,
                        sbuffer,
                    )

                    # Reset the frame buffers
                    pccds, bccds, mccds, nframes = [], [], [], []

        if len(mccds):
            # out of loop now. Finish processing any remaining
            # frames.
            results = processor(pccds, bccds, mccds, nframes)

            # write out results to the log file
            alerts = logfile.write_results(results)

            # print out any accumulated alert messages
            if len(alerts):
                print("\n".join(alerts))

            update_plots(
                results,
                rfile,
                implot,
                lplot,
                imdev,
                lcdev,
                pccd,
                ccds,
                msub,
                nx,
                iset,
                plo,
                phi,
                ilo,
                ihi,
                xlo,
                xhi,
                ylo,
                yhi,
                lpanel,
                xpanel,
                ypanel,
                tpanel,
                spanel,
                tkeep,
                lbuffer,
                xbuffer,
                ybuffer,
                tbuffer,
                sbuffer,
            )

            print("reduce finished")
