import torch

from modules_forge.supported_preprocessor import Preprocessor, PreprocessorParameter
from modules_forge.shared import add_supported_preprocessor
from backend.sampling.sampling_function import sampling_function_inner
from backend import attention


def sdp(q, k, v, transformer_options):
    if q.shape[0] == 0:
        return q

    return attention.attention_function(q, k, v, heads=transformer_options["n_heads"], mask=None)


def adain(x, target_std, target_mean):
    if x.shape[0] == 0:
        return x

    std, mean = torch.std_mean(x, dim=(2, 3), keepdim=True, correction=0)
    return (((x - mean) / std) * target_std) + target_mean


def zero_cat(a, b, dim):
    if a.shape[0] == 0:
        return b
    if b.shape[0] == 0:
        return a
    return torch.cat([a, b], dim=dim)


class PreprocessorReference(Preprocessor):
    def __init__(self, name, use_attn=True, use_adain=True, priority=0):
        super().__init__()
        self.name = name
        self.use_attn = use_attn
        self.use_adain = use_adain
        self.sorting_priority = priority
        self.tags = ['Reference']
        self.slider_resolution = PreprocessorParameter(visible=False)
        self.slider_1 = PreprocessorParameter(label='Style Fidelity', value=0.5, minimum=0.0, maximum=1.0, step=0.01, visible=True)
        self.show_control_mode = False
        self.corp_image_with_a1111_mask_when_in_img2img_inpaint_tab = False
        self.do_not_need_model = True

        self.is_recording_style = False
        self.recorded_attn1 = {}
        self.recorded_h = {}

    def process_before_every_sampling(self, process, cond, mask, *args, **kwargs):
        unit = kwargs['unit']
        weight = float(unit.weight)
        style_fidelity = float(unit.threshold_a)
        start_percent = float(unit.guidance_start)
        end_percent = float(unit.guidance_end)

        if process.sd_model.is_sdxl:
            style_fidelity = style_fidelity ** 3.0  # sdxl is very sensitive to reference so we lower the weights

        vae = process.sd_model.forge_objects.vae
        # This is a powerful VAE with integrated memory management, bf16, and tiled fallback.

        latent_image = vae.encode(cond.movedim(1, -1))
        latent_image = process.sd_model.forge_objects.vae.first_stage_model.process_in(latent_image)

        gen_seed = process.seeds[0] + 1
        gen_cpu = torch.Generator().manual_seed(gen_seed)

        unet = process.sd_model.forge_objects.unet.clone()
        sigma_max = unet.model.predictor.percent_to_sigma(start_percent)
        sigma_min = unet.model.predictor.percent_to_sigma(end_percent)

        self.recorded_attn1 = {}
        self.recorded_h = {}

        def conditioning_modifier(model, x, timestep, uncond, cond, cond_scale, model_options, seed):
            sigma = timestep[0].item()
            if not (sigma_min <= sigma <= sigma_max):
                return model, x, timestep, uncond, cond, cond_scale, model_options, seed

            self.is_recording_style = True

            xt = latent_image.to(x) + torch.randn(x.size(), dtype=x.dtype, generator=gen_cpu).to(x) * sigma
            sampling_function_inner(model, xt, timestep, uncond, cond, 1, model_options, seed)

            self.is_recording_style = False

            return model, x, timestep, uncond, cond, cond_scale, model_options, seed

        def block_proc(h, flag, transformer_options):
            if not self.use_adain:
                return h

            if flag != 'after':
                return h

            location = transformer_options['block']

            sigma = transformer_options["sigmas"][0].item()
            if not (sigma_min <= sigma <= sigma_max):
                return h

            channel = int(h.shape[1])
            minimal_channel = 1500 - 1000 * weight

            if channel < minimal_channel:
                return h

            if self.is_recording_style:
                self.recorded_h[location] = torch.std_mean(h, dim=(2, 3), keepdim=True, correction=0)
                return h
            else:
                cond_indices = transformer_options['cond_indices']
                uncond_indices = transformer_options['uncond_indices']
                cond_or_uncond = transformer_options['cond_or_uncond']
                r_std, r_mean = self.recorded_h[location]

                h_c = h[cond_indices]
                h_uc = h[uncond_indices]

                o_c = adain(h_c, r_std, r_mean)
                o_uc_strong = h_uc
                o_uc_weak = adain(h_uc, r_std, r_mean)
                o_uc = o_uc_weak + (o_uc_strong - o_uc_weak) * style_fidelity

                recon = []
                for cx in cond_or_uncond:
                    if cx == 0:
                        recon.append(o_c)
                    else:
                        recon.append(o_uc)

                o = torch.cat(recon, dim=0)
                return o

        def attn1_proc(q, k, v, transformer_options):
            if not self.use_attn:
                return sdp(q, k, v, transformer_options)

            sigma = transformer_options["sigmas"][0].item()
            if not (sigma_min <= sigma <= sigma_max):
                return sdp(q, k, v, transformer_options)

            location = (transformer_options['block'][0], transformer_options['block'][1],
                        transformer_options['block_index'])

            channel = int(q.shape[2])
            minimal_channel = 1500 - 1280 * weight

            if channel < minimal_channel:
                return sdp(q, k, v, transformer_options)

            if self.is_recording_style:
                self.recorded_attn1[location] = (k, v)
                return sdp(q, k, v, transformer_options)
            else:
                cond_indices = transformer_options['cond_indices']
                uncond_indices = transformer_options['uncond_indices']
                cond_or_uncond = transformer_options['cond_or_uncond']

                q_c = q[cond_indices]
                q_uc = q[uncond_indices]

                k_c = k[cond_indices]
                k_uc = k[uncond_indices]

                v_c = v[cond_indices]
                v_uc = v[uncond_indices]

                k_r, v_r = self.recorded_attn1[location]

                o_c = sdp(q_c, zero_cat(k_c, k_r, dim=1), zero_cat(v_c, v_r, dim=1), transformer_options)
                o_uc_strong = sdp(q_uc, k_uc, v_uc, transformer_options)
                o_uc_weak = sdp(q_uc, zero_cat(k_uc, k_r, dim=1), zero_cat(v_uc, v_r, dim=1), transformer_options)
                o_uc = o_uc_weak + (o_uc_strong - o_uc_weak) * style_fidelity

                recon = []
                for cx in cond_or_uncond:
                    if cx == 0:
                        recon.append(o_c)
                    else:
                        recon.append(o_uc)

                o = torch.cat(recon, dim=0)
                return o

        unet.add_block_modifier(block_proc)
        unet.add_conditioning_modifier(conditioning_modifier)
        unet.set_model_replace_all(attn1_proc, 'attn1')

        process.sd_model.forge_objects.unet = unet

        return cond, mask


add_supported_preprocessor(PreprocessorReference(
    name='reference_only',
    use_attn=True,
    use_adain=False,
    priority=100
))

add_supported_preprocessor(PreprocessorReference(
    name='reference_adain',
    use_attn=False,
    use_adain=True
))

add_supported_preprocessor(PreprocessorReference(
    name='reference_adain+attn',
    use_attn=True,
    use_adain=True
))