import warnings
import numpy as np
import soundfile as sf
import pywt

# https://pywavelets.readthedocs.io/en/latest/index.html
def denoise(in_wav: str, out_wav: str):
	info = sf.info(in_wav)  # getting info of the audio
	rate = info.samplerate

	warnings.simplefilter('ignore')
	warnings.simplefilter('error', RuntimeWarning)
	warnings.simplefilter('error', UserWarning)

	with sf.SoundFile(out_wav, "w", samplerate=rate, channels=info.channels) as of:
		for block in sf.blocks(in_wav, int(rate * info.duration * 0.10)):
			try:
				# Check fixes "UserWarning: Level value of 2 is too high: all coefficients will experience boundary effects."
				# All zero blocks seem safe to ignore since there's no visible spectrogram difference and the audio sounds slightly better.
				# The only concern is that longer pauses may be cut shorter.
				if not np.all(block == 0):
					# Set axis=0 for mono audio
					axis = 0 if info.channels == 1 else -1
					coefficients = pywt.wavedec(block, 'db4', mode='per', level=2, axis=axis)

					# getting variance of the input signal
					sigma = mad(coefficients[- 1])

					# VISU Shrink thresholding by applying the universal threshold proposed by Donoho and Johnstone
					thresh = sigma * np.sqrt(2 * np.log(len(block)))

					# thresholding using the noise threshold generated
					coefficients[1:] = (pywt.threshold(i, value=thresh, mode='soft') for i in coefficients[1:])

					# getting the clean signal as in original form and writing to the file
					clean = pywt.waverec(coefficients, 'db4', mode='per', axis=axis)
					of.write(clean)
				else:
					of.write(block)
			except RuntimeWarning:
				# Caused by "RuntimeWarning: invalid value encountered in divide"
				# more than likely b/c the block is mostly quiet anyway. Write it as-is.
				of.write(block)
			except UserWarning:
				# With the check above, ending here means it's similar to the RuntimeWarning.
				# Therefore write the block as-is.
				of.write(block)