"""Scheduler, sampler, and latent-noise helpers for Z-Image inference."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import numpy as np
import torch
from scipy.stats import beta

from diffusion_cli.config import LATENT_CHANNELS, LATENT_DOWNSCALE
from diffusion_cli.errors import DiffusionCliError


@dataclass(frozen=True)
class ModelSamplingDiscreteFlow:
    """Discrete flow sigma conversion used by the workflow sampling node."""

    shift: float = 3.0
    multiplier: float = 1.0
    timesteps: int = 1000

    def sigma(self, timestep: torch.Tensor) -> torch.Tensor:
        """Convert normalized timesteps to shifted flow sigmas."""

        t = timestep / self.multiplier
        if self.shift == 1.0:
            return t
        return self.shift * t / (1.0 + (self.shift - 1.0) * t)

    def timestep(self, sigma: torch.Tensor) -> torch.Tensor:
        """Convert shifted flow sigmas back to model timesteps."""

        return sigma * self.multiplier

    def sigmaTable(self, device=None, dtype=None) -> torch.Tensor:
        """Return the 1,000-step shifted flow table used by ComfyUI."""

        indices = torch.arange(
            1,
            self.timesteps + 1,
            device=device,
            dtype=dtype or torch.float32,
        )
        timesteps = indices / self.timesteps * self.multiplier
        return self.sigma(timesteps)

    def percentToSigma(self, percent: float) -> torch.Tensor:
        """Convert a denoise percentage to shifted flow sigma."""

        if percent <= 0.0:
            return torch.tensor(1.0)
        if percent >= 1.0:
            return torch.tensor(0.0)
        timestep = torch.tensor(1.0 - percent) * self.multiplier
        return self.sigma(timestep)


def latentShape(
    batch_size: int,
    height: int,
    width: int,
) -> tuple[int, int, int, int]:
    """Return the Z-Image latent tensor shape for image dimensions."""

    return (
        batch_size,
        LATENT_CHANNELS,
        height // LATENT_DOWNSCALE,
        width // LATENT_DOWNSCALE,
    )


def buildInitialNoise(
    batch_size: int,
    height: int,
    width: int,
    seed: int,
    device,
    dtype,
) -> torch.Tensor:
    """Build deterministic Gaussian latent noise on the selected device."""

    generator = torch.Generator(device=device)
    generator.manual_seed(seed)
    return torch.randn(
        latentShape(batch_size, height, width),
        generator=generator,
        device=device,
        dtype=dtype,
    )


def betaScheduler(
    steps: int,
    sampling: ModelSamplingDiscreteFlow | None = None,
    device: object | None = None,
    dtype: object | None = None,
) -> torch.Tensor:
    """Return ComfyUI beta scheduler sigmas with a final zero sigma."""

    if steps < 1:
        raise ValueError("steps must be at least 1")

    sampling = sampling or ModelSamplingDiscreteFlow()
    table = sampling.sigmaTable(device=device, dtype=dtype)
    total_timesteps = len(table) - 1
    percentiles = 1.0 - np.linspace(0.0, 1.0, steps, endpoint=False)
    timesteps = np.rint(beta.ppf(percentiles, 0.6, 0.6) * total_timesteps)

    values = []
    last_t = -1
    for timestep in timesteps:
        if timestep != last_t:
            values.append(table[int(timestep)])
        last_t = timestep

    sigmas = torch.stack(values) if values else table[:0]
    sigmas = torch.cat(
        [sigmas, torch.zeros(1, device=sigmas.device, dtype=sigmas.dtype)]
    )
    return sigmas


def sigmaToHalfLogSnr(sigma: torch.Tensor) -> torch.Tensor:
    """Convert sigma to half-log-SNR for flow samplers."""

    return torch.logit(sigma).neg()


def halfLogSnrToSigma(half_log_snr: torch.Tensor) -> torch.Tensor:
    """Convert half-log-SNR to sigma for flow samplers."""

    return half_log_snr.neg().sigmoid()


def offsetFirstSigmaForSnr(
    sigmas: torch.Tensor,
    sampling: ModelSamplingDiscreteFlow | None = None,
    percent_offset: float = 1e-4,
) -> torch.Tensor:
    """Offset an initial sigma of 1 to avoid infinite flow log-SNR."""

    if len(sigmas) <= 1 or sigmas[0] < 1:
        return sigmas
    sampling = sampling or ModelSamplingDiscreteFlow()
    sigmas = sigmas.clone()
    sigmas[0] = sampling.percentToSigma(percent_offset).to(
        device=sigmas.device,
        dtype=sigmas.dtype,
    )
    return sigmas


def eiHPhi1(h: torch.Tensor) -> torch.Tensor:
    """Compute h * phi_1(h) for exponential integrator updates."""

    return torch.expm1(h)


def eiHPhi2(h: torch.Tensor) -> torch.Tensor:
    """Compute h * phi_2(h) for exponential integrator updates."""

    return (torch.expm1(h) - h) / h


def defaultNoiseSampler(
    latent: torch.Tensor,
    seed: int | None = None,
) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor]:
    """Return a deterministic Gaussian noise sampler for stochastic steps."""

    if seed is None:
        generator = None
    else:
        seed_value = seed + 1 if latent.device == torch.device("cpu") else seed
        generator = torch.Generator(device=latent.device)
        generator.manual_seed(seed_value)

    def sample(_sigma, _sigma_next):
        return torch.randn(
            latent.size(),
            dtype=latent.dtype,
            layout=latent.layout,
            device=latent.device,
            generator=generator,
        )

    return sample


def cfgDenoise(model, latent, sigma, conditioning, cfg: float) -> torch.Tensor:
    """Run classifier-free guidance for one denoising call."""

    positive = model.forward(latent, sigma, conditioning.positive)
    negative = model.forward(latent, sigma, conditioning.negative)
    return negative + (positive - negative) * cfg


def _checkFinite(latent: torch.Tensor, context: str) -> None:
    if not torch.all(torch.isfinite(latent)):
        raise DiffusionCliError(
            f"SEEDS-3 sampler produced NaN or Inf {context}"
        )


@torch.no_grad()
def sampleSeeds3(
    model,
    latent: torch.Tensor,
    sigmas: torch.Tensor,
    conditioning,
    cfg: float,
    *,
    eta: float = 1.0,
    s_noise: float = 1.0,
    seed: int | None = None,
    noise_sampler=None,
    r_1: float = 1.0 / 3.0,
    r_2: float = 2.0 / 3.0,
    callback=None,
) -> torch.Tensor:
    """Run the SEEDS-3 solver over a beta sigma schedule."""

    if len(sigmas) < 2:
        raise ValueError("sigmas must contain at least two values")

    x = latent.float()
    sigmas = sigmas.to(device=x.device, dtype=torch.float32)
    sigmas = offsetFirstSigmaForSnr(sigmas)
    noise_sampler = noise_sampler or defaultNoiseSampler(x, seed=seed)
    s_in = torch.ones(x.shape[0], device=x.device, dtype=torch.float32)
    inject_noise = eta > 0 and s_noise > 0

    for i in range(len(sigmas) - 1):
        sigma = sigmas[i]
        sigma_next = sigmas[i + 1]
        denoised = cfgDenoise(model, x, sigma * s_in, conditioning, cfg)
        _checkFinite(denoised, f"after denoising step {i}")
        if callback is not None:
            callback({
                "x": x,
                "i": i,
                "sigma": sigma,
                "sigma_hat": sigma,
                "denoised": denoised,
            })

        if sigma_next == 0:
            x = denoised
            _checkFinite(x, f"at final step {i}")
            continue

        lambda_s = sigmaToHalfLogSnr(sigma)
        lambda_t = sigmaToHalfLogSnr(sigma_next)
        h = lambda_t - lambda_s
        h_eta = h * (eta + 1)
        lambda_s_1 = torch.lerp(lambda_s, lambda_t, r_1)
        lambda_s_2 = torch.lerp(lambda_s, lambda_t, r_2)
        sigma_s_1 = halfLogSnrToSigma(lambda_s_1)
        sigma_s_2 = halfLogSnrToSigma(lambda_s_2)

        alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
        alpha_s_2 = sigma_s_2 * lambda_s_2.exp()
        alpha_t = sigma_next * lambda_t.exp()

        x_2 = (
            sigma_s_1 / sigma
            * (-r_1 * h * eta).exp()
            * x
            - alpha_s_1 * eiHPhi1(-r_1 * h_eta) * denoised
        )
        if inject_noise:
            sde_noise = (
                (-2 * r_1 * h * eta).expm1().neg().sqrt()
                * noise_sampler(sigma, sigma_s_1)
            )
            x_2 = x_2 + sde_noise * sigma_s_1 * s_noise
        denoised_2 = cfgDenoise(
            model,
            x_2,
            sigma_s_1 * s_in,
            conditioning,
            cfg,
        )

        a3_2 = r_2 / r_1 * eiHPhi2(-r_2 * h_eta)
        a3_1 = eiHPhi1(-r_2 * h_eta) - a3_2
        x_3 = (
            sigma_s_2 / sigma
            * (-r_2 * h * eta).exp()
            * x
            - alpha_s_2 * (a3_1 * denoised + a3_2 * denoised_2)
        )
        if inject_noise:
            segment_factor = (r_1 - r_2) * h * eta
            sde_noise = sde_noise * segment_factor.exp()
            sde_noise = sde_noise + (
                segment_factor.mul(2).expm1().neg().sqrt()
                * noise_sampler(sigma_s_1, sigma_s_2)
            )
            x_3 = x_3 + sde_noise * sigma_s_2 * s_noise
        denoised_3 = cfgDenoise(
            model,
            x_3,
            sigma_s_2 * s_in,
            conditioning,
            cfg,
        )

        b3 = eiHPhi2(-h_eta) / r_2
        b1 = eiHPhi1(-h_eta) - b3
        x = (
            sigma_next / sigma
            * (-h * eta).exp()
            * x
            - alpha_t * (b1 * denoised + b3 * denoised_3)
        )
        if inject_noise:
            segment_factor = (r_2 - 1) * h * eta
            sde_noise = sde_noise * segment_factor.exp()
            sde_noise = sde_noise + (
                segment_factor.mul(2).expm1().neg().sqrt()
                * noise_sampler(sigma_s_2, sigma_next)
            )
            x = x + sde_noise * sigma_next * s_noise

        _checkFinite(x, f"after step {i}")

    return x


def sampleLatents(
    model,
    conditioning,
    *,
    batch_size: int,
    height: int,
    width: int,
    seed: int,
    steps: int,
    cfg: float,
    device,
    dtype,
) -> torch.Tensor:
    """Build initial noise and run the default Z-Image sampler."""

    latent = buildInitialNoise(
        batch_size,
        height,
        width,
        seed,
        device,
        torch.float32,
    )
    sigmas = betaScheduler(steps, device=device, dtype=torch.float32)
    # Keep the first standalone milestone deterministic; the stochastic
    # SEEDS-3 path is more sensitive to numeric parity with ComfyUI internals.
    return sampleSeeds3(
        model,
        latent,
        sigmas,
        conditioning,
        cfg,
        eta=0.0,
        s_noise=0.0,
        seed=seed,
    )