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_PupilLyotAlignment(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 lyot. 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')

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

        ###############################
        # SETUP FOR LYOT CENTER CHECK #
        ###############################

        sci_filt_w1 = ins_s.get("SCI_FILT_W1")
        sci_filt_w2 = ins_s.get("SCI_FILT_W2")
        
        align_ins_s = {
            "CAL_MIRROR_DEP" : "IN",
            "CAL_FF_LAMP": "ON",
            #"LAMP_TEL_SIM" : "OFF",
            "SCI_FILT_W1" : "ND2",
            "SCI_FILT_W2" : "HOLE",
            "SHUTTER" : "IN",
        }

        ins_s.update(align_ins_s)

        orig_det_s = {
            "NGROUPS":"1",
            "NREADS":"1",
            "NDROPS":"0",
            "NDIT": "1",
            "NCOADDS":"1",
            "READOUT":det_s["READOUT"]
        }

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

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

        #########################
        # DETERMINE LYOT CENTER #
        #########################
        if self.paragraph("DETERMINE LYOT CENTER"):

            logger.info("Determine LYOT center")
            
            which_lyot = {
                "FQPM_LYOT_STOP": 'FQPM',
                "GAUSS_LYOT_STOP": 'GAUSS',
                "SP_LYOT_STOP": "SP1",
                }
            used_lyot = ins_s["LYOT_GRISM_W"]
            our_lyot = which_lyot[used_lyot]
            
            saly.AddACoro(our_lyot)
            
            if our_lyot == 'FQPM':
                saly.FQPM.LS.LoadImage(fitsDir+self.filenames[-1])
                # saly.FQPM.LS.removeHotPixels(verbose=True)
                saly.FQPM.LS.Center(method='isophote', level=15000, display=False)
                lyot_counts = saly.FQPM.LS.median_counts()
                lyotCX = (saly.FQPM.LS.xc).n
                lyotCY = (saly.FQPM.LS.yc).n
            elif our_lyot == 'GAUSS':
                saly.GAUSS.LS.LoadImage(fitsDir+self.filenames[-1])
                # saly.GAUSS.LS.removeHotPixels(verbose=True)
                saly.GAUSS.LS.Center(method='isophote', level=15000, display=False)
                lyot_counts = saly.GAUSS.LS.median_counts()
                lyotCX = (saly.GAUSS.LS.xc).n # It's an ufloat(2, 0.25)
                lyotCY = (saly.GAUSS.LS.yc).n # .n is the number, .s is the error
            elif our_lyot == 'SP1':
                saly.SP1.LS.LoadImage(fitsDir+self.filenames[-1])
                # saly.SP1.LS.removeHotPixels(verbose=True)
                saly.SP1.LS.Center(method='isophote', level=15000, display=False)
                lyot_counts = saly.SP1.LS.median_counts()
                lyotCX = (saly.SP1.LS.xc).n
                lyotCY = (saly.SP1.LS.yc).n
            
            logger.warning("CHECK: Lyot center values are {}, {}".format(lyotCX, lyotCY))

        thrs_counts = 20000
        if lyot_counts < thrs_counts:
            logger.warning("Lyot counts ({}) below threashold value ({}) ... Aborting template".format(lyot_counts, thrs_counts))
            self.abort()
        
        #####################################
        # PREPARE SETUP FOR LBT orientation #
        #####################################

        logger.info("Setup to retrive LBT orientation: move coro masks to hole")

        align_ins_s = {
            "LYOT_GRISM_W": "HOLE",
            "CAL_MIRROR_DEP" : "OUT",
            "CAL_FF_LAMP": "OFF",
            "SCI_FILT_W1" : sci_filt_w1,
            "SCI_FILT_W2" : sci_filt_w2,
            "SHUTTER" : "OUT",
        }

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

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

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

        if counts < thrs_counts:
            logger.warning("Pupil counts ({}) below threashold value ({}) ... Aborting template".format(counts, thrs_counts))
            self.abort()
        
        ##############################
        # ITERATIONS TO ALIGN PUPILS #
        ##############################

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

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

        LBTCXarr = []
        LBTCYarr = []

        internal_circle = True
        rmin, rmax = 35, 45 if internal_circle else 335, 345

        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")
                    pass
                
                ##############################
                # DETERMINE LBT PUPIL CENTER #
                ##############################
                
                logger.info("Determine LBT center")
                
                saly.Pup.LoadImage(fitsDir+self.filenames[-1])
                # saly.Pup.removeHotPixels(verbose=True)
                
                # saly.Pup.LoadTemplate("SCICAM")
                # saly.Pup.ZumTemplate(shift=(-3500,-1700),scale=4.0)
	        # saly.PupilCenter(method='FLANN', display=False, distance_ratio=0.5)
                saly.PupilCenter(method='Hough', rmin=rmin, rmax=rmax, hthresh=(100,150), display=False)
                # LBTCX, LBTCY = saly.Pup.xc, saly.Pup.yc
                LBTCX, LBTCY = (saly.Pup.xc).n, (saly.Pup.yc).n
                
                logger.warning("CHECK: LBT center are {}, {}".format(LBTCX, LBTCY))
                
                ####################
                # GET CURRENT DROT #
                ####################
                
                inDict = {"DROT":-9999.9}
                unit = "DEGREE"
                
                self.status, outDict = obsController.getPosition(unit, inDict)
                logger.warning("DEROT is: {} degrees".format(outDict["DROT"]))
                
                ########################
                # COMPUTE DISPLACEMENT #
                ########################
                
                logger.info("Calculate displacement between orientation and lyot centers")
                
                deltaX =  LBTCX - lyotCX
                deltaY =  LBTCY - lyotCY
                
                logger.info("Deltas are {} {}".format(deltaX, deltaY))
                
                thetaInTT = ufun.transform_offset(disp=[deltaX, deltaY],
                                                  thetaRot=outDict["DROT"])
                
                logger.info("Apply this displacementon INBEAM_TT: {}".format(thetaInTT))
                
                #########################
                # GET CURRENT INBEAM_TT #
                #########################
                
                logger.info("Take current INBEAMTT positions")
                
                outDict = {}
                
                self.status, ttx = obsController.getPosition("STEPS", {"INBEAM_TT_X": -9999})
                self.status, tty = obsController.getPosition("STEPS", {"INBEAM_TT_Y": -9999})
                
                outDict.update(ttx)
                outDict.update(tty)
                
                logger.warning("INBEAM_TT are: {} steps".format(outDict))
                
                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_n = {k: outDict.get(k, 0) + tt_ins_rel.get(k, 0) for k in set(outDict) & set(tt_ins_rel)}
                    keys_values = tt_ins_n.items()
                    tt_ins_s = {k:str(v) for (k,v) in keys_values}
                
                    logger.info("Waiting setup to apply displacement...")
                    gen_base.main(ins_s=tt_ins_s, **kwargs)
                    logger.info("Waiting setup to apply displacement... DONE")

        ###################
        # END OF THE LOOP #
        ###################
        if self.paragraph("END OF THE LOOP"):

            LBTCXarr.append(lyotCX)
            LBTCYarr.append(lyotCY)
            
            print(LBTCXarr)
            print(LBTCYarr)
            
            logger.warning("CHECK: LBT pupil center: [{:.2f}, {:.2f}]".format(saly.Pup.xc, saly.Pup.yc))
            logger.warning("CHECK: Lyot center: [{:.2f}, {:.2f}]".format(lyotCX, lyotCY))
            logger.warning("CHECK: Displacement: deltaX: {} deltaY: {}".format(deltaX, deltaY))

        #######################################################
        # REMOVING PUPIL LENS DEPLOYER AND PUT BACK LYOT STOP #
        #######################################################
        if self.paragraph("REMOVING PUPIL LENS DEPLOYER and PUT BACK LYOT STOP AT END OF PROCEDURE"):
            obsController.setupInstrument({"PUPIL_LENS_DEP":"OUT" , "LYOT_GRISM_W":used_lyot})
            logger.info("Pupil Lens is now OUT")
            logger.info("Lyot Stop 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")
            
            # if not converged:
            #     logger.error("Pupil alignment did not converged")
            #     self.status = uvars.EXIT_FAILURE
            #     raise RuntimeError("Pupil alignment did not converged")

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