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

from SHARKNIR_gen_base import SHARKNIR_gen_base

class SHARKNIR_cal_tec_NCPA_PD(SHINSApp):

    @exception.handler_for_run
    def run(self, ins_s=None, det_s=None, tel_s=None, rtc_s=None, **kwargs):
        '''Calculate NCPA with PD using star.
        '''

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

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

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

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

        gen_base.main(ins_s=ins_s, det_s=det_s, expose=False)
        
        dit = float(det_s["DIT"])

        ##########################
        # OPEN REAL-TIME TT LOOP #
        ##########################
        if self.paragraph("OPEN REAL-TIME TT LOOP"):
            logger.info("TT open loop")
            self.status = obsController.RTCTTOpenLoop()
            ufun.error_handler(self.status, "RTCTTOpenLoop")
            logger.debug(self.status)
            logger.info("Tip tilt loop open on TECCAM")
        
        self.master_array = [] # the list of images CDS-subtracted and meaned over NREADS=10

        # Apo positions are shared between focus and defocus loops
        several_apos = ["HOLE3", "Lens_PD_neg", "Lens_PD_pos"]

        #############################
        # First LOOP (FOCUSED STAR) #
        #############################

        # several_drops = ["0", "85", "85"] # ARGOS
        # several_drops = ["0", "25", "25"] # 24 Tau
        # several_drops = ["13", "189", "189"] # Eta Tau
        # several_factors =  [1, 13.5, 13.5] # Luca
        # several_factors =  [1, 25, 25] # HD49197
        several_factors =  [1, 50, 50] 
        
        logger.info("***Starting Focused star")

        #for apo, drop in zip(several_apos, several_drops):
        for apo, factor in zip(several_apos, several_factors):

            ins_s = {
                "APODIZER_W": apo,
            }
            
            #det_s["NDROPS"] = drop
            det_s_dropped = ufun.calc_ramp_settings(dit*factor, "128x128", ngmax=2, keep_ngmax=True)
            
            logger.info("****Start focus setup and image")
            self.status, self.filename, self.other = gen_base.main(ins_s, det_s_dropped, ngmax=uvars.Tech_Image_NGMAX, **kwargs)
            logger.info("****End focus setup and image")

            if self.paragraph("Focused {}".format(apo)):
                
                self.master_array = ufun.ncpa_masterize(fitsDir, self.filename, self.master_array)

#        ##################
#        # PAUSE / RESUME #
#        ##################

#        self.paused = True
#        logger.warning("Resume when ready")
#        self.check_pause_or_abort()

        ############################
        # APPLY 3mm OFFSET ON SOUL #
        ############################

        deltaZ = -3
        logger.info("Moving SOUL with an offset of {} mm ...".format(deltaZ))
        self.obsController.OffsetZAO(+deltaZ, "LEFT")
        logger.info("Moving SOUL with an offset of {} mm ... DONE".format(deltaZ))


        
        ################################
        # Second LOOP (DEFOCUSED STAR) #
        ################################

        # several_drops = ["21", "200", "35"] # ARGOS
        # several_drops = ["4", "65", "9"] # 24 Tau
        # several_drops = ["30", "492", "68"] # Eta Tau
        # several_factors =  [2.3, 35.5, 5] # Luca
        #several_factors =  [5, 35, 5] # HD49197
        several_factors =  [10, 70, 10] 

        logger.warning("Starting Defocused star")
        
        #for apo, drop in zip(several_apos, several_drops):
        for apo, factor in zip(several_apos, several_factors):

            ins_s = {
                "APODIZER_W": apo,
            }
            
            #det_s["NDROPS"] = drop
            det_s_dropped = ufun.calc_ramp_settings(dit*factor, "128x128", ngmax=2, keep_ngmax=True)
            
            logger.info("****Start defocus setup and image")
            self.status, self.filename, self.other = gen_base.main(ins_s, det_s_dropped, ngmax=uvars.Tech_Image_NGMAX)
            logger.info("****End defocus setup and image")
            
            if self.paragraph("Defocused {}".format(apo)):
                
                self.master_array = ufun.ncpa_masterize(fitsDir, self.filename, self.master_array)

        ########################################
        # ONERA CODE in IDL, using IDL bridges #
        ########################################            
        if self.paragraph("ONERA processing"):
            wvl = ufun.determine_filter_wavelength(sci1 = ins_s["SCI_FILT_W1"] , sci2 = ins_s["SCI_FILT_W2"])
            zernikes = ufun.call_onera(self.master_array, wavelength=wvl)

            # Print Strehl
                    
            saly = aly(camera='HAWAII H-2RG')
            saly.PSFNoCoro.LoadImage(self.master_array[0])
            strehl = saly.PSFNoCoro.getStrehl(0.0145, method='esposito',
                                          wavelengthInUm=wvl, Dpup=8.22*0.98, eps=0.1,
                                          verbose=False, returnSR=True)
            
            logger.info("Strehl ratio: {:.2}".format(strehl))

        ###################
        # Calculate modes #
        ###################
        if self.paragraph("Calculate modes modes"):

            logger.info("Found these zernikes:")
            logger.info(str(zernikes))

            modes = ufun.zernikes_to_modes(zernikes)

            logger.info("Function zernike_to_modes returns:")
            logger.info(str(modes))

 
        ############################
        # MOVING BACK SOUL TO -3mm #
        ############################

        
        logger.info("Moving SOUL back where at the begining of template ...")
        self.obsController.OffsetZAO(-deltaZ, "LEFT")
        logger.info("Moving SOUL back where at the begining of template ... DONE")

        
        ##################
        # PAUSE / RESUME #
        ##################

        self.paused = True
        logger.warning("Resume if you agree to apply; Quit otherwise.")
        print("(R)esume, (Q)uit")
        
        import sys
        from select import select
        
        timeout = 2
               
        answer = None
        print("Enter something:"),
        self.paused = True
        
        
        while self.paused==True and self.aborted==False:

            rlist, _, _ = select([sys.stdin], [], [], timeout)
            if rlist:
                answer = sys.stdin.readline()
                print(answer)
            else:
                print("still waiting...")

            if answer:
                if answer.lower().startswith('r'):
                    self.paused = False
                    pass
                if answer.lower().startswith('q'):
                    self.paused = False
                    self.aborted = True
                    self.abort()
                    return
                else:
                    pass

            
        ###############
        # Apply Modes #
        ###############
        
        reserved = [0,0]
        zeros = [0,0,0,0]

        logger.info("Upload a mode vector to internal DM")
        self.status = obsController.RTCTTModeUpload(reserved+zeros+modes, absolute=False)
        logger.info("Waiting setup to apply displacement... DONE")

         
