# Third Party Imports
import torch
import onnxruntime as ort

# Local Imports
from vocal_isolation.models.mdx_net import Conv_TDF_net_trimm

# Global Variables
from vocal_isolation.constants import EXECUTION_PROVIDER_LIST, COMPUTATION_DEVICE
from vocal_isolation.constants import ONNX_MODEL_PATH, PRETRAINED_MODELS_DIRECTORY


class KimVocal:
    def __init__(self):
        self.models = [
            Conv_TDF_net_trimm(   
                ONNX_MODEL_PATH, use_onnx=True, target_name='vocals', 
                L=11, l=3, g=48, bn=8, bias=False, 
                dim_f=11, dim_t=8
            )
        ]

    def demix_both(self, music_tensor, sample_rate):
        """
        Isolating vocals AND instrumental using an ONNX model. 
        Assuming the audio is loaded correctly at 41000hz samplerate.

        Args:
            music_tensor (torch.Tensor): Input tensor.
            model (torch.nn): Model used for inferring.

        Returns:
            torch.Tensor: Output tensor after passing through the network.
        """
        number_of_samples = music_tensor.shape[1]
        vocals_tensor = None

        # * Extracting vocals
        overlap = self.models[0].overlap
        chunk_size = self.models[0].chunk_size
        gen_size = chunk_size - 2 * overlap
        pad_size = gen_size - number_of_samples % gen_size
        # Along the column dimensions (used for features), we pad the left and right side of the mix to keep the integrity of the whole tensor
        # overlap is added to ensure there's overlap between chunks.
        mix_padded = torch.cat([torch.zeros(2, overlap), music_tensor, torch.zeros(2, pad_size+overlap)], 1)

        ort_session = ort.InferenceSession(f'{PRETRAINED_MODELS_DIRECTORY}/{self.models[0].target_name}.onnx', providers=EXECUTION_PROVIDER_LIST)

        # process one chunk at a time (batch_size=1)
        demixed_chunks = []
        i = 0
        while i < number_of_samples + pad_size:                
            chunk = mix_padded[:, i : i + chunk_size]
            x = self.models[0].stft(chunk.unsqueeze(0).to(COMPUTATION_DEVICE))                
            with torch.no_grad():
                x = torch.tensor(ort_session.run(None, {'input': x.cpu().numpy()})[0]) 
            x = self.models[0].stft.inverse(x).squeeze(0)
            x = x[...,overlap:-overlap]
            demixed_chunks.append(x)
            i += gen_size

        vocals_tensor = torch.cat(demixed_chunks, -1)[...,:-pad_size].cpu()

        # Subtract vocals output from the input mix for the remaining models
        music_minus_vocals_tensor = music_tensor - vocals_tensor

        # Returning two tensors.
        return music_minus_vocals_tensor, vocals_tensor

    def demix_vocals(self, music_tensor, sample_rate):
        """
        Isolating vocals using an ONNX model. 
        Assuming the audio is loaded correctly at 41000hz samplerate.

        Args:
            music_tensor (torch.Tensor): Input tensor.
            model (torch.nn): Model used for inferring.

        Returns:
            torch.Tensor: Output tensor after passing through the network.
        """
        number_of_samples = music_tensor.shape[1]
        overlap = self.models[0].overlap

        # Calculate chunk_size and gen_size based on the sample rate
        chunk_size = self.models[0].chunk_size
        gen_size = chunk_size - 2 * overlap
        pad_size = gen_size - number_of_samples % gen_size
        mix_padded = torch.cat(
            [torch.zeros(2, overlap), music_tensor, torch.zeros(2, pad_size + overlap)],
            1,
        )

        # Start running the session for the model
        ort_session = ort.InferenceSession(
            ONNX_MODEL_PATH, providers=EXECUTION_PROVIDER_LIST
        )

        # process one chunk at a time (batch_size=1)
        demixed_chunks = []
        i = 0
        while i < number_of_samples + pad_size:
            # Computation
            chunk = mix_padded[:, i : i + chunk_size]
            x = self.models[0].stft(chunk.unsqueeze(0).to(COMPUTATION_DEVICE))
            with torch.no_grad():
                x = torch.tensor(ort_session.run(None, {"input": x.cpu().numpy()})[0])
            x = self.models[0].stft.inverse(x).squeeze(0)
            x = x[..., overlap:-overlap]
            demixed_chunks.append(x)
            i += gen_size

        vocals_output = torch.cat(demixed_chunks, -1)[..., :-pad_size].cpu()

        return vocals_output