
from shinsapp import SHINSApp
from util.shinsLogger import logger
from util import variables as uvars
from util import exception
from util import functions as ufun
from alignment.allineamento import Align_the_coronagraph as aly 
from alignment.utils import read_image
from SHARKNIR_gen_base import SHARKNIR_gen_base
from scipy.optimize import minimize_scalar as minimize
import numpy as np

class SHARKNIR_gen_focus(SHINSApp):
    """calibration
    """

    @exception.handler_for_run
    def run(self, ins_s=None, det_s=None, tel_s=None, rtc_s=None, **kwargs):

        logger.info(__name__)
        obsController = self.obsController
        self.status = uvars.EXIT_SUCCESS

        self.gen_base = SHARKNIR_gen_base();        gen_base = self.gen_base

        fitsDir = uvars.fitsDir + obsController.getCurrentDataDirectory() +'/'

        saly = aly(camera='HAWAII H-2RG')
        saly.detector.setStackType('CDS')
        saly.detector.setCrop(1)

        #############################################
        # START AND WAIT INSTRUMENT AND SASHA SETUP #
        #############################################

        gen_base.main(det_s=det_s, expose=False)

        #######################
        # LOOP FOR EVERY STEP #
        #######################

        # For example, val = 2.2
        steps = int(ins_s["STEPS"])
        val = float(ins_s["AMOUNT"])
        n_steps = -(steps//2)*val      ########################-(steps-1)//2*val

        metric_array = []

        # [-4.4, -2.2,  0.0, +2.2, +4.4]
        pos_array = [val*s for s in range(-steps//2+1, steps//2+1)]

        is_absolute = False

        if is_absolute:
            # # [-4.4, -2.2,  0.0, +2.2, +4.4]
            base_array = list(pos_array)  # copy
        else:
            # # [-4.4, +2.2, +2.2, +2.2, +2.2]
            base_array = [val for s in range(steps)]
            base_array[0] = n_steps


        for position in base_array :

            ###########################################
            # APPLY NEW FOCUS VALUE AND TAKE AN IMAGE #
            ###########################################
            if self.paragraph("Focus step: {}. ".format(position)):

                #########################
                # APPLY NEW FOCUS VALUE #
                #########################
                logger.info("Waiting SOUL z movement to complete ...")
                self.obsController.OffsetZAO(position, "LEFT")
                logger.info("Waiting SOUL z movement to complete ... DONE")

                ##################
                # EXPOSING SASHA #
                ##################

                logger.info("Waiting camera exposure to complete ...")
                self.status, self.filenames = obsController.sashaExpose()
                logger.info("Waiting camera exposure to complete ... DONE")
                logger.info("FITS files are "+str(self.filenames) )

            #####################
            # MEASURE THE IMAGE #
            #####################
            if self.paragraph("MEASURE THE IMAGE"):

                im_arr = []
                for image_name in self.filenames:
                    logger.debug("reading {}".format(image_name))
                    img_ = read_image(fitsDir+image_name, saly.detector)
                    im_arr.append(img_)

                logger.debug("Making mean of {} images".format(len(im_arr)))
                meaned = np.mean(np.asarray(im_arr), 0)

                logger.debug("Loading")
                saly.PSFNoCoro.LoadImage(meaned)
                saly.PSFNoCoro.removeHotPixels(verbose=True)
                logger.debug("Getting Metric")

                if rtc_s["METRIC"] == "Strehl":

                    metric = saly.PSFNoCoro.getStrehl(0.0145, method='esposito', wavelengthInUm=1.091,
                                                      Dpup=8.22*0.98, eps=0.0, verbose=False,
                                                      doNotRecompDLPsf=True, returnSR=True)

                elif rtc_s["METRIC"] == "EE_frac":

                    radius = 2.35 * 0.7 # Depends on wavelength and pixel scale. 70% of lampda/D
                    metric = saly.PSFNoCoro.getEncircledEnergy(radius=radius, oversample=5, box=20,
                                                                    verbose=False, display=False,
                                                                    returnEE=True)

                metric_array.append(metric)
                logger.info("{} is {}".format(rtc_s["METRIC"], metric))


        ########################################
        # FITTING FILES FOR THIS THROUGH FOCUS #
        ########################################
        if self.paragraph("FITTING FILES FOR THIS THROUGH FOCUS"):

            logger.info("This is the array of measured {}  for each image:".format(rtc_s["METRIC"]))
            logger.info("{}".format(metric_array))

            # Get best width
            # First we polyfit the strhels:
            # pos_array is the amplitude vector: [-0.1, -0.08, #-0.06, -0.04, -0.02, 0.0, 0.02, 0.04, 0.06, 0.08, 0.1]),

            pol_degree = 2

            logger.info("Fitting a {} degree polynomial to that array!".format(pol_degree))

            fit  = np.polyfit(pos_array, metric_array, pol_degree)
            y_fit = np.poly1d(fit)

            # Then we calculate the local maximum of the polynomial.
            # x_best is the best amplitude that we want to pass to the DM

            minimization = minimize(lambda x: -1 * y_fit(x), bounds=(np.array(pos_array).min(),
                                                                     np.array(pos_array).max()),
                                                                     method='bounded')
            best_metric = minimization.x

            logger.info("Base array: {}".format(base_array))
            logger.info("Pos array:  {}".format(pos_array))
            logger.info("Best {} is at x={} ".format(rtc_s["METRIC"], best_metric))

        ##################
        # PAUSE / RESUME #
        ##################

        # self.paused = True
        # logger.warning("Resume when ready")
        # self.check_pause_or_abort()

        # #self.paused = True
        # logger.warning("Resume if you agree to apply; Quit otherwise.")
        # print("(R)esume, (Q)uit")

        # import sys
        # from select import select

        # timeout = 2

        # answer = None
        # print("Enter something:"),
        # self.paused = True


        # while self.paused==True and self.aborted==False:

        #     rlist, _, _ = select([sys.stdin], [], [], timeout)
        #     if rlist:
        #         answer = sys.stdin.readline()
        #         print(answer)
        #     else:
        #         print("still waiting...")

        #     if answer:
        #         if answer.lower().startswith('r'):
        #             self.paused = False
        #             pass
        #         if answer.lower().startswith('q'):
        #             self.paused = False
        #             self.aborted = True
        #             self.abort()
        #             return
        #         else:
        #             pass




        ##############################
        # APPLY OPTIMIZED FOCUS VALUE#
        ##############################
        if self.paragraph("APPLY OPTIMIZED FOCUS"):
            if is_absolute:
                pass
            else:
                #logger.warning("RELATIVE APPLYING NOT YET IMPLEMENTED!")            
                logger.info("Coming back to relative zero, moving by {} steps".format(n_steps))
                self.obsController.OffsetZAO(n_steps, "LEFT")
                logger.info("SOUL is now at relative zero position")

        logger.warning("APPLYING THIS TO SOUL NOW: {}".format(best_metric))

        logger.info("Waiting SOUL z movement to complete ...")
        self.obsController.OffsetZAO(best_metric, "LEFT")
        logger.info("Waiting SOUL z movement to complete ... DONE")
