import torch

from .protocols import DiffusionStepProtocol
from ..utils import to_velocity


class EulerDiffusionStep(DiffusionStepProtocol):
    """
    First-order Euler method for diffusion sampling.
    Takes a single step from the current noise level (sigma) to the next by
    computing velocity from the denoised prediction and applying: sample + velocity * dt.
    """

    def step(
        self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int
    ) -> torch.Tensor:
        sigma = sigmas[step_index]
        sigma_next = sigmas[step_index + 1]
        dt = sigma_next - sigma
        velocity = to_velocity(sample, sigma, denoised_sample)

        return (sample.to(torch.float32) + velocity.to(torch.float32) * dt).to(sample.dtype)


class EulerCFGPPDiffusionStep(DiffusionStepProtocol):
    """Rectified-flow Euler CFG++ step used by the official LTX-2.3 upscaler workflow."""

    uses_unconditional = True

    def step(
        self,
        sample: torch.Tensor,
        denoised_sample: torch.Tensor,
        sigmas: torch.Tensor,
        step_index: int,
        unconditional_denoised_sample: torch.Tensor,
    ) -> torch.Tensor:
        sigma = sigmas[step_index].to(torch.float32)
        sigma_next = sigmas[step_index + 1].to(torch.float32)
        if sigma_next == 0:
            return denoised_sample.to(sample.dtype)

        alpha = 1.0 - sigma
        alpha_next = 1.0 - sigma_next
        noise = (sample.float() - alpha * unconditional_denoised_sample.float()) / sigma
        return (alpha_next * denoised_sample.float() + sigma_next * noise).to(sample.dtype)


class EulerAncestralDiffusionStep(DiffusionStepProtocol):
    """Rectified-flow ancestral Euler step used by LTX-2.5 Distilled."""

    def __init__(self, eta: float = 1.0, s_noise: float = 1.0) -> None:
        self.eta = eta
        self.s_noise = s_noise

    def step(
        self,
        sample: torch.Tensor,
        denoised_sample: torch.Tensor,
        sigmas: torch.Tensor,
        step_index: int,
        noise: torch.Tensor | None = None,
    ) -> torch.Tensor:
        sigma = sigmas[step_index].to(torch.float32)
        sigma_next = sigmas[step_index + 1].to(torch.float32)
        if sigma_next == 0:
            return denoised_sample.to(sample.dtype)
        if self.eta > 0 and noise is None:
            raise ValueError("EulerAncestralDiffusionStep requires noise when eta > 0")

        downstep_ratio = 1.0 + (sigma_next / sigma - 1.0) * self.eta
        sigma_down = sigma_next * downstep_ratio
        sigma_down_ratio = sigma_down / sigma
        x_next = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised_sample.float()
        if self.eta > 0:
            alpha_next = 1.0 - sigma_next
            alpha_down = 1.0 - sigma_down
            renoise_coeff = (sigma_next.square() - sigma_down.square() * alpha_next.square() / alpha_down.square()).clamp(min=0).sqrt()
            x_next = alpha_next / alpha_down * x_next + noise.float() * self.s_noise * renoise_coeff
        return x_next.to(sample.dtype)


class Res2sDiffusionStep(DiffusionStepProtocol):
    """
    Second-order diffusion step for the res_2s HQ sampler.
    """

    @staticmethod
    def get_sde_coeff(
        sigma_next: torch.Tensor,
        sigma_up: torch.Tensor | None = None,
        sigma_down: torch.Tensor | None = None,
        sigma_max: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if sigma_down is not None:
            alpha_ratio = (1 - sigma_next) / (1 - sigma_down)
            sigma_up = (sigma_next**2 - sigma_down**2 * alpha_ratio**2).clamp(min=0) ** 0.5
        elif sigma_up is not None:
            sigma_up = sigma_up.clamp(max=sigma_next * 0.9999)
            sigmax = sigma_max if sigma_max is not None else torch.ones_like(sigma_next)
            sigma_signal = sigmax - sigma_next
            sigma_residual = (sigma_next**2 - sigma_up**2).clamp(min=0) ** 0.5
            alpha_ratio = sigma_signal + sigma_residual
            sigma_down = sigma_residual / alpha_ratio
        else:
            alpha_ratio = torch.ones_like(sigma_next)
            sigma_down = sigma_next
            sigma_up = torch.zeros_like(sigma_next)

        sigma_up = torch.nan_to_num(sigma_up if sigma_up is not None else torch.zeros_like(sigma_next), 0.0)
        sigma_down = sigma_down.clone()
        nan_mask = torch.isnan(sigma_down)
        sigma_down[nan_mask] = sigma_next[nan_mask].to(sigma_down.dtype)
        alpha_ratio = torch.nan_to_num(alpha_ratio, 1.0)

        return alpha_ratio, sigma_down, sigma_up

    def step(
        self,
        sample: torch.Tensor,
        denoised_sample: torch.Tensor,
        sigmas: torch.Tensor,
        step_index: int,
        noise: torch.Tensor,
    ) -> torch.Tensor:
        sigma = sigmas[step_index]
        sigma_next = sigmas[step_index + 1]
        alpha_ratio, sigma_down, sigma_up = self.get_sde_coeff(sigma_next, sigma_up=sigma_next * 0.5)
        if torch.any(sigma_up == 0) or torch.any(sigma_next == 0):
            return denoised_sample

        eps_next = (sample - denoised_sample) / (sigma - sigma_next)
        denoised_next = sample - sigma * eps_next
        x_noised = alpha_ratio * (denoised_next + sigma_down * eps_next) + sigma_up * noise
        return x_noised.to(denoised_sample.dtype)