import sys
import os
import time

import numpy as np
import matplotlib.pyplot as plt
from astropy.time import Time

from trm import cline
from trm.cline import Cline

import hipercam as hcam
from hipercam import spooler, defect

__all__ = [
    "filtid",
]

##################################################################
#
# filtid -- attempts to identify the filters in use from sky flats
#
##################################################################


def filtid(args=None):
    """``filtid [source] (run first last [twait tmax] |
    flist) trim ([ncol nrow]) ccdref maxlev plot``

    Computes the ratio of the mean values in each CCD versus those of a reference CCD.
    This is supposed to be applied to flat fields. The hope is that this might be
    characteristic of the filter combination. Obviously requires > 1 CCD, i.e. not
    ULTRASPEC.

    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; set = 0 to always
           try to get the most recent frame (if it has changed).  For data
           from the |hiper| server, a negative number tries to get a frame not
           quite at the end.  i.e. -10 will try to get 10 from the last
           frame. This is mainly to sidestep a difficult bug with the
           acquisition system.

        last : int [if source ends 's' or 'l']
           last frame to reduce. 0 to just continue until the end.

        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 [if source starts with 'u']
           True to trim columns and/or rows off the edges of windows nearest
           the readout which can sometimes 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)

        ccdref : str
           The reference CCD (usually choose the g-band one)

        bias : str
           bias frame to subtract (required)

        lower : list of floats
           Lower limits to the mean count level for a flat to be
           included (after bias subtraction).  Should be the same
           number as the number of CCDs, and will be assumed to be in
           the same order. Separate with spaces. Prevents low exposure
           data from being included.

        upper : list of floats
           Upper limits to the mean count level for a flat to be
           included.  Should be the same number as the selected CCDs,
           and will be assumed to be in the same order. Use 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.

        plot: bool
           Plot the fit or not.

    .. Note::

       This is currently adapted specifically for ULTRACAM data

    """

    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("trim", Cline.GLOBAL, Cline.PROMPT)
        cl.register("ncol", Cline.GLOBAL, Cline.HIDE)
        cl.register("nrow", Cline.GLOBAL, 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("ccdref", Cline.LOCAL, Cline.PROMPT)
        cl.register("bias", Cline.LOCAL, Cline.PROMPT)
        cl.register("lower", Cline.LOCAL, Cline.PROMPT)
        cl.register("upper", Cline.LOCAL, Cline.PROMPT)
        cl.register("plot", 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 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 consider", 1, 1)
            last = cl.get_value("last", "last frame to consider", 0, 0)
            if last and last < first+2:
                raise hcam.HipercamError("Must consider at least 2 exposures, and preferably more")

            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)

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

        if len(ccdinf) == 1:
            raise hcam.HipercamError("Only one CCD; this routine only works for > 1 CCD")

        ccds = list(ccdinf.keys())
        ccdref = cl.get_value("ccdref", "the reference CCD for the ratios", ccds[1], lvals=ccds)

        bias = cl.get_value("bias", "bias frame", cline.Fname("bias", hcam.HCAM))

        # read the bias frame
        bias = hcam.MCCD.read(bias)
        if len(bias) != len(ccds):
            raise ValueError('bias has an incompatible number of CCDs')

        # need to check that the default has the right number of items, if not,
        # over-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,),
        )
        lowers = {k:v for k,v in zip(ccds,lowers)}

        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,),
        )
        uppers = {k:v for k,v in zip(ccds,uppers)}

        plot = cl.get_value("plot", "plot the results?", True)

    ################################################################
    #
    # all the inputs have now been obtained. Get on with doing stuff

    total_time = 0.
    xdata, ydata = {}, {}

    # only want to look at CCDs other than ccdref
    ccds.remove(ccdref)

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

        # 'spool' is an iterable source of MCCDs
        nframe = 0
        for mccd in spool:

            if server_or_local:
                # Handle the waiting game ...
                give_up, try_again, total_time = spooler.hang_about(
                    mccd, twait, tmax, total_time
                )

                if give_up:
                    print("rtplot stopped")
                    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)

            if nframe == 0:
                # crop the bias on the first frame only
                bias = bias.crop(mccd)

            # bias subtraction
            mccd -= bias

            # indicate progress
            tstamp = Time(mccd.head["TIMSTAMP"], format="isot", precision=3)
            print(
                f"{mccd.head.get('NFRAME',nframe+1)}, utc= {tstamp.iso} " +
                f"({'ok' if mccd.head.get('GOODTIME', True) else 'nok'})"
            )

            # add in data if it is in range
            xd  = mccd[ccdref].mean()
            if xd > lowers[ccdref] and xd < uppers[ccdref]:
                for cnam in ccds:
                    yd = mccd[cnam].mean()
                    if yd > lowers[cnam] and yd < uppers[cnam]:
                        if cnam in xdata:
                            xdata[cnam].append(xd)
                            ydata[cnam].append(yd)
                        else:
                            xdata[cnam] = [xd]
                            ydata[cnam] = [yd]

            if last and nframe == last:
                break

            # update the frame number
            nframe += 1

    print()
    invalid = True
    for cnam in ccds:
        if cnam in xdata:
            xdata[cnam] = np.array(xdata[cnam])
            ydata[cnam] = np.array(ydata[cnam])
            print(
                f'{cnam}-vs-{ccdref}: found {len(xdata[cnam])} valid points ranging from' +
                f' {xdata[cnam].min():.1f} to {xdata[cnam].max():.1f} (CCD {ccdref}), ' +
                f' {ydata[cnam].min():.1f} to {ydata[cnam].max():.1f} (CCD {cnam})'
            )
            invalid = False

    if invalid:
        raise hcam.HipercamError(f"No valid points found at all")

    print()
    if plot:
        fig,axes = plt.subplots(len(ccds),1,sharex=True)

        for cnam, ax in zip(ccds, axes):
            if cnam in xdata:
                xds = xdata[cnam]
                ratios = ydata[cnam] / xds
                mratio = ratios.mean()
                ax.plot(xds,ratios,'.g')
                ax.plot([xds.min(),xds.max()],[mratio,mratio],'r--')
                ax.set_ylabel(f'Ratio CCD {cnam} / CCD {ccdref}')
                if cnam == ccds[-1]:
                    ax.set_xlabel(f'CCD = {ccdref} counts/pixel')
                print(f'{cnam}-vs-{ccdref} = {mratio:.4f}')
        plt.show()

    else:
        for cnam in ccds:
            if cnam in xdata:
                ratios = ydata[cnam] / xdata[cnam]
                mratio = ratios.mean()
                print(f'{cnam}-vs-{ccdref} = {mratio:.4f}')

