import torch
import numpy as np
import gradio as gr
from librosa import resample
import logging


SAMPLE_RATE = 16000


def _convert_input(arr):
    if "int" in str(arr.dtype):
        max_val = np.iinfo(arr.dtype).max
        arr = (arr.astype(np.float64) / max_val).astype(np.float32)
    if len(arr.shape) == 1:
        pass
    elif len(arr.shape) == 2 and arr.shape[1] in (1, 2):
        arr = arr.mean(axis=1)
    else:
        raise ValueError("wrong input shape")
    return arr


def _convert_output(arr):
    return np.vstack([arr / 2, arr / 2]).T


#########
# Model #
#########


from denoiser import pretrained


class Args():
    def __init__(self, model_path=None, dns48=None, dns64=None, master64=None, valentini_nc=None):
        self.model_path = model_path
        self.dns48 = dns48
        self.dns64 = dns64
        self.master64 = master64
        self.valentini_nc = valentini_nc


model = pretrained.get_model(Args(master64=True))


def apply_model(arr, wet=0.99):
    """using facebook model"""
    arr_enh = model(torch.from_numpy(arr).view(1, -1))
    arr_enh = (arr_enh / max(arr_enh.abs().max().item(), 1)).cpu().detach().numpy()[0]
    arr_enh = (1 - wet) * arr + wet * arr_enh
    return arr_enh


##########
# Server #
##########


def main(audio, audio_rec):#, amount):
    if audio is None and audio_rec is None:
        return "no audio defined" 
    elif audio is None:
        sr, arr = audio_rec
    else:
        sr, arr = audio
    arr = _convert_input(arr)
    arr = resample(arr, sr, SAMPLE_RATE)
    arr_enh = apply_model(arr)#, wet=amount)
    arr_enh = _convert_output(arr_enh)
    return SAMPLE_RATE, arr_enh


iface = gr.Interface(
    fn=main, inputs=[
        gr.inputs.Audio("upload", label="upload a file", optional=True),
        gr.inputs.Audio("microphone", label="create a recording", optional=True),
        #gr.inputs.Slider(0.0, 1.0, step=0.01, default=0.99, label="Aggressiveness")
    ], outputs="audio",
    examples=[["samples/restaurant.wav", ""]],#, 0.99]],
    allow_screenshot=False, allow_flagging=True, server_name="0.0.0.0", server_port=7863
)
iface.launch(
    ssl=('../cert/383fff33778762e8.crt', '../cert/383fff33778762e8.key'),
)