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_DBI(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)

        fiber_test = False
        if fiber_test:
            logger.debug("Fiber TEST mode enabled")
            test_ins_s = {
                # "STEPS"                : "3",
                "CAL_FIBER_DEP"        : "FOCUS_IN",
                "SHUTTER"              : "IN",
                "CAL_FIBER_FOCUS_LAMP" : "ON",
                "SCI_FILT_W1"          : "NB_H2",
                "SCI_FILT_W2"          : "ND3",
                "ADC1_PosAng"          : "min_disp",
                "ADC2_PosAng"          : "min_disp",
                "ND_FILT_W"            : "OD1_3",
                "APODIZER_W"           : "HOLE",
                "DB_FILT_W"            : "WINDOW"
            }
            
            # test_det_s = {
            #     "DIT"                  : "0.3"

            # }
            
            ins_s.update(test_ins_s)
            # det_s.update(test_det_s)
            
        #############################################
        # START AND WAIT INSTRUMENT AND SASHA SETUP #
        #############################################
        
        gen_base.main(ins_s=ins_s, 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_r = []
        opt_ee_array_l = []
        
        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 = []
                metric_array_dbi = [] # questa lista contiene 2 ee per ogni pos 
                # [-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)

                        # Split meaned image in two
                        n_col = np.round(meaned.shape[1])
                        psfs_dbi = [meaned[:,:n_col/2], meaned[:,n_col/2:]]
                        
                        # Compute metric for the the two sides separately
                        
                        for psf in psfs_dbi: 
                            logger.debug("Loading {} Image".format("LEFT" if i==0 else "RIGHT"))
                            saly.PSFNoCoro.LoadImage(psf)
                            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_dbi.append(metric)
                            
                                                
                        logger.info("LEFT {} is {}".format(rtc_s["METRIC"], metric_array_dbi[-2]))
                        logger.info("RIGHT {} is {}".format(rtc_s["METRIC"], metric_array_dbi[-1]))


                ###############################
                # FITTING FILES FOR THIS MODE #
                ###############################
                if self.paragraph("FITTING FILES FOR THIS MODE"):
                    
                    best_metric_arr = np.zeros(2)
                    
                    for i in range(2):
                        metric_array = metric_array_dbi[i::2] # Quando i=0 prende gli indici pari (canale di sx), quando i=1 i dispari (canale di dx)
                        which_image = "LEFT" if i==0 else "RIGHT"
                        logger.info("This is the array of measured {}  for each {} image:".format(rtc_s["METRIC"],which_image))
                        logger.info("{}".format(metric_array))

                        # saving the EE with previous best mode applied and current mode not applied yet
                        if which_image == "LEFT" :
                            opt_ee_array_l.append(metric_array[len(metric_array)/2])
                        elif which_image == "RIGHT" :
                            opt_ee_array_r.append(metric_array[len(metric_array)/2])

                        # Get best width
                        # First we polyfit the metric:
                        # 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, CovMat  = np.polyfit(pos_array, metric_array, pol_degree, full=False, cov='scaled')# highest power first!
                        
                        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')
                        print(minimization.x)
                        best_metric_arr[i] = minimization.x
                    logger.info("Best metrics are at x={} (left) and x={} (right) ".format(best_metric_arr[0], best_metric_arr[1]))

                    best_metric = np.mean(best_metric_arr) # prendiamo la media??
                    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 "+str(rtc_s["METRIC"])+" for every mode:"+str(final_array))
        logger.warning(final_array)


        ##################
        # 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)
                
            # Split meaned image in two
            n_col = np.round(meaned.shape[1])
            psfs_dbi = [meaned[:,:n_col/2], meaned[:,n_col/2:]]
                        
            # Compute metric for the the two sides separately
                        
            for psf in psfs_dbi: 
                logger.debug("Loading {} Image".format("LEFT" if i==0 else "RIGHT"))
                saly.PSFNoCoro.LoadImage(psf)
                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_dbi.append(metric)
                            
                                                
            logger.info("LEFT {} is {}".format(rtc_s["METRIC"], metric_array_dbi[-2]))
            logger.info("RIGHT {} is {}".format(rtc_s["METRIC"], metric_array_dbi[-1]))

        # saving the final EE  
        opt_ee_array_l.append(metric_array_dbi[-2])
        logger.debug("Evolution of Encircled Energy on LEFT PSF after application of calculated optimal mode:")
        logger.debug(opt_ee_array_l)
        opt_ee_array_r.append(metric_array_dbi[-1])
        logger.debug("Evolution of Encircled Energy on RIGHT PSF after application of calculated optimal mode:")
        logger.debug(opt_ee_array_r)

        obsController.unsetFlags(8,0)
        fast_mode_enabled = False
        logger.info("Trial and Error Fast Mode DISABLED")

        self.tpl_output = {
            uvars.tpl_output_names["trial_dbi_ee_l"] : opt_ee_array_l[-1],
            uvars.tpl_output_names["trial_dbi_ee_r"] : opt_ee_array_r[-1]
        }
