import argparse
import os
import secrets
import threading
from pathlib import Path

from PIL import Image


ROOT = Path(__file__).resolve().parent
MODEL_ID = "google/ddpm-celebahq-256"
MODEL_SIZES = {MODEL_ID: 256, "google/ddpm-cifar10-32": 32}
CHECKPOINTS = (0, 50, 100, 200, 350, 500, 700, 1000)

# Gradio may resume a generator on a different worker thread between yields.
# A primitive Lock can be released there; RLock requires the acquiring thread.
_MODEL_LOCK = threading.Lock()
_CACHE_LOCK = threading.Lock()
_PIPELINES = {}
_PIPELINE = None
_TORCH = None
_DEVICE = None


def load_model(device=None, model_id=MODEL_ID):
    """Load and cache the diffusion pipeline, returning it and its device."""
    global _PIPELINE, _TORCH, _DEVICE

    if model_id not in MODEL_SIZES:
        raise ValueError(f"Unsupported model: {model_id}")
    with _CACHE_LOCK:
        try:
            import torch
            from diffusers import DDPMPipeline
        except ImportError as error:
            raise RuntimeError(
                "Diffusers dependencies are missing from this Python environment. "
                "Install them with `python -m pip install -r requirements.txt`. "
                f"Import detail: {error}"
            ) from error

        if device is None:
            if torch.cuda.is_available():
                device = "cuda"
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                device = "mps"
            else:
                device = "cpu"
        device = str(device)
        key = (model_id, device)
        if key not in _PIPELINES:
            print(f"Loading {model_id} on device: {device}")
            pipeline = DDPMPipeline.from_pretrained(
                model_id, cache_dir=Path(os.environ.get("MODEL_CACHE_DIR", ROOT / "models")))
            if device == "cuda" and os.environ.get("SPACES_ZERO_GPU", "").lower() in {"1", "true", "yes"}:
                # ZeroGPU packs CUDA aliases at startup. Give it owned CPU
                # storage rather than tensors backed by mapped checkpoint files.
                for component in pipeline.components.values():
                    if isinstance(component, torch.nn.Module):
                        component._apply(lambda tensor: tensor.clone())
            pipeline.to(device)
            _PIPELINES[key] = pipeline
        _PIPELINE = _PIPELINES[key]
        _TORCH = torch
        _DEVICE = device
        return _PIPELINE, _TORCH, _DEVICE


def is_model_loaded():
    return _PIPELINE is not None


def sample_to_image(sample):
    """Convert a model sample tensor into a PIL image."""
    image = sample[0].detach().float().cpu().clamp(-1, 1)
    image = ((image + 1) * 127.5).round().byte().permute(1, 2, 0).numpy()
    return Image.fromarray(image, mode="RGB")


def generate_images(seed, checkpoints=CHECKPOINTS, initial_image=None, device=None, model_id=MODEL_ID):
    """Yield every remaining-step count and an image at requested checkpoints."""
    with _MODEL_LOCK:
        pipeline, torch, device = load_model(device=device, model_id=model_id)
        scheduler = pipeline.scheduler
        training_steps = scheduler.config.num_train_timesteps
        scheduler.set_timesteps(training_steps, device=device)
        # MPS has no torch.Generator implementation. Keep its random stream
        # on CPU; Diffusers also accepts this generator for scheduler noise.
        random_device = "cpu" if device == "mps" else device
        generator = torch.Generator(device=random_device).manual_seed(seed)
        unet = pipeline.unet
        sample_size = unet.config.sample_size
        if isinstance(sample_size, (tuple, list)):
            height, width = sample_size
        else:
            height = width = sample_size
        if initial_image is None:
            sample = torch.randn(
                (1, unet.config.in_channels, height, width),
                generator=generator,
                device=random_device,
                dtype=unet.dtype,
            ).to(device) * scheduler.init_noise_sigma
        else:
            import numpy as np

            pixels = np.asarray(initial_image.convert("RGB").resize((width, height))).copy()
            sample = torch.from_numpy(pixels).permute(2, 0, 1).unsqueeze(0)
            sample = sample.to(device=device, dtype=unet.dtype) / 127.5 - 1

        yield (
            training_steps,
            sample_to_image(sample) if training_steps in checkpoints else None,
        )

        with torch.inference_mode():
            for index, timestep in enumerate(scheduler.timesteps):
                model_output = unet(sample, timestep).sample
                sample = scheduler.step(
                    model_output,
                    timestep,
                    sample,
                    generator=generator,
                ).prev_sample
                remaining_steps = training_steps - index - 1
                image = sample_to_image(sample) if remaining_steps in checkpoints else None
                yield remaining_steps, image


def generate_image(seed, model_id=MODEL_ID):
    """Generate and return the final image for a seed."""
    final_image = None
    for _, image in generate_images(seed, checkpoints=(0,), model_id=model_id):
        if image is not None:
            final_image = image
    if final_image is None:
        raise RuntimeError("The diffusion pipeline did not produce an image.")
    return final_image


def main():
    parser = argparse.ArgumentParser(description="Generate an unconditional DDPM image.")
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed (defaults to a randomly generated seed)",
    )
    parser.add_argument("--model", choices=list(MODEL_SIZES), default=MODEL_ID)
    args = parser.parse_args()
    seed = args.seed if args.seed is not None else secrets.randbelow(2**31)
    if seed < 0 or seed >= 2**63:
        parser.error("seed must be between 0 and 2^63 - 1")

    print(f"Generating with seed: {seed}")
    image = generate_image(seed, model_id=args.model)

    import matplotlib.pyplot as plt

    plt.imshow(image)
    plt.title(f"{args.model} — seed {seed}")
    plt.axis("off")
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()