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

from SHARKNIR_gen_base import SHARKNIR_gen_base

class SHARKNIR_gen_tec_PupilApoAlignment(SHINSApp):
    
    @exception.handler_for_run
    def run(self, ins_s=None, det_s=None, tel_s=None, rtc_s=None, **kwargs):
        '''template to align Pupil on apodizer. Steps:
        . instrument setup START and WAIT
        . setup SASHA
        . expose SASHA
        . retrieve FITS file
        . retrieve apo mask position
        . instrument setup START and WAIT
        . retrieve pupil position and move tip/tilt mirror (iterate until convergence)
        . save residual displacement X,Y
        . instrument setup START and WAIT (for rot alignment)
        . setup SASHA
        . expose SASHA
        . retrieve FITS file
        '''
        logger.info(__name__)

        self.status = uvars.EXIT_SUCCESS
        time_out = 120000        
        obsController = self.obsController
        
        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')

        # logger.info("Control if SHARK-NIR is authorized at the LEFT side of LBT ...")
        # try:
        #     self.status = self.obsController.OffsetXYAO(0,0,"LEFT")
        #     logger.info("SHARK-NIR is authorized")
        # except Exception as e:
        #     if "Not authorized" in str(e):
        #         logger.error("SHARK-NIR is NOT authorized ... Aborting template")
        #         self.abort()
        #     else:
        #         logger.error(e)
        #         #self.abort()
                
        ########################################
        # PREPARE SETUP FOR ROTATION ALIGNMENT #
        ########################################

        sci1 = ins_s["SCI_FILT_W1"]
        sci2 = ins_s["SCI_FILT_W2"]
        
        align_ins_s = {
            "CAL_FIBER_DEP" : "FOCUS_IN",
            "CAL_FIBER_FOCUS_LAMP" : "ON",
            "SCI_FILT_W1" : "HOLE",
            "SCI_FILT_W2" : "NB_ContH",
            "SHUTTER" : "IN",
        }

        ins_s.update(align_ins_s)

        apo_det_s = {"DIT" : "5", "READOUT": det_s["READOUT"] }

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

        logger.info("Waiting setup to retrive coro mask image ...")
        self.status, self.filenames, _ = gen_base.main(ins_s=ins_s, det_s=apo_det_s, ngmax = uvars.Tech_Image_NGMAX)
        logger.info("Waiting setup to retrive coro mask image ... DONE")
        
        ###############################################
        # DETERMINE APO SPIDER CENTER AND ORIENTATION #
        ###############################################
        if self.paragraph("DETERMINE APO SPIDER CENTER AND ORIENTATION"):

            logger.info("Determine orientation of APO spider")
            
            which_apo = {
                'SP1_APO' : 'SP1',
                "SP2a_APO": 'SP2a',
                "SP2b_APO": 'SP2b',
                }

            used_apo = ins_s["APODIZER_W"]
            our_apo = which_apo[used_apo]  #ins_s["APODIZER_W"]]            

            saly.AddACoro(our_apo)
            
            if our_apo == 'SP1':
                saly.SP1.Apo.LoadImage(fitsDir+self.filenames[-1])
                # saly.SP1.Apo.removeHotPixels(verbose=True)
                saly.SP1.Apo.LoadTemplate("SCICAM_IntFib_NoBaffle")
                saly.SP1.Apo.Center(method='FLANN', display=False)
                apoOrientation = saly.SP1.Apo.orientation + saly.SP1.Apo.tpl_orientation
                apoCX = saly.SP1.Apo.xc
                apoCY = saly.SP1.Apo.yc
                hole_chosen = "HOLE"
            elif  our_apo == 'SP2a':
                saly.SP2a.Apo.LoadImage(fitsDir+self.filenames[-1])
                # saly.SP2a.Apo.removeHotPixels(verbose=True)
                saly.SP2a.Apo.LoadTemplate("SCICAM_IntFib_NoBaffle")
                saly.SP2a.Apo.Center(method='FLANN', display=False)
                apoOrientation = saly.SP2a.Apo.orientation + saly.SP2a.Apo.tpl_orientation
                apoCX = saly.SP2a.Apo.xc
                apoCY = saly.SP2a.Apo.yc
                hole_chosen = "HOLE"
            elif  our_apo == 'SP2b':
                saly.SP2b.Apo.LoadImage(fitsDir+self.filenames[-1])
                # saly.SP2b.Apo.removeHotPixels(verbose=True)
                saly.SP2b.Apo.LoadTemplate("SCICAM_IntFib_NoBaffle")
                saly.SP2b.Apo.Center(method='FLANN', display=False)
                apoOrientation = saly.SP2b.Apo.orientation + saly.SP2b.Apo.tpl_orientation
                apoCX = saly.SP2b.Apo.xc
                apoCY = saly.SP2b.Apo.yc
                hole_chosen = "HOLE"
            
            logger.warning("CHECK: Apo orientation and center values are {}, {}, {}".format(apoOrientation, apoCX, apoCY))

        #####################################
        # PREPARE SETUP FOR LBT orientation #
        #####################################
        
        logger.info("Setup to retrive apodizer mask center: move coro masks to hole")
        
        align_ins_s = {
            "APODIZER_W" : hole_chosen,
            "CAL_FIBER_DEP" : "NO_FIBERS",
            "CAL_FIBER_FOCUS_LAMP" : "OFF",
            #"LAMP_TEL_SIM" : "ON",
            "SCI_FILT_W1" : sci1,
            "SCI_FILT_W2" : sci2,
            "SHUTTER" : "OUT",
        }
    
        #det_s["NREADS"] = "1"

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

        self.status, self.filenames, _ = gen_base.main(ins_s=align_ins_s, det_s=det_s, ngmax = uvars.Tech_Image_NGMAX)

        #########################
        # COUNTS IN PUPIL IMAGE #
        #########################

        saly.detector.setStackType('last')
        saly.Pup.LoadImage(fitsDir+self.filenames[-1])
        counts = saly.Pup.median_counts()
        logger.info("Median counts in the pupil: {}".format(counts))
        saly.detector.setStackType('CDS')

        thrs_counts = 20000
        if counts < thrs_counts:
            logger.warning("Pupil counts ({}) below threashold value ({}) ... Aborting template".format(counts, thrs_counts))
            self.abort()
            
        ####################################
        # DETERMINE LBT SPIDER ORIENTATION #
        ####################################
        if self.paragraph("DETERMINE LBT SPIDER ORIENTATION"):

            logger.info("Determine orientation of LBT spider")
            
            saly.Pup.LoadImage(fitsDir+self.filenames[-1])
            # saly.Pup.removeHotPixels(verbose=True)
            saly.Pup.LoadTemplate("SCICAM_OnSky")
            # saly.Pup.ZumTemplate(shift=(-3500,-1700),scale=4.0, display=False)
            saly.PupilCenter(method='FLANN', display=False, distance_ratio=0.7)
            if not saly.Pup.orientation:
                logger.error("FLANN FAILED: Returned None value")
                self.status = uvars.EXIT_FAILURE
                return
            LBTOrientation = saly.Pup.orientation + saly.Pup.tpl_orientation
            
            logger.warning("CHECK: LBT spider orientation is {}!".format(LBTOrientation))
            
            ####################
            # GET CURRENT DROT #
            ####################
            
            logger.debug("Retrieve position from derotator")
            
            inDict = {"DROT":-9999.9}
            unit = "DEGREE"
            self.status, outDict = obsController.getPosition(unit, inDict)
            logger.warning(self.status)
            
            deltaRot = apoOrientation - LBTOrientation
            newPosAng = deltaRot + outDict["DROT"]
            
            logger.warning("outDict:{}".format(outDict))
            logger.warning("CHECK: deltaRot: {}".format(deltaRot))
            logger.warning("CHECK: newPosAng: {}".format(newPosAng))
            
            ##########################
            # SETUP DEROT TO ABS POS #
            ##########################
            
            logger.debug("Setup the derotator on new position angle")
            
            align_ins_s.clear()
            align_ins_s = {"DROT_PosAng": str(newPosAng)}
            
            logger.info("Waiting setup to set new pos ang ...")
            gen_base.main(ins_s=align_ins_s, **kwargs)
            logger.warning(self.status)
            logger.info("Waiting setup to set new pos ang ... DONE")

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

        converged = False
        thrsConv = 1.0 # pixels. The requirement was this:  660px * 0.5% = 3.3px. 
        thrsRot = 0.5 # Rotation threshold in degrees
        nIteration = 0
        maxIteration = 5

        LBTCXarr = []
        LBTCYarr = []
        quasi_final_arr = []            
        
        while nIteration < maxIteration and not converged:
            
            if self.paragraph("ITERATION {} of {}".format(nIteration, maxIteration)):
  
                if True: #nIteration > 0:
                    #if condition set True in order to take image and evaluate pupil, as requested by Luca and Daniele
                    
                    ##################
                    # 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 PUPIL #
                ##############################
                
                logger.info("Calculate LBT pupil center") # and orientation")
                
                saly.Pup.LoadImage(fitsDir+self.filenames[-1])
                # saly.Pup.removeHotPixels(verbose=True)
                # saly.Pup.LoadTemplate("SCICAM")
                saly.PupilCenter(method='Hough', rmin=35, rmax=45, hthresh=(100,150), display=False)
                
                LBTCX, LBTCY = saly.Pup.xc.n, saly.Pup.yc.n
                
                LBTCXarr.append(LBTCX)
                LBTCYarr.append(LBTCY)
                
                logger.warning("CHECK: LBT pupil orientation and center are {}, {}, {}".format(saly.Pup.orientation, LBTCX, LBTCY))
                #logger.warning("CHECK: LBT spider orientation is {}!".format(saly.Pup.orientation + saly.Pup.tpl_orientation))
                
                ####################
                # GET CURRENT DROT #
                ####################
                
                inDict = {"DROT":-9999.9}
                unit = "DEGREE"
                
                self.status, outDict = obsController.getPosition(unit, inDict)
                logger.warning(self.status)
                
                ########################
                # COMPUTE DISPLACEMENT #
                ########################
                
                logger.info("Calculate displacement between orientation and apo centers")
                
                deltaX = LBTCX - apoCX
                deltaY = LBTCY - apoCY
                
                thetaInTT = ufun.transform_offset(disp=[deltaX, deltaY],
                                                  thetaRot=outDict["DROT"])
                
                logger.info("Apply this displacementon INBEAM_TT:")
                print(thetaInTT)
                
                #########################
                # GET CURRENT INBEAM_TT #
                #########################
                
                logger.info("Take current INBEAMTT positions")
                
                inDict = { "INBEAM_TT_X" : -9999, "INBEAM_TT_Y": -9999 }
                unit = "STEPS"
                
                self.status, outDict = obsController.getPosition(unit, inDict)
                logger.warning(self.status)
                
                tt_ins_rel = {
                    "INBEAM_TT_X": thetaInTT[0],
                    "INBEAM_TT_Y": thetaInTT[1]
                }
                
                ######################
                # 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.debug("CHECK: Converged: {}".format(converged))
                
                nIteration += 1
                
                if not converged and nIteration != maxIteration:
                
                    ##################################
                    # SUM AND APPLY THE DISPLACEMENT #
                    ##################################
                    
                    logger.info("Sum values")
                    

                    tt_ins_s = {k: str(outDict.get(k, 0) + tt_ins_rel.get(k, 0)) for k in set(outDict)}

                    logger.debug("**********************************")
                    logger.debug(outDict)
                    logger.debug("+")
                    logger.debug(tt_ins_rel)
                    logger.debug("=")
                    logger.debug(tt_ins_s)
                    logger.debug("**********************************")
                    
                    logger.info("Waiting setup to apply displacement...")
                    gen_base.main(ins_s=tt_ins_s, **kwargs)
                    logger.warning(self.status)
                    logger.info("Waiting setup to apply displacement... DONE")

                    
        ## Following paragraph is commented due to request to remove the rotation after the pupil alignment 
        # ###############################
        # # ITERATIONS TO ROTATE PUPILS #
        # ###############################
        
        # # this loop goes up to convergence or to a max num of
        # # iteration, after which it fails

        # converged = False
        # thrsRot = 0.5 # Rotation threshold in degrees
        # nIteration = 0
        # maxIteration = 5
        
        # logger.info("Now iterate to rotate pupil")
                    
        # 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 taken for XY alignment")
        #             pass
                
        #         ################################
        #         # CHECK CENTERING AND CLOCKING #
        #         ################################
                
        #         logger.info("Calculate FINAL LBT Pupil centering and clocking")
                
        #         saly.Pup.LoadImage(fitsDir+self.filenames[-1])
        #         # saly.Pup.removeHotPixels(verbose=True)
        #         # saly.Pup.LoadTemplate("SCICAM")
        #         saly.PupilCenter(method='FLANN', display=False, distance_ratio=0.7)
        #         LBTCX, LBTCY = saly.Pup.xc, saly.Pup.yc            
        #         deltaX = LBTCX - apoCX
        #         deltaY = LBTCY - apoCY
                
        #         quasi_final_orientation = saly.Pup.orientation + saly.Pup.tpl_orientation
                
        #         quasi_final_arr.append(quasi_final_orientation)
                
        #         logger.warning("CHECK: LBT pupil orientation: {:.2f}".format(quasi_final_orientation))
        #         logger.warning("CHECK: LBT pupil center: [{:.2f}, {:.2f}]".format(saly.Pup.xc, saly.Pup.yc))
        #         logger.warning("CHECK: Apo orientation: {:.2f}".format(apoOrientation))
        #         logger.warning("CHECK: Apo center: [{:.2f}, {:.2f}]".format(apoCX, apoCY))
        #         logger.debug("CHECK: Displacement: deltaX: {} deltaY: {}".format(deltaX, deltaY))
                
                
        #         ####################
        #         # GET CURRENT DROT #
        #         ####################
                
        #         logger.debug("Retrieve position from derotator")
                
        #         inDict = {"DROT":-9999.9}
        #         unit = "DEGREE"
        #         self.status, outDict = obsController.getPosition(unit, inDict)
        #         logger.warning(self.status)
                
        #         deltaRot = apoOrientation - quasi_final_orientation
        #         newPosAng = deltaRot + outDict["DROT"]
                
        #         logger.warning("outDict:{}".format(outDict))
        #         logger.warning("CHECK: deltaRot: {}".format(deltaRot))
        #         logger.warning("CHECK: newPosAng: {}".format(newPosAng))
                
        #         ######################
        #         # CHECK IF CONVERGED #
        #         ######################
                
        #         logger.debug("Check if {} < {}".format(deltaRot, thrsRot))
        #         converged = True if (np.abs(deltaRot) <= thrsRot) else False
                
        #         logger.debug("CHECK: Converged: {}".format(converged))
                
        #         nIteration += 1
                
        #         if not converged:
                
        #             ##########################
        #             # SETUP DEROT TO ABS POS #
        #             ##########################
                    
        #             logger.debug("Setup the derotator on new position angle")
                    
        #             align_ins_s.clear()
        #             align_ins_s = {"DROT_PosAng": str(newPosAng)}
                    
        #             logger.info("Waiting setup to set new pos ang ...")
        #             gen_base.main(ins_s=align_ins_s, **kwargs)
        #             logger.warning(self.status)
        #             logger.info("Waiting setup to set new pos ang ... DONE")


        ##################################
        # CHECK CENTERING AND CLOCKING !!!
        ##################################
        if self.paragraph("CHECK CENTERING AND CLOCKING"):

            logger.info("Calculate FINAL LBT Pupil centering and clocking")
            
            saly.Pup.LoadImage(fitsDir+self.filenames[-1])
            # saly.Pup.removeHotPixels(verbose=True)
            # saly.Pup.LoadTemplate("SCICAM")
            saly.PupilCenter(method='Hough', rmin=35, rmax=45, hthresh=(100,150), display=False)
            LBTCX, LBTCY = saly.Pup.xc.n, saly.Pup.yc.n
            deltaX = LBTCX - apoCX
            deltaY = LBTCY - apoCY
            
            # final_orientation = saly.Pup.orientation + saly.Pup.tpl_orientation
            
            logger.warning("CHECK: here below, the arrays with iterations of LBTCX, LBTCY, quasi_final_orientation")
            
            LBTCXarr.append(LBTCX)
            LBTCYarr.append(LBTCY)
            # quasi_final_arr.append(final_orientation)
            
            print(LBTCXarr)
            print(LBTCYarr)
            print(quasi_final_arr)
            
            logger.warning("CHECK: LBT pupil center: [{:.2f}, {:.2f}]".format(saly.Pup.xc, saly.Pup.yc))
            logger.warning("CHECK WARNING: LBT pupil orientation NOT yet computated")#logger.warning("CHECK: LBT pupil orientation: {:.2f}".format(final_orientation))
            logger.warning("CHECK: Apo orientation: {:.2f}".format(apoOrientation))
            logger.warning("CHECK: Apo center: [{:.2f}, {:.2f}]".format(apoCX, apoCY))
            logger.warning("CHECK: Displacement: deltaX: {} deltaY: {}".format(deltaX, deltaY))
            
        ######################################################
        # REMOVING PUPIL LENS DEPLOYER AND PUT BACK APO MASK #
        ######################################################
        if self.paragraph("REMOVING PUPIL LENS DEPLOYER and PUT BACK APO AT END OF PROCEDURE"):
            obsController.setupInstrument({"PUPIL_LENS_DEP":"OUT" , "APODIZER_W":used_apo})
            logger.info("Pupil Lens is now OUT")
            logger.info("Apodizer Mask has been put back")
            
        ##################################
        # SAVE RESIDUAL DISPLACEMENT X,Y #
        ##################################
        # .... to HDF5 file
        logger.warning("Saving residual displacement X, Y on HDF5 file NOT YET IMPLEMENTED")
        
        # saly.PSFNoCoro.calculate_bkg()
        
        self.tpl_output = {
            uvars.tpl_output_names["pupil_counts"] : counts,
            uvars.tpl_output_names["pupil_max"]    : np.max(saly.Pup.image),
            uvars.tpl_output_names["converged"]    : converged,
            uvars.tpl_output_names["threshold"]    : thrsConv,
            uvars.tpl_output_names["pupil_dx"]     : np.round(deltaX,3),
            uvars.tpl_output_names["pupil_dy"]     : np.round(deltaY,3)
        }
