import gradio as gr

import tensorflow as tf
import tensorflow_hub as hub

import numpy as np
import matplotlib.pyplot as plt
import librosa
from librosa import display as librosadisplay

import logging
import math
import statistics
import sys

from IPython.display import Audio, Javascript
from scipy.io import wavfile

from base64 import b64decode

import music21
from pydub import AudioSegment
from IPython.core.display import display, HTML, Javascript
import json, random

EXPECTED_SAMPLE_RATE = 16000
MAX_ABS_INT16 = 32768.0
A4 = 440
C0 = A4 * pow(2, -4.75)
note_names = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]

MAX_ABS_INT16 = 32768.0

def plot_stft(x, sample_rate, show_black_and_white=False):
  x_stft = np.abs(librosa.stft(x, n_fft=2048))
  fig, ax = plt.subplots()
  fig.set_size_inches(20, 10)
  x_stft_db = librosa.amplitude_to_db(x_stft, ref=np.max)
  if(show_black_and_white):
    librosadisplay.specshow(data=x_stft_db, y_axis='log', 
                             sr=sample_rate, cmap='gray_r')
  else:
    librosadisplay.specshow(data=x_stft_db, y_axis='log', sr=sample_rate)

  plt.colorbar(format='%+2.0f dB')
  return fig

def showScore(score):
    xml = open(score.write('musicxml')).read()
    showMusicXML(xml)
    
def showMusicXML(score):
    xml = open(score.write('musicxml')).read()
    DIV_ID = "OSMD_div"
    display(HTML('<div id="'+DIV_ID+'">loading OpenSheetMusicDisplay</div>'))
    script = """
    var div_id = {{DIV_ID}};
    function loadOSMD() { 
        return new Promise(function(resolve, reject){
            if (window.opensheetmusicdisplay) {
                return resolve(window.opensheetmusicdisplay)
            }
            // OSMD script has a 'define' call which conflicts with requirejs
            var _define = window.define // save the define object 
            window.define = undefined // now the loaded script will ignore requirejs
            var s = document.createElement( 'script' );
            s.setAttribute( 'src', "https://cdn.jsdelivr.net/npm/opensheetmusicdisplay@0.7.6/build/opensheetmusicdisplay.min.js" );
            //s.setAttribute( 'src', "/custom/opensheetmusicdisplay.js" );
            s.onload=function(){
                window.define = _define
                resolve(opensheetmusicdisplay);
            };
            document.body.appendChild( s ); // browser will try to load the new script tag
        }) 
    }
    loadOSMD().then((OSMD)=>{
        window.openSheetMusicDisplay = new OSMD.OpenSheetMusicDisplay(div_id, {
          drawingParameters: "compacttight"
        });
        openSheetMusicDisplay
            .load({{data}})
            .then(
              function() {
                openSheetMusicDisplay.render();
              }
            );
    })
    """.replace('{{DIV_ID}}',DIV_ID).replace('{{data}}',json.dumps(xml))
    display(Javascript(script))
    return

def output2hz(pitch_output):
  # Constants taken from https://tfhub.dev/google/spice/2
  PT_OFFSET = 25.58
  PT_SLOPE = 63.07
  FMIN = 10.0;
  BINS_PER_OCTAVE = 12.0;
  cqt_bin = pitch_output * PT_SLOPE + PT_OFFSET;
  return FMIN * 2.0 ** (1.0 * cqt_bin / BINS_PER_OCTAVE)

def quantize_predictions(group, ideal_offset):
  # Group values are either 0, or a pitch in Hz.
  non_zero_values = [v for v in group if v != 0]
  zero_values_count = len(group) - len(non_zero_values)

  # Create a rest if 80% is silent, otherwise create a note.
  if zero_values_count > 0.8 * len(group):
    # Interpret as a rest. Count each dropped note as an error, weighted a bit
    # worse than a badly sung note (which would 'cost' 0.5).
    return 0.51 * len(non_zero_values), "Rest"
  else:
    # Interpret as note, estimating as mean of non-rest predictions.
    h = round(
        statistics.mean([
            12 * math.log2(freq / C0) - ideal_offset for freq in non_zero_values
        ]))
    octave = h // 12
    n = h % 12
    note = note_names[n] + str(octave)
    # Quantization error is the total difference from the quantized note.
    error = sum([
        abs(12 * math.log2(freq / C0) - ideal_offset - h)
        for freq in non_zero_values
    ])
    return error, note


def get_quantization_and_error(pitch_outputs_and_rests, predictions_per_eighth,
                               prediction_start_offset, ideal_offset):
  # Apply the start offset - we can just add the offset as rests.
  pitch_outputs_and_rests = [0] * prediction_start_offset + \
                            pitch_outputs_and_rests
  # Collect the predictions for each note (or rest).
  groups = [
      pitch_outputs_and_rests[i:i + predictions_per_eighth]
      for i in range(0, len(pitch_outputs_and_rests), predictions_per_eighth)
  ]

  quantization_error = 0

  notes_and_rests = []
  for group in groups:
    error, note_or_rest = quantize_predictions(group, ideal_offset)
    quantization_error += error
    notes_and_rests.append(note_or_rest)

  return quantization_error, notes_and_rests

def convert_audio_for_model(user_file, output_file='converted_audio_file.wav'):
  audio = AudioSegment.from_file(user_file)
  audio = audio.set_frame_rate(EXPECTED_SAMPLE_RATE).set_channels(1)
  audio.export(output_file, format="wav")
  return output_file

def hz2offset(freq):
  # This measures the quantization error for a single note.
    if freq == 0:  # Rests always have zero error.
      return None
  # Quantized note.
    h = round(12 * math.log2(freq / C0))
    return 12 * math.log2(freq / C0) - h

def greet(uploaded_file_name):

  converted_audio_file = convert_audio_for_model(uploaded_file_name)
  sample_rate, audio_samples = wavfile.read(converted_audio_file, 'rb')
  audio_samples = audio_samples / float(MAX_ABS_INT16)

  model = hub.load("https://tfhub.dev/google/spice/2")
  model_output = model.signatures["serving_default"](tf.constant(audio_samples, tf.float32))

  pitch_outputs = model_output["pitch"]
  uncertainty_outputs = model_output["uncertainty"]

  # 'Uncertainty' basically means the inverse of confidence.
  confidence_outputs = 1.0 - uncertainty_outputs

  confidence_outputs = list(confidence_outputs)
  pitch_outputs = [ float(x) for x in pitch_outputs]

  indices = range(len (pitch_outputs))
  confident_pitch_outputs = [ (i,p)  
    for i, p, c in zip(indices, pitch_outputs, confidence_outputs) if  c >= 0.9  ]
  confident_pitch_outputs_x, confident_pitch_outputs_y = zip(*confident_pitch_outputs)

  pitch_outputs_and_rests = [
    output2hz(p) if c >= 0.9 else 0
    for i, p, c in zip(indices, pitch_outputs, confidence_outputs)
  ]
  
  offsets = [hz2offset(p) for p in pitch_outputs_and_rests if p != 0]

  ideal_offset = statistics.mean(offsets)
  
  best_error = float("inf")
  best_notes_and_rests = None
  best_predictions_per_note = None

  for predictions_per_note in range(20, 65, 1):
    for prediction_start_offset in range(predictions_per_note):

      error, notes_and_rests = get_quantization_and_error(
          pitch_outputs_and_rests, predictions_per_note,
          prediction_start_offset, ideal_offset)

      if error < best_error:      
        best_error = error
        best_notes_and_rests = notes_and_rests
        best_predictions_per_note = predictions_per_note

  # At this point, best_notes_and_rests contains the best quantization.
  # Since we don't need to have rests at the beginning, let's remove these:
  while best_notes_and_rests[0] == 'Rest':
    best_notes_and_rests = best_notes_and_rests[1:]
  # Also remove silence at the end.
  while best_notes_and_rests[-1] == 'Rest':
    best_notes_and_rests = best_notes_and_rests[:-1]

  sc = music21.stream.Score()
  # Adjust the speed to match the actual singing.
  bpm = 60 * 60 / best_predictions_per_note
  print ('bpm: ', bpm)
  a = music21.tempo.MetronomeMark(number=bpm)
  sc.insert(0,a)

  for snote in best_notes_and_rests:   
      d = 'half'
      if snote == 'Rest':      
        sc.append(music21.note.Rest(type=d))
      else:
        sc.append(music21.note.Note(snote, type=d))

  converted_audio_file_as_midi = converted_audio_file[:-4] + '.mid'
  fp = sc.write('midi', fp=converted_audio_file_as_midi)

  wav_from_created_midi = converted_audio_file_as_midi.replace(' ', '_') + "_midioutput.wav"
  #!timidity $converted_audio_file_as_midi -Ow -o $wav_from_created_midi

  #return Audio(wav_from_created_midi)

  # ------- PLOT 1 -------
  fig1 = plt.figure()
  plt.plot(audio_samples)

  # ------- PLOT 2 -------
  fig2, ax = plt.subplots()
  fig2.set_size_inches(90, 50)
  plt.plot(pitch_outputs, label='pitch')
  plt.plot(confidence_outputs, label='confidence')
  plt.legend(loc="lower right")

  # ------- PLOT 3 -------
  x = audio_samples / MAX_ABS_INT16
  sample_rate = EXPECTED_SAMPLE_RATE
  show_black_and_white=False

  x_stft = np.abs(librosa.stft(x, n_fft=2048))
  fig3, ax1 = plt.subplots()
  fig3.set_size_inches(20, 10)
  x_stft_db = librosa.amplitude_to_db(x_stft, ref=np.max)
  if(show_black_and_white):
    librosadisplay.specshow(data=x_stft_db, y_axis='log', 
                             sr=sample_rate, cmap='gray_r')
  else:
    librosadisplay.specshow(data=x_stft_db, y_axis='log', sr=sample_rate)

  # -------PLOT 4 -------
  fig4, ax2 = plt.subplots()
  fig4.set_size_inches(20, 10)
  ax2.set_ylim([0, 1])
  plt.scatter(confident_pitch_outputs_x, confident_pitch_outputs_y, )
  plt.scatter(confident_pitch_outputs_x, confident_pitch_outputs_y, c="r")

  # ------- PLOT 5 -------
  x = audio_samples / MAX_ABS_INT16
  sample_rate = EXPECTED_SAMPLE_RATE
  show_black_and_white=True

  x_stft = np.abs(librosa.stft(x, n_fft=2048))
  fig5, ax3 = plt.subplots()
  fig5.set_size_inches(20, 10)
  x_stft_db = librosa.amplitude_to_db(x_stft, ref=np.max)
  if(show_black_and_white):
    librosadisplay.specshow(data=x_stft_db, y_axis='log', 
                             sr=sample_rate, cmap='gray_r')
  else:
    librosadisplay.specshow(data=x_stft_db, y_axis='log', sr=sample_rate)

  confident_pitch_values_hz = [ output2hz(p) for p in confident_pitch_outputs_y ]  
  plt.scatter(confident_pitch_outputs_x, confident_pitch_values_hz, c="r")

  return fig1,fig2,fig3,fig4,fig5,uploaded_file_name,sc.show(fp)


#audio = gr.inputs.Audio(source="upload",type='filepath')

audio = gr.inputs.Audio(source="upload",type='filepath')

out = gr.outputs.Audio(type="auto", label='Salida')
fig1 = gr.outputs.Plot(type="auto")
fig2 = gr.outputs.Plot(type="auto")
fig3 = gr.outputs.Plot(type="auto")
fig4 = gr.outputs.Plot(type="auto")
fig5 = gr.outputs.Plot(type="auto")
out2 = gr.outputs.Audio(type="auto", label='Salida')


iface = gr.Interface(fn=greet, inputs=audio, outputs=[fig1, fig2, fig3, fig4, fig5, fig6, out, out2])
iface.launch(debug=True)