import copy
import random
import numpy as np
import args_manager
import modules.flags
import modules.config
import modules.flags as flags
import modules.constants as constants
import modules.hack_async_worker_class as worker

class Fooocus():
    def __init__(self) -> None:
        self._preset_parameters()
        self._inpaint_parameters()
        self._define_models()
        self._load_generator()
        
    def _seed_set(self, seed = 0, seed_random = True):
        if seed_random:
            self.image_seed = random.randint(constants.MIN_SEED, constants.MAX_SEED)
        else:
            try:
                seed_value = int(seed)
                if constants.MIN_SEED <= seed_value <= constants.MAX_SEED:
                    self.image_seed = seed_value
            except ValueError:
                pass
            self.image_seed = random.randint(constants.MIN_SEED, constants.MAX_SEED)

    def _preset_parameters(self):
        '''
        We keep the default parameters of Fooocus unchanged!
        '''
        self.output_format = modules.config.default_output_format # png
        self.style_selections = copy.deepcopy(modules.config.default_styles) # Fooocus style unchanged:Fooocus V2+Enhanec+Sharp
        self.performance_selection = modules.config.default_performance
        self.aspect_ratios_selection = modules.config.default_aspect_ratio

        self.current_tab = 'inpaint'
        self.input_image_checkbox = True
        self.advanced_checkbox = modules.config.default_advanced_checkbox
        self.uov_method = flags.disabled
        self.uov_input_image = None

        # unchanged diffusion setting
        self.dev_mode = False
        self.guidance_scale = modules.config.default_cfg_scale
        self.sharpness = modules.config.default_sample_sharpness

        self.adm_scaler_end = 0.3
        self.adm_scaler_positive = 1.5
        self.adm_scaler_negative = 0.8
        self.refiner_swap_method = flags.refiner_swap_method
        self.adaptive_cfg = modules.config.default_cfg_tsnr
        self.clip_skip = modules.config.default_clip_skip
        self.sampler_name = modules.config.default_sampler
        self.scheduler_name = modules.config.default_scheduler
        self.vae_name = modules.config.default_vae

        self.generate_image_grid = False
        self.overwrite_step = modules.config.default_overwrite_step
        self.overwrite_switch = modules.config.default_overwrite_switch
        self.overwrite_width = -1
        self.overwrite_height = -1
        self.overwrite_vary_strength = -1
        self.overwrite_upscale_strength = -1
        self.disable_preview = modules.config.default_black_out_nsfw
        self.disable_intermediate_results = flags.Performance.has_restricted_features(modules.config.default_performance)
        self.disable_seed_increment = False
        self.read_wildcards_in_order = False

        self.black_out_nsfw = modules.config.default_black_out_nsfw
        self.save_metadata_to_images = modules.config.default_save_metadata_to_images
        self.metadata_scheme = modules.config.default_metadata_scheme

        self.debugging_cn_preprocessor = False
        self.skipping_cn_preprocessor = False
        self.mixing_image_prompt_and_vary_upscale = False
        self.mixing_image_prompt_and_inpaint = False
        self.controlnet_softness = 0.25
        self.canny_low_threshold = 64
        self.canny_high_threshold = 128

        freeu_enabled = False
        freeu_b1 = 1.01
        freeu_b2 = 1.02
        freeu_s1 = 0.99
        freeu_s2 = 0.95
        self.freeu_ctrls = [freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2]

        self.image_number = 1 
        self.default_prompt = modules.config.default_prompt
        self.negative_prompt = modules.config.default_prompt_negative

        self.inpaint_additional_prompt = ''
        self.outpaint_selections = []
        self.outpaint_extend_times = 0.0
        self.inpaint_ctrls = [False,False,'None',0,0.0,False,False,0]
        
        self.ip_ctrls = []
        for _ in range(flags.controlnet_image_count):
            ip_image = None
            # WHAT'S THIS
            self.ip_ctrls.append(ip_image)
            default_end, default_weight = flags.default_parameters[flags.default_ip]
            ip_stop = default_end
            self.ip_ctrls.append(ip_stop)
            ip_weight = default_weight
            self.ip_ctrls.append(ip_weight)
            ip_type = flags.default_ip
            self.ip_ctrls.append(ip_type)

    def _upsample_parameteres(self):
        # argue mode
        self.current_tab = 'uov'
        self.uov_method = 'Upscale (2x)'
        
        self.inpaint_additional_prompt = ''
        self.outpaint_selections = []
        self.outpaint_extend_times = 0.0
        self.inpaint_ctrls = [False,False,'None',0,0.0,False,False,0]

    def _inpaint_parameters(self, 
                            inpaint_additional_prompt = '',
                            outpaint_selections = ['Left', 'Right', 'Top', 'Bottom'],
                            outpaint_extend_times = 0.4,
                            inpaint_mode = modules.flags.inpaint_option_default
        ):
        # mode change
        assert inpaint_mode in modules.flags.inpaint_options

        # argue mode
        self.current_tab = 'inpaint'
        self.uov_method = flags.disabled

        debugging_inpaint_preprocessor = False

        if inpaint_mode == modules.flags.inpaint_option_detail:
            self.inpaint_additional_prompt = inpaint_additional_prompt + 'Sharp objects, Clear objects, optimize all blur objects'
            self.outpaint_selections = []
            self.outpaint_extend_times = 0.0
            inpaint_disable_initial_latent = False
            inpaint_engine = 'None'
            inpaint_strength = 0.3
            inpaint_respective_field = 0.0

        elif inpaint_mode == modules.flags.inpaint_option_modify:
            self.inpaint_additional_prompt = inpaint_additional_prompt
            self.outpaint_selections = []
            self.outpaint_extend_times = 0.0
            inpaint_disable_initial_latent = True
            inpaint_engine = modules.config.default_inpaint_engine_version
            inpaint_strength = 1.0
            inpaint_respective_field = 0.0

        else:
            self.inpaint_additional_prompt = ''
            self.outpaint_selections = outpaint_selections
            self.outpaint_extend_times = outpaint_extend_times
            inpaint_disable_initial_latent = False
            inpaint_engine = modules.config.default_inpaint_engine_version
            inpaint_strength = 1.0
            inpaint_respective_field = 0.618

        inpaint_erode_or_dilate = 0
        invert_mask_checkbox = False
        inpaint_mask_upload_checkbox = False

        self.inpaint_ctrls = [debugging_inpaint_preprocessor, inpaint_disable_initial_latent, inpaint_engine,
                              inpaint_strength, inpaint_respective_field,
                              inpaint_mask_upload_checkbox, invert_mask_checkbox, inpaint_erode_or_dilate]

    def _define_models(self):
        # define models
        self.base_model = modules.config.default_base_model_name
        self.refiner_model = modules.config.default_refiner_model_name
        self.refiner_switch = modules.config.default_refiner_switch
        self.lora_ctrls = []
        for i, (enabled, filename, weight) in enumerate(modules.config.default_loras):
            lora_enabled = enabled
            lora_model = filename
            lora_weight = weight
            self.lora_ctrls += [lora_enabled, lora_model, lora_weight]

    def _generate_images(self, task: worker.AsyncTask):
        import ldm_patched.modules.model_management as model_management

        with model_management.interrupt_processing_mutex:
            model_management.interrupt_processing = False
        # outputs=[progress_html, progress_window, progress_gallery, gallery]
        if len(task.args) == 0:
            return 
        imgs = self.worker(task)
        return imgs

    def _set_ctrls(self, 
                   prompt, 
                   negative_prompt,
                   # for inpainting/outpainting
                   inpaint_input_image=None,
                   inpaint_mask_image=None,
                   # for upsampling
                   upsample_input_image=None):
        
        prompt = prompt + self.default_prompt
        negative_prompt = negative_prompt + self.negative_prompt
        currentTask = worker.AsyncTask(args=[])

        inpaint_input_image = {'image':inpaint_input_image,
                               'mask':inpaint_mask_image}
        inpaint_mask_image = None

        ctrls = [currentTask, self.generate_image_grid]
        ctrls += [
            prompt, negative_prompt, self.style_selections,
            self.performance_selection, self.aspect_ratios_selection, self.image_number, self.output_format, self.image_seed,
            self.read_wildcards_in_order, self.sharpness, self.guidance_scale
        ]
        ctrls += [self.base_model, self.refiner_model, self.refiner_switch] + self.lora_ctrls
        ctrls += [self.input_image_checkbox, self.current_tab]
        ctrls += [self.uov_method, upsample_input_image]
        ctrls += [self.outpaint_selections,self.outpaint_extend_times, inpaint_input_image, self.inpaint_additional_prompt, inpaint_mask_image]
        ctrls += [self.disable_preview, self.disable_intermediate_results, self.disable_seed_increment, self.black_out_nsfw]
        ctrls += [self.adm_scaler_positive, self.adm_scaler_negative, self.adm_scaler_end, self.adaptive_cfg, self.clip_skip]
        ctrls += [self.sampler_name, self.scheduler_name, self.vae_name]
        ctrls += [self.overwrite_step, self.overwrite_switch, self.overwrite_width, self.overwrite_height, self.overwrite_vary_strength]
        ctrls += [self.overwrite_upscale_strength, self.mixing_image_prompt_and_vary_upscale, self.mixing_image_prompt_and_inpaint]
        ctrls += [self.debugging_cn_preprocessor, self.skipping_cn_preprocessor, self.canny_low_threshold, self.canny_high_threshold]
        ctrls += [self.refiner_swap_method, self.controlnet_softness]

        ctrls += self.freeu_ctrls
        ctrls += self.inpaint_ctrls

        if not args_manager.args.disable_metadata:
            ctrls += [self.save_metadata_to_images, self.metadata_scheme]
        ctrls += self.ip_ctrls
        return ctrls

    def _get_task(self,args):
        args.pop(0)
        return worker.AsyncTask(args=args)

    def _load_generator(self):
        self.worker = worker.Hack_worker()

    def _refine(self,
                image_number = 1,
                prompt = '', 
                negative_prompt = '',
                outpaint_selections=[],
                outpaint_extend_times=0.0,
                origin_image = None,
                mask_image = None,
                seed = None):
        '''
        origin_image numpy HW3 0-255
        mask_image numpy HW3 inpaint area be 255,255,255 else be 0,0,0
        '''
        # input check
        if np.amax(origin_image) < 1.1:
            origin_image = (origin_image*255).astype(np.uint8)
        if np.amax(mask_image) < 1.1:
            mask_image = (mask_image*255).astype(np.uint8)
        if mask_image.ndim < 3:
            mask_image = mask_image[:,:,None].repeat(3,axis=-1)
        
        # set seed
        self.image_number = image_number
        if seed is None: self._seed_set()
        else: self._seed_set(seed,seed_random=False)
        # conduct inpainting
        self._inpaint_parameters(outpaint_selections=outpaint_selections,
                                 outpaint_extend_times=outpaint_extend_times,
                                 inpaint_mode = modules.flags.inpaint_option_detail)
        ctrls = self._set_ctrls(prompt, negative_prompt, origin_image, mask_image)
        currentTask = self._get_task(ctrls)
        output = self._generate_images(currentTask)
        output = [np.array(o)/255. for o in output]
        return output

    def _upsample(self,
                 image :np.array = None,
                 image_number = 1,
                 prompt = '', 
                 negative_prompt = '',
                 control_strength = 0.382,
                 seed = None):
        # input check
        if np.amax(image) < 1.1:
            image = (image*255).astype(np.uint8)
        # set seed
        if seed is None: self._seed_set()
        else: self._seed_set(seed,seed_random=False)
        self.image_number=image_number
        # conduct inpainting
        self._upsample_parameteres()
        self.overwrite_upscale_strength = control_strength
        ctrls = self._set_ctrls(prompt, negative_prompt, upsample_input_image=image)
        currentTask = self._get_task(ctrls)
        output = self._generate_images(currentTask)
        output = [np.array(o)/255. for o in output]
        return output

    def __call__(self,
                 image_number = 1,
                 prompt = '', 
                 negative_prompt = '',
                 outpaint_selections = ['Left', 'Right', 'Top', 'Bottom'],
                 outpaint_extend_times = 0.4,
                 origin_image = None,
                 mask_image = None,
                 seed = None):
        '''
        origin_image numpy HW3 0-255 / 0-1
        mask_image numpy HW(3) inpaint area be 255,255(,255) / 1,1(,1) else be 0,0(,0)
        '''
        # input check
        if np.amax(origin_image) < 1.1:
            origin_image = (origin_image*255).astype(np.uint8)
        if np.amax(mask_image) < 1.1:
            mask_image = (mask_image*255).astype(np.uint8)
        if mask_image.ndim < 3:
            mask_image = mask_image[:,:,None].repeat(3,axis=-1)
            
        # set seed
        if seed is None: self._seed_set()
        else: self._seed_set(seed,seed_random=False)
        self.image_number=image_number
        # conduct inpainting
        self._inpaint_parameters(outpaint_selections=outpaint_selections,
                                 outpaint_extend_times=outpaint_extend_times,
                                 inpaint_mode = modules.flags.inpaint_option_default)
        ctrls = self._set_ctrls(prompt, negative_prompt, origin_image, mask_image)
        currentTask = self._get_task(ctrls)
        output = self._generate_images(currentTask)
        output = [np.array(o)/255. for o in output]
        return output