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

from scipy.optimize import minimize_scalar as minimize

from SHARKNIR_gen_base import SHARKNIR_gen_base


class SHARKNIR_cal_tec_NCPA_trial(SHINSApp):

    @exception.handler_for_run
    def run(self, ins_s=None, det_s=None, tel_s=None, rtc_s=None, **kwargs):
        '''Check that BCU has been initialized, then send specific command
        to CRED2 camera
        '''

        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, ngmax = uvars.Tech_Image_NGMAX)

        if self.paragraph("RTC SETUP"):

            ##################
            # SHOW RTC SETUP #
            ##################
            
            logger.debug("Enumerating RTC setup")
            for element, value in rtc_s.items():
                logger.debug("Element : % 15s  -  Value % 15s    " % (element, value))
            
            ##############
            # SAVE MODES #
            ##############
            
            modes = [ float(rtc_s["VECTOR_{}".format(i)]) for i in range(1, 39) ]
                        
            logger.warning("Modes: {}".format(modes))
            
            # ##################################
            # # Retrieving RTC centroid origin #
            # ##################################
            
            # logger.info("Read centroid origin on RTC...")
            # self.status, [cx, cy] = obsController.RTCTTReadSingle(uvars.CentroidOrigin, 2, 'dsp')
            # logger.info("Read centroid origin on RTC... Done")
            
            # # self.status, rtc_filenames, self.other = gen_tec_CloseLoop.main()
            # # cx,cy = self.other
            
            # cx_orig = cx
            # cy_orig = cy
            
            # logger.info("Retrieved center: {} {}".format(cx, cy))
            
            ###############################################
            # CHECK FOR ABSOLUTE OR RELATIVE DISPLACEMENT #
            ###############################################
            
            if "ABSOLUTE" in rtc_s:
                #logger.warning("ABSOLUTE keyword found")
                if (rtc_s["ABSOLUTE"] == "true") : # or rtc_s["ABSOLUTE"]=="0"):
                    is_absolute = True
                    logger.warning("Absolute movement")
                else:
                    is_absolute = False
                    logger.warning("Relative movement")
            else:
                is_absolute = False
                logger.warning("Relative movement")

        #######################
        # LOOP FOR EVERY MODE #
        #######################

        steps = int(ins_s["STEPS"])

        base_original = [0]*len(modes)
        base = list(base_original) # copy
        final_array = [0]*len(modes)
        opt_ee_array = []
        
        fast_mode_enabled = False
        
        #for idx, val in enumerate(modes):
        for idx, val in reversed(list(enumerate(modes))):
            if val != 0:
                count = 0 # Iteration count for every non-zero mode
                # For example, modes[1] = 2.2

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

                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)]

                if is_absolute:
                    # # [-4.4, -2.2,  0.0, +2.2, +4.4]
                    base_array = list(pos_array)  # copy
                else:
                    # Reset base to act only on the concerned
                    base = list(base_original) # copy
                    # logger.warning("Resetting to {} ".format(base))
                    # # [-4.4, +2.2, +2.2, +2.2, +2.2]
                    base_array = [val for s in range(steps)]
                    base_array[0] = -((steps-1)//2)*val

                
                for b in base_array :
                    base[idx] = b
                    count += 1
                    if self.paragraph("Loop mode {} and step {}  ({}/{})".format(idx+1, b, count, len(base_array))):

                        # For example:
                        # [0, -4.4, 0, 0]
                        # [0, -2.2, 0, 0]
                        # [0, 0, 0, 0]
                        # [0, 2.2, 0, 0]
                        # [0, 4.4, 0, 0]
                        
                        logger.info("Base is: {}".format(str(base)))
                        
                        # ############################
                        # # Pass centroid to the RTC #
                        # ############################
                        
                        # logger.info("Change centroid origin on the technical camera")
                        # self.status = obsController.RTCTTSetCentroidPos(cx, cy)
                        # ufun.error_handler(self.status, "RTCTTSetCentroidPos")
                        
                        # logger.info("TECCAM CentroidOrigin is now {}, {}".format(cx, cy))
                        
                        #############################
                        # APPLY THE DM DISPLACEMENT #
                        #############################
                        
                        logger.info("Upload a mode vector to internal DM")
                        self.status = obsController.RTCTTModeUpload(base, absolute=is_absolute)
                        logger.info("Waiting setup to apply displacement... 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) )

                        if(not fast_mode_enabled):
                            obsController.setFlags(8,0)
                            logger.info("Trial and Error Fast Mode ENABLED")
                            fast_mode_enabled = True

                    #####################
                    # 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":

                            sci_filt_1, sci_filt_2 = ufun.get_status_by_name(os=obsController, names=["SCI_FILT_W1","SCI_FILT_W2"])
                            print(sci_filt_1)
                            print(sci_filt_2)
                            wvl = ufun.determine_filter_wavelength(sci1 = sci_filt_1, sci2 = sci_filt_2)
                            print(wvl)
                            radius = (3.432 * 0.7)*wvl/uvars.filters_wavelength["NB_H2"] # Depends on wavelength and pixel scale. 70% of lambda/D @1.593um and with LS.
                            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 MODE #
                ###############################
                if self.paragraph("FITTING FILES FOR THIS MODE"):

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

                    # saving the EE with previous best mode applied and current mode not applied yet
                    opt_ee_array.append(metric_array[len(metric_array)/2])
                    
                    # 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 = 4
                    
                    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
                    final_array[idx] = best_metric
                    
                    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))

                ###############################
                # APPLY OPTIMIZED MODES TO DM #
                ###############################
                if self.paragraph("APPLY OPTIMIZED MODES TO DM"):

                    if is_absolute:
                        base[idx] = best_metric # <--- Optimized value
                    else:
                        logger.warning("RELATIVE APPLYING NOT YET IMPLEMENTED!")
                    
                        come_back = list(base_original)
                        n_steps = -(steps-1)//2*val
                        logger.info("Coming back to relative zero, moving by {} steps".format(n_steps))
                        come_back[idx] = n_steps
                        logger.warning("APPLYING THIS TO DM NOW: {}".format(come_back))
                        self.status = obsController.RTCTTModeUpload(come_back, absolute=is_absolute)
                        logger.info("Coming back to relative zero... DONE")
                        base[idx] = best_metric # <--- Optimized value
                    
                    logger.warning("APPLYING THIS TO DM NOW: {}".format(base))
                    logger.info("Upload tip/tilt signal to internal DM")
                    self.status = obsController.RTCTTModeUpload(base, absolute=is_absolute)
                    logger.info("Waiting setup to apply displacement... DONE")
                    
                    logger.warning("Terminated this list of displacements for this mode: {}".format(str(pos_array)))

        logger.warning("Best base after taking best fwhm for every mode: "+str(final_array))
        logger.warning(final_array)


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

        if self.paragraph("TAKING FINAL TRIAL AND ERROR IMAGE"):
            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":
                
                metric = saly.PSFNoCoro.getEncircledEnergy(radius=radius, oversample=5, box=20,
                                                                verbose=False, display=False,
                                                                returnEE=True)
            
            # metric_array.append(metric)
            logger.info("FINAL {} is {}".format(rtc_s["METRIC"], metric))

        # saving the final EE  
        opt_ee_array.append(metric)
        logger.debug("Evolution of Encircled Energy after application of calculated optimal mode:")
        logger.debug(opt_ee_array)
        
        obsController.unsetFlags(8,0)
        fast_mode_enabled = False
        logger.info("Trial and Error Fast Mode DISABLED")

        opt_ee_array = [np.round(ee, 3) for ee in opt_ee_array]
        self.tpl_output = {
            uvars.tpl_output_names["final_ee"]      : metric,
            uvars.tpl_output_names["trial_ee_evol"] : opt_ee_array
        }
