#!/usr/bin/env python3
"""
rec7_encode.py — real audio -> M3 states, codes, renderer condition.

    python rec7_encode.py song.mp3 --m3 /path/to/minimax_music3 --out song_rec7.pt

Writes a .pt with:
    h       [T, 4096] float16   the states (what the LM would have produced)
    codes   [T, 8]    int16     c0 (16384-way) + c1..c7 (1024-way each)
    cond    [T, 32768] float16  optional (--cond): the renderer condition
    z       [128, L]  float16   optional (--z): the Flow-VAE latents
25 frames per second. Works on the whole file (any length).
"""
import argparse
import sys
from pathlib import Path

import torch

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("audio", type=Path)
    ap.add_argument("--m3", type=Path, required=True,
                    help="MiniMax-Music3 checkpoint directory")
    ap.add_argument("--weights", type=Path, default=HERE / "weights")
    ap.add_argument("--out", type=Path, default=None)
    ap.add_argument("--seconds", type=float, default=0,
                    help="0 = whole file")
    ap.add_argument("--cond", action="store_true",
                    help="also write the 32768-d renderer condition")
    ap.add_argument("--z", action="store_true", help="also write latents")
    a = ap.parse_args()

    import torchaudio
    from rec7_model import (load_rec7, load_dav, lm_semantic_tables,
                            encode_audio, read_states, states_to_streams)
    dev = "cuda"
    wav, sr = torchaudio.load(str(a.audio))
    if sr != 44100:
        wav = torchaudio.transforms.Resample(sr, 44100)(wav)
    if wav.shape[0] == 1:
        wav = wav.repeat(2, 1)
    wav = wav[:2].float()
    if a.seconds > 0:
        wav = wav[:, : int(a.seconds * 44100)]

    dav = load_dav(a.m3, dev)
    z = encode_audio(dav, wav, dev)
    del dav
    torch.cuda.empty_cache()

    trunk, head = load_rec7(a.weights, dev)
    h = read_states(trunk, head, z, dev)
    del trunk, head
    torch.cuda.empty_cache()

    from diffusers.models.transformers.minimax_music3_rvq_depth_decoder \
        import MiniMaxMusic3RVQDepthDecoder
    depth = MiniMaxMusic3RVQDepthDecoder.from_pretrained(
        str(a.m3 / "rvq_depth_decoder"), torch_dtype=torch.bfloat16
    ).to(dev).eval()
    W_head, W_embed = lm_semantic_tables(a.m3, dev)
    codes, cond = states_to_streams(h, depth, W_head, W_embed, dev)

    out = {"h": h.cpu().to(torch.float16), "codes": codes.to(torch.int16),
           "frame_rate": 25, "source": str(a.audio)}
    if a.cond:
        out["cond"] = cond.cpu().to(torch.float16)
    if a.z:
        out["z"] = z.to(torch.float16)
    p = a.out or a.audio.with_suffix(".rec7.pt")
    torch.save(out, p)
    print(f"{p}: {h.shape[0]} frames ({h.shape[0] / 25:.1f}s)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())