from modules import scripts, script_callbacks, shared
import aesthetic_clip
import gradio as gr

aesthetic = aesthetic_clip.AestheticCLIP()
aesthetic_imgs_components = []


class AestheticScript(scripts.Script):
    def title(self):
        return "Aesthetic embeddings"

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

    def ui(self, is_img2img):
        aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative = aesthetic_clip.create_ui()

        self.infotext_fields = [
            (aesthetic_lr, "Aesthetic LR"),
            (aesthetic_weight, "Aesthetic weight"),
            (aesthetic_steps, "Aesthetic steps"),
            (aesthetic_imgs, "Aesthetic embedding"),
            (aesthetic_slerp, "Aesthetic slerp"),
            (aesthetic_imgs_text, "Aesthetic text"),
            (aesthetic_text_negative, "Aesthetic text negative"),
            (aesthetic_slerp_angle, "Aesthetic slerp angle"),
        ]

        aesthetic_imgs_components.append(aesthetic_imgs)

        return [aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative]

    def process(self, p, aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative):
        aesthetic.set_aesthetic_params(p, float(aesthetic_lr), float(aesthetic_weight), int(aesthetic_steps), aesthetic_imgs, aesthetic_slerp, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative)


def on_model_loaded(sd_model):
    aesthetic.process_tokens = sd_model.cond_stage_model.process_tokens
    sd_model.cond_stage_model.process_tokens = aesthetic


def on_script_unloaded():
    cond_stage_model = shared.sd_model.cond_stage_model
    if type(cond_stage_model.process_tokens) == aesthetic_clip.AestheticCLIP:
        cond_stage_model.process_tokens = cond_stage_model.process_tokens.process_tokens


def on_ui_tabs():
    with gr.Blocks(analytics_enabled=False) as aesthetic_interface:
        with gr.Row().style(equal_height=False):
            with gr.Column(variant='panel'):
                gr.HTML(value="Create an aesthetic embedding out of any number of images")

                new_embedding_name_ae = gr.Textbox(label="Name")
                process_src_ae = gr.Textbox(label='Source directory')
                batch_ae = gr.Slider(minimum=1, maximum=1024, step=1, label="Batch size", value=256)

                with gr.Row():
                    with gr.Column(scale=3):
                        gr.HTML(value="")

                    with gr.Column():
                        create_embedding_ae = gr.Button(value="Create images embedding", variant='primary')

            with gr.Column():
                output = gr.Text(value="", show_label=False)

        dropdown_components = aesthetic_imgs_components.copy()

        def generate_embs(*args):
            res = aesthetic_clip.generate_imgs_embd(*args)

            aesthetic_clip.update_aesthetic_embeddings()
            updates = [gr.Dropdown.update(choices=sorted(aesthetic_clip.aesthetic_embeddings.keys())) for _ in range(len(dropdown_components))]

            return [*updates, res]

        create_embedding_ae.click(
            fn=generate_embs,
            inputs=[
                new_embedding_name_ae,
                process_src_ae,
                batch_ae
            ],
            outputs=[
                *dropdown_components,
                output
            ]
        )

        aesthetic_imgs_components.clear()

    return [(aesthetic_interface, "Create aesthetic embedding", "aesthetic_interface")]


script_callbacks.on_script_unloaded(on_script_unloaded)
script_callbacks.on_model_loaded(on_model_loaded)
script_callbacks.on_ui_tabs(on_ui_tabs)