import modules.scripts as scripts
import gradio as gr
import os
import cv2
import numpy as np
from PIL import Image

from modules import images
from modules.processing import process_images, Processed
from modules.processing import Processed
from modules.shared import opts, cmd_opts, state
from scipy.ndimage import binary_erosion
from scripts.weshop_util import image_segment_by_mask

class Script(scripts.Script):  

# The title of the script. This is what will be displayed in the dropdown menu.
    def title(self):
        return "ChooseControlnetImage"


# Determines when the script should be shown in the dropdown menu via the 
# returned value. As an example:
# is_img2img is True if the current tab is img2img, and False if it is txt2img.
# Thus, return is_img2img to only show the script on the img2img tab.

    def show(self, is_img2img):
        return scripts.AlwaysVisible

# How the script's is displayed in the UI. See https://gradio.app/docs/#components
# for the different UI components you can use and how to create them.
# Most UI components can return a value, such as a boolean for a checkbox.
# The returned values are passed to the run method as parameters.

    def ui(self, is_img2img):
        with gr.Accordion("Chose Controlnet Image" , open=False):
            controlnet_1 = gr.Dropdown(['0' , '1' , '2'] , label="Controlnet-0 image" , info="controlnet image, 0:original image , 1:mask image , 2: masked original image")
            controlnet_2 = gr.Dropdown(['0' , '1' , '2'] , label="Controlnet-1 image" , info="controlnet image, 0:original image , 1:mask image , 2: masked original image")
            controlnet_3 = gr.Dropdown(['0' , '1' , '2'] , label="Controlnet-2 image" , info="controlnet image, 0:original image , 1:mask image , 2: masked original image")
            controlnet_4 = gr.Dropdown(['0' , '1' , '2'] , label="Controlnet-3 image" , info="controlnet image, 0:original image , 1:mask image , 2: masked original image")
            controlnet_5 = gr.Dropdown(['0' , '1' , '2'] , label="Controlnet-4 image" , info="controlnet image, 0:original image , 1:mask image , 2: masked original image")
            return [controlnet_1,controlnet_2,controlnet_3,controlnet_4,controlnet_5]

    def before_process(self, p, *args):
        pass

    def process(self, p, *args):
        """
        This function is called before processing begins for AlwaysVisible scripts.
        You can modify the processing object (p) here, inject hooks, etc.
        args contains all values returned by components from ui()
        """

        controlnet_input_images = []
        for conf in args:
            input_image = None
            if conf is None :
                pass
            elif conf == "0" :
                pass
            elif conf == "1" :
                input_image = p.image_mask
                pass
            elif conf == "2" :
                if p.init_images is not None and len(p.init_images) > 0 and p.image_mask is not None:
                    image = np.array(p.init_images[0])
                    image_mask = np.array(p.image_mask)
                    segment_result = image_segment_by_mask(image , image_mask)
                    input_image = Image.fromarray(segment_result)
            controlnet_input_images.append(input_image)

        setattr(p , 'control_net_input_image' , controlnet_input_images)

    def postprocess(self, p, processed, *args):
        pass

# This is where the additional processing is implemented. The parameters include
# self, the model object "p" (a StableDiffusionProcessing class, see
# processing.py), and the parameters returned by the ui method.
# Custom functions can be defined here, and additional libraries can be imported 
# to be used in processing. The return value should be a Processed object, which is
# what is returned by the process_images method.

    # def run(self, p, angle, hflip, vflip, overwrite):
    #     pass
