<reponame>jiwoncpark/baobab
import os, sys
from datetime import datetime
import warnings
from importlib import import_module
from addict import Dict
import json
from collections import OrderedDict 
import lenstronomy.SimulationAPI.ObservationConfig as obs_cfg

class BaobabConfig:
    """Nested dictionary representing the configuration for Baobab data generation
    
    """
    def __init__(self, user_cfg):
        """
        Parameters
        ----------
        user_cfg : dict or Dict
            user-defined configuration
        
        """
        self.__dict__ = Dict(user_cfg)
        if not hasattr(self, 'out_dir'):
            # Default out_dir path if not specified
            self.out_dir = os.path.join(self.destination_dir, '{:s}_{:s}_prior={:s}_seed={:d}'.format(self.name, self.train_vs_val, self.bnn_prior_class, self.seed))
        self.out_dir = os.path.abspath(self.out_dir)
        if not hasattr(self, 'checkpoint_interval'):
            self.checkpoint_interval = max(100, self.n_data // 100)
        self.get_survey_info(self.survey_info, self.psf.type)
        self.interpret_magnification_cfg()
        self.interpret_kinematics_cfg()
        self.log_filename = datetime.now().strftime("log_%m-%d-%Y_%H:%M_baobab.json")
        self.log_path = os.path.join(self.out_dir, self.log_filename)

    @classmethod
    def from_file(cls, user_cfg_path):
        """Alternative constructor that accepts the path to the user-defined configuration python file
        Parameters
        ----------
        user_cfg_path : str or os.path object
            path to the user-defined configuration python file
        """
        dirname, filename = os.path.split(os.path.abspath(user_cfg_path))
        module_name, ext = os.path.splitext(filename)
        sys.path.insert(0, dirname)
        if ext == '.py':
            #user_cfg_file = map(__import__, module_name)
            #user_cfg = getattr(user_cfg_file, 'cfg')
            user_cfg_script = import_module(module_name)
            user_cfg = getattr(user_cfg_script, 'cfg').deepcopy()
            return cls(user_cfg)
        elif ext == '.json':
            with open(user_cfg_path, 'r') as f:
                user_cfg_str = f.read()
            user_cfg = Dict(json.loads(user_cfg_str)).deepcopy()
            return cls(user_cfg)
        else:
            raise NotImplementedError("This extension is not supported.")

    def export_log(self):
        """Export the baobab log to the current working directory

        """
        with open(self.log_path, 'w') as f:
            json.dump(self.__dict__, f)
            print("Exporting baobab log to {:s}".format(self.log_path))

    def interpret_magnification_cfg(self):
        if 'agn_light' not in self.components:
            if len(self.bnn_omega.magnification.frac_error_sigma) != 0: # non-empty dictionary
                warnings.warn("`bnn_omega.magnification.frac_error_sigma` field is ignored as the images do not contain AGN.")
                self.bnn_omega.magnification.frac_error_sigma = 0.0
        else:
            if 'magnification' not in self.bnn_omega:
                self.bnn_omega.magnification.frac_error_sigma = 0.0
            elif self.bnn_omega.magnification is None:
                self.bnn_omega.magnification.frac_error_sigma = 0.0
            
        if ('magnification' not in self.bnn_omega) and 'agn_light' in self.components:
            self.bnn_omega.magnification.frac_error_sigma = 0.0

    def interpret_kinematics_cfg(self):
        """Validate the kinematics config

        """
        kinematics_cfg = self.bnn_omega.kinematics_cfg
        if kinematics_cfg.anisotropy_model == 'analytic':
            warnings.warn("Since velocity dispersion computation is analytic, any entry other than `sampling_number` in `kinematics.numerics_kwargs` will be ignored.")

    def get_survey_info(self, survey_info, psf_type):
        """Fetch the camera and instrument information corresponding to the survey string identifier

        """
        sys.path.insert(0, obs_cfg.__path__[0])
        survey_module = import_module(survey_info['survey_name'])
        survey_class = getattr(survey_module, survey_info['survey_name'])
        coadd_years = survey_info['coadd_years'] if 'coadd_years' in survey_info else None

        self.survey_object_dict = OrderedDict()
        for bp in survey_info['bandpass_list']:
            survey_object = survey_class(band=bp, psf_type=psf_type, coadd_years=coadd_years)
            # Overwrite ObservationConfig PSF type with user-configured PSF type
            if hasattr(self, 'psf'):
                survey_object.obs['psf_type'] = self.psf.type
            if survey_object.obs['psf_type'] == 'PIXEL':
                if hasattr(self, 'psf'):
                    if hasattr(self.psf, 'psf_kernel_size'):
                        survey_object.psf_kernel_size = self.psf.kernel_size
                    else:
                        raise ValueError("Observation dictionary must specify PSF kernel size if psf_type is PIXEL.")
                    if hasattr(self.psf, 'which_psf_maps'):
                        survey_object.which_psf_maps = self.psf.which_psf_maps
                    else:
                        raise ValueError("Observation dictionary must specify indices of PSF kernel maps if psf_type is PIXEL.")
                else:
                    raise ValueError("User must supply PSF kwargs in the Baobab config if PSF type is PIXEL.")
            else: # 'GAUSSIAN'
                survey_object.psf_kernel_size = None
                survey_object.which_psf_maps = None
            # Override default survey specs with user-specified kwargs
            survey_object.camera.update(survey_info['override_camera_kwargs'])
            survey_object.obs.update(survey_info['override_obs_kwargs'])
            self.survey_object_dict[bp] = survey_object
        # Camera dict is same across bands, so arbitrarily take the last band
        self.instrument = survey_object.camera

    def get_noise_kwargs(self,bandpass):
        """
        Return the noise kwargs defined in the babobab config, e.g. for passing to the noise model for online data augmentation

        Returns
        -------
            (dict): A dict containing the noise kwargs to be passed to the noise
                model.
            (str): The bandpass to pull the noise information for

        """
        # Go through the baobab config and pull out the noise kwargs one by one.
        noise_kwargs = {}
        noise_kwargs.update(self.instrument)
        noise_kwargs.update(self.survey_object_dict[bandpass].kwargs_single_band())
        return noise_kwargs

import os, sys
from datetime import datetime
import warnings
from importlib import import_module
from addict import Dict
import json
from collections import OrderedDict 
import lenstronomy.SimulationAPI.ObservationConfig as obs_cfg

class BaobabConfig:
    """Nested dictionary representing the configuration for Baobab data generation
    
    """
    def __init__(self, user_cfg):
        """
        Parameters
        ----------
        user_cfg : dict or Dict
            user-defined configuration
        
        """
        self.__dict__ = Dict(user_cfg)
        if not hasattr(self, 'out_dir'):
            # Default out_dir path if not specified
            self.out_dir = os.path.join(self.destination_dir, '{:s}_{:s}_prior={:s}_seed={:d}'.format(self.name, self.train_vs_val, self.bnn_prior_class, self.seed))
        self.out_dir = os.path.abspath(self.out_dir)
        if not hasattr(self, 'checkpoint_interval'):
            self.checkpoint_interval = max(100, self.n_data // 100)
        self.get_survey_info(self.survey_info, self.psf.type)
        self.interpret_magnification_cfg()
        self.interpret_kinematics_cfg()
        self.log_filename = datetime.now().strftime("log_%m-%d-%Y_%H:%M_baobab.json")
        self.log_path = os.path.join(self.out_dir, self.log_filename)

    @classmethod
    def from_file(cls, user_cfg_path):
        """Alternative constructor that accepts the path to the user-defined configuration python file
        Parameters
        ----------
        user_cfg_path : str or os.path object
            path to the user-defined configuration python file
        """
        dirname, filename = os.path.split(os.path.abspath(user_cfg_path))
        module_name, ext = os.path.splitext(filename)
        sys.path.insert(0, dirname)
        if ext == '.py':
            #user_cfg_file = map(__import__, module_name)
            #user_cfg = getattr(user_cfg_file, 'cfg')
            user_cfg_script = import_module(module_name)
            user_cfg = getattr(user_cfg_script, 'cfg').deepcopy()
            return cls(user_cfg)
        elif ext == '.json':
            with open(user_cfg_path, 'r') as f:
                user_cfg_str = f.read()
            user_cfg = Dict(json.loads(user_cfg_str)).deepcopy()
            return cls(user_cfg)
        else:
            raise NotImplementedError("This extension is not supported.")

    def export_log(self):
        """Export the baobab log to the current working directory

        """
        with open(self.log_path, 'w') as f:
            json.dump(self.__dict__, f)
            print("Exporting baobab log to {:s}".format(self.log_path))

    def interpret_magnification_cfg(self):
        if 'agn_light' not in self.components:
            if len(self.bnn_omega.magnification.frac_error_sigma) != 0: # non-empty dictionary
                warnings.warn("`bnn_omega.magnification.frac_error_sigma` field is ignored as the images do not contain AGN.")
                self.bnn_omega.magnification.frac_error_sigma = 0.0
        else:
            if 'magnification' not in self.bnn_omega:
                self.bnn_omega.magnification.frac_error_sigma = 0.0
            elif self.bnn_omega.magnification is None:
                self.bnn_omega.magnification.frac_error_sigma = 0.0
            
        if ('magnification' not in self.bnn_omega) and 'agn_light' in self.components:
            self.bnn_omega.magnification.frac_error_sigma = 0.0

    def interpret_kinematics_cfg(self):
        """Validate the kinematics config

        """
        kinematics_cfg = self.bnn_omega.kinematics_cfg
        if kinematics_cfg.anisotropy_model == 'analytic':
            warnings.warn("Since velocity dispersion computation is analytic, any entry other than `sampling_number` in `kinematics.numerics_kwargs` will be ignored.")

    def get_survey_info(self, survey_info, psf_type):
        """Fetch the camera and instrument information corresponding to the survey string identifier

        """
        sys.path.insert(0, obs_cfg.__path__[0])
        survey_module = import_module(survey_info['survey_name'])
        survey_class = getattr(survey_module, survey_info['survey_name'])
        coadd_years = survey_info['coadd_years'] if 'coadd_years' in survey_info else None

        self.survey_object_dict = OrderedDict()
        for bp in survey_info['bandpass_list']:
            survey_object = survey_class(band=bp, psf_type=psf_type, coadd_years=coadd_years)
            # Overwrite ObservationConfig PSF type with user-configured PSF type
            if hasattr(self, 'psf'):
                survey_object.obs['psf_type'] = self.psf.type
            if survey_object.obs['psf_type'] == 'PIXEL':
                if hasattr(self, 'psf'):
                    if hasattr(self.psf, 'psf_kernel_size'):
                        survey_object.psf_kernel_size = self.psf.kernel_size
                    else:
                        raise ValueError("Observation dictionary must specify PSF kernel size if psf_type is PIXEL.")
                    if hasattr(self.psf, 'which_psf_maps'):
                        survey_object.which_psf_maps = self.psf.which_psf_maps
                    else:
                        raise ValueError("Observation dictionary must specify indices of PSF kernel maps if psf_type is PIXEL.")
                else:
                    raise ValueError("User must supply PSF kwargs in the Baobab config if PSF type is PIXEL.")
            else: # 'GAUSSIAN'
                survey_object.psf_kernel_size = None
                survey_object.which_psf_maps = None
            # Override default survey specs with user-specified kwargs
            survey_object.camera.update(survey_info['override_camera_kwargs'])
            survey_object.obs.update(survey_info['override_obs_kwargs'])
            self.survey_object_dict[bp] = survey_object
        # Camera dict is same across bands, so arbitrarily take the last band
        self.instrument = survey_object.camera

    def get_noise_kwargs(self,bandpass):
        """
        Return the noise kwargs defined in the babobab config, e.g. for passing to the noise model for online data augmentation

        Returns
        -------
            (dict): A dict containing the noise kwargs to be passed to the noise
                model.
            (str): The bandpass to pull the noise information for

        """
        # Go through the baobab config and pull out the noise kwargs one by one.
        noise_kwargs = {}
        noise_kwargs.update(self.instrument)
        noise_kwargs.update(self.survey_object_dict[bandpass].kwargs_single_band())
        return noise_kwargs
Nested dictionary representing the configuration for Baobab data generation Parameters         ----------         user_cfg : dict or Dict             user-defined configuration # Default out_dir path if not specified Alternative constructor that accepts the path to the user-defined configuration python file         Parameters         ----------         user_cfg_path : str or os.path object             path to the user-defined configuration python file #user_cfg_file = map(__import__, module_name) #user_cfg = getattr(user_cfg_file, 'cfg') Export the baobab log to the current working directory # non-empty dictionary Validate the kinematics config Fetch the camera and instrument information corresponding to the survey string identifier # Overwrite ObservationConfig PSF type with user-configured PSF type # 'GAUSSIAN' # Override default survey specs with user-specified kwargs # Camera dict is same across bands, so arbitrarily take the last band Return the noise kwargs defined in the babobab config, e.g. for passing to the noise model for online data augmentation          Returns         -------             (dict): A dict containing the noise kwargs to be passed to the noise                 model.             (str): The bandpass to pull the noise information for # Go through the baobab config and pull out the noise kwargs one by one.