"""Portable reference example for the exact native12-codebook CFG backend.

pip install torch transformers safetensors huggingface_hub soundfile numpy
Use the versions/dependencies documented by the pinned model repositories.
No API key is embedded. Original and rank4 use different base-loading paths.
"""
import argparse,importlib.util,json
from pathlib import Path

def module(path):
    s=importlib.util.spec_from_file_location('published_native_model_api',path);m=importlib.util.module_from_spec(s);s.loader.exec_module(m);return m

def main():
    p=argparse.ArgumentParser(description=__doc__);p.add_argument('--model',choices=['original15','rank4'],required=True)
    p.add_argument('--recipe-take',type=Path,required=True,help='One take object selected from the public scenario JSON')
    group=p.add_mutually_exclusive_group();group.add_argument('--reference-wav',type=Path)
    group.add_argument('--reference-codes',type=Path,help='Published references/codes.npz for the exact fixed reference codes')
    p.add_argument('--output',type=Path,required=True);a=p.parse_args()
    import numpy as np,soundfile as sf,torch
    from huggingface_hub import snapshot_download
    from native_cfg import generate
    if a.output.exists():raise FileExistsError(a.output)
    spec={'original15':('laion/Humaneness-Voice-Large-1.5-Alpha','f5fdc4f06db63b1d755252997f0e8135121db264'),
          'rank4':('laion/Humaneness-Voice-Large-1.5-RealSpeech-LoRAs','fed33e554f094fcdc89fc0124aed2f9d42018714')}
    repo,revision=spec[a.model];folder=Path(snapshot_download(repo,revision=revision))
    api=module(folder/'infer.py')
    model,proc,base=api.load(directory=folder,**({'rank':4} if a.model=='rank4' else {}))
    base.load_codec(proc,'OpenMOSS-Team/MOSS-Audio-Tokenizer-v2')
    row=json.loads(a.recipe_take.read_text());reference=None
    if row['model']!=a.model:raise ValueError('Take and model differ')
    if bool(a.reference_wav or a.reference_codes)!=row['reference']:raise ValueError('Respect the recorded reference/no-reference arm')
    if a.reference_codes:
        reference=torch.as_tensor(np.load(a.reference_codes)[row['language']],dtype=torch.long)
        if reference.ndim!=2 or reference.shape[1]!=12 or not 63<=len(reference)<=187:raise ValueError('Invalid reference-code shape/window')
    if a.reference_wav:
        wave,sr=sf.read(a.reference_wav,dtype='float32',always_2d=True)
        if len(wave)/sr<5.04:raise ValueError('At least5.04seconds of one speaker required')
        wave=wave[:int(14.96*sr)]
        with torch.inference_mode():reference=proc.encode_audios_from_wav([torch.from_numpy(wave.T.copy())],sr,n_vq=12)[0][:187]
    frames=round(row['requested_seconds']*12.5)
    kwargs=dict(text=row['text'],frames=frames,language=row['language'],reference=reference)
    cond=base.build_user(proc,prompt=row['prompt'],**kwargs);neutral=base.build_user(proc,prompt=row['neutral_prompt'],**kwargs)
    seeds=[1701,1702,1703,1704,1705]
    if row['seed'] not in seeds:raise ValueError('Unknown fixed search seed')
    codes,timing=generate(model,proc,cond,neutral,seeds,row['cfg'],min(375,frames+45))
    audio,sr=base.decode(proc,codes[seeds.index(row['seed'])]);a.output.parent.mkdir(parents=True,exist_ok=True);sf.write(a.output,audio,sr)
    print(json.dumps(timing))
if __name__=='__main__':main()
