
from shinsapp import SHINSApp
from util.shinsLogger import logger
from alignment.allineamento import Align_the_coronagraph as aly
from alignment.utils import read_image
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_cal_Flux_MaskRemoval(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:

        . Read from OS the current scientific configuration of the instrument (CORO_SLIT_W, SCI_FILT_W1-2)
        
        . Set the CORO_SLIT_W to WINDOW to remove the mask
        . Set the ND filter specified by the user, leave untouched the scientific filter in use
        . Instrument Setup and SASHA Setup
        . Expose SASHA
        . Retrive FITS file
        
        . Put back the mask and filters as in the previous configuration
        . Instrument Setup without taking image

        . End of Template
        
        '''
        
        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

        # #########################
        # # OPEN INTERNAL TT LOOP #
        # #########################
        # logger.info("Open internal TT loop ...")
        # self.status = obsController.RTCTTOpenLoop()
        # logger.info("Open internal TT loop ... DONE")

        prev_setup = ufun.get_status_by_name(obsController, ["CORO_SLIT_W", "SCI_FILT_W1", "SCI_FILT_W2"])
        coro_mask = prev_setup[0]
        init_sci1 = prev_setup[1]
        init_sci2 = prev_setup[2]
        logger.debug("Coro Mask is in pos {}".format(coro_mask))
        logger.debug("Sci Filt 1 is in pos {}".format(init_sci1))
        logger.debug("Sci Filt 2 is in pos {}".format(init_sci2))

        sci , sci_pos , nd = "", "", ""
        if ins_s["SCI_FILT_W1"][0:2] == "ND":
            nd  = "SCI_FILT_W1"
            sci = "SCI_FILT_W2"
            sci_pos = init_sci2

        elif ins_s["SCI_FILT_W2"][0:2] == "ND":
            nd  = "SCI_FILT_W2"
            sci = "SCI_FILT_W1"
            sci_pos = init_sci1

        else:
            logger.warning("NO ND FILTER SPECIFIED!")

        ins_s.update({sci : sci_pos})
        ins_s.update({"CORO_SLIT_W" : "WINDOW"})
        
        # flux_ins_s = {
        #     "CORO_SLIT_W" : "WINDOW",
        #     sci           : sci_pos ,
        #     nd            : ins_s[nd]
        # }

        #ins_s["CORO_SLIT_W"] = "WINDOW"
        logger.debug("Saving current configuration and putting the coro wheel in {}".format(ins_s['CORO_SLIT_W']))
        
        #################
        # Setup + image #
        #################

        logger.info("Waiting setup to acquire image without the mask ...")
        self.status, self.filenames, _ = gen_base.main(ins_s=ins_s, det_s=det_s, **kwargs)
        logger.info("Waiting setup to acquire image without the mask ... DONE")

        # ins_s["CORO_SLIT_W"] = coro_mask

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

        ins_s.update({
            "CORO_SLIT_W" : coro_mask,
            "SCI_FILT_W1" : init_sci1,
            "SCI_FILT_W2" : init_sci2
        })
        
        logger.info("Waiting setup to restore original configuration ...")
        self.status, _, __ = gen_base.main(ins_s=ins_s, expose=False)
        logger.info("Waiting setup to restore original configuration ... DONE")

        # ##########################
        # # CLOSE INTERNAL TT LOOP #
        # ##########################
        # logger.info("Close internal TT loop ...")
        # self.status = obsController.RTCTTCloseLoop()
        # logger.info("Close internal TT loop ... DONE")

        saly = aly(camera='HAWAII H-2RG')
        saly.detector.setCrop(3)
        saly.detector.setStackType('CDS')
        im_arr = []
        for image_name in self.filenames:
            logger.debug("reading {}".format(image_name))
            img_ = read_image(fitsDir+image_name, saly.detector)
            im_arr.append(img_)
        meaned = np.mean(np.asarray(im_arr), 0)
        saly.PSFNoCoro.LoadImage(meaned)
        saly.PSFNoCoro.removeHotPixels(verbose=True)
        #saly.PSFNoCoro.calculate_bkg()
        
        self.tpl_output = {
            #uvars.tpl_output_names["sasha_flux"] : np.round(saly.PSFNoCoro.flux, 2),
            uvars.tpl_output_names["sasha_max"]  : np.max(saly.PSFNoCoro.image),
            uvars.tpl_output_names["sasha_min"]  : np.min(saly.PSFNoCoro.image)
        }
