import torch
import k_diffusion.sampling
from k_diffusion.sampling import to_d, default_noise_sampler, get_ancestral_step
from tqdm.auto import trange
from modules import sd_samplers, sd_samplers_common, sd_samplers_kdiffusion

_SAMPLER_REGISTERED = False


@torch.no_grad()
def sample_kohaku_lonyu_yog(model, x, sigmas, extra_args=None, callback=None,
                            disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'),
                            s_noise=1., noise_sampler=None, eta=1.):
    """
    Kohaku_LoNyu_Yog Sampler - Geometric Second-Order Method
    """
    extra_args = {} if extra_args is None else extra_args
    s_in = x.new_ones([x.shape[0]])
    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler

    steps_total = len(sigmas) - 1
    halfway_point = steps_total // 2

    for i in trange(steps_total, disable=disable, desc="Kohaku Sampling"):
        gamma = min(s_churn / steps_total, 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
        sigma_hat = sigmas[i] * (gamma + 1)

        if gamma > 0:
            eps = torch.randn_like(x) * s_noise
            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5

        denoised = model(x, sigma_hat * s_in, **extra_args)
        d = to_d(x, sigma_hat, denoised)

        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
        dt = sigma_down - sigmas[i]

        if i <= halfway_point:
            x_antipode = -x

            denoised2 = model(x_antipode, sigma_hat * s_in, **extra_args)
            d2 = to_d(x_antipode, sigma_hat, denoised2)

            v_down = (d + d2) / 2
            x_closer = x + v_down * dt

            denoised3 = model(x_closer, sigma_hat * s_in, **extra_args)
            d3 = to_d(x_closer, sigma_hat, denoised3)

            real_d = (d + d3) / 2
            x = x + real_d * dt

            if sigma_up > 0:
                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
        else:
            x = x + d * dt
            if sigma_up > 0:
                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up

        if callback is not None:
            callback({
                'x': x,
                'i': i,
                'sigma': sigmas[i],
                'sigma_hat': sigma_hat,
                'denoised': denoised
            })

    return x


def register_kohaku_sampler():
    """Регистрирует Kohaku_LoNyu_Yog сэмплер в WebUI"""
    global _SAMPLER_REGISTERED

    if _SAMPLER_REGISTERED:
        return

    if any(s.name == 'Kohaku_LoNyu_Yog' for s in sd_samplers.all_samplers):
        _SAMPLER_REGISTERED = True
        return

    if not hasattr(k_diffusion.sampling, 'sample_kohaku_lonyu_yog'):
        setattr(k_diffusion.sampling, 'sample_kohaku_lonyu_yog', sample_kohaku_lonyu_yog)

    sampler_data = sd_samplers_common.SamplerData(
        name='Kohaku_LoNyu_Yog',
        constructor=lambda model: sd_samplers_kdiffusion.KDiffusionSampler('sample_kohaku_lonyu_yog', model),
        aliases=['kohaku', 'lonyu'],
        options={'second_order': True}
    )

    sd_samplers.all_samplers.append(sampler_data)
    sd_samplers.all_samplers_map = {x.name: x for x in sd_samplers.all_samplers}

    _SAMPLER_REGISTERED = True
    print("✓ Kohaku_LoNyu_Yog sampler registered successfully!")