import os
import torch

import gradio as gr

from modules import script_callbacks, scripts, shared


class CLIPBlend:
    def __init__(self):
        self.process_tokens = None

    def __call__(self, remade_batch_tokens, multipliers):
        if shared.opts.clip_blend_enabled:
            ratio1 = shared.opts.clip_blend_ratio1
            ratio2 = shared.opts.clip_blend_ratio2

            CLIP_original = shared.opts.CLIP_stop_at_last_layers

            shared.opts.CLIP_stop_at_last_layers = 1
            z1 = self.process_tokens(remade_batch_tokens, multipliers) # type: ignore
            shared.opts.CLIP_stop_at_last_layers = 2
            z2 = self.process_tokens(remade_batch_tokens, multipliers) # type: ignore

            shared.opts.CLIP_stop_at_last_layers = CLIP_original

            blend1 = z1*(1.-ratio1)+z2*ratio1

            if shared.opts.clip_blend_concat:
                blend2 = z1*(1.-ratio2)+z2*ratio2
                return torch.hstack([blend1,blend2])
            return blend1

        z = self.process_tokens(remade_batch_tokens, multipliers) # type: ignore
        return z

clip_blend = CLIPBlend()

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

def on_script_unloaded():
    cond_stage_model = shared.sd_model.cond_stage_model # type: ignore
    if type(cond_stage_model.process_tokens) == CLIPBlend:
        cond_stage_model.process_tokens = cond_stage_model.process_tokens.process_tokens

def on_ui_settings():
    section = ('clip_blend', 'CLIP Blend')
    shared.opts.add_option('clip_blend_enabled', shared.OptionInfo(
        False, 'Enable CLIP Blend', section=section
    ))
    shared.opts.add_option("clip_blend_ratio1", shared.OptionInfo(
        0, "Blend Ratio (R1) - The result is calculated as Z1 * (1 - R1) + Z2 * R1. Set to 0 to get CLIP skip=1; set to 1 to get CLIP skip=2", gr.Slider, {"minimum": -1, "maximum": 2, "step": 0.1}, section=section
    ))
    shared.opts.add_option("clip_blend_concat", shared.OptionInfo(
        False, "Concatenate second CLIP Blend", section=section
    ))
    shared.opts.add_option("clip_blend_ratio2", shared.OptionInfo(
        1, "Blend Ratio (R2) - The result is calculated as Z1 * (1 - R2) + Z2 * R2. Set to 0 to get CLIP skip=1; set to 1 to get CLIP skip=2", gr.Slider, {"minimum": -1, "maximum": 2, "step": 0.1}, section=section
    ))

def format_bool(x):
    return not x.lower().strip() in ['0', 'false']

def apply_override(p, opt, x, is_bool=False):
    getattr(p, 'override_settings')[opt] = format_bool(x) if is_bool else x

def apply_label(label, x, is_bool=False):
    return f"{label + (': ' if label else '')}{str(format_bool(x)) if is_bool else x}"

for scriptDataTuple in scripts.scripts_data:
    if os.path.basename(scriptDataTuple.path) == "xyz_grid.py":
        xy_grid = scriptDataTuple.module
        blend_enabled = xy_grid.AxisOption(
            "[CLIP Blend] Blend Enabled",
            str,
            lambda p, field, _: apply_override(p, 'clip_blend_enabled', field, is_bool=True),
            format_value=lambda p, opt, x: apply_label('CLIP Blend Enabled', x, is_bool=True)
        )
        xy_grid.axis_options.extend([blend_enabled])

        ratio1 = xy_grid.AxisOption(
            "[CLIP Blend] Blend Ratio (R1)",
            float,
            lambda p, field, _: apply_override(p, 'clip_blend_ratio1', field, is_bool=False)
        )
        xy_grid.axis_options.extend([ratio1])

        concat_enabled = xy_grid.AxisOption(
            "[CLIP Blend] Concatenate Second Blend",
            str,
            lambda p, field, _: apply_override(p, 'clip_blend_concat', field, is_bool=True),
            format_value=lambda p, opt, x: apply_label('Concatenated Second CLIP Blend', x, is_bool=True)
        )
        xy_grid.axis_options.extend([concat_enabled])

        ratio2 = xy_grid.AxisOption(
            "[CLIP Blend] Blend Ratio (R2)",
            float,
            lambda p, field, _: apply_override(p, 'clip_blend_ratio2', field, is_bool=False)
        )
        xy_grid.axis_options.extend([ratio2])

script_callbacks.on_script_unloaded(on_script_unloaded)
script_callbacks.on_model_loaded(on_model_loaded)
script_callbacks.on_ui_settings(on_ui_settings)