# separation_utils.py
import numpy as np
import os
import librosa
import soundfile as sf
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.special import rel_entr
import nussl
import types

# Required to patch mask validation for nussl

def _validate_mask_patched(self, mask_):
    assert isinstance(mask_, np.ndarray), 'Mask must be a numpy array!'
    if mask_.dtype == bool:
        return mask_
    mask_ = mask_ > 0.5
    if not np.all(np.logical_or(mask_, np.logical_not(mask_))):
        raise ValueError('All mask entries must be 0 or 1.')
    return mask_

nussl.core.masks.binary_mask.BinaryMask._validate_mask = types.MethodType(
    _validate_mask_patched, nussl.core.masks.binary_mask.BinaryMask)

# Separation methods
def Repet(mix):
    return nussl.separation.primitive.Repet(mix)( )

def Repet_Sim(mix):
    return nussl.separation.primitive.RepetSim(mix)( )

def Two_DFT(mix):
    return nussl.separation.primitive.FT2D(mix)( )

# Audio metrics
def calculate_psnr(clean_signal, separated_signal):
    min_length = min(len(clean_signal), len(separated_signal))
    clean_signal = clean_signal[:min_length]
    separated_signal = separated_signal[:min_length]
    mse = np.mean((clean_signal - separated_signal) ** 2)
    if mse == 0:
        return float('inf')
    max_val = np.max(np.abs(clean_signal))
    return 10 * np.log10((max_val ** 2) / mse)

def calculate_melspectrogram_kl_divergence(clean_signal, separated_signal, sr):
    clean_mel = compute_mel_spectrogram(clean_signal, sr)
    separated_mel = compute_mel_spectrogram(separated_signal, sr)
    clean_mel_norm = clean_mel / np.sum(clean_mel)
    separated_mel_norm = separated_mel / np.sum(separated_mel)
    return np.sum(rel_entr(np.clip(clean_mel_norm, 1e-10, None), np.clip(separated_mel_norm, 1e-10, None)))

def compute_mel_spectrogram(signal, sr, n_fft=2048, hop_length=512, n_mels=128):
    return librosa.feature.melspectrogram(
        y=signal, sr=sr, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels, power=2.0
    )

# Main function used in Gradio

def process_audio(file_path):
    signal = nussl.AudioSignal(file_path)
    mix_signal, sr1 = librosa.load(file_path, sr=None)

    ft2d_bg, ft2d_fg = Two_DFT(signal)
    repet_bg, repet_fg = Repet(signal)
    rsim_bg, rsim_fg = Repet_Sim(signal)

    output_file1 = "output_foreground_2dft.wav"
    output_file2 = "output_foreground_repet.wav"
    output_file3 = "output_foreground_rsim.wav"

    ft2d_fg.write_audio_to_file(output_file1)
    repet_fg.write_audio_to_file(output_file2)
    rsim_fg.write_audio_to_file(output_file3)

    output_snr1 = calculate_psnr(signal.audio_data, ft2d_fg.audio_data)
    output_snr2 = calculate_psnr(signal.audio_data, repet_fg.audio_data)
    output_snr3 = calculate_psnr(signal.audio_data, rsim_fg.audio_data)

    output_kl1 = calculate_melspectrogram_kl_divergence(signal.audio_data, ft2d_fg.audio_data, sr1)
    output_kl2 = calculate_melspectrogram_kl_divergence(signal.audio_data, repet_fg.audio_data, sr1)
    output_kl3 = calculate_melspectrogram_kl_divergence(signal.audio_data, rsim_fg.audio_data, sr1)

    return output_file1, output_snr1, output_kl1, output_file2, output_snr2, output_kl2, output_file3, output_snr3, output_kl3