# coding=utf-8
# Remove voice-over
import logging
import argparse
import subprocess
import librosa
import os
import torch
import soundfile as sf
import numpy as np


# Using the embedded version of Python can also correctly import the utils module.
# current_dir = os.path.dirname(os.path.abspath(__file__))
# sys.path.append(current_dir)

from third_party.MusicSourceSeparationTraining.utils import demix, load_config, normalize_audio, denormalize_audio, draw_spectrogram
from third_party.MusicSourceSeparationTraining.utils import prefer_target_instrument, apply_tta, load_start_checkpoint
from third_party.MusicSourceSeparationTraining.models.bs_roformer import BSRoformer
import warnings

warnings.filterwarnings("ignore")

model_base_dir = "pretrained/remove_vo/checkpoints"
MODEL_PATHS = {"bs_roformer": [f"{model_base_dir}/model_bs_roformer_ep_317_sdr_12.9755.ckpt", f"{model_base_dir}/model_bs_roformer_ep_317_sdr_12.9755.yaml"]}


class Step3:
    def __init__(self, model_type="bs_roformer"):
        model_path, config_path = MODEL_PATHS[model_type]
        
        self.log = logging.getLogger(self.__class__.__name__)
        self.log.setLevel(logging.INFO)
        self.device = 'cpu'
        if torch.cuda.is_available():
            self.device = 'cuda'
        elif torch.backends.mps.is_available():
            self.device = 'mps'
        else:
            self.log.warning('CUDA/MPS are not available, running on CPU')
        
        self.model_type = model_type

        # self.model, self.config = get_model_from_config(model_type, config_path)
        self.config = load_config(model_type, config_path)
        self.model = BSRoformer(**dict(self.config.model))
        args = argparse.Namespace()
        args.start_check_point = model_path
        args.model_type = model_type
        args.lora_checkpoint = ''
        load_start_checkpoint(args, self.model, type_='inference')
        self.model = self.model.to(self.device)
        self.sample_rate = getattr(self.config.audio, 'sample_rate', 44100)

        
    def run(self,
            input_audio_path,
            temp_store_dir,  # for remove result dir
            output_dir,  # for final dir
            disable_detailed_pbar: bool=False,
            use_tta: bool= False,
            extract_instrumental: bool=True,
            codec="wav",
            subtype="FLOAT",
            draw_spectro=0,
            ):
        
        # self.log.info("Step3: Remove voice-over from audio.")
        
        os.makedirs(output_dir, exist_ok=True)
        
        if disable_detailed_pbar:
            detailed_pbar = False
        else:
            detailed_pbar = True

        instruments = prefer_target_instrument(self.config)[:]
        
        mix, sr = librosa.load(input_audio_path, sr=self.sample_rate, mono=False)
        # If mono audio we must adjust it depending on model
        if len(mix.shape) == 1:
            mix = np.expand_dims(mix, axis=0)
            if 'num_channels' in self.config.audio:
                if self.config.audio['num_channels'] == 2:
                    print(f'Convert mono track to stereo...')
                    mix = np.concatenate([mix, mix], axis=0)

        mix_orig = mix.copy()
        if 'normalize' in self.config.inference:
            if self.config.inference['normalize'] is True:
                mix, norm_params = normalize_audio(mix)

        waveforms_orig = demix(self.config, self.model, mix, self.device, model_type=self.model_type, pbar=detailed_pbar)
        if use_tta:
            waveforms_orig = apply_tta(self.config, self.model, mix, waveforms_orig, self.device, self.model_type)

        if extract_instrumental:
            instr = 'vocals' if 'vocals' in instruments else instruments[0]
            waveforms_orig['instrumental'] = mix_orig - waveforms_orig[instr]
            if 'instrumental' not in instruments:
                instruments.append('instrumental')

        file_name = os.path.splitext(os.path.basename(input_audio_path))[0].replace(".step1", "")
        temp_output_dir = os.path.join(temp_store_dir, file_name)
        os.makedirs(temp_output_dir, exist_ok=True)

        for instr in instruments:
            estimates = waveforms_orig[instr]
            if 'normalize' in self.config.inference:
                if self.config.inference['normalize'] is True:
                    estimates = denormalize_audio(estimates, norm_params)

            output_path = os.path.join(temp_output_dir, f"{instr}.{codec}")
            sf.write(output_path, estimates.T, sr, subtype=subtype)
            if draw_spectro > 0:
                output_img_path = os.path.join(temp_output_dir, f"{instr}.jpg")
                draw_spectrogram(estimates.T, sr, draw_spectro, output_img_path)


        instrumental_file = os.path.join(temp_output_dir, 'instrumental.wav')
        step3_audio_path = f"{output_dir}/{file_name}.step3.wav"
        subprocess.run(['cp', instrumental_file, step3_audio_path])

        self.log.info(f"The voice-over has been removed, and the audio is saved in {step3_audio_path}")
        self.log.info("Finish Step3 successfully.\n")
        return step3_audio_path