# %%
"""Provisional Roman instrument and observational settings."""
import lenstronomy.Util.util as util
import os
import astropy.io.fits as pyfits

# magnitude_zero_point: table 1 from https://iopscience.iop.org/article/10.3847/1538-4357/aac08b/pdf
# ccd_gain found right under table 1 in paper
# seeing, read_noise, pixel_scale (labelled as plate_scale on website): https://roman.gsfc.nasa.gov/science/WFI_technical.html
# sky brightness calculated using count rates per pixel at minimum Zodiacal light given in website above
# For microlensing survey mode: exposure time, number of exposures for F146 and F087: table 1 with the mission design of WFIRST Cycle 7 from https://iopscience.iop.org/article/10.3847/1538-4365/aafb69/meta#apjsaafb69t1fnd
# For wide area survey mode: exposure time and number of exposures for relevant filters set as given in https://roman.gsfc.nasa.gov/high_latitude_wide_area_survey.html


__all__ = ['Roman']

F062_band_obs = {'sky_brightness': 23.19,
                 'magnitude_zero_point': 26.56,
                 'seeing': 0.058}

F087_band_obs = {'sky_brightness': 22.93,
                 'magnitude_zero_point': 26.30,
                 'seeing': 0.073}

F106_band_obs = {'sky_brightness': 22.99,
                 'magnitude_zero_point': 26.44,
                 'seeing': 0.087}

F129_band_obs = {'sky_brightness': 22.99,
                 'magnitude_zero_point': 26.40,
                 'seeing': 0.105}

F158_band_obs = {'sky_brightness': 23.10,
                 'magnitude_zero_point': 26.43,
                 'seeing': 0.127}

F184_band_obs = {'sky_brightness': 23.22,
                 'magnitude_zero_point': 25.95,
                 'seeing': 0.151}

F146_band_obs = {'sky_brightness': 22.03,
                 'magnitude_zero_point': 26.65,
                 'seeing': 0.105}


"""
:keyword sky_brightness: sky brightness (in magnitude per square arcseconds in units of electrons)
:keyword magnitude_zero_point: magnitude in which 1 count (e-) per second per arcsecond square is registered
:keyword seeing: Full-Width-at-Half-Maximum (FWHM) of PSF
:keyword psf_type: string, type of PSF ('GAUSSIAN' supported)
"""


class Roman(object):
    """
    class contains Roman instrument and observation configurations
    """

    def __init__(self, band='F062', psf_type='GAUSSIAN', survey_mode='wide_area'):
        """

        :param band: string, 'F062', 'F087', 'F106', 'F129', 'F158' , 'F184' or 'F146' supported. Determines obs dictionary.
        :param psf_type: string, type of PSF ('GAUSSIAN', 'PIXEL' supported).
        """
        
        if band == 'F062':
            self.obs = F062_band_obs
        elif band == 'F087':
            self.obs = F087_band_obs
        elif band == 'F106':
            self.obs = F106_band_obs
        elif band == 'F129':
            self.obs = F129_band_obs
        elif band == 'F158':
            self.obs = F158_band_obs
        elif band == 'F184':
            self.obs = F184_band_obs
        elif band == 'F146':
            self.obs = F146_band_obs
        else:
            raise ValueError("band %s not supported! Choose 'F062', 'F087', 'F106', 'F129', 'F158' , 'F184' or 'F146'" % band) 

        if survey_mode == 'wide_area':
            # the number of exposures is given per sector
            # a full pass of the High Latitude Wide Area Survey is 155 sectors
            exp_per_tile = 0

            if band in ['F106','F158','F184', 'F062']:
                exp_per_tile = 3
            elif band == 'F129':
                exp_per_tile = 4
            else:
                raise ValueError("band %s is not supported with the wide_area survey mode! Choose 'F106', 'F062, 'F158', 'F184' or F129" % band)
            
            self.obs.update({'exposure_time': 146, 'num_exposures': 32*exp_per_tile})
        elif survey_mode == 'microlensing':
            if band == 'F146':
                # These are the exposure times and number of exposures for the primary filter, F146
                self.obs.update({'exposure_time': 46.8, 'num_exposures': 41000})
            elif band == 'F087':
                # These are the exposure times and number of exposures for the secondary filter, F087
                self.obs.update({'exposure_time': 286., 'num_exposures': 860})
            else:
                raise ValueError("band %s is not supported with the microlensing survey mode! Choose 'F146' or 'F087'" % band)
        else:
            raise ValueError("survey mode %s not supported! Choose 'wide_area' or 'microlensing'" % survey_mode)

        if psf_type == 'PIXEL':
            import lenstronomy
            module_path = os.path.dirname(lenstronomy.__file__)
            psf_filename = os.path.join(module_path, 'SimulationAPI/ObservationConfig/PSF_models/{}.fits'.format(band))
            kernel = pyfits.getdata(psf_filename)
            self.obs.update({'psf_type': 'PIXEL','kernel_point_source': kernel})
        elif psf_type == 'GAUSSIAN':
            self.obs.update({'psf_type': 'GAUSSIAN'})
        else:
            raise ValueError("psf_type %s not supported!" % psf_type)

        self.camera = {'read_noise': 15.5,
                       'pixel_scale': 0.11, 
                       'ccd_gain': 1,}
        """
        :keyword read_noise: std of noise generated by read-out (in units of electrons)
        :keyword pixel_scale: scale (in arcseconds) of pixels
        :keyword ccd_gain: electrons/ADU (analog-to-digital unit).
        """

    def kwargs_single_band(self):
        """

        :return: merged kwargs from camera and obs dicts
        """
        kwargs = util.merge_dicts(self.camera, self.obs)
        return kwargs
