"""Module for separating audio sources using MDX architecture models."""

import os
import platform
import torch
import onnx
import onnxruntime as ort
import numpy as np
import onnx2torch
from tqdm import tqdm

from PolUVR.separator.uvr_lib_v5 import spec_utils
from PolUVR.separator.uvr_lib_v5.stft import STFT
from PolUVR.separator.common_separator import CommonSeparator


class MDXSeparator(CommonSeparator):
    """
    MDXSeparator is responsible for separating audio sources using MDX models.
    It initializes with configuration parameters and prepares the model for separation tasks.
    """

    def __init__(self, common_config, arch_config):
        # Any configuration values which can be shared between architectures should be set already in CommonSeparator,
        # e.g. user-specified functionality choices (self.output_single_stem) or common model parameters (self.primary_stem_name)
        super().__init__(config=common_config)

        # Initializing user-configurable parameters, passed through with an mdx_from the CLI or Separator instance

        # Pick a segment size to balance speed, resource use, and quality:
        # - Smaller sizes consume less resources.
        # - Bigger sizes consume more resources, but may provide better results.
        # - Default size is 256. Quality can change based on your pick.
        self.segment_size = arch_config.get("segment_size", 256)

        # Ensure segment_size is within the range [32, 4000]
        if self.segment_size < 32:
            self.logger.warning(f"segment_size {self.segment_size} is less than the minimum allowed value of 32. Setting segment_size to 32.")
            self.segment_size = 32
        elif self.segment_size > 4000:
            self.logger.warning(f"segment_size {self.segment_size} is greater than the maximum allowed value of 4000. Setting segment_size to 4000.")
            self.segment_size = 4000

        # This option controls the amount of overlap between prediction windows.
        #  - Higher values can provide better results, but will lead to longer processing times.
        #  - For Non-MDX23C models: You can choose between 0.001-0.999
        self.overlap = arch_config.get("overlap", 0.25)

        # Ensure overlap is within the range [0.001, 0.99]
        if self.overlap < 0.001:
            self.logger.warning(f"overlap {self.overlap} is less than the minimum allowed value of 0.001. Setting overlap to 0.001.")
            self.overlap = 0.001
        elif self.overlap > 0.999:
            self.logger.warning(f"overlap {self.overlap} is greater than the maximum allowed value of 0.999. Setting overlap to 0.999.")
            self.overlap = 0.999

        # Number of batches to be processed at a time.
        # - Higher values mean more RAM usage but slightly faster processing times.
        # - Lower values mean less RAM usage but slightly longer processing times.
        # - Batch size value has no effect on output quality.
        # BATCH_SIZE = ('1', ''2', '3', '4', '5', '6', '7', '8', '9', '10')
        self.batch_size = arch_config.get("batch_size", 1)

        # Ensure batch_size is within the range [1, 16]
        if self.batch_size < 1:
            self.logger.warning(f"batch_size {self.batch_size} is less than the minimum allowed value of 1. Setting batch_size to 1.")
            self.batch_size = 1
        elif self.batch_size > 16:
            self.logger.warning(f"batch_size {self.batch_size} is greater than the maximum allowed value of 16. Setting batch_size to 16.")
            self.batch_size = 16

        # hop_length is equivalent to the more commonly used term "stride" in convolutional neural networks
        # In machine learning, particularly in the context of convolutional neural networks (CNNs),
        # the term "stride" refers to the number of pixels by which we move the filter across the input image.
        # Strides are a crucial component in the convolution operation, a fundamental building block of CNNs used primarily in the field of computer vision.
        # Stride is a parameter that dictates the movement of the kernel, or filter, across the input data, such as an image.
        # When performing a convolution operation, the stride determines how many units the filter shifts at each step.
        # The choice of stride affects the model in several ways:
        # Output Size: A larger stride will result in a smaller output spatial dimension.
        # Computational Efficiency: Increasing the stride can decrease the computational load.
        # Field of View: A higher stride means that each step of the filter takes into account a wider area of the input image.
        #   This can be beneficial when the model needs to capture more global features rather than focusing on finer details.
        self.hop_length = arch_config.get("hop_length", 1024)

        # Ensure hop_length is within the range [32, 2048]
        if self.hop_length < 32:
            self.logger.warning(f"hop_length {self.hop_length} is less than the minimum allowed value of 32. Setting hop_length to 32.")
            self.hop_length = 32
        elif self.hop_length > 2048:
            self.logger.warning(f"hop_length {self.hop_length} is greater than the maximum allowed value of 2048. Setting hop_length to 2048.")
            self.hop_length = 2048

        # If enabled, model will be run twice to reduce noise in output audio.
        self.enable_denoise = arch_config.get("enable_denoise")

        self.logger.debug(f"MDX arch params: batch_size={self.batch_size}, segment_size={self.segment_size}")
        self.logger.debug(f"MDX arch params: overlap={self.overlap}, hop_length={self.hop_length}, enable_denoise={self.enable_denoise}")

        # Initializing model-specific parameters from model_data JSON
        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.logger.debug(f"MDX arch params: compensate={self.compensate}, dim_f={self.dim_f}, dim_t={self.dim_t}, n_fft={self.n_fft}")
        self.logger.debug(f"MDX arch params: config_yaml={self.config_yaml}")

        # In UVR, these variables are set but either aren't useful or are better handled in PolUVR.
        # Leaving these comments explaining to help myself or future developers understand why these aren't in PolUVR.

        # "chunks" is not actually used for anything in UVR...
        # self.chunks = 0

        # "adjust" is hard-coded to 1 in UVR, and only used as a multiplier in run_model, so it does nothing.
        # self.adjust = 1

        # "hop" is hard-coded to 1024 in UVR. We have a "hop_length" parameter instead
        # self.hop = 1024

        # "margin" maps to sample rate and is set from the GUI in UVR (default: 44100). We have a "sample_rate" parameter instead.
        # self.margin = 44100

        # "dim_c" is hard-coded to 4 in UVR, seems to be a parameter for the number of channels, and is only used for checkpoint models.
        # We haven't implemented support for the checkpoint models here, so we're not using it.
        # self.dim_c = 4

        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):
        """
        Load the model into memory from file on disk, initialize it with config from the model data,
        and prepare for inferencing using hardware accelerated Torch device.
        """
        self.logger.debug("Loading ONNX model for inference...")

        if self.segment_size == self.dim_t:
            ort_session_options = ort.SessionOptions()
            if self.log_level > 10:
                ort_session_options.log_severity_level = 3
            else:
                ort_session_options.log_severity_level = 0

            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]
            self.logger.debug("Model loaded successfully using ONNXruntime inferencing session.")
        else:
            if platform.system() == 'Windows':
                onnx_model = onnx.load(self.model_path)
                self.model_run = onnx2torch.convert(onnx_model)
            else:
                self.model_run = onnx2torch.convert(self.model_path)
   
            self.model_run.to(self.torch_device).eval()
            self.logger.warning("Model converted from onnx to pytorch due to segment size not matching dim_t, processing may be slower.")

    def separate(self, audio_file_path, custom_output_names=None):
        """
        Separates the audio file into primary and secondary sources based on the model's configuration.
        It processes the mix, demixes it into sources, normalizes the sources, and saves the output files.

        Args:
            audio_file_path (str): The path to the audio file to be processed.
            custom_output_names (dict, optional): Custom names for the output files. Defaults to None.

        Returns:
            list: A list of paths to the output files generated by the separation process.
        """
        self.audio_file_path = audio_file_path
        self.audio_file_base = os.path.splitext(os.path.basename(audio_file_path))[0]

        # Prepare the mix for processing
        self.logger.debug(f"Preparing mix for input audio file {self.audio_file_path}...")
        mix = self.prepare_mix(self.audio_file_path)

        self.logger.debug("Normalizing mix before demixing...")
        mix = spec_utils.normalize(wave=mix, max_peak=self.normalization_threshold, min_peak=self.amplification_threshold)

        # Start the demixing process
        source = self.demix(mix)
        self.logger.debug("Demixing completed.")

        # In UVR, the source is cached here if it's a vocal split model, but we're not supporting that yet

        # Initialize the list for output files
        output_files = []
        self.logger.debug("Processing output files...")

        # Normalize and transpose the primary source if it's not already an array
        if not isinstance(self.primary_source, np.ndarray):
            self.logger.debug("Normalizing primary source...")
            self.primary_source = spec_utils.normalize(wave=source, max_peak=self.normalization_threshold, min_peak=self.amplification_threshold).T

        # Process the secondary source if not already an array
        if not isinstance(self.secondary_source, np.ndarray):
            self.logger.debug("Producing secondary source: demixing in match_mix mode")
            raw_mix = self.demix(mix, is_match_mix=True)

            if self.invert_using_spec:
                self.logger.debug("Inverting secondary stem using spectogram as invert_using_spec is set to True")
                self.secondary_source = spec_utils.invert_stem(raw_mix, source)
            else:
                self.logger.debug("Inverting secondary stem by subtracting of transposed demixed stem from transposed original mix")
                self.secondary_source = mix.T - source.T

        # Save and process the secondary stem if needed
        if not self.output_single_stem or self.output_single_stem.lower() == self.secondary_stem_name.lower():
            self.secondary_stem_output_path = self.get_stem_output_path(self.secondary_stem_name, custom_output_names)

            self.logger.info(f"Saving {self.secondary_stem_name} stem to {self.secondary_stem_output_path}...")
            self.final_process(self.secondary_stem_output_path, self.secondary_source, self.secondary_stem_name)
            output_files.append(self.secondary_stem_output_path)

        # Save and process the primary stem if needed
        if not self.output_single_stem or self.output_single_stem.lower() == self.primary_stem_name.lower():
            self.primary_stem_output_path = self.get_stem_output_path(self.primary_stem_name, custom_output_names)

            if not isinstance(self.primary_source, np.ndarray):
                self.primary_source = source.T

            self.logger.info(f"Saving {self.primary_stem_name} stem to {self.primary_stem_output_path}...")
            self.final_process(self.primary_stem_output_path, self.primary_source, self.primary_stem_name)
            output_files.append(self.primary_stem_output_path)

        # Not yet implemented from UVR features:
        # self.process_vocal_split_chain(secondary_sources)
        # self.logger.debug("Vocal split chain processed.")

        return output_files

    def initialize_model_settings(self):
        """
        This function sets up the necessary parameters for the model, like the number of frequency bins (n_bins), the trimming size (trim),
        the size of each audio chunk (chunk_size), and the window function for spectral transformations (window).
        It ensures that the model is configured with the correct settings for processing the audio data.
        """
        self.logger.debug("Initializing model settings...")

        # n_bins is half the FFT size plus one (self.n_fft // 2 + 1).
        self.n_bins = self.n_fft // 2 + 1

        # trim is half the FFT size (self.n_fft // 2).
        self.trim = self.n_fft // 2

        # chunk_size is the hop_length size times the segment size minus one
        self.chunk_size = self.hop_length * (self.segment_size - 1)

        # gen_size is the chunk size minus twice the trim size
        self.gen_size = self.chunk_size - 2 * self.trim

        self.stft = STFT(self.logger, self.n_fft, self.hop_length, self.dim_f, self.torch_device)

        self.logger.debug(f"Model input params: n_fft={self.n_fft} hop_length={self.hop_length} dim_f={self.dim_f}")
        self.logger.debug(f"Model settings: n_bins={self.n_bins}, trim={self.trim}, chunk_size={self.chunk_size}, gen_size={self.gen_size}")

    def initialize_mix(self, mix, is_ckpt=False):
        """
        After prepare_mix segments the audio, initialize_mix further processes each segment.
        It ensures each audio segment is in the correct format for the model, applies necessary padding,
        and converts the segments into tensors for processing with the model.
        This step is essential for preparing the audio data in a format that the neural network can process.
        """
        # Log the initialization of the mix and whether checkpoint mode is used
        self.logger.debug(f"Initializing mix with is_ckpt={is_ckpt}. Initial mix shape: {mix.shape}")

        # Ensure the mix is a 2-channel (stereo) audio signal
        if mix.shape[0] != 2:
            error_message = f"Expected a 2-channel audio signal, but got {mix.shape[0]} channels"
            self.logger.error(error_message)
            raise ValueError(error_message)

        # If in checkpoint mode, process the mix differently
        if is_ckpt:
            self.logger.debug("Processing in checkpoint mode...")
            # Calculate padding based on the generation size and trim
            pad = self.gen_size + self.trim - (mix.shape[-1] % self.gen_size)
            self.logger.debug(f"Padding calculated: {pad}")
            # Add padding at the beginning and the end of the mix
            mixture = np.concatenate((np.zeros((2, self.trim), dtype="float32"), mix, np.zeros((2, pad), dtype="float32")), 1)
            # Determine the number of chunks based on the mixture's length
            num_chunks = mixture.shape[-1] // self.gen_size
            self.logger.debug(f"Mixture shape after padding: {mixture.shape}, Number of chunks: {num_chunks}")
            # Split the mixture into chunks
            mix_waves = [mixture[:, i * self.gen_size : i * self.gen_size + self.chunk_size] for i in range(num_chunks)]
        else:
            # If not in checkpoint mode, process normally
            self.logger.debug("Processing in non-checkpoint mode...")
            mix_waves = []
            n_sample = mix.shape[1]
            # Calculate necessary padding to make the total length divisible by the generation size
            pad = self.gen_size - n_sample % self.gen_size
            self.logger.debug(f"Number of samples: {n_sample}, Padding calculated: {pad}")
            # Apply padding to the mix
            mix_p = np.concatenate((np.zeros((2, self.trim)), mix, np.zeros((2, pad)), np.zeros((2, self.trim))), 1)
            self.logger.debug(f"Shape of mix after padding: {mix_p.shape}")

            # Process the mix in chunks
            i = 0
            while i < n_sample + pad:
                waves = np.array(mix_p[:, i : i + self.chunk_size])
                mix_waves.append(waves)
                self.logger.debug(f"Processed chunk {len(mix_waves)}: Start {i}, End {i + self.chunk_size}")
                i += self.gen_size

        # Convert the list of wave chunks into a tensor for processing on the specified device
        mix_waves_tensor = torch.tensor(mix_waves, dtype=torch.float32).to(self.torch_device)
        self.logger.debug(f"Converted mix_waves to tensor. Tensor shape: {mix_waves_tensor.shape}")

        return mix_waves_tensor, pad

    def demix(self, mix, is_match_mix=False):
        """
        Demixes the input mix into its constituent sources. If is_match_mix is True, the function adjusts the processing
        to better match the mix, affecting chunk sizes and overlaps. The demixing process involves padding the mix,
        processing it in chunks, applying windowing for overlaps, and accumulating the results to separate the sources.
        """
        self.logger.debug(f"Starting demixing process with is_match_mix: {is_match_mix}...")
        self.initialize_model_settings()

        # Preserves the original mix for later use.
        # In UVR, this is used for the pitch fix and VR denoise processes, which aren't yet implemented here.
        org_mix = mix
        self.logger.debug(f"Original mix stored. Shape: {org_mix.shape}")

        # Initializes a list to store the separated waveforms.
        tar_waves_ = []

        # Handling different chunk sizes and overlaps based on the matching requirement.
        if is_match_mix:
            # Sets a smaller chunk size specifically for matching the mix.
            chunk_size = self.hop_length * (self.segment_size - 1)
            # Sets a small overlap for the chunks.
            overlap = 0.02
            self.logger.debug(f"Chunk size for matching mix: {chunk_size}, Overlap: {overlap}")
        else:
            # Uses the regular chunk size defined in model settings.
            chunk_size = self.chunk_size
            # Uses the overlap specified in the model settings.
            overlap = self.overlap
            self.logger.debug(f"Standard chunk size: {chunk_size}, Overlap: {overlap}")

        # Calculates the generated size after subtracting the trim from both ends of the chunk.
        gen_size = chunk_size - 2 * self.trim
        self.logger.debug(f"Generated size calculated: {gen_size}")

        # Calculates padding to make the mix length a multiple of the generated size.
        pad = gen_size + self.trim - ((mix.shape[-1]) % gen_size)
        # Prepares the mixture with padding at the beginning and the end.
        mixture = np.concatenate((np.zeros((2, self.trim), dtype="float32"), mix, np.zeros((2, pad), dtype="float32")), 1)
        self.logger.debug(f"Mixture prepared with padding. Mixture shape: {mixture.shape}")

        # Calculates the step size for processing chunks based on the overlap.
        step = int((1 - overlap) * chunk_size)
        self.logger.debug(f"Step size for processing chunks: {step} as overlap is set to {overlap}.")

        # Initializes arrays to store the results and to account for overlap.
        result = np.zeros((1, 2, mixture.shape[-1]), dtype=np.float32)
        divider = np.zeros((1, 2, mixture.shape[-1]), dtype=np.float32)

        # Initializes counters for processing chunks.
        total = 0
        total_chunks = (mixture.shape[-1] + step - 1) // step
        self.logger.debug(f"Total chunks to process: {total_chunks}")

        # Processes each chunk of the mixture.
        for i in tqdm(range(0, mixture.shape[-1], step)):
            total += 1
            start = i
            end = min(i + chunk_size, mixture.shape[-1])
            self.logger.debug(f"Processing chunk {total}/{total_chunks}: Start {start}, End {end}")

            # Handles windowing for overlapping chunks.
            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))
                self.logger.debug("Window applied to the chunk.")

            # Zero-pad the chunk to prepare it for processing.
            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)

            # Converts the chunk to a tensor for processing.
            mix_part = torch.tensor([mix_part_], dtype=torch.float32).to(self.torch_device)
            # Splits the chunk into smaller batches if necessary.
            mix_waves = mix_part.split(self.batch_size)
            total_batches = len(mix_waves)
            self.logger.debug(f"Mix part split into batches. Number of batches: {total_batches}")

            with torch.no_grad():
                # Processes each batch in the chunk.
                batches_processed = 0
                for mix_wave in mix_waves:
                    batches_processed += 1
                    self.logger.debug(f"Processing mix_wave batch {batches_processed}/{total_batches}")

                    # Runs the model to separate the sources.
                    tar_waves = self.run_model(mix_wave, is_match_mix=is_match_mix)

                    # Applies windowing if needed and accumulates the results.
                    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]

        # Normalizes the results by the divider to account for overlap.
        self.logger.debug("Normalizing result by dividing result by divider.")
        tar_waves = result / divider
        tar_waves_.append(tar_waves)

        # Reshapes the results to match the original dimensions.
        tar_waves_ = np.vstack(tar_waves_)[:, :, self.trim : -self.trim]
        tar_waves = np.concatenate(tar_waves_, axis=-1)[:, : mix.shape[-1]]

        # Extracts the source from the results.
        source = tar_waves[:, 0:None]
        self.logger.debug(f"Concatenated tar_waves. Shape: {tar_waves.shape}")

        # TODO: In UVR, pitch changing happens here. Consider implementing this as a feature.

        # Compensates the source if not matching the mix.
        if not is_match_mix:
            source *= self.compensate
            self.logger.debug("Match mix mode; compensate multiplier applied.")

        # TODO: In UVR, VR denoise model gets applied here. Consider implementing this as a feature.

        self.logger.debug("Demixing process completed.")
        return source

    def run_model(self, mix, is_match_mix=False):
        """
        Processes the input mix through the model to separate the sources.
        Applies STFT, handles spectrum modifications, and runs the model for source separation.
        """
        # Applying the STFT to the mix. The mix is moved to the specified device (e.g., GPU) before processing.
        # self.logger.debug(f"Running STFT on the mix. Mix shape before STFT: {mix.shape}")
        spek = self.stft(mix.to(self.torch_device))
        self.logger.debug(f"STFT applied on mix. Spectrum shape: {spek.shape}")

        # Zeroing out the first 3 bins of the spectrum. This is often done to reduce low-frequency noise.
        spek[:, :, :3, :] *= 0
        # self.logger.debug("First 3 bins of the spectrum zeroed out.")

        # Handling the case where the mix needs to be matched (is_match_mix = True)
        if is_match_mix:
            # self.logger.debug("Match mix mode is enabled. Converting spectrum to NumPy array.")
            spec_pred = spek.cpu().numpy()
            self.logger.debug("is_match_mix: spectrum prediction obtained directly from STFT output.")
        else:
            # If denoising is enabled, the model is run on both the negative and positive spectrums.
            if self.enable_denoise:
                # Assuming spek is a tensor and self.model_run can process it directly
                spec_pred_neg = self.model_run(-spek)  # Ensure this line correctly negates spek and runs the model
                spec_pred_pos = self.model_run(spek)
                # Ensure both spec_pred_neg and spec_pred_pos are tensors before applying operations
                spec_pred = (spec_pred_neg * -0.5) + (spec_pred_pos * 0.5)  # [invalid-unary-operand-type]
                self.logger.debug("Model run on both negative and positive spectrums for denoising.")
            else:
                spec_pred = self.model_run(spek)
                self.logger.debug("Model run on the spectrum without denoising.")

        # Applying the inverse STFT to convert the spectrum back to the time domain.
        result = self.stft.inverse(torch.tensor(spec_pred).to(self.torch_device)).cpu().detach().numpy()
        self.logger.debug(f"Inverse STFT applied. Returning result with shape: {result.shape}")

        return result
