def estimate_beats(audio_path, onset_type, plot=False):
    """Compute beat positions using either a spectral flux or a machine learned onset novelty function,
    followed by computing a tempogram and PLP.

    When onset_type = 'spectral_flux', use librosa.onset.onset strength to compute the novelty function.
    When onset_type = 'machine_learning', use madmom.features.beats.RNNBeatProcessor() to compute the
    novelty function.

    When `plot=True`, plot the onset novelty function and the estimated beats as vertical
    lines on the same plot. Your plot should include all axis labels and the time stamps
    should be in seconds.

    Parameters
    ----------
    audio_path : str
        Path to input audio file
    onset_type : str
        One of 'spectral_flux' or 'machine_learning'
    plot : bool
        If True, plots the onset novelty curve and the estimated beat positions.

    Returns
    -------
    beat_times : 1-d np.array
        Array of time stamps of the estimated beats in seconds.

    """


    #Load audio
    y, sr = librosa.load(audio_path)

    #Spectrogram for reference:
#     D = np.abs(librosa.stft(y))

    if onset_type == 'spectral_flux':
        #Onset envelope:
        onset_env = librosa.onset.onset_strength(y=y, sr=sr)
        times = librosa.times_like(onset_env)
#         print('onset shape1:', onset_env.shape)

        #PLP function
        plp_function = librosa.beat.plp(onset_envelope=onset_env, sr=sr)
        beats_plp = np.flatnonzero(librosa.util.localmax(plp_function))


        if (plot == True):
            plt.figure(figsize=(20,10))
            ax1 = plt.subplot(2, 1, 1)
#             plt.axis('tight')
#             librosa.display.specshow(librosa.amplitude_to_db(D, ref=np.max), y_axis='log', x_axis='time')
#             plt.title('Power spectrogram')

            #Plot onset
            plt.subplot(2, 1, 1, sharex=ax1)
            plt.ylabel('Onset & Est Beats')
            plt.plot(times, librosa.util.normalize(onset_env))

            #Est beats:

#             times = librosa.times_like(onset_env, sr=sr)
            plt.vlines(times[beats_plp], 0, 1, alpha=0.5, color='r', linestyle='--', label='PLP Beats')


        return times[beats_plp]





    if onset_type == 'machine_learning':
        # use madmom.features.beats.RNNBeatProcessor() to compute the machine learned novelty curve
        #Onset envelope:
        proc = madmom.features.beats.RNNBeatProcessor()
        onset_env = proc(audio)
        print('onset shape:', onset_env.shape)
#         onset_env = librosa.samples_to_time(onset_env, sr)
#         beat_samples = librosa.frames_to_samples(onset_env, 100)
#         print('beat_samples shape', beat_samples.shape)
#         times = librosa.times_like(onset_env, sr)#sr doesn't make a difference really
        times = librosa.samples_to_time(onset_env, 480000)
        print('time shape:', times.shape)

#         times_linspace = np.linspace(0, ( len(audio)/sr ) )
#         print('linspace len: ', len(times_linspace))

        #PLP function
        plp_function = librosa.beat.plp(onset_envelope=onset_env, sr=sr)
        beats_plp = np.flatnonzero(librosa.util.localmax(plp_function))


        if plot is True:
            plt.figure(figsize=(20,5))
            plt.plot(times)
            plt.vlines(times[beats_plp], 0, 1, alpha=0.5, color='r', linestyle='--', label='Beats [PLP]')
            plt.legend(frameon=True, framealpha=0.75)
            plt.title('Machine learning')

        return times[beats_plp]



    pass