import os
import sys
import onnx
import torch
import platform
import warnings
import onnx2torch

import numpy as np
import onnxruntime as ort

from tqdm import tqdm

sys.path.append(os.getcwd())

from main.library.uvr5_lib import spec_utils
from main.library.uvr5_lib.common_separator import CommonSeparator

warnings.filterwarnings("ignore")

class MDXSeparator(CommonSeparator):
    def __init__(self, common_config, arch_config):
        super().__init__(config=common_config)
        self.segment_size = arch_config.get("segment_size")
        self.overlap = arch_config.get("overlap")
        self.batch_size = arch_config.get("batch_size", 1)
        self.hop_length = arch_config.get("hop_length")
        self.enable_denoise = arch_config.get("enable_denoise")
        self.compensate = self.model_data["compensate"]
        self.dim_f = self.model_data["mdx_dim_f_set"]
        self.dim_t = 2 ** self.model_data["mdx_dim_t_set"]
        self.n_fft = self.model_data["mdx_n_fft_scale_set"]
        self.config_yaml = self.model_data.get("config_yaml", None)
        self.load_model()
        self.n_bins = 0
        self.trim = 0
        self.chunk_size = 0
        self.gen_size = 0
        self.stft = None
        self.primary_source = None
        self.secondary_source = None
        self.audio_file_path = None
        self.audio_file_base = None

    def load_model(self):
        if self.segment_size == self.dim_t:
            ort_session_options = ort.SessionOptions()
            ort_session_options.log_severity_level = 3
            ort_inference_session = ort.InferenceSession(self.model_path, providers=self.onnx_execution_provider, sess_options=ort_session_options)
            self.model_run = lambda spek: ort_inference_session.run(None, {"input": spek.cpu().numpy()})[0]
        else:
            self.model_run = onnx2torch.convert(onnx.load(self.model_path)) if platform.system() == 'Windows' else onnx2torch.convert(self.model_path)
            self.model_run.to(self.torch_device).eval()

    def separate(self, audio_file_path):
        self.audio_file_path = audio_file_path
        self.audio_file_base = os.path.splitext(os.path.basename(audio_file_path))[0]
        mix = self.prepare_mix(self.audio_file_path)
        mix = spec_utils.normalize(wave=mix, max_peak=self.normalization_threshold)
        source = self.demix(mix)
        output_files = []

        if not isinstance(self.primary_source, np.ndarray):
            self.primary_source = spec_utils.normalize(wave=source, max_peak=self.normalization_threshold).T

        if not isinstance(self.secondary_source, np.ndarray):
            raw_mix = self.demix(mix, is_match_mix=True)

            if self.invert_using_spec:
                self.secondary_source = spec_utils.invert_stem(raw_mix, source)
            else:
                self.secondary_source = mix.T - source.T

        if not self.output_single_stem or self.output_single_stem.lower() == self.secondary_stem_name.lower():
            self.secondary_stem_output_path = os.path.join(f"{self.audio_file_base}_({self.secondary_stem_name})_{self.model_name}.{self.output_format.lower()}")
            self.final_process(self.secondary_stem_output_path, self.secondary_source, self.secondary_stem_name)
            output_files.append(self.secondary_stem_output_path)

        if not self.output_single_stem or self.output_single_stem.lower() == self.primary_stem_name.lower():
            self.primary_stem_output_path = os.path.join(f"{self.audio_file_base}_({self.primary_stem_name})_{self.model_name}.{self.output_format.lower()}")
            if not isinstance(self.primary_source, np.ndarray): self.primary_source = source.T

            self.final_process(self.primary_stem_output_path, self.primary_source, self.primary_stem_name)
            output_files.append(self.primary_stem_output_path)

        return output_files

    def initialize_model_settings(self):
        self.n_bins = self.n_fft // 2 + 1
        self.trim = self.n_fft // 2
        self.chunk_size = self.hop_length * (self.segment_size - 1)
        self.gen_size = self.chunk_size - 2 * self.trim
        self.stft = STFT(self.n_fft, self.hop_length, self.dim_f, self.torch_device)

    def initialize_mix(self, mix, is_ckpt=False):
        if is_ckpt:
            pad = self.gen_size + self.trim - (mix.shape[-1] % self.gen_size)
            mixture = np.concatenate((np.zeros((2, self.trim), dtype="float32"), mix, np.zeros((2, pad), dtype="float32")), 1)

            num_chunks = mixture.shape[-1] // self.gen_size
            mix_waves = [mixture[:, i * self.gen_size : i * self.gen_size + self.chunk_size] for i in range(num_chunks)]
        else:
            mix_waves = []
            n_sample = mix.shape[1]

            pad = self.gen_size - n_sample % self.gen_size
            mix_p = np.concatenate((np.zeros((2, self.trim)), mix, np.zeros((2, pad)), np.zeros((2, self.trim))), 1)

            i = 0
            while i < n_sample + pad:
                mix_waves.append(np.array(mix_p[:, i : i + self.chunk_size]))
                i += self.gen_size

        mix_waves_tensor = torch.tensor(mix_waves, dtype=torch.float32).to(self.torch_device)
        return mix_waves_tensor, pad

    def demix(self, mix, is_match_mix=False):
        self.initialize_model_settings()
        tar_waves_ = []

        if is_match_mix:
            chunk_size = self.hop_length * (self.segment_size - 1)
            overlap = 0.02
        else:
            chunk_size = self.chunk_size
            overlap = self.overlap

        gen_size = chunk_size - 2 * self.trim
        mixture = np.concatenate((np.zeros((2, self.trim), dtype="float32"), mix, np.zeros((2, gen_size + self.trim - ((mix.shape[-1]) % gen_size)), dtype="float32")), 1)
        step = int((1 - overlap) * chunk_size)

        result = np.zeros((1, 2, mixture.shape[-1]), dtype=np.float32)
        divider = np.zeros((1, 2, mixture.shape[-1]), dtype=np.float32)
        total = 0

        for i in tqdm(range(0, mixture.shape[-1], step), ncols=100, unit="f"):
            total += 1
            start = i
            end = min(i + chunk_size, mixture.shape[-1])

            chunk_size_actual = end - start
            window = None

            if overlap != 0:
                window = np.hanning(chunk_size_actual)
                window = np.tile(window[None, None, :], (1, 2, 1))

            mix_part_ = mixture[:, start:end]
            
            if end != i + chunk_size:
                pad_size = (i + chunk_size) - end
                mix_part_ = np.concatenate((mix_part_, np.zeros((2, pad_size), dtype="float32")), axis=-1)

            mix_waves = torch.tensor([mix_part_], dtype=torch.float32).to(self.torch_device).split(self.batch_size)

            with torch.no_grad():
                batches_processed = 0
                
                for mix_wave in mix_waves:
                    batches_processed += 1
                    tar_waves = self.run_model(mix_wave, is_match_mix=is_match_mix)

                    if window is not None:
                        tar_waves[..., :chunk_size_actual] *= window
                        divider[..., start:end] += window
                    else: divider[..., start:end] += 1

                    result[..., start:end] += tar_waves[..., : end - start]

        tar_waves = result / divider
        tar_waves_.append(tar_waves)
        tar_waves = np.concatenate(np.vstack(tar_waves_)[:, :, self.trim : -self.trim], axis=-1)[:, : mix.shape[-1]]

        source = tar_waves[:, 0:None]

        if not is_match_mix:
            source *= self.compensate

        return source

    def run_model(self, mix, is_match_mix=False):
        spek = self.stft(mix.to(self.torch_device))
        spek[:, :, :3, :] *= 0

        if is_match_mix:
            spec_pred = spek.cpu().numpy()
        else:
            if self.enable_denoise:
                spec_pred_neg = self.model_run(-spek)  
                spec_pred_pos = self.model_run(spek)
                spec_pred = (spec_pred_neg * -0.5) + (spec_pred_pos * 0.5)
            else:
                spec_pred = self.model_run(spek)

        result = self.stft.inverse(torch.tensor(spec_pred).to(self.torch_device)).cpu().detach().numpy()
        return result

class STFT:
    def __init__(self, n_fft, hop_length, dim_f, device):
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.dim_f = dim_f
        self.device = device
        self.hann_window = torch.hann_window(window_length=self.n_fft, periodic=True)

    def __call__(self, input_tensor):
        is_non_standard_device = not input_tensor.device.type in ["cuda", "cpu"]

        if is_non_standard_device: input_tensor = input_tensor.cpu()

        batch_dimensions = input_tensor.shape[:-2]
        channel_dim, time_dim = input_tensor.shape[-2:]

        permuted_stft_output = torch.stft(input_tensor.reshape([-1, time_dim]), n_fft=self.n_fft, hop_length=self.hop_length, window=self.hann_window.to(input_tensor.device), center=True, return_complex=False).permute([0, 3, 1, 2])
        final_output = permuted_stft_output.reshape([*batch_dimensions, channel_dim, 2, -1, permuted_stft_output.shape[-1]]).reshape([*batch_dimensions, channel_dim * 2, -1, permuted_stft_output.shape[-1]])

        if is_non_standard_device: final_output = final_output.to(self.device)
        return final_output[..., : self.dim_f, :]

    def pad_frequency_dimension(self, input_tensor, batch_dimensions, channel_dim, freq_dim, time_dim, num_freq_bins):
        return torch.cat([input_tensor, torch.zeros([*batch_dimensions, channel_dim, num_freq_bins - freq_dim, time_dim]).to(input_tensor.device)], -2)

    def calculate_inverse_dimensions(self, input_tensor):
        channel_dim, freq_dim, time_dim = input_tensor.shape[-3:]

        return input_tensor.shape[:-3], channel_dim, freq_dim, time_dim, self.n_fft // 2 + 1

    def prepare_for_istft(self, padded_tensor, batch_dimensions, channel_dim, num_freq_bins, time_dim):
        permuted_tensor = padded_tensor.reshape([*batch_dimensions, channel_dim // 2, 2, num_freq_bins, time_dim]).reshape([-1, 2, num_freq_bins, time_dim]).permute([0, 2, 3, 1])

        return permuted_tensor[..., 0] + permuted_tensor[..., 1] * 1.0j

    def inverse(self, input_tensor):
        is_non_standard_device = not input_tensor.device.type in ["cuda", "cpu"]
        if is_non_standard_device: input_tensor = input_tensor.cpu()

        batch_dimensions, channel_dim, freq_dim, time_dim, num_freq_bins = self.calculate_inverse_dimensions(input_tensor)
        final_output = torch.istft(self.prepare_for_istft(self.pad_frequency_dimension(input_tensor, batch_dimensions, channel_dim, freq_dim, time_dim, num_freq_bins), batch_dimensions, channel_dim, num_freq_bins, time_dim), n_fft=self.n_fft, hop_length=self.hop_length, window=self.hann_window.to(input_tensor.device), center=True).reshape([*batch_dimensions, 2, -1])

        if is_non_standard_device: final_output = final_output.to(self.device)
        return final_output