import os
import sys
import traceback
import logging

import numpy as np

from infer.lib.audio import load_audio
from model import rmvpe, fp16

logging.getLogger("numba").setLevel(logging.WARNING)


class FeatureInput(object):
    def __init__(self, exp_dir, samplerate=16000, hop_size=160):
        self.exp_dir = exp_dir
        self.logfile = open("%s/extract_f0_feature.log" % exp_dir, "a+")
        self.fs = samplerate
        self.hop = hop_size
        self.is_half = fp16

        self.f0_bin = 256
        self.f0_max = 1100.0
        self.f0_min = 50.0
        self.f0_mel_min = 1127 * np.log(1 + self.f0_min / 700)
        self.f0_mel_max = 1127 * np.log(1 + self.f0_max / 700)

    def println(self, strr):
        print(strr)
        self.logfile.write("%s\n" % strr)
        self.logfile.flush()

    def compute_f0(self, path, f0_method):
        x = load_audio(path, self.fs)
        # p_len = x.shape[0] // self.hop
        if f0_method == "rmvpe":
            if hasattr(self, "model_rmvpe") == False:
                self.model_rmvpe = rmvpe
            f0 = self.model_rmvpe.infer_from_audio(x, thred=0.03)
        return f0

    def coarse_f0(self, f0):
        f0_mel = 1127 * np.log(1 + f0 / 700)
        f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - self.f0_mel_min) * (
            self.f0_bin - 2
        ) / (self.f0_mel_max - self.f0_mel_min) + 1

        # use 0 or 1
        f0_mel[f0_mel <= 1] = 1
        f0_mel[f0_mel > self.f0_bin - 1] = self.f0_bin - 1
        f0_coarse = np.rint(f0_mel).astype(int)
        assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (
            f0_coarse.max(),
            f0_coarse.min(),
        )
        return f0_coarse

    def go(self, paths, f0_method):
        if len(paths) == 0:
            self.println("no-f0-todo")
        else:
            self.println("todo-f0-%s" % len(paths))
            n = max(len(paths) // 5, 1)  # 每个进程最多打印5条
            for idx, (inp_path, opt_path1, opt_path2) in enumerate(paths):
                try:
                    if idx % n == 0:
                        self.println(
                            "f0ing,now-%s,all-%s,-%s" % (idx, len(paths), inp_path)
                        )
                    if (
                        os.path.exists(opt_path1 + ".npy") == True
                        and os.path.exists(opt_path2 + ".npy") == True
                    ):
                        continue
                    featur_pit = self.compute_f0(inp_path, f0_method)
                    np.save(
                        opt_path2,
                        featur_pit,
                        allow_pickle=False,
                    )  # nsf
                    coarse_pit = self.coarse_f0(featur_pit)
                    np.save(
                        opt_path1,
                        coarse_pit,
                        allow_pickle=False,
                    )  # ori
                except:
                    self.println(
                        "f0fail-%s-%s-%s" % (idx, inp_path, traceback.format_exc())
                    )

    def run(self):
        inp_root = "%s/1_16k_wavs" % (self.exp_dir)
        opt_root1 = "%s/2a_f0" % (self.exp_dir)
        opt_root2 = "%s/2b-f0nsf" % (self.exp_dir)
        os.makedirs(opt_root1, exist_ok=True)
        os.makedirs(opt_root2, exist_ok=True)

        paths = []
        for name in sorted(list(os.listdir(inp_root))):
            inp_path = "%s/%s" % (inp_root, name)
            if "spec" in inp_path:
                continue
            opt_path1 = "%s/%s" % (opt_root1, name)
            opt_path2 = "%s/%s" % (opt_root2, name)
            paths.append([inp_path, opt_path1, opt_path2])

        self.go(paths, "rmvpe")


if __name__ == "__main__":
    now_dir = os.getcwd()
    sys.path.append(now_dir)

    n_part = int(sys.argv[1])
    i_part = int(sys.argv[2])
    os.environ["CUDA_VISIBLE_DEVICES"] = sys.argv[3]
    exp_dir = sys.argv[4]

    featureInput = FeatureInput(exp_dir)
    featureInput.run()