import galsim
import galsim.config.stamp as stamp
import logging
import os
import numpy as np

import injection_pipeline.Grid as grid

class AddOnImageBuilder(galsim.config.image_scattered.ScatteredImageBuilder):
    
    def setup(self, config, base, image_num, obj_num, ignore, logger):
        ignore = ignore + ['initial_image']
        return super(AddOnImageBuilder, self).setup(config, base, image_num, obj_num, ignore, logger)
    
    def addNoise(self, image, config, base, image_num, obj_num, current_var, logger):
        super(AddOnImageBuilder, self).addNoise(image, config, base, image_num, obj_num, current_var, logger)

        initial_image_name = galsim.config.ParseValue(config, 'initial_image', base, str)[0]
        initial_image = galsim.fits.read(initial_image_name)
        image += initial_image


galsim.config.RegisterImageType('AddOnImage', AddOnImageBuilder())

class InjectImageBuilder(AddOnImageBuilder):

    def setup(self, config, base, image_num, obj_num, ignore, logger):
        extra_ignore = ignore + ['tile_list', 'geom_file', 'tile_dir', 'config_dir', 'psf_dir',
                                 'version', 'run_name', 'bands', 'n_objects', 'n_realizations',
                                 'object_density', 'inj_objs_only', 'pos_sampling', 'realizations',
                                 'extinct_objs', 'rotate_objs']
        for key in config:
            if 'N_' in key:
                extra_ignore.append(key)
        full_xsize, full_ysize = super(InjectImageBuilder, self).setup(config, base, image_num, obj_num, extra_ignore, logger)

        # config = parse_inject_image_inputs(config, base)

        return full_xsize, full_ysize

    def addNoise(self, image, config, base, image_num, obj_num, current_var, logger):
        try:
            ioo = config['inj_objs_only']
            if (type(ioo) is bool) and (ioo is True):
                return super(AddOnImageBuilder, self).addNoise(image,
                                                               config,
                                                               base,
                                                               image_num,
                                                               obj_num,
                                                               current_var,
                                                               logger)
            elif (isinstance(ioo, dict)) and (ioo['value'] is True):
                # Still want to use existing image if changed to be BKG
                if (ioo['noise']) and ('BKG' in ioo['noise']):
                    return super(InjectImageBuilder, self).addNoise(image,
                                                                    config,
                                                                    base,
                                                                    image_num,
                                                                    obj_num,
                                                                    current_var,
                                                                    logger)
                else:
                    return super(AddOnImageBuilder, self).addNoise(image,
                                                                   config,
                                                                   base,
                                                                   image_num,
                                                                   obj_num,
                                                                   current_var,
                                                                   logger)
            else:
                # Default is to add on top of initial images
                return super(InjectImageBuilder, self).addNoise(image,
                                                                config,
                                                                base,
                                                                image_num,
                                                                obj_num,
                                                                current_var,
                                                                logger)

        except KeyError:
            # Default is to add on top of initial images
            return super(InjectImageBuilder, self).addNoise(image,
                                                            config,
                                                            base,
                                                            image_num,
                                                            obj_num,
                                                            current_var,
                                                            logger)

galsim.config.RegisterImageType('InjectImage', InjectImageBuilder())
