from __future__ import annotations

import json
import secrets
from itertools import pairwise
from pathlib import Path

import gradio as gr
import torch

from audio_dit import CHANNELS, FRAMES, MAX_TAGS, load_model
from same_l_decoder import SAMPLE_RATE, load_decoder

DEVICE = ("cuda" if torch.cuda.is_available() else
          "mps" if torch.backends.mps.is_available() else "cpu")
DTYPE = torch.float32 if DEVICE == "cpu" else torch.float16
AUTOCAST = torch.bfloat16 if DEVICE == "cuda" else torch.float16

STEPS = 50
DEFAULT_CFG = 3.5
TIME_SHIFT = 2.0

TAG_GROUPS = json.loads((Path(__file__).parent / "tags.json").read_text(encoding="utf-8"))
METADATA_TAGS = TAG_GROUPS["metadata"]
OTHER_TAGS = TAG_GROUPS["other"]
SYNTHETIC_TAG_ID = METADATA_TAGS["synthetic"]
TAGS = METADATA_TAGS | OTHER_TAGS
del TAGS["synthetic"]  # Source conditioning is added automatically, not shown in the UI.

TIMESTEPS = [1 - i / STEPS for i in range(STEPS + 1)]
TIMESTEPS = [TIME_SHIFT * t / (1 + (TIME_SHIFT - 1) * t) for t in TIMESTEPS]

model, latent_mean, latent_std = load_model(DEVICE)
decoder = load_decoder(device=DEVICE, dtype=DTYPE)


@torch.no_grad()
def generate(tag_names: list[str], cfg_scale: float, seed: int):
    if not tag_names:
        raise gr.Error("Pick at least one tag.")
    seed = secrets.randbits(63) if seed == -1 else seed
    torch.manual_seed(seed)
    ids = [TAGS[name] for name in tag_names]
    if all(name in METADATA_TAGS for name in tag_names):
        if len(ids) == MAX_TAGS:
            raise gr.Error(f"Choose at most {MAX_TAGS - 1} metadata tags.")
        ids.append(SYNTHETIC_TAG_ID)
    tags = torch.tensor([ids + [0] * (MAX_TAGS - len(ids))], device=DEVICE)
    null = torch.zeros_like(tags)
    x = torch.randn(1, CHANNELS, FRAMES, device=DEVICE)

    for i, (t, t_next) in enumerate(pairwise(TIMESTEPS)):
        # Guidance alternates between the shallow path-drop branch and a full
        # unconditional pass as the weak model to steer away from.
        weak = (model.shallow, model)[i % 2]
        with torch.autocast(DEVICE, dtype=AUTOCAST, enabled=DEVICE != "cpu"):
            cond = model(x, t, tags).float()
            uncond = weak(x, t, null).float()
        x = x + (t_next - t) * (uncond + cfg_scale * (cond - uncond))

    audio = decoder.decode((x * latent_std + latent_mean).to(DTYPE))
    return SAMPLE_RATE, audio[0].float().clamp(-1, 1).T.cpu().numpy()


with gr.Blocks(title="Audio DiT") as demo:
    gr.Markdown("# Audio DiT\nChoose 1-8 tags.")
    tag_box = gr.Dropdown(choices=list(TAGS), value=[], multiselect=True,
                          max_choices=MAX_TAGS, label="Tags", filterable=True)
    cfg_scale = gr.Slider(2.0, 6.0, value=DEFAULT_CFG, step=0.5,
                          label="CFG scale")
    seed = gr.Number(value=-1, precision=0, minimum=-1,
                     label="Seed (-1 = random)")
    button = gr.Button("Generate", variant="primary")
    audio_out = gr.Audio(label="Output", interactive=False)
    button.click(generate, inputs=[tag_box, cfg_scale, seed], outputs=audio_out,
                 concurrency_limit=1)

if __name__ == "__main__":
    demo.queue().launch()