"""Movetoy: a tiny from-scratch text-to-video model, running on a ZeroGPU Space (CPU fallback for local use)."""
import os
import tempfile

import gradio as gr
import imageio
import numpy as np
import torch
from huggingface_hub import hf_hub_download

try:
    import spaces                      # ZeroGPU: a GPU is attached only while a decorated call runs
    gpu = spaces.GPU(duration=lambda prompt, seed, steps, guidance, seconds, progress=None: 10 + 4 * int(seconds))
except ImportError:                    # running locally without the `spaces` package
    def gpu(fn):
        return fn

from tiny_video.model import load_checkpoint, null_tokens, to_uint8, tokenizer_for

torch.set_num_threads(max(1, os.cpu_count() or 1))
MODEL_REPO = os.environ.get("MOVETOY_REPO", "Kompaktion/Movetoy")
DEVICE = "cuda" if (torch.cuda.is_available() or "spaces" in globals()) else "cpu"
model, _ = load_checkpoint(hf_hub_download(MODEL_REPO, "movetoy.pt"), "cpu")
model = model.to(DEVICE).eval()
tokenizer = tokenizer_for(model.cfg)
SCALE = 8      # 32x32 native -> 256x256 display, nearest-neighbour (no interpolation)
FPS = 8
WINDOW = model.cfg.frames   # 16 frames = 2 s: everything the model was trained to produce in one go
OVERLAP = 4                 # frames carried from one window into the next when stitching longer videos


@torch.inference_mode()
def sample_window(cond, steps, guidance, gen, known=None, shift=2.0):
    """Movetoy's Euler sampler for one 16-frame window. If `known` [3,k,H,W] is given, the first k frames are
    pinned to it at every step (re-noised to the current time), so the window continues the previous one."""
    cfg = model.cfg
    dev = next(model.parameters()).device
    x = torch.randn(1, 3, WINDOW, cfg.size, cfg.size, generator=gen).to(dev)   # CPU generator: same video on any device
    if known is not None:
        k = known.shape[1]
        eps_k = torch.randn(1, 3, k, cfg.size, cfg.size, generator=gen).to(dev)
    u = torch.linspace(1, 0, steps + 1, device=dev)
    ts = shift * u / (1 + (shift - 1) * u)
    for i in range(steps):
        if known is not None:
            x[:, :, :k] = (1 - ts[i]) * known + ts[i] * eps_k
        with torch.autocast(dev.type, dtype=torch.bfloat16, enabled=dev.type == "cuda"):
            v = model(torch.cat([x, x]), ts[i].expand(2), cond).float()
        vc, vu = v.chunk(2)
        x = x + (ts[i + 1] - ts[i]) * (vu + guidance * (vc - vu))
    if known is not None:
        x[:, :, :k] = known
    return x.clamp(-1, 1)[0]                                   # [3, 16, H, W]


@gpu
def generate(prompt, seed, steps, guidance, seconds, progress=gr.Progress()):
    """Text -> video. Up to 2 s is one native pass; longer videos stitch 16-frame windows (12 new frames each)."""
    prompt = (prompt or "").strip() or "a red circle and blue circle merge into a purple circle"
    seed = int(seed) if seed is not None and int(seed) >= 0 else int(np.random.randint(0, 2**31 - 1))
    gen = torch.Generator().manual_seed(seed)
    tokens = tokenizer([prompt])
    cond = torch.cat([tokens, null_tokens(tokens, model.cfg.text)]).to(next(model.parameters()).device)
    target = int(min(15, max(2, float(seconds))) * FPS)
    n_windows = 1 + max(0, -(-(target - WINDOW) // (WINDOW - OVERLAP)))
    clip = sample_window(cond, int(steps), float(guidance), gen)
    for w in range(1, n_windows):
        progress(w / n_windows, desc=f"window {w + 1}/{n_windows}")
        nxt = sample_window(cond, int(steps), float(guidance), gen, known=clip[:, -OVERLAP:])
        clip = torch.cat([clip, nxt[:, OVERLAP:]], dim=1)
    frames = to_uint8(clip[:, :target].unsqueeze(0))[0].numpy()   # [T, 32, 32, 3]
    frames = frames.repeat(SCALE, axis=1).repeat(SCALE, axis=2)
    path = os.path.join(tempfile.mkdtemp(), "movetoy.mp4")
    imageio.mimwrite(path, list(frames), fps=FPS, codec="libx264", quality=9, macro_block_size=1)
    note = f"seed {seed} · {len(frames)} frames ({len(frames) / FPS:.1f} s)"
    if n_windows > 1:
        note += f" · stitched from {n_windows} windows of 16 frames"
    return path, note


EXAMPLES = [
    ["a red circle and blue circle merge into a purple circle", 424243, 32, 3.0, 2],
    ["a yellow circle and red circle merge into an orange circle", 7, 32, 3.0, 2],
    ["a big green star falls down at sunset", 11, 32, 3.0, 2],
    ["a pink heart spins in the night sky", 5, 32, 3.0, 2],
    ["a small blue square grows bigger on a white background", 3, 32, 3.0, 2],
    ["a purple dog dances on the moon", 1, 32, 3.0, 2],
]

with gr.Blocks(title="Movetoy") as demo:
    gr.Markdown(
        "# Movetoy\n"
        "A **51M-parameter text-to-video model trained from scratch** (random weights, no pretrained parts, "
        "no downloaded data) on one RTX 4070. It makes 2-second, 16-frame **32×32** clips of coloured shapes, "
        "shown here 8× bigger. It reads *any* text (bytes + hashed words, no vocabulary), but it only *understands* "
        "its toy world: colours, shapes, motions, falls, spins, and two-object merges with colour mixing. "
        "Words it never saw get you a vibe-based blob.\n\n"
        "Runs on ZeroGPU: a 2-second clip takes about a second of GPU. Longer clips (up to 15 s) are stitched from "
        "2-second windows that the model continues one after another (it was only ever trained on 2 s), "
        "so they drift in funny ways.\n\n"
        "**How much training did this take?** ~382,000 procedurally generated captioned clips "
        "(23,874 steps × batch 16; every clip unique, rendered on the fly) — about **212 hours** of 32×32 video, "
        "~112 minutes on one RTX 4070 — just to learn coloured shapes moving and merging. "
        "And it still gets the direction wrong half the time."
    )
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", value="a red circle and blue circle merge into a purple circle", lines=2)
            with gr.Row():
                seed = gr.Number(label="Seed (-1 = random)", value=-1, precision=0)
                steps = gr.Slider(8, 64, value=32, step=1, label="Sampling steps")
                guidance = gr.Slider(1.0, 6.0, value=3.0, step=0.5, label="Guidance")
            seconds = gr.Slider(2, 15, value=2, step=1, label="Length in seconds (over 2 s = stitched, slow on CPU)")
            go = gr.Button("Generate", variant="primary")
        with gr.Column():
            out = gr.Video(label="Movetoy output (32×32, upscaled)", autoplay=True, loop=True)
            used = gr.Markdown()
    go.click(generate, [prompt, seed, steps, guidance, seconds], [out, used], api_name="generate")
    prompt.submit(generate, [prompt, seed, steps, guidance, seconds], [out, used], api_name=False)
    gr.Examples(EXAMPLES, [prompt, seed, steps, guidance, seconds], [out, used], fn=generate, cache_examples=False)
    gr.Markdown(f"Model: [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO})")

demo.queue(max_size=32).launch()