#!/usr/bin/env python3
"""YuE2 song generation on Apple Silicon with MLX.

    python generate.py --model ~/Desktop/Yue2-3B-MLX --style "indie pop, warm vocal" \
        --lyrics-file lyrics.txt --cot off --seed 831001 --out song.wav

Ports the yue2_infer protocol: symbolic ABC planning (cot=melody/full), semantic
codec AR decode with CFG, 32-step midpoint flow matching, tiled VAE decode.
"""
from __future__ import annotations

import argparse
import base64
import json
import math
import sys
import time
import unicodedata
import wave
from dataclasses import dataclass
from pathlib import Path

import mlx.core as mx
import numpy as np

from yue2_model import KVCache, Yue2Model, load_model
from yue2_vae import load_vae

EOD = 151643
ABC_START, ABC_END = 151847, 151848
MUSIC_START, MUSIC_END = 151851, 151852
CODEC_OFFSET, CODEC_SIZE = 151853, 32768
VOCAB, CONTEXT = 184704, 24576
SAMPLE_RATE = 48000
INSTRUCTIONS = {
    "off": "Generate music with codec tokens from the given conditions.",
    "melody": "Generate a melody-only ABC transcription without chord symbols, then generate music with codec tokens from the given conditions.",
    "full": "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions.",
}


class Tokenizer:
    """The frozen Qwen text/ABC BPE (not the audio codec)."""

    def __init__(self, merge_file: Path):
        import tiktoken
        ranks = {base64.b64decode(t): int(r) for t, r in
                 (line.split() for line in Path(merge_file).read_bytes().splitlines() if line)}
        specials = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "<R>", "<S>", "<X>", "<mask>", "<sep>"]
        specials += [f"<extra_{i}>" for i in range(200)]
        specials[204:206] = ["<abc>", "</abc>"]
        pattern = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
        self._enc = tiktoken.Encoding("YuE2", pat_str=pattern, mergeable_ranks=ranks,
                                      special_tokens={s: i + len(ranks) for i, s in enumerate(specials)})

    def encode(self, text):
        return self._enc.encode_ordinary(unicodedata.normalize("NFC", text))

    def decode(self, ids):
        return self._enc.decode([int(i) for i in ids if 0 <= i < self._enc.n_vocab], errors="replace")


@dataclass(frozen=True)
class Sampling:
    temperature: float = 1.0
    top_p: float = 0.95
    top_k: int = 100
    repetition_penalty: float = 1.2
    penalty_window: int = 50
    min_tokens: int = 200
    max_tokens: int = 9000


def request_text(style, lyrics, cot):
    return f"{INSTRUCTIONS[cot]}\n[Tags]\n{style}\n[Lyrics]\n{lyrics}\n"


def token_prefix(tok, style, lyrics, cot, abc_ids=None):
    base = [EOD] + tok.encode(request_text(style, lyrics, cot))
    if cot == "off":
        return base + [ABC_START, ABC_END, MUSIC_START]
    if abc_ids is None:
        return base + [ABC_START]
    return base + [ABC_START] + list(abc_ids) + [ABC_END, MUSIC_START]


def negative_prefix(tok, cot, abc_ids):
    base = [EOD] + tok.encode(INSTRUCTIONS[cot])
    if cot == "off":
        return base + [MUSIC_START]
    return base + [ABC_START] + list(abc_ids) + [ABC_END, MUSIC_START]


# ── AR sampling ──────────────────────────────────────────────────────────────

def allowed_mask(phase, dtype):
    idx = mx.arange(VOCAB)
    end = ABC_END if phase == "abc" else MUSIC_END
    ok = (idx < EOD) if phase == "abc" else ((idx >= CODEC_OFFSET) & (idx < CODEC_OFFSET + CODEC_SIZE))
    return mx.where(ok | (idx == end), 0.0, -mx.inf).astype(dtype)


def distribution(logits, s: Sampling, history, step, phase, allowed, legacy_off):
    # Historical cot=off arithmetic stays in BF16; symbolic/CFG paths upcast.
    scores = logits if legacy_off else logits.astype(mx.float32)
    end = ABC_END if phase == "abc" else MUSIC_END
    scores = scores + allowed
    if step < s.min_tokens:
        scores[end] = -mx.inf
    recent = history[-s.penalty_window:]
    if s.repetition_penalty != 1.0 and recent:
        freq = mx.zeros((VOCAB,), scores.dtype).at[mx.array(recent)].add(1)
        alpha = mx.power(mx.array(s.repetition_penalty, scores.dtype), freq)
        scores = mx.where(scores < 0, scores * alpha, scores / alpha)
    if s.temperature == 0:
        return scores
    if s.temperature != 1:
        scores = scores / s.temperature
    kth = VOCAB - min(s.top_k, VOCAB)
    threshold = mx.partition(scores, kth)[kth]
    scores = mx.where(scores < threshold, -mx.inf, scores)
    if s.top_p < 1:
        order = mx.argsort(-scores)
        values = scores[order]
        probs = mx.softmax(values.astype(mx.float32), axis=-1)
        removed = (mx.cumsum(probs) - probs) > s.top_p
        removed = removed & (mx.arange(VOCAB) >= (3 if legacy_off else 1))
        scores = mx.put_along_axis(scores, order, mx.where(removed, -mx.inf, values), axis=0)
    return scores


def generate_tokens(model: Yue2Model, prefix, s: Sampling, seed, phase, negative=None, cfg_scale=1.0,
                    legacy_off=False, on_token=None):
    """Return (ids, truncated). Two KV caches when CFG is active (cond / uncond)."""
    if len(prefix) + s.max_tokens > CONTEXT or (negative and len(negative) + s.max_tokens > CONTEXT):
        raise ValueError("Prefix + generation budget exceeds the 24576 context")
    if cfg_scale != 1 and negative is None:
        raise ValueError("CFG requires a negative prefix")
    n_layers = len(model.model.layers)
    caches = [[KVCache() for _ in range(n_layers)] for _ in range(2 if cfg_scale != 1 else 1)]
    cond = model.ar_step(mx.array([prefix]), caches[0])
    uncond = model.ar_step(mx.array([negative]), caches[1]) if cfg_scale != 1 else None
    end = ABC_END if phase == "abc" else MUSIC_END
    allowed = allowed_mask(phase, cond.dtype if legacy_off else mx.float32)
    key = mx.random.key(seed)
    history, eos = [], False
    for step in range(s.max_tokens):
        logits = cond if cfg_scale == 1 else uncond + cfg_scale * (cond - uncond)
        scores = distribution(logits[0], s, history, step, phase, allowed, legacy_off)
        if s.temperature == 0:
            token = int(mx.argmax(scores).item())
        else:
            key, sub = mx.random.split(key)
            token = int(mx.random.categorical(scores.astype(mx.float32), key=sub).item())
        if on_token is not None:
            on_token(phase, token)
        if token == end:
            eos = True
            break
        history.append(token)
        if step + 1 < s.max_tokens:
            nxt = mx.array([[token]])
            cond = model.ar_step(nxt, caches[0])
            if uncond is not None:
                uncond = model.ar_step(nxt, caches[1])
    return history, not eos


# ── NAR flow matching ────────────────────────────────────────────────────────

def chunk_ranges(frames, prefix_tokens, context=CONTEXT):
    size = min((context - prefix_tokens - 3) // 2, CONTEXT)
    if frames < 1 or size < 1:
        raise ValueError("Empty codec or prefix leaves no acoustic context")
    return [(a, min(a + size, frames)) for a in range(0, frames, size)]


def _logit(t):
    return max(-20.0, min(20.0, math.log(t / (1 - t)))) if 0 < t < 1 else (20.0 if t >= 1 else -20.0)


def synthesize(model: Yue2Model, prefix, codec, seed, steps=32, noise=None, on_progress=None):
    """Return latents [frames,64] float32 via midpoint ODE from t=1 → 0, solved per context chunk."""
    if noise is None:
        noise = mx.random.normal((len(codec), 64), key=mx.random.key(seed))
    dt = 1.0 / steps
    out = []
    for a, b in chunk_ranges(len(codec), len(prefix)):
        ar_tokens = prefix + [c + CODEC_OFFSET for c in codec[a:b]] + [MUSIC_END]
        ar_cache = model.nar_prefill(ar_tokens)
        state = noise[a:b].astype(mx.bfloat16)
        for step in range(steps):
            t = 1.0 - step * dt
            v1 = model.nar_velocity(state, _logit(t), ar_cache, len(ar_tokens))
            mid = state - v1 * (dt / 2)
            state = state - model.nar_velocity(mid, _logit(t - dt / 2), ar_cache, len(ar_tokens)) * dt
            mx.eval(state)
            if on_progress is not None:
                on_progress(step + 1, steps)
        out.append(state.astype(mx.float32))
    return mx.concatenate(out)


# ── Pipeline ─────────────────────────────────────────────────────────────────

class Yue2Pipeline:
    def __init__(self, path, log=print):
        self.path, self.log = Path(path), log
        gen = json.loads((self.path / "yue2_generation_config.json").read_text())
        self.abc_sampling, self.semantic_sampling = Sampling(**gen["abc"]), Sampling(**gen["semantic"])
        self.ode_steps = gen.get("ode_steps", 32)
        self.tokenizer = Tokenizer(self.path / "qwen.tiktoken")
        self.vae = load_vae(self.path)  # small; load first so a bad file fails before minutes of AR work
        self.model = load_model(self.path)

    def _progress(self, label):
        count, start = [0], time.perf_counter()

        def on_token(_phase, _token):
            count[0] += 1
            if count[0] % 200 == 0:
                self.log(f"[{label}] {count[0]} tokens, {count[0] / (time.perf_counter() - start):.1f} tok/s")
        return on_token

    def __call__(self, style, lyrics, cot="full", seed=831001, abc=None, cfg_scale=None,
                 semantic_sampling=None, steps=None, on_latents=None):
        if cot not in INSTRUCTIONS:
            raise ValueError("cot must be off, melody or full")
        if abc is not None and cot == "off":
            raise ValueError("External ABC requires cot=melody/full")
        tok = self.tokenizer
        # 1. symbolic plan
        abc_ids, abc_text = [], None
        if cot != "off":
            if abc is not None:
                abc_ids, abc_text = tok.encode(abc), abc
            else:
                self.log("[plan] generating ABC score")
                abc_ids, truncated = generate_tokens(self.model, token_prefix(tok, style, lyrics, cot),
                                                     self.abc_sampling, seed, "abc", on_token=self._progress("abc"))
                abc_text = tok.decode(abc_ids)
                if truncated:
                    self.log("[plan] ABC hit max_tokens")
        prefix = token_prefix(tok, style, lyrics, cot, abc_ids)
        # 2. semantic codec tokens
        guidance = (1.01 if cot == "off" else 1.0) if cfg_scale is None else cfg_scale
        negative = negative_prefix(tok, cot, abc_ids) if guidance != 1 else None
        self.log(f"[semantic] prefix {len(prefix)} tokens, cfg {guidance}")
        ids, truncated = generate_tokens(self.model, prefix, semantic_sampling or self.semantic_sampling, seed,
                                         "semantic", negative, guidance, legacy_off=cot == "off",
                                         on_token=self._progress("semantic"))
        if truncated:
            self.log("[semantic] hit max_tokens")
        codec = [t - CODEC_OFFSET for t in ids]
        if not codec:
            raise RuntimeError("Semantic stage produced no codec tokens")
        # 3. acoustic latents
        self.log(f"[nar] {len(codec)} frames ({len(codec) * 1920 / SAMPLE_RATE:.1f}s), {steps or self.ode_steps} midpoint steps")
        latents = synthesize(self.model, prefix, codec, seed, steps or self.ode_steps,
                             on_progress=lambda i, n: i % 8 == 0 and self.log(f"[nar] step {i}/{n}"))
        if on_latents is not None:
            on_latents(latents)
        # 4. waveform
        self.log("[vae] decoding")
        return self.decode(latents), {"abc": abc_text, "codec": codec, "latents": latents, "prefix": prefix}

    def decode(self, latents):
        return mx.clip(self.vae.decode_tiled(latents), -1, 1)


def write_wav(path, audio, sample_rate=SAMPLE_RATE):
    pcm = (np.clip(np.array(audio, dtype=np.float32), -1, 1) * 32767).astype("<i2")
    with wave.open(str(path), "wb") as f:
        f.setnchannels(pcm.shape[1])
        f.setsampwidth(2)
        f.setframerate(sample_rate)
        f.writeframes(pcm.tobytes())


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--model", type=Path, required=True, help="MLX package dir from convert.py")
    parser.add_argument("--style", required=True)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--lyrics")
    group.add_argument("--lyrics-file", type=Path)
    parser.add_argument("--cot", default="full", choices=list(INSTRUCTIONS))
    parser.add_argument("--abc-file", type=Path, help="external ABC score (cot=melody/full)")
    parser.add_argument("--seed", type=int, default=831001)
    parser.add_argument("--cfg-scale", type=float)
    parser.add_argument("--steps", type=int, help="NAR midpoint steps (default from generation config)")
    parser.add_argument("--max-semantic-tokens", type=int, help="cap song length in codec frames (40 ms each)")
    parser.add_argument("--out", type=Path, default=Path("yue2.wav"))
    parser.add_argument("--decode-latents", type=Path, help="skip generation; decode a saved <out>.latents.npy")
    args = parser.parse_args()

    lyrics = args.lyrics if args.lyrics is not None else args.lyrics_file.read_text()
    abc = args.abc_file.read_text() if args.abc_file else None
    log = lambda msg: print(msg, file=sys.stderr, flush=True)
    pipe = Yue2Pipeline(args.model, log=log)
    sampling = None
    if args.max_semantic_tokens:
        sampling = Sampling(**{**pipe.semantic_sampling.__dict__, "max_tokens": args.max_semantic_tokens,
                               "min_tokens": min(pipe.semantic_sampling.min_tokens, args.max_semantic_tokens)})
    start = time.perf_counter()
    latents_path = args.out.with_suffix(".latents.npy")
    if args.decode_latents:
        audio, info = pipe.decode(mx.array(np.load(args.decode_latents))), {"abc": None}
    else:
        audio, info = pipe(args.style, lyrics, cot=args.cot, seed=args.seed, abc=abc, cfg_scale=args.cfg_scale,
                           semantic_sampling=sampling, steps=args.steps,
                           on_latents=lambda z: np.save(latents_path, np.array(z)))
    write_wav(args.out, audio)
    if info["abc"] is not None:
        args.out.with_suffix(".abc").write_text(info["abc"])
    log(f"[done] {args.out} {audio.shape[0] / SAMPLE_RATE:.1f}s in {time.perf_counter() - start:.0f}s")


if __name__ == "__main__":
    main()