# $Header: https://oberon.oapd.inaf.it/shins/trunk/shins/src/scripts/templates/tpl/SHARKNIR_di_mon_MaskRepeatability.py 3997 2022-07-07 21:32:29Z depa $

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

from SHARKNIR_gen_base import SHARKNIR_gen_base

class SHARKNIR_di_mon_MaskRepeatability(SHINSApp):
    """Tests the stability of each mask

    """

    @exception.handler_for_run
    def run(self, ins_s=None, det_s=None, tel_s=None, rtc_s=None, **kwargs):

        logger.info(__name__)
        obsController = self.obsController
        self.status = uvars.EXIT_SUCCESS

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

        self.gen_base = SHARKNIR_gen_base();        gen_base = self.gen_base

        saly = aly(camera='HAWAII H-2RG')
        saly.detector.setCrop(3)

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

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

        ##########################
        # Apply RTC shape (flat) #
        ##########################
        if self.paragraph("Apply RTC shape (flat)"):
        
            #        0, 0, tilt, tip, other modes
            flatten = [ 0 ]*39
            logger.info("Flattening the DM")
            self.status = obsController.RTCTTModeUpload(flatten, absolute=True)
            logger.info("Waiting setup to apply displacement... DONE")

        ##################
        # LOOP APO MASKS #
        ##################

        several_apos = ["SP1_APO", "SP2a_APO", "SP2b_APO"]
        det_s = {"NREADS": "3",
                 "READOUT":"Center",
                 }

        logger.info("***** Looping apo masks")
        apoArr = []
        for apo in several_apos:
                    
            #################
            # Setup + image #
            #################

            ins_s = {"APODIZER_W": apo,
                     "CAL_FF_LAMP": "OFF",
                     "CAL_FIBER_DEP": "FOCUS_IN",
                     "CAL_FIBER_FOCUS_LAMP": "ON",
                     "CAL_MIRROR_DEP": "OUT",
                     "CORO_SLIT_W": "WINDOW",
                     "LYOT_GRISM_W": "HOLE",
                     "PUPIL_LENS_DEP": "IN",
                     "SCI_FILT_W1": "HOLE",
                     "SCI_FILT_W2": "NB_ContH",
                     }

            logger.info("**** Apo {} setup and image".format(apo))
            self.status, self.filenames, self.other = gen_base.main(ins_s, det_s)

            ########################
            # DETERMINE APO CENTER #
            ########################
            if self.paragraph("DETERMINE APO CENTER"):
            
                which_apo = {
                    'SP1_APO' : 'SP1',
                    "SP2a_APO": 'SP2a',
                    "SP2b_APO": 'SP2b'
                }            
                
                our_apo = which_apo[apo]
                
                saly.AddACoro(our_apo)
                
                if apo == 'SP1_APO':
                    saly.SP1.Apo.LoadImage(fitsDir+self.filenames[-1])
                    saly.SP1.Apo.LoadTemplate("SCICAM_IntFib_NoBaffle")
                    saly.SP1.Apo.Center(method='FLANN', display=False)
                    apoCX = saly.SP1.Apo.xc
                    apoCY = saly.SP1.Apo.yc
                elif  apo == 'SP2a_APO':
                    saly.SP2a.Apo.LoadImage(fitsDir+self.filenames[-1])
                    saly.SP2a.Apo.LoadTemplate("SCICAM_IntFib_NoBaffle")
                    saly.SP2a.Apo.Center(method='FLANN', display=False)
                    apoCX = saly.SP2a.Apo.xc
                    apoCY = saly.SP2a.Apo.yc
                elif  apo == 'SP2b_APO':
                    saly.SP2b.Apo.LoadImage(fitsDir+self.filenames[-1])
                    saly.SP2b.Apo.LoadTemplate("SCICAM_IntFib_NoBaffle")
                    saly.SP2b.Apo.Center(method='FLANN', display=False)
                    apoCX = saly.SP2b.Apo.xc
                    apoCY = saly.SP2b.Apo.yc
                
                apoArr.append([apo, apoCX, apoCY])

        logger.info("Apo array:")
        logger.info(apoArr)
        
        ###################
        # LOOP LYOT MASKS #
        ###################

        several_lyots = ["GAUSS_LYOT_STOP", "FQPM_LYOT_STOP", "SP_LYOT_STOP"]
        det_s = {"NREADS": "1",
                 # "READOUT":"Center", # same as apo
                 }

        logger.info("***** Looping lyot masks")
        lyotArr = []
        for lyot in several_lyots:
                    
            #################
            # Setup + image #
            #################

            ins_s = {"APODIZER_W": "HOLE",
                     "CAL_FF_LAMP": "ON",
                     "CAL_FIBER_DEP": "NO_FIBERS",
                     "CAL_FIBER_FOCUS_LAMP": "OFF",
                     "CAL_MIRROR_DEP": "IN",
                     #"CORO_SLIT_W": "WINDOW",
                     "LYOT_GRISM_W": lyot,
                     #"PUPIL_LENS_DEP": "IN", # same as apo
                     "SCI_FILT_W1": "ND2",
                     "SCI_FILT_W2": "HOLE",
                     }

            logger.info("**** Lyot {} setup and image".format(lyot))
            self.status, self.filenames, self.other = gen_base.main(ins_s, det_s)

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

                which_lyot = {
                    "FQPM_LYOT_STOP": 'FQPM',
                    "GAUSS_LYOT_STOP": 'GAUSS',
                    "SP_LYOT_STOP": "SP1",
                    }
                
                our_lyot = which_lyot[lyot]
                
                saly.AddACoro(our_lyot)
                
                if our_lyot == 'FQPM':
                    saly.FQPM.LS.LoadImage(fitsDir+self.filenames[-1])
                    saly.FQPM.LS.Center(method='isophote', level=15000, display=False)
                    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.Center(method='isophote', level=15000, display=False)
                    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.Center(method='isophote', level=15000, display=False)
                    lyotCX = (saly.SP1.LS.xc).n
                    lyotCY = (saly.SP1.LS.yc).n
                
                lyotArr.append([lyot, lyotCX, lyotCY])

        logger.info("Lyot array:")
        logger.info(lyotArr)
        
        ###################
        # LOOP CORO MASKS #
        ###################

        several_coros = ["FQPM", "GAUSS", "SP1_FPM", "SP2a_FPM",
                         "SP2b_FPM", "SLIT1", "SLIT2"]

        det_s = {"NREADS": "4",
                 "READOUT":"128x128",
                 }

        logger.info("***** Looping coro masks")
        coroArr = []
        for coro in several_coros:
                    
            #################
            # Setup + image #
            #################

            ins_s = {# "APODIZER_W": "HOLE",          # same as lyot
                     #"CAL_FF_LAMP": "ON",            # same as lyot
                     #"CAL_FIBER_DEP": "NO_FIBERS",   # same as lyot
                     #"CAL_FIBER_FOCUS_LAMP": "OFF",  # same as lyot
                     #"CAL_MIRROR_DEP": "IN",         # same as lyot
                     "CORO_SLIT_W": coro,
                     "LYOT_GRISM_W": "HOLE",
                     "PUPIL_LENS_DEP": "OUT",
                     "SCI_FILT_W1": "BB_H",
                     #"SCI_FILT_W2": "HOLE",          # same as lyot
                     }

            logger.info("**** Coro {} setup and image".format(coro))
            self.status, self.filenames, self.other = gen_base.main(ins_s, det_s)

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

                if coro == 'GAUSS':
                    saly.detector.setStackType('CDS')
                else:
                    saly.detector.setStackType("last")
                
                if coro == 'FQPM':
                    mask = saly.FQPM.FPM
                    mask.LoadImage(fitsDir+self.filenames[-1])
                    mask.Center(method='Cross-Auto', display=False, tcrop=5, bthreshold=150)
                elif coro == 'GAUSS':
                    mask = saly.GAUSS.FPM
                    mask.LoadImage(fitsDir+self.filenames[-1])
                    mask.Center(method='GaussFit', box=50, display=False)
                elif coro == 'SP1_FPM':
                    mask = saly.SP1.FPM
                    mask.LoadImage(fitsDir+self.filenames[-1])
                    mask.Center(method='isophote', display=False, level=15000)
                elif coro == "SP2a_FPM":
                    mask = saly.SP2a.FPM
                    mask.LoadImage(fitsDir+self.filenames[-1])
                    mask.Center(method='Hough', rmin=20, rmax=25, hthresh=(75,150), display=False)
                elif coro == "SP2b_FPM":
                    mask = saly.SP2b.FPM
                    mask.LoadImage(fitsDir+self.filenames[-1])
                    mask.Center(method='Hough', rmin=20, rmax=25, hthresh=(75,150), display=False)
                elif coro == "SLIT1":
                    saly.AddACoro("SLIT_1")
                    mask = saly.SLIT_1.FPM
                    mask.LoadImage(fitsDir+self.filenames[-1])
                    mask.Center(method='contours', bthreshold=80, borderLen=5, display=False)
                elif coro == "SLIT2":
                    saly.AddACoro("SLIT_2")
                    mask = saly.SLIT_2.FPM
                    mask.LoadImage(fitsDir+self.filenames[-1])
                    mask.Center(method='contours', bthreshold=80, borderLen=5, display=False)
                else:
                    logger.error("Bad Coro")
                    
                if type(mask.xc) == float or type(mask.xc) == np.float64 :
                    coroCX = mask.xc
                    coroCY = mask.yc
                elif mask.xc == None or mask.yc == None:
                    coroCX = mask.xc
                    coroCY = mask.yc                
                else:
                    coroCX = (mask.xc).n
                    coroCY = (mask.yc).n
                
                coroArr.append([coro, coroCX, coroCY])

        logger.info("Coro array:")
        logger.info(coroArr)
        
        ############################
        # TAKE PSF REFERENCE IMAGE #
        ############################

        det_s = {"NREADS": "1",
                 #"READOUT":"128x128", # same as coro 
                 }


        ins_s = {"APODIZER_W": "HOLE",
                 "CAL_FF_LAMP": "OFF",
                 "CAL_FIBER_DEP": "FOCUS_IN",
                 "CAL_FIBER_FOCUS_LAMP": "ON",
                 "CAL_MIRROR_DEP": "OUT",
                 "CORO_SLIT_W": "WINDOW",
                 # "LYOT_GRISM_W": "HOLE", # same as coro
                 # "PUPIL_LENS_DEP": "OUT", # same as coro
                 "SCI_FILT_W1": "NB_H2",
                 "SCI_FILT_W2": "ND3",
        }

        logger.info("**** PSF reference setup and image")
        self.status, self.filenames, self.other = gen_base.main(ins_s, det_s)
        
        ############################
        # RETRIEVE POSITION OF PSF #
        ############################
        if self.paragraph("RETRIEVE POSITION OF PSF"):
        
            logger.info("Retrieve position of target with algorithm")
            logger.info("Crop: {}".format(saly.detector.crop))
            
            saly.detector.setStackType('CDS')
            
            saly.PSFNoCoro.LoadImage(fitsDir+self.filenames[-1])
            print(fitsDir+self.filenames[-1])
            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))
            
            ##########################
            # Summary of the results #
            ##########################
            
            logger.info("****************Summary********************")
            logger.info("Apo array, Lyot array, Coro array, PSF center")
            logger.info(apoArr)
            logger.info(lyotArr)
            logger.info(coroArr)

