from __future__ import annotations

import json
import secrets
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
MODEL_DTYPE = (torch.bfloat16 if DEVICE == "cuda" else
               torch.float16 if DEVICE == "mps" else torch.float32)
AUTOCAST = MODEL_DTYPE

STEPS = 50
DEFAULT_CFG = 4.0
TIME_SHIFT = 3.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"]
AESTHETIC_TAG_ID = OTHER_TAGS["aesthetic"]
TAGS = METADATA_TAGS | OTHER_TAGS
del TAGS["synthetic"]
del TAGS["aesthetic"]

VISIBLE_MAX_TAGS = 6
SAMPLE_MP3S = {
    path.stem: path
    for path in sorted((Path(__file__).parent / "samples").glob("*.mp3"))
}

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


def sampling_timesteps(device, dtype):
    timesteps = torch.linspace(1.0, 0.0, STEPS + 1, device=device, dtype=dtype)[:-1]
    return TIME_SHIFT * timesteps / (1 + (TIME_SHIFT - 1) * timesteps)


@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 len(ids) > VISIBLE_MAX_TAGS:
        raise gr.Error(f"Choose at most {VISIBLE_MAX_TAGS} tags.")
    if all(name in METADATA_TAGS for name in tag_names):
        ids.append(SYNTHETIC_TAG_ID)
    else:
        ids.append(AESTHETIC_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)
    timesteps = sampling_timesteps(x.device, x.dtype)

    for i in range(STEPS):
        t = timesteps[i]
        t_next = (timesteps[i + 1] if i + 1 < STEPS else
                  torch.tensor(0.0, device=x.device, dtype=x.dtype))
        # The model predicts clean latents. Guidance alternates between the
        # shallow path-drop branch and a full unconditional pass.
        weak = (model.shallow, model)[i % 2]
        with torch.autocast(DEVICE, dtype=AUTOCAST, enabled=DEVICE != "cpu"):
            cond_x0 = model(x, t, tags).float()
            uncond_x0 = weak(x, t, null).float()
        cond = (x - cond_x0) / t
        uncond = (x - uncond_x0) / t
        x = x - (t - t_next) * (uncond + cfg_scale * (cond - uncond))

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


def load_sample(tag_names: list[str]):
    return SAMPLE_MP3S[tag_names[0]]


with gr.Blocks(title="Audio DiT") as demo:
    gr.Markdown(
        f"# Audio DiT\nChoose 1-{VISIBLE_MAX_TAGS} tags to generate a song, "
        "or click a sample song to load it."
    )
    tag_box = gr.Dropdown(choices=list(TAGS), value=[], multiselect=True,
                          max_choices=VISIBLE_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", format="mp3", interactive=False)
    gr.Examples(
        examples=[[[tag]] for tag in SAMPLE_MP3S],
        inputs=tag_box,
        outputs=audio_out,
        fn=load_sample,
        cache_examples=False,
        run_on_click=True,
        example_labels=list(SAMPLE_MP3S),
        label="Sample songs",
    )
    button.click(generate, inputs=[tag_box, cfg_scale, seed], outputs=audio_out,
                 concurrency_limit=1)

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