import onnxruntime
import librosa
import numpy as np
import soundfile


def load_audio_fast(path, target_sr):
    # 1. Coba torchaudio (sangat cepat, ~11ms)
    try:
        import torchaudio
        wav, sr = torchaudio.load(path)
        if sr != target_sr:
            import torchaudio.transforms as T
            resampler = T.Resample(sr, target_sr)
            wav = resampler(wav)
        if wav.shape[0] > 1:
            wav = wav.mean(dim=0)
        wav = wav.squeeze().numpy()
        return wav, target_sr
    except Exception:
        pass

    # 2. Coba pydub (cepat, ~80ms)
    try:
        from pydub import AudioSegment
        audio_seg = AudioSegment.from_file(path)
        audio_seg = audio_seg.set_frame_rate(target_sr).set_channels(1)
        wav = np.array(audio_seg.get_array_of_samples(), dtype=np.float32) / 32768.0
        return wav, target_sr
    except Exception:
        pass

    # 3. Fallback ke librosa asli
    return librosa.load(path, sr=target_sr, mono=True)

class ContentVec:
    def __init__(self, vec_path="pretrained/vec-768-layer-12.onnx", device=None):
        print("load model(s) from {}".format(vec_path))
        import onnxruntime as ort
        opts = ort.SessionOptions()
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        opts.intra_op_num_threads = 2
        
        if device == "cpu" or device is None:
            providers = ["CPUExecutionProvider"]
        elif device == "cuda":
            providers = [
                ("CUDAExecutionProvider", {
                    "device_id": 0,
                    "arena_extend_strategy": "kNextPowerOfTwo",
                    "cudnn_conv_algo_search": "EXHAUSTIVE",
                    "do_copy_in_default_stream": True,
                }),
                "CPUExecutionProvider"
            ]
        elif device == "dml":
            providers = ["DmlExecutionProvider"]
        else:
            raise RuntimeError("Unsportted Device")
        self.model = ort.InferenceSession(vec_path, sess_options=opts, providers=providers)

    def __call__(self, wav):
        return self.forward(wav)

    def forward(self, wav):
        feats = wav
        if feats.ndim == 2:  # double channels
            feats = feats.mean(-1)
        assert feats.ndim == 1, feats.ndim
        feats = np.expand_dims(np.expand_dims(feats, 0), 0)
        onnx_input = {self.model.get_inputs()[0].name: feats}
        logits = self.model.run(None, onnx_input)[0]
        return logits.transpose(0, 2, 1)


class RMVPEF0Predictor:
    def __init__(self, model_path="rmvpe.pt", is_half=False, device="cpu", sampling_rate=40000):
        import torch
        from rmvpe import RMVPE
        self.model = RMVPE(model_path, is_half=is_half, device=device)
        self.sampling_rate = sampling_rate

    def interpolate_f0(self, f0):
        data = np.reshape(f0, (f0.size, 1))
        vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
        vuv_vector[data > 0.0] = 1.0
        vuv_vector[data <= 0.0] = 0.0

        ip_data = data
        frame_number = data.size
        last_value = 0.0
        for i in range(frame_number):
            if data[i] <= 0.0:
                j = i + 1
                for j in range(i + 1, frame_number):
                    if data[j] > 0.0:
                        break
                if j < frame_number - 1:
                    if last_value > 0.0:
                        step = (data[j] - data[i - 1]) / float(j - i)
                        for k in range(i, j):
                            ip_data[k] = data[i - 1] + step * (k - i + 1)
                    else:
                        for k in range(i, j):
                            ip_data[k] = data[j]
                else:
                    for k in range(i, frame_number):
                        ip_data[k] = last_value
            else:
                ip_data[i] = data[i]
                last_value = data[i]

        return ip_data[:, 0], vuv_vector[:, 0]

    def compute_f0(self, wav16k, p_len):
        # Input 'wav16k' sudah pasti berada di 16000 Hz karena di-bypass pada tingkat atas
        f0 = self.model.infer_from_audio(wav16k, thred=0.03)
        
        # Resize f0 to match p_len perfectly using np.interp (sama dengan resize_f0 di Dio)
        source = np.array(f0)
        source[source < 0.001] = np.nan
        target = np.interp(
            np.arange(0, len(source) * p_len, len(source)) / p_len,
            np.arange(0, len(source)),
            source,
        )
        res = np.nan_to_num(target)
        
        # Lakukan interpolasi agar kontinu (menghindari suara robotik & glitch pitch)
        return self.interpolate_f0(res)[0]


def get_f0_predictor(f0_predictor, hop_length, sampling_rate, **kargs):
    device = kargs.get("device", "cpu")
    if f0_predictor == "pm":
        from lib.infer_pack.modules.F0Predictor.PMF0Predictor import PMF0Predictor

        f0_predictor_object = PMF0Predictor(
            hop_length=hop_length, sampling_rate=sampling_rate
        )
    elif f0_predictor == "harvest":
        from lib.infer_pack.modules.F0Predictor.HarvestF0Predictor import (
            HarvestF0Predictor,
        )

        f0_predictor_object = HarvestF0Predictor(
            hop_length=hop_length, sampling_rate=sampling_rate
        )
    elif f0_predictor == "dio":
        from lib.infer_pack.modules.F0Predictor.DioF0Predictor import DioF0Predictor

        f0_predictor_object = DioF0Predictor(
            hop_length=hop_length, sampling_rate=sampling_rate
        )
    elif f0_predictor == "rmvpe":
        is_half = kargs.get("is_half", False)
        f0_predictor_object = RMVPEF0Predictor(
            model_path="rmvpe.pt",
            is_half=is_half,
            device=device,
            sampling_rate=sampling_rate
        )
    else:
        raise Exception("Unknown f0 predictor")
    return f0_predictor_object


class OnnxRVC:
    def __init__(
        self,
        model_path,
        sr=40000,
        hop_size=512,
        vec_path="vec-768-layer-12",
        device="cpu",
    ):
        vec_path = f"pretrained/{vec_path}.onnx"
        self.vec_model = ContentVec(vec_path, device)
        import onnxruntime as ort
        opts = ort.SessionOptions()
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        opts.intra_op_num_threads = 2
        
        if device == "cpu" or device is None:
            providers = ["CPUExecutionProvider"]
        elif device == "cuda":
            providers = [
                ("CUDAExecutionProvider", {
                    "device_id": 0,
                    "arena_extend_strategy": "kNextPowerOfTwo",
                    "cudnn_conv_algo_search": "EXHAUSTIVE",
                    "do_copy_in_default_stream": True,
                }),
                "CPUExecutionProvider"
            ]
        elif device == "dml":
            providers = ["DmlExecutionProvider"]
        else:
            raise RuntimeError("Unsportted Device")
        self.model = ort.InferenceSession(model_path, sess_options=opts, providers=providers)
        self.sampling_rate = sr
        self.hop_size = hop_size
        self.device = device

    def forward(self, hubert, hubert_length, pitch, pitchf, ds, rnd):
        onnx_input = {
            self.model.get_inputs()[0].name: hubert,
            self.model.get_inputs()[1].name: hubert_length,
            self.model.get_inputs()[2].name: pitch,
            self.model.get_inputs()[3].name: pitchf,
            self.model.get_inputs()[4].name: ds,
            self.model.get_inputs()[5].name: rnd,
        }
        return (self.model.run(None, onnx_input)[0] * 32767).astype(np.int16)

    def inference(
        self,
        raw_path,
        sid,
        f0_method="dio",
        f0_up_key=0,
        pad_time=0.5,
        cr_threshold=0.02,
        rmvpe_fp16=False,
    ):
        f0_min = 50
        f0_max = 1100
        f0_mel_min = 1127 * np.log(1 + f0_min / 700)
        f0_mel_max = 1127 * np.log(1 + f0_max / 700)
        f0_predictor = get_f0_predictor(
            f0_method,
            hop_length=self.hop_size,
            sampling_rate=self.sampling_rate,
            threshold=cr_threshold,
            device=self.device,
            is_half=rmvpe_fp16,
        )
        if f0_method == "rmvpe":
            wav16k, sr = load_audio_fast(raw_path, 16000)
            org_length = int(len(wav16k) * (self.sampling_rate / 16000))
            if len(wav16k) / 16000 > 50.0:
                raise RuntimeError("Reached Max Length")
        else:
            wav, sr = load_audio_fast(raw_path, self.sampling_rate)
            org_length = len(wav)
            if org_length / sr > 50.0:
                raise RuntimeError("Reached Max Length")
            wav16k, _ = load_audio_fast(raw_path, 16000)

        hubert = self.vec_model(wav16k)
        hubert = np.repeat(hubert, 2, axis=2).transpose(0, 2, 1).astype(np.float32)
        hubert_length = hubert.shape[1]

        if f0_method == "rmvpe":
            pitchf = f0_predictor.compute_f0(wav16k, hubert_length)
        else:
            pitchf = f0_predictor.compute_f0(wav, hubert_length)
        pitchf = pitchf * 2 ** (f0_up_key / 12)
        pitch = pitchf.copy()
        f0_mel = 1127 * np.log(1 + pitch / 700)
        f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * 254 / (
            f0_mel_max - f0_mel_min
        ) + 1
        f0_mel[f0_mel <= 1] = 1
        f0_mel[f0_mel > 255] = 255
        pitch = np.rint(f0_mel).astype(np.int64)

        pitchf = pitchf.reshape(1, len(pitchf)).astype(np.float32)
        pitch = pitch.reshape(1, len(pitch))
        ds = np.array([sid]).astype(np.int64)

        rnd = np.random.randn(1, 192, hubert_length).astype(np.float32)
        hubert_length = np.array([hubert_length]).astype(np.int64)

        out_wav = self.forward(hubert, hubert_length, pitch, pitchf, ds, rnd).squeeze()
        out_wav = np.pad(out_wav, (0, 2 * self.hop_size), "constant")
        return out_wav[0:org_length]