import os,sys,traceback
device=sys.argv[1]
n_part=int(sys.argv[2])
i_part=int(sys.argv[3])
if len(sys.argv) == 5:
    exp_dir=sys.argv[4]
else:
    i_gpu=sys.argv[4]
    exp_dir=sys.argv[5]
    os.environ["CUDA_VISIBLE_DEVICES"]=str(i_gpu)

import torch
import torch.nn.functional as F
import soundfile as sf
import numpy as np
from fairseq import checkpoint_utils

f = open("%s/extract_f0_feature.log"%exp_dir, "a+")
def printt(strr):
    print(strr)
    f.write("%s\n" % strr)
    f.flush()
printt(sys.argv)
model_path = "hubert_base.pt"

printt(exp_dir)
wavPath = "%s/1_16k_wavs"%exp_dir
outPath = "%s/3_feature256"%exp_dir
os.makedirs(outPath,exist_ok=True)
# wave must be 16k, hop_size=320
def readwave(wav_path, normalize=False):
    wav, sr = sf.read(wav_path)
    assert sr == 16000
    feats = torch.from_numpy(wav).float()
    if feats.dim() == 2:  # double channels
        feats = feats.mean(-1)
    assert feats.dim() == 1, feats.dim()
    if normalize:
        with torch.no_grad():
            feats = F.layer_norm(feats, feats.shape)
    feats = feats.view(1, -1)
    return feats
# HuBERT model
printt("load model(s) from {}".format(model_path))
models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
    [model_path],
    suffix="",
)
model = models[0]
model = model.to(device)
printt("move model to "+device)
if device != "cpu": model = model.half()
model.eval()

todo=sorted(list(os.listdir(wavPath)))[i_part::n_part]
n = max(1,len(todo) // 10)  # 最多打印十条
if(len(todo)==0):printt("no-feature-todo")
else:
    printt("all-feature-%s"%len(todo))
    for idx,file in enumerate(todo):
        try:
            if file.endswith(".wav"):
                wav_path = "%s/%s"%(wavPath,file)
                out_path = "%s/%s"%(outPath,file.replace("wav","npy"))

                if(os.path.exists(out_path)):continue

                feats = readwave(wav_path, normalize=saved_cfg.task.normalize)
                padding_mask = torch.BoolTensor(feats.shape).fill_(False)
                inputs = {
                    "source": feats.half().to(device) if device != "cpu" else feats.to(device),
                    "padding_mask": padding_mask.to(device),
                    "output_layer": 9,  # layer 9
                }
                with torch.no_grad():
                    logits = model.extract_features(**inputs)
                    feats = model.final_proj(logits[0])

                feats = feats.squeeze(0).float().cpu().numpy()
                if(np.isnan(feats).sum()==0):
                    np.save(out_path, feats, allow_pickle=False)
                else:
                    printt("%s-contains nan"%file)
                if (idx % n == 0):printt("now-%s,all-%s,%s,%s"%(len(todo),idx,file,feats.shape))
        except:
            printt(traceback.format_exc())
    printt("all-feature-done")