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

# Other templates
from SHARKNIR_gen_tec_CloseLoop import SHARKNIR_gen_tec_CloseLoop
from SHARKNIR_gen_base import SHARKNIR_gen_base

class SHARKNIR_gen_tec_PSFAlignment_tiptilt_lab(SHINSApp):

    @exception.handler_for_run
    def run(self, ins_s=None, det_s=None, tel_s=None, rtc_s=None, **kwargs):
        '''template to align PSF on masks. Steps:

        . Instrument Setup with coro and SASHA Setup
        . Expose SASHA
        . Retrive FITS file
        
        . Determine coro mask center
        
        . Instrument and SASHA Setup without coro
        . Retrive RTC centroid origin
        . Expose SASHA
        
        . Iterations loop to align PSF:

        . SASHA image if niter > 0, else recycle
	. Retrive position of PSF meaning images making mean of images and compute displacement
	. if converged:
		.exit
	. else 
		. Transform deltas in relative tiptilt
		. Load tiptilt modes to the DM (iterate until convergence)
                . repeat

        
        . Putting back the mask
        . End of Template
        
        '''
        
        logger.info(__name__)
        self.status = uvars.EXIT_SUCCESS
        obsController = self.obsController

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

        self.gen_base = SHARKNIR_gen_base();        gen_base = self.gen_base
        gen_tec_CloseLoop = SHARKNIR_gen_tec_CloseLoop()
        
        saly = aly(camera='HAWAII H-2RG')
        saly.detector.setCrop(3)

        ######################
        # Reset RTC tiptilt  #
        ######################
        #        0, 0, tilt, tip
        modes = [0, 0, 0, 0] # Tiptilt to zero
        self.status = obsController.RTCTTModeUpload(modes, absolute=True)
        logger.info("Waiting setup to apply displacement... DONE")
    
        ######################################
        # PREPARE SETUP TO CHECK CORO CENTER #
        ######################################

        sci_filt1 = ins_s["SCI_FILT_W1"] # save for later
        sci_filt2 = ins_s["SCI_FILT_W2"]
        db_filt_w = ins_s["DB_FILT_W"]
        apodizer_w = ins_s["APODIZER_W"]
        coro_slit_w = ins_s["CORO_SLIT_W"]
        lyot_grism_w = ins_s["LYOT_GRISM_W"]            
            
        align_ins_s = {
            "SHUTTER": "IN",
            "DB_FILT_W": "WINDOW",
            "APODIZER_W": "HOLE",
            "LYOT_GRISM_W": "HOLE",
            "CAL_FF_LAMP": "ON",
            "CAL_MIRROR_DEP": "IN",
            "SCI_FILT_W1": "BB_H",
            "SCI_FILT_W2" : "HOLE",
            "CAL_FIBER_DEP": "NO_FIBERS" # NO_FIBERS in the REF at the telescope
        }

        ins_s.update(align_ins_s)

        our_det_s = {
            "NCOADDS" : "1",
            "NGROUPS" : "2",
            "NREADS"  : "1",
            "READOUT" : det_s["READOUT"], ### in ref
            "NDIT"    : "1",
            "NDROPS"  : "2",
            "SAVE"    : det_s["SAVE"], ### in ref
            }

        #################
        # Setup + image #
        #################

        logger.info("Waiting setup to retrive coro mask center ...")
        self.status, self.filenames, _ = gen_base.main(ins_s=ins_s, det_s=our_det_s, **kwargs)
        logger.info("Waiting setup to retrive coro mask center ... DONE")

        time.sleep(1)
        
        #########################
        # DETERMINE CORO CENTER #
        #########################
        if self.paragraph("DETERMINE CORO CENTER"):
        
            which_coro = {
                "FQPM": "FQPM",
                "GAUSS": "GAUSS",
                "SP1_FPM": "SP1",
                "SP2a_FPM": "SP2a",
                "SP2b_FPM": "SP2b",
                "SLIT1": "SLIT_1",
                "SLIT2": "SLIT_2",
                "WINDOW": "WINDOW"
            }
            
            our_coro = which_coro[coro_slit_w]
            
            saly.AddACoro(our_coro)
            
            logger.info("Crop: {}, Stack type: {}".format(saly.detector.crop, saly.detector.stackType))
            
            if our_coro == 'FQPM':
                saly.FQPM.FPM.LoadImage(fitsDir+self.filenames[-1])
                # saly.FQPM.FPM.removeHotPixels(verbose=True)
                saly.FQPM.FPM.Center(method='auto', display=False, tcrop=5, bthreshold=100)
                coroCX = (saly.FQPM.FPM.xc).n
                coroCY = (saly.FQPM.FPM.yc).n
            elif our_coro == 'GAUSS':
                saly.detector.setStackType('CDS')
                saly.GAUSS.FPM.LoadImage(fitsDir+self.filenames[-1])
                # saly.GAUSS.FPM.removeHotPixels(verbose=True)
                saly.GAUSS.FPM.Center(method='auto', box=50, display=False)
                coroCX = (saly.GAUSS.FPM.xc).n
                coroCY = (saly.GAUSS.FPM.yc).n
            elif our_coro == 'SP1':
                saly.SP1.FPM.LoadImage(fitsDir+self.filenames[-1])
                # saly.SP1.FPM.removeHotPixels(verbose=True)
                saly.SP1.FPM.Center(method='auto', display=False, level=15000)
                coroCX = (saly.SP1.FPM.xc).n
                coroCY = (saly.SP1.FPM.yc).n
            elif our_coro == 'SP2a':
                saly.SP2a.FPM.LoadImage(fitsDir+self.filenames[-1])
                # saly.SP2a.FPM.removeHotPixels(verbose=True)
                saly.SP2a.FPM.Center(method='Hough', rmin=20, rmax=25, hthresh=(75,150), display=False)
                coroCX = (saly.SP2a.FPM.xc).n
                coroCY = (saly.SP2a.FPM.yc).n
            elif our_coro == 'SP2b':
                saly.SP2b.FPM.LoadImage(fitsDir+self.filenames[-1])
                # saly.SP2b.FPM.removeHotPixels(verbose=True)
                saly.SP2b.FPM.Center(method='Hough', rmin=20, rmax=25, hthresh=(75,150), display=False)
                coroCX = (saly.SP2b.FPM.xc).n
                coroCY = (saly.SP2b.FPM.yc).n
            elif our_coro == 'SLIT_1':
                saly.SLIT_1.FPM.LoadImage(fitsDir+self.filenames[-1])
                # saly.SLIT_1.FPM.removeHotPixels(verbose=True)
                saly.SLIT_1.FPM.Center(method='contours', bthreshold=80, borderLen=5, display=False)
                coroCX = saly.SLIT_1.FPM.xc
                coroCY = saly.SLIT_1.FPM.yc
            elif our_coro == 'SLIT_2':
                saly.SLIT_2.FPM.LoadImage(fitsDir+self.filenames[-1])
                # saly.SLIT_2.FPM.removeHotPixels(verbose=True)
                saly.SLIT_2.FPM.Center(method='contours', bthreshold=80, borderLen=5, display=False)
                coroCX = saly.SLIT_2.FPM.xc
                coroCY = saly.SLIT_2.FPM.yc
            else:
                # CORO IS WINDOW
                coroCX = 77 # center of the CCD in this case
                coroCY = 84 # center of the CCD in this case
                
            
            logger.warning("CHECK: Coro center values are {}, {}".format(coroCX, coroCY))

        ##############################
        # PREPARE SETUP WITHOUT CORO #
        ##############################

        logger.info("Setup to retrive PSF center: move coro masks to window")

        align_ins_s = {
            #"SHUTTER": "OUT", # commented because it's lab version
            "DB_FILT_W": "WINDOW",
            "CORO_SLIT_W": "WINDOW",
            "APODIZER_W": apodizer_w, #corresponding_apo,
            "CAL_FF_LAMP": "OFF",
            "CAL_MIRROR_DEP": "OUT",
            "LYOT_GRISM_W": lyot_grism_w,
            "SCI_FILT_W1": sci_filt1,
            "SCI_FILT_W2": sci_filt2,
            "CAL_FIBER_DEP": "FOCUS_IN" # NO_FIBERS in REF
        }

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

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

        # ##############
        # # Close loop #
        # ##############

        # self.status, rtc_filenames, _ = gen_tec_CloseLoop.main(rtc_s=rtc_s)
        
        ##################################
        # 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")

        ##################
        # EXPOSING SASHA #
        ##################
        if self.paragraph("DETERMINE PSF CENTER"):

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

        ###########################
        # ITERATIONS TO ALIGN PSF #
        ###########################

        # this loop goes up to convergence or to a max num of
        # iteration, after which it fails

        converged = False
        thrsConv = 0.2 # pixels.
        nIteration = 0
        maxIteration = 5

        psfCXarr = []
        psfCYarr = []

        tilt = 0
        tip = 0
                
        saly.detector.setStackType('CDS')

        while nIteration < maxIteration and not converged:
            
            if self.paragraph("ITERATION {} of {}".format(nIteration, maxIteration)):

                if nIteration > 0:
                
                    ##################
                    # EXPOSING SASHA #
                    ##################
                
                    logger.info("Waiting camera exposure to complete ...")
                    self.status, self.filenames = obsController.sashaExpose()
                    logger.info("FITS files are "+str(self.filenames) )
                
                else:
                    logger.debug("Use last image of LBT taken for rotation alignment")
                    pass
                
                ############################
                # RETRIEVE POSITION OF PSF #
                ############################
                logger.info("Retrieve position of target with algorithm")
                logger.info("Crop: {}, Stack type: {}".format(saly.detector.crop, saly.detector.stackType))
                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)
        
                saly.PSFNoCoro.LoadImage(meaned)
                saly.PSFNoCoro.removeHotPixels(verbose=True)
                saly.PSFNoCoro.Center(model='Gaussian2D', box=20, display=False)
                psfCX = saly.PSFNoCoro.xc
                psfCY = saly.PSFNoCoro.yc
                
                psfCXarr.append(psfCX)
                psfCYarr.append(psfCY)
                
                logger.info("CHECK: PSF centers on SCICAM are: {} {}".format(psfCX, psfCY))
                
                ########################
                # COMPUTE DISPLACEMENT #
                ########################
                
                logger.info("Calculate displacement between psf and coro centers")
                
                deltaX = psfCX - coroCX
                deltaY = psfCY - coroCY
                
                logger.info("Deltas between PSF position and CORO center are: deltaX: {} deltaY: {}".format(deltaX, deltaY))
                
                logger.info("Compute displacement " \
                            " and convert to tilt/tip signal")
                
                # 100nm tilt = 20um on focal plane.
                tilt_i,tip_i = ufun.transform_offset(disp=[deltaX, deltaY], conv="PSF")
                tilt += tilt_i
                tip += tip_i
                modes = [0, 0, tilt, tip]
                
                logger.warning("Tilt: {}, tip: {}".format(tilt, tip))
                
                ######################
                # CHECK IF CONVERGED #
                ######################
                
                logger.debug("CHECK: Displacement: deltaX: {} deltaY: {}".format(deltaX, deltaY))
                
                converged = True if ( (np.abs(deltaX) <= thrsConv) and (np.abs(deltaY) <= thrsConv) ) else False
                
                logger.info("CHECK if delstas are < {}: {}".format(thrsConv, converged))
                
                nIteration += 1
                
                if not converged and nIteration != maxIteration:
                
                    #############################
                    # APPLY THE DM DISPLACEMENT #
                    #############################
                
                    logger.info("Upload tilt/tip signal to internal DM")
                    self.status = obsController.RTCTTModeUpload(modes, absolute=True)
                    logger.info("Waiting setup to apply displacement... DONE")


        # ##############
        # # Close loop #
        # ##############

        # self.status, rtc_filenames, [cx, cy] = gen_tec_CloseLoop.main(rtc_s=rtc_s)

        # #logger.warning("END OF PSFALIGNMENT {} {}".format(cx, cy))

        # ############################################
        # # Set RTC new centroid to Coro Mask Center #
        # ############################################
                                                          # NON VOGLIAMO CHE IL CENTROIDE CAMBI SUL CENTRO DELLA MASCHERA, AI101 ANNULLATA
        # logger.info("Set RTC new centroid to Coro Mask Center...")
        # self.status = obsController.RTCTTSetCentroidPos(coroCX, coroCY)
        # ufun.error_handler(self.status, "RTCTTSetCentroidPos")
        # logger.info("Set RTC new centroid to Coro Mask Center... DONE")
        
        #########################
        # Putting back the mask #
        #########################

        ins_s = {
            "DB_FILT_W": db_filt_w,
            "CORO_SLIT_W": coro_slit_w,
        }

        logger.info("Waiting setup to retrive coro mask center ...")
        gen_base.main(ins_s=ins_s, expose=False, **kwargs)
        logger.info("Waiting setup to retrive coro mask center ... DONE")


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

        # ############################
        # # RETRIEVE POSITION OF PSF #
        # ############################
        # logger.info("Retrieve position of target with algorithm")
        # logger.info("Crop: {}, Stack type: {}".format(saly.detector.crop, saly.detector.stackType))

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

        # saly.PSFNoCoro.LoadImage(meaned)
        # saly.PSFNoCoro.removeHotPixels(verbose=True)
        # saly.PSFNoCoro.Center(model='Gaussian2D', box=20, display=False)
        # psfCX = saly.PSFNoCoro.xc
        # psfCY = saly.PSFNoCoro.yc
        # logger.info("CHECK: PSF centers on SCICAM are: {} {}".format(psfCX, psfCY))
        
        
        saly.PSFNoCoro.calculate_bkg()
        
        self.tpl_output = {
            uvars.tpl_output_names["sasha_flux"] : np.round(saly.PSFNoCoro.flux, 2),
            # uvars.tpl_output_names["sasha_bkg"]  : saly.PSFNoCoro.background,
            uvars.tpl_output_names["sasha_max"]  : np.max(saly.PSFNoCoro.image),
            uvars.tpl_output_names["sasha_min"]  : np.min(saly.PSFNoCoro.image),
            uvars.tpl_output_names["converged"]  : converged,
            uvars.tpl_output_names["threshold"]  : thrsConv,
            uvars.tpl_output_names["psf_dx"]     : np.round(deltaX,3),
            uvars.tpl_output_names["psf_dy"]     : np.round(deltaY,3)
        }
