"""Flow-matching sampling for Krea 2 Turbo and Raw checkpoints."""

from __future__ import annotations

import math
import warnings

import torch

from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.krea2_model import Krea2Conditioning

KREA2_LATENT_DOWNSCALE = 8
KREA2_PATCH_SIZE = 2


def roundUp(value: int, multiple: int, name: str) -> int:
    """Round a dimension up and issue the required visible warning."""

    if value <= 0:
        raise DiffusionCliError(f"{name} must be positive: got {value}")
    aligned = ((value + multiple - 1) // multiple) * multiple
    if aligned != value:
        warnings.warn(
            f"Krea {name}={value} is not a multiple of {multiple}; "
            f"rounding up to {aligned}",
            UserWarning,
            stacklevel=2,
        )
    return aligned


def prepareConditioning(
    latent: torch.Tensor,
    conditioning: Krea2Conditioning,
    patch_size: int = KREA2_PATCH_SIZE,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Patchify a latent and build image/text positions and masks."""

    if latent.ndim != 4:
        raise DiffusionCliError("Krea sampler expects an NCHW latent")
    batch, channels, height, width = latent.shape
    if height % patch_size or width % patch_size:
        raise DiffusionCliError("Krea latent is not patch aligned")
    image_height = height // patch_size
    image_width = width // patch_size
    image = latent.reshape(
        batch,
        channels,
        image_height,
        patch_size,
        image_width,
        patch_size,
    ).permute(0, 2, 4, 1, 3, 5).reshape(
        batch,
        image_height * image_width,
        channels * patch_size * patch_size,
    )
    image_positions = torch.zeros(
        batch,
        image_height,
        image_width,
        3,
        device=latent.device,
        dtype=torch.float32,
    )
    image_positions[..., 1] = torch.arange(
        image_height,
        device=latent.device,
    )[None, :, None]
    image_positions[..., 2] = torch.arange(
        image_width,
        device=latent.device,
    )[None, None, :]
    text_positions = torch.zeros(
        batch,
        conditioning.hidden_states.shape[1],
        3,
        device=latent.device,
        dtype=torch.float32,
    )
    positions = torch.cat(
        (text_positions, image_positions.reshape(batch, -1, 3)),
        dim=1,
    )
    mask = torch.cat(
        (
            conditioning.attention_mask.to(device=latent.device, dtype=torch.bool),
            torch.ones(
                batch,
                image.shape[1],
                device=latent.device,
                dtype=torch.bool,
            ),
        ),
        dim=1,
    )
    return image, positions, mask


def timestepMu(
    image_token_count: int,
    *,
    min_resolution: int = 256,
    max_resolution: int = 1280,
    y1: float = 0.5,
    y2: float = 1.15,
) -> float:
    """Interpolate the Raw timestep shift from image-token resolution."""

    minimum = (min_resolution // KREA2_LATENT_DOWNSCALE) ** 2
    maximum = (max_resolution // KREA2_LATENT_DOWNSCALE) ** 2
    if maximum == minimum:
        return y1
    slope = (y2 - y1) / (maximum - minimum)
    return slope * image_token_count + (y1 - slope * minimum)


def timesteps(
    image_token_count: int,
    steps: int,
    *,
    y1: float = 0.5,
    y2: float = 1.15,
    sigma: float = 1.0,
    mu: float | None = None,
) -> list[float]:
    """Return a resolution-shifted descending flow-matching grid."""

    if steps < 1:
        raise DiffusionCliError(f"Steps must be at least 1: got {steps}")
    if mu is None:
        mu = timestepMu(image_token_count, y1=y1, y2=y2)
    values = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64)
    exp_mu = math.exp(mu)
    shifted = exp_mu / (
        exp_mu + (1.0 / values - 1.0).pow(sigma)
    )
    shifted[-1] = 0.0
    return shifted.tolist()


def buildInitialNoise(
    batch_size: int,
    channels: int,
    height: int,
    width: int,
    seed: int,
    *,
    device,
    dtype,
) -> torch.Tensor:
    """Create independent per-image Gaussian noise using ``seed + index``."""

    if batch_size < 1:
        raise DiffusionCliError("Batch size must be at least 1")
    values = []
    for index in range(batch_size):
        generator = torch.Generator(device=device).manual_seed(seed + index)
        values.append(
            torch.randn(
                1,
                channels,
                height // KREA2_LATENT_DOWNSCALE,
                width // KREA2_LATENT_DOWNSCALE,
                device=device,
                dtype=dtype,
                generator=generator,
            )
        )
    return torch.cat(values, dim=0)


@torch.no_grad()
def sampleKrea2(
    model,
    conditioning: Krea2Conditioning,
    *,
    negative_conditioning: Krea2Conditioning | None = None,
    batch_size: int,
    height: int,
    width: int,
    seed: int,
    steps: int,
    cfg: float,
    device,
    dtype,
    latent_channels: int = 16,
    patch_size: int = 2,
    mu: float | None = None,
    shift_y1: float = 0.5,
    shift_y2: float = 1.15,
    progress_callback=None,
) -> torch.Tensor:
    """Run Krea Euler flow integration and return an NCHW latent."""

    if cfg < 0:
        raise DiffusionCliError(f"CFG must be non-negative: got {cfg}")
    height = roundUp(height, KREA2_LATENT_DOWNSCALE * patch_size, "height")
    width = roundUp(width, KREA2_LATENT_DOWNSCALE * patch_size, "width")
    if conditioning.hidden_states.shape[0] != batch_size:
        if conditioning.hidden_states.shape[0] == 1:
            conditioning = Krea2Conditioning(
                conditioning.hidden_states.expand(
                    batch_size,
                    -1,
                    -1,
                    -1,
                ),
                conditioning.attention_mask.expand(batch_size, -1),
            )
        else:
            raise DiffusionCliError("Krea conditioning batch mismatch")
    if cfg > 0 and negative_conditioning is None:
        raise DiffusionCliError(
            "Krea CFG requires negative conditioning when cfg is nonzero"
        )
    if cfg > 0 and negative_conditioning.hidden_states.shape[0] == 1:
        negative_conditioning = Krea2Conditioning(
            negative_conditioning.hidden_states.expand(
                batch_size,
                -1,
                -1,
                -1,
            ),
            negative_conditioning.attention_mask.expand(batch_size, -1),
        )
    elif cfg > 0 and negative_conditioning.hidden_states.shape[0] != batch_size:
        raise DiffusionCliError("Krea negative conditioning batch mismatch")
    latent = buildInitialNoise(
        batch_size,
        latent_channels,
        height,
        width,
        seed,
        device=device,
        dtype=dtype,
    )
    image_tokens = (height // KREA2_LATENT_DOWNSCALE // patch_size) * (
        width // KREA2_LATENT_DOWNSCALE // patch_size
    )
    schedule = timesteps(
        image_tokens,
        steps,
        y1=shift_y1,
        y2=shift_y2,
        mu=mu,
    )
    for index, (current, following) in enumerate(
        zip(schedule[:-1], schedule[1:])
    ):
        time = torch.full(
            (batch_size,),
            current,
            device=device,
            dtype=dtype,
        )
        conditional = model(latent, time, conditioning)
        if cfg > 0:
            unconditional = model(latent, time, negative_conditioning)
            velocity = conditional + cfg * (conditional - unconditional)
        else:
            velocity = conditional
        if not torch.isfinite(velocity).all():
            raise DiffusionCliError(
                f"Krea diffusion output contains NaN or Inf at step {index}"
            )
        latent = latent + (following - current) * velocity
        if not torch.isfinite(latent).all():
            raise DiffusionCliError(
                f"Krea latent contains NaN or Inf at step {index}"
            )
        if progress_callback is not None:
            progress_callback(index + 1, steps)
    return latent


# Compatibility aliases make the sampling boundary easy to discover.
sample = sampleKrea2
sampleLatents = sampleKrea2
prepare = prepareConditioning