import os 

os.system("wget https://csteinmetz1.github.io/steerable-nafx/models/compressor_full.pt")
os.system("wget https://csteinmetz1.github.io/steerable-nafx/models/reverb_full.pt")
os.system("wget https://csteinmetz1.github.io/steerable-nafx/models/amp_full.pt")
os.system("wget https://csteinmetz1.github.io/steerable-nafx/models/delay_full.pt")
os.system("wget https://csteinmetz1.github.io/steerable-nafx/models/synth2synth_full.pt")

import sys
import math
import torch
import librosa.display
import auraloss
import torchaudio
import numpy as np
import scipy.signal
from tqdm.notebook import tqdm
from time import sleep
import pyloudnorm as pyln
import gradio as gr

def measure_rt60(h, fs=1, decay_db=30, rt60_tgt=None):
    """
    Analyze the RT60 of an impulse response.
    Args:
        h (ndarray): The discrete time impulse response as 1d array.
        fs (float, optional): Sample rate of the impulse response. (Default: 48000)
        decay_db (float, optional): The decay in decibels for which we actually estimate the time. (Default: 60)
        rt60_tgt (float, optional): This parameter can be used to indicate a target RT60. (Default: None)
    Returns:
        est_rt60 (float): Estimated RT60.
    """

    h = np.array(h)
    fs = float(fs)

    # The power of the impulse response in dB
    power = h ** 2
    energy = np.cumsum(power[::-1])[::-1]  # Integration according to Schroeder

    try:
        # remove the possibly all zero tail
        i_nz = np.max(np.where(energy > 0)[0])
        energy = energy[:i_nz]
        energy_db = 10 * np.log10(energy)
        energy_db -= energy_db[0]

        # -5 dB headroom
        i_5db = np.min(np.where(-5 - energy_db > 0)[0])
        e_5db = energy_db[i_5db]
        t_5db = i_5db / fs

        # after decay
        i_decay = np.min(np.where(-5 - decay_db - energy_db > 0)[0])
        t_decay = i_decay / fs

        # compute the decay time
        decay_time = t_decay - t_5db
        est_rt60 = (60 / decay_db) * decay_time
    except:
        est_rt60 = np.array(0.0)

    return est_rt60
    
def causal_crop(x, length: int):
    if x.shape[-1] != length:
        stop = x.shape[-1] - 1
        start = stop - length
        x = x[..., start:stop]
    return x

class FiLM(torch.nn.Module):
    def __init__(
        self,
        cond_dim,  # dim of conditioning input
        num_features,  # dim of the conv channel
        batch_norm=True,
    ):
        super().__init__()
        self.num_features = num_features
        self.batch_norm = batch_norm
        if batch_norm:
            self.bn = torch.nn.BatchNorm1d(num_features, affine=False)
        self.adaptor = torch.nn.Linear(cond_dim, num_features * 2)

    def forward(self, x, cond):

        cond = self.adaptor(cond)
        g, b = torch.chunk(cond, 2, dim=-1)
        g = g.permute(0, 2, 1)
        b = b.permute(0, 2, 1)

        if self.batch_norm:
            x = self.bn(x)  # apply BatchNorm without affine
        x = (x * g) + b  # then apply conditional affine

        return x

class TCNBlock(torch.nn.Module):
  def __init__(self, in_channels, out_channels, kernel_size, dilation, cond_dim=0, activation=True):
    super().__init__()
    self.conv = torch.nn.Conv1d(
        in_channels, 
        out_channels, 
        kernel_size, 
        dilation=dilation, 
        padding=0, #((kernel_size-1)//2)*dilation,
        bias=True)
    if cond_dim > 0:
      self.film = FiLM(cond_dim, out_channels, batch_norm=False)
    if activation:
      #self.act = torch.nn.Tanh()
      self.act = torch.nn.PReLU()
    self.res = torch.nn.Conv1d(in_channels, out_channels, 1, bias=False)

  def forward(self, x, c=None):
    x_in = x
    x = self.conv(x)
    if hasattr(self, "film"):
      x = self.film(x, c)
    if hasattr(self, "act"):
      x = self.act(x)
    x_res = causal_crop(self.res(x_in), x.shape[-1])
    x = x + x_res

    return x

class TCN(torch.nn.Module):
  def __init__(self, n_inputs=1, n_outputs=1, n_blocks=10, kernel_size=13, n_channels=64, dilation_growth=4, cond_dim=0):
    super().__init__()
    self.kernel_size = kernel_size
    self.n_channels = n_channels
    self.dilation_growth = dilation_growth
    self.n_blocks = n_blocks
    self.stack_size = n_blocks

    self.blocks = torch.nn.ModuleList()
    for n in range(n_blocks):
      if n == 0:
        in_ch = n_inputs
        out_ch = n_channels
        act = True
      elif (n+1) == n_blocks:
        in_ch = n_channels
        out_ch = n_outputs
        act = True
      else:
        in_ch = n_channels
        out_ch = n_channels
        act = True
      
      dilation = dilation_growth ** n
      self.blocks.append(TCNBlock(in_ch, out_ch, kernel_size, dilation, cond_dim=cond_dim, activation=act))

  def forward(self, x, c=None):
    for block in self.blocks:
      x = block(x, c)

    return x
  
  def compute_receptive_field(self):
    """Compute the receptive field in samples."""
    rf = self.kernel_size
    for n in range(1, self.n_blocks):
        dilation = self.dilation_growth ** (n % self.stack_size)
        rf = rf + ((self.kernel_size - 1) * dilation)
    return rf
    
# setup the pre-trained models
model_comp = torch.load("compressor_full.pt", map_location="cpu").eval()
model_verb = torch.load("reverb_full.pt", map_location="cpu").eval()
model_amp = torch.load("amp_full.pt", map_location="cpu").eval()
model_delay = torch.load("delay_full.pt", map_location="cpu").eval()
model_synth = torch.load("synth2synth_full.pt", map_location="cpu").eval()



def inference(aud, effect_type,gain_dB,c0,c1,mix,width,max_length,stereo,tail):
  x_p, sample_rate = torchaudio.load(aud)
  
  effect_type = effect_type #@param ["Compressor", "Reverb", "Amp", "Analog Delay", "Synth2Synth"]
  gain_dB = gain_dB #@param {type:"slider", min:-24, max:24, step:0.1}
  c0 = c0 #@param {type:"slider", min:-10, max:10, step:0.1}
  c1 = c1 #@param {type:"slider", min:-10, max:10, step:0.1}
  mix = mix #@param {type:"slider", min:0, max:100, step:1}
  width = width #@param {type:"slider", min:0, max:100, step:1}
  max_length = max_length#@param {type:"slider", min:5, max:120, step:1}
  stereo = stereo #@param {type:"boolean"}
  tail = tail #@param {type:"boolean"}
  
  # select model type
  if effect_type == "Compressor":
    pt_model = model_comp
  elif effect_type == "Reverb":
    pt_model = model_verb
  elif effect_type == "Amp":
    pt_model = model_amp
  elif effect_type == "Analog Delay":
    pt_model = model_delay
  elif effect_type == "Synth2Synth":
    pt_model = model_synth
  
  # measure the receptive field
  pt_model_rf = pt_model.compute_receptive_field()
  
  # crop input signal if needed
  max_samples = int(sample_rate * max_length)
  x_p_crop = x_p[:,:max_samples]
  chs = x_p_crop.shape[0]
  
  # if mono and stereo requested
  if chs == 1 and stereo:
    x_p_crop = x_p_crop.repeat(2,1)
    chs = 2
  
  # pad the input signal
  front_pad = pt_model_rf-1
  back_pad = 0 if not tail else front_pad
  x_p_pad = torch.nn.functional.pad(x_p_crop, (front_pad, back_pad))
  
  # design highpass filter
  sos = scipy.signal.butter(
      8, 
      20.0, 
      fs=sample_rate, 
      output="sos", 
      btype="highpass"
  )
  
  # compute linear gain 
  gain_ln = 10 ** (gain_dB / 20.0)
  
  # process audio with pre-trained model
  with torch.no_grad():
    y_hat = torch.zeros(x_p_crop.shape[0], x_p_crop.shape[1] + back_pad)
    for n in range(chs):
      if n == 0:
        factor = (width*5e-3)
      elif n == 1:
        factor = -(width*5e-3)
      c = torch.tensor([float(c0+factor), float(c1+factor)]).view(1,1,-1)
      y_hat_ch = pt_model(gain_ln * x_p_pad[n,:].view(1,1,-1), c)
      y_hat_ch = scipy.signal.sosfilt(sos, y_hat_ch.view(-1).numpy())
      y_hat_ch = torch.tensor(y_hat_ch)
      y_hat[n,:] = y_hat_ch
  
  # pad the dry signal 
  x_dry = torch.nn.functional.pad(x_p_crop, (0,back_pad))
  
  # normalize each first
  y_hat /= y_hat.abs().max()
  x_dry /= x_dry.abs().max()
  
  # mix
  mix = mix/100.0
  y_hat = (mix * y_hat) + ((1-mix) * x_dry)
  
  # remove transient
  y_hat = y_hat[...,8192:]
  y_hat /= y_hat.abs().max()
  
  torchaudio.save("output.mp3", y_hat.view(chs,-1), sample_rate, compression=320.0)
  return "output.mp3"
  
title = "Steerable nafx"
description = """Gradio demo for Steerable discovery of neural audio effects. To use it, simply upload your audio, or click one of the examples to load them. Read more at the links below. 
Now set the audio effect parameters. Here are some more insights into the controls:

effect_type - Choose from one of the pre-trained models.
gain_dB - Adjust the input gain. This can have a big effect since the effects are very nonlinear.
c0 and c1 - These are the effect controls which will adjust perceptual aspects of the effect, depending on the effect type. Very large values will often result in more extreme effects.
mix - Control the wet/dry mix of the effect.
width - Increase stereo width of the effect.
max_length - If you uploaded a very long file this will truncate it.
stereo - Convert mono input to stereo output.
tail - If checked, we will also compute the effect tail (nice for reverbs)."""
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.02926' target='_blank'>Steerable discovery of neural audio effects</a> | <a href='https://github.com/csteinmetz1/steerable-nafx' target='_blank'>Github Repo</a></p><center><img src='https://visitor-badge.glitch.me/badge?page_id=akhaliq_steerablenafx' alt='visitor badge'></center>"

gr.Interface(
    inference, 
    [gr.inputs.Audio(type="filepath", label="Input"),gr.inputs.Dropdown(choices=["Compressor", "Reverb", "Amp", "Analog Delay", "Synth2Synth"], type="value", default="Analog Delay", label="Effect Type"),gr.inputs.Slider(minimum=-24, maximum=24, step=1, default=-24, label="gain dB"),gr.inputs.Slider(minimum=-10, maximum=10, step=0.1, default=-1.4, label="c0"),
    gr.inputs.Slider(minimum=-10, maximum=10, step=1, default=3, label="c1"),
    gr.inputs.Slider(minimum=0, maximum=100, step=1, default=70, label="mix"),
    gr.inputs.Slider(minimum=0, maximum=100, step=1, default=50, label="width"),
    gr.inputs.Slider(minimum=5, maximum=120, step=1, default=30, label="max length"),
    gr.inputs.Checkbox(default=True, label="stereo"),
    gr.inputs.Checkbox(default=True, label="tail")], 
    gr.outputs.Audio(type="file", label="Output"),
    title=title,
    description=description,
    article=article,
    examples=[['vocals.wav',"Analog Delay",-24,-1.4,3,70,50,30,True,True]],
    enable_queue=True
    ).launch(debug=True)