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 functions as ufun
from util import variables as uvars
import Basda
import numpy as np
import time

# Other templates
from SHARKNIR_gen_base import SHARKNIR_gen_base

class SHARKNIR_gen_tec_SoulAlignment(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:
 
        . Wait for Instrument and SASHA Setup

        . Convergence loop
  
        . Takes SASHA image

        . determine PSF centroid

        . check if convergence.

        - if converge, exit
        - else,
          - apply offset on SOUL with OffsetXYAO func 
          -repeat

        . save residual displacement

        '''
        
        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
        
        saly = aly(camera='HAWAII H-2RG')
        saly.detector.setStackType('CDS')
        saly.detector.setCrop(5)

        #############################################
        # START AND WAIT INSTRUMENT AND SASHA SETUP #
        #############################################
        fiber_test = False
        if fiber_test:
            test_ins_s = {
                "CAL_FIBER_DEP"        : "FOCUS_IN",
                "CAL_FIBER_FOCUS_LAMP" : "ON",
                "SCI_FILT_W1"          : "NB_H2",
                "SCI_FILT_W2"          : "ND3",
                "ADC1_PosAng"          : "min_disp",
                "ADC2_PosAng"          : "min_disp",
                "ND_FILT_W"            : "OD1_3",
                "APODIZER_W"           : "HOLE",
                "STEPS"                : "1"
            }

            test_det_s = {
                "DIT"                  : "1",
                "READOUT"              : "128x128" 
            }

            ins_s.update(test_ins_s)
            det_s.update(test_det_s)
        gen_base.main(ins_s=ins_s, det_s=det_s, expose=False, ngmax = uvars.Tech_Image_NGMAX)
        ###########################
        # ITERATIONS TO ALIGN PSF #
        ###########################

        # this loop goes up to convergence or to a max num of
        # iteration, after which it fails
        
        if self.paragraph("Definition of Target Position and Iteration Parameters"):
            converged = False
            thrsConv = 2 # pixels.  
            nIteration = 0
            maxIteration = int(ins_s["STEPS"]) #5
            logger.info("Definition of threshold convergence: {} [px] and Max Iterations: {} ".format(thrsConv, maxIteration)) 

            psfCXarr = []
            psfCYarr = []

            cor_x , cor_y = uvars.SOUL_TARGET_POS
            logger.info("Target positions coordinates in Full_Image are ({} , {})".format(cor_x, cor_y))
            cx, cy = ufun.convert_readout(cx=cor_x , cy=cor_y , readout=det_s["READOUT"])
            logger.info("TARGET POSITION coordinates in SCICAM {} configuration for the SOUL Alignment: ({} , {})".format(det_s["READOUT"], cx, cy))
        
        if fiber_test:
            t_cx, t_cy = ufun.to_teccam(cx, cy)
            print("center of convergence in TECCAM coordinates : {} , {}".format(t_cx, t_cy))

        while nIteration < maxIteration and not converged:
            nIteration += 1
            if self.paragraph("ITERATION {} of {}".format(nIteration, maxIteration)):
            
                ##################
                # 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))
                                
                saly.PSFNoCoro.LoadImage(fitsDir+self.filenames[-1])
                saly.PSFNoCoro.removeHotPixels(min_size=5, 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( round(psfCX,2), round(psfCY,2) ) )
                
                ########################
                # COMPUTE DISPLACEMENT #
                ########################
                
                logger.info("Calculate displacement between orientation and lyot centers")
                
                deltaX = psfCX - cx
                deltaY = psfCY - cy
                
                logger.info("DELTAs between PSF position and target center are:                deltaX: {}  ,    deltaY: {}".format(round(deltaX), round(deltaY) ) )
                
                ######################
                # CHECK IF CONVERGED #
                ######################
                
                converged = True if ( (np.abs(deltaX) <= thrsConv) and (np.abs(deltaY) <= thrsConv) ) else False
                
                logger.info("CHECK if deltas are < {}: {}".format(thrsConv, converged))
                
                

                if not converged: #and nIteration < maxIteration :
                
                    ####################
                    # GET CURRENT DROT #
                    ####################
                    
                    inDict = {"DROT":-9999.9}
                    unit = "DEGREE"
                    
                    self.status, drotDict = obsController.getPosition(unit, inDict)
                    drotPos = drotDict["DROT"]
                    
                    ###########################
                    # Transform to SOUL delta #
                    ###########################
                    step_x, step_y = ufun.transform_offset(disp = [deltaX,deltaY], thetaRot = drotPos , conv = "SOUL") if(not fiber_test) else ufun.to_teccam(deltaX, deltaY)

                    #step_x, step_y = ufun.to_soul(deltaX, deltaY)
                    logger.info("DELTAs in SOUL coordinates to be applied :                   {}  ,  {}".format(round(step_x,3), round(step_y,3) ) )
                   
                                    
                    ##############################
                    # Apply Offset in XY on SOUL #
                    ##############################
                
                    # if fiber_test:
                    #     self.status = obsController.RTCTTSetCentroidPos(step_x, step_y)
                    if (np.abs(step_x) > 2 or np.abs(step_y) > 2):
                        logger.error("Steps calculated are too large")
                        break
                    else:
                        logger.info("Applying offset in XY to SOUL...")
                        self.status = obsController.OffsetXYAO(step_x , step_y , "LEFT")
                        # self.status = ufun.soul_offsetXY_step(os=obsController, Dx=step_x, Dy=step_y)
                        logger.debug('Status is {}'.format(self.status))
                        logger.info("Applying offset in XY to SOUL... DONE")


        #############################################
        # EXPOSING SASHA TO EVALUATE FINAL POSITION #
        #############################################
        
        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))
                                
        saly.PSFNoCoro.LoadImage(fitsDir+self.filenames[-1])
        saly.PSFNoCoro.removeHotPixels(min_size=5, 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)
        deltaX = psfCX - cx
        deltaY = psfCY - cy
        converged = True if ( (np.abs(deltaX) <= thrsConv) and (np.abs(deltaY) <= thrsConv) ) else False
        print(psfCXarr)
        print(psfCYarr)
        if converged:
            logger.info("REACHED CONVERGENCE")
        else:
            logger.warning("SOUL ALIGNMENT PROCEDURE TERMINATED WITHOUT CONVERGENCE")
        logger.warning("END of SoulAlignment at Position{} , {}".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["pupil_dx"]   : np.round(deltaX,3),
            uvars.tpl_output_names["pupil_dy"]   : np.round(deltaY,3)
        }
