#!/usr/bin/env python3
"""
rec7_render.py — "cover" a real recording through M3.

Reads the recording with rec7, builds the renderer condition, and runs
the M3 pipeline with that condition injected in place of the language
model's states. The LM is skipped entirely; timbre, guitar and
performance come from the recording.

    python rec7_render.py song.mp3 --m3 /path/to/minimax_music3 \
        --lyrics lyrics.txt --out cover.flac

Caption and lyrics are still fed to the pipeline (it expects them), but
the states drive the audio, so they matter little. 30 steps, guidance
as the pipeline defaults.
"""
import argparse
import sys
from pathlib import Path

import numpy as np
import torch

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


def render_with_condition(pipe, cond, caption, lyrics, steps=30, seed=7):
    """Run the M3 pipeline with cond [T, 32768] injected as the
    semantic step's frame hiddens. Flat [1, T, 8*4096] is what the
    condition encoder expects."""
    from diffusers.modular_pipelines.minimax_music3.encoders import (
        MiniMaxMusic3SemanticGenerationStep as SemStep)
    orig = SemStep.__call__
    T = cond.shape[0]
    inject = cond.reshape(1, T, 8 * H_DIM).to("cuda", torch.bfloat16)

    def patched(self, components, state):
        bs = self.get_block_state(state)
        bs.frame_hiddens = inject
        self.set_block_state(state, bs)
        return components, state

    SemStep.__call__ = patched
    try:
        audio = pipe(prompt=caption, lyrics=lyrics, audio_duration=T / 25.0,
                     num_inference_steps=steps,
                     generator=torch.Generator("cuda").manual_seed(seed),
                     output="audios")[0]
    finally:
        SemStep.__call__ = orig
    arr = np.asarray(audio, dtype=np.float32)
    return arr[None, :] if arr.ndim == 1 else arr


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("audio", type=Path)
    ap.add_argument("--m3", type=Path, required=True)
    ap.add_argument("--weights", type=Path, default=HERE / "weights")
    ap.add_argument("--lyrics", type=Path, default=None)
    ap.add_argument("--caption", type=str, default="acoustic song")
    ap.add_argument("--seconds", type=float, default=60.0,
                    help="length to cover (0 = whole file)")
    ap.add_argument("--steps", type=int, default=30)
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--out", type=Path, default=None)
    a = ap.parse_args()
    for k in ("audio", "m3", "weights", "out", "lyrics"):
        v = getattr(a, k, None)
        if isinstance(v, Path):
            setattr(a, k, v.expanduser().resolve())
    if not a.audio.exists():
        print(f"audio not found: {a.audio}")
        return 1

    import soundfile as sf
    import torchaudio
    from rec7_model import (load_rec7, load_dav, lm_semantic_tables,
                            encode_audio, read_states, states_to_streams)
    from diffusers import ModularPipeline
    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
    trunk, head = load_rec7(a.weights, dev)
    h = read_states(trunk, head, z, dev)
    del trunk, head
    torch.cuda.empty_cache()

    pipe = ModularPipeline.from_pretrained(str(a.m3))
    try:
        pipe.load_components(
            names=["tokenizer", "rvq_depth_decoder", "condition_encoder",
                   "transformer", "scheduler", "vocoder", "guider"],
            dtype=torch.bfloat16)
    except TypeError:
        pipe.load_components(dtype=torch.bfloat16)
    pipe.to(dev)
    W_head, W_embed = lm_semantic_tables(a.m3, dev)
    _codes, cond = states_to_streams(h, pipe.rvq_depth_decoder,
                                     W_head, W_embed, dev)

    # the pipeline's text encoder refuses empty lyrics; the injected
    # states drive the audio, so a placeholder is fine when none are given
    lyr = a.lyrics.read_text(encoding="utf-8").strip() if a.lyrics else ""
    if not lyr:
        lyr = "[verse]\nla la la la la\nla la la la la\n"
    arr = render_with_condition(pipe, cond, a.caption, lyr, a.steps, a.seed)
    if arr.shape[0] == 1:
        arr = np.repeat(arr, 2, axis=0)
    p = a.out or a.audio.with_suffix(".cover.flac")
    sf.write(str(p), arr.T, 44100, subtype="PCM_24")
    print(f"{p}: {arr.shape[-1] / 44100:.1f}s")
    return 0


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