#!/usr/bin/env python3
"""
RVC 模型推理测试。纯 CPU 模式。
两个音频：参考音频 (提取昔涟音色) + 输入音频 (要转换的声音)
"""
import os, sys, time, warnings
warnings.filterwarnings('ignore')

# 添加 RVC 到路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'rvc'))

import torch
import librosa
import soundfile as sf
import numpy as np
from infer.modules.vc.modules import VC

# === 配置 ===
MODEL_PATH = r"D:\Users\Aska\Documents\G_2333333.pth"
HUBERT_PATH = r"D:\Project\Code\Uni\Cyrene\scripts\voice\rvc\assets\hubert\hubert_base.pt"
REF_AUDIO = r"D:\Project\Code\Uni\Cyrene-Voice-Model\data\cyrene_group01\VoBanks31_0078_004ab065.wav"  # 昔涟参考
INPUT_AUDIO = REF_AUDIO  # 测试: 用同一文件验证模型能正确重建
OUTPUT_AUDIO = r"D:\Users\Aska\Documents\rvc_test_output.wav"

print("=" * 50)
print("RVC 推理测试 (CPU)")
print("=" * 50)

# 检查文件
for label, path in [("模型", MODEL_PATH), ("HuBERT", HUBERT_PATH), ("参考音频", REF_AUDIO)]:
    if not os.path.exists(path):
        print(f"[ERROR] {label} 不存在: {path}")
        sys.exit(1)
    print(f"[OK] {label}: {os.path.getsize(path)/1024/1024:.1f} MB")

# 加载参考音频
print(f"\n[1/3] 加载参考音频...")
ref_audio, ref_sr = librosa.load(REF_AUDIO, sr=40000, mono=True)
print(f"  采样率: {ref_sr}Hz, 时长: {len(ref_audio)/ref_sr:.1f}s")

# 加载输入音频
print(f"[2/3] 加载输入音频...")
input_audio, input_sr = librosa.load(INPUT_AUDIO, sr=40000, mono=True)
print(f"  采样率: {input_sr}Hz, 时长: {len(input_audio)/input_sr:.1f}s")

# 推理
print(f"[3/3] RVC 推理中...")
t0 = time.time()

# 路径中有中文，先复制到临时路径
import tempfile, shutil
tmp_dir = tempfile.mkdtemp()
tmp_model = os.path.join(tmp_dir, "model.pth")
tmp_ref = os.path.join(tmp_dir, "ref.wav")
tmp_input = os.path.join(tmp_dir, "input.wav")
shutil.copy(MODEL_PATH, tmp_model)
sf.write(tmp_ref, ref_audio, 40000)
sf.write(tmp_input, input_audio, 40000)

# 保存当前目录
old_cwd = os.getcwd()
rvc_dir = os.path.join(os.path.dirname(__file__), 'rvc')
os.chdir(rvc_dir)
# RVC 需要 assets/hubert/hubert_base.pt
os.makedirs('assets/hubert', exist_ok=True)
if not os.path.exists('assets/hubert/hubert_base.pt'):
    import shutil as _sh
    _sh.copy(HUBERT_PATH, 'assets/hubert/hubert_base.pt')

try:
    # 初始化 VC pipeline
    vc = VC()
    vc.get_vc(tmp_model, device="cpu", use_jit=False)

    # 转换
    output, output_sr = vc.vc_single(
        sid=0,
        input_audio_path=tmp_input,
        f0_up_key=0,          # 不改变音高
        f0_file=None,
        f0_method="rmvpe",
        file_index="",         # 不使用 index
        file_index2="",
        index_rate=0,
        filter_radius=3,
        resample_sr=40000,
        rms_mix_rate=0.25,
        protect=0.33,
    )

    elapsed = time.time() - t0
    print(f"\n  完成! 耗时: {elapsed:.1f}s")
    print(f"  输出采样率: {output_sr}Hz, 时长: {len(output)/output_sr:.1f}s")

    # 保存
    sf.write(OUTPUT_AUDIO, output, output_sr)
    print(f"  已保存: {OUTPUT_AUDIO}")

finally:
    os.chdir(old_cwd)
    shutil.rmtree(tmp_dir, ignore_errors=True)

print(f"\n试听: {OUTPUT_AUDIO}")