import essentia.standard as ess
import json
import numpy as np
from scipy.signal import firwin, lfilter, freqz
import math

DEF_PARAMS = {}
DEF_PARAMS['audio_samplerate'] = 44100
DEF_PARAMS['framerate'] = 30
DEF_PARAMS['hop_size'] = int(DEF_PARAMS['audio_samplerate'] / DEF_PARAMS['framerate'])
DEF_PARAMS['window_size'] = DEF_PARAMS['hop_size'] * 4

def extractAllFeatures(audio, params):
    """Extract all dynamic features using Essentia's standard extractor.

    Args:
        audio (essentia audio) - loaded audio file
        params (dict) - Dict of analysis parameters (req: framerate, hop_size,
                        window_size)

    Returns:
        features (dictionary) - Dict containing the results of the analysis
    """
    pool = ess.Extractor(
        sampleRate=params['audio_samplerate'],
        lowLevelFrameSize=params['window_size'],
        lowLevelHopSize=params['hop_size'],
        rhythm=False
    )(audio)
    params['num_frames'] = math.floor(len(audio) / params['hop_size'] + 2)
    
    features = {}
    for descriptor in pool.descriptorNames():
        val = pool[descriptor]
        
        if type(val) is np.ndarray and val.shape[0] == params['num_frames']:
            features[descriptor] = pool[descriptor]
    return features, params

def smoothPool(features, params, cutoff, N=201):
    """Apply a lowpass filter to each feature of the pool for smoothing.

    Args:
        pool (essentia Pool) - Pool containing the results of the analysis
        cutoff (float) - cutoff frequency ([0. 1.]])

    Returns:
        features (dictionary) - Dict containing the results of the analysis
    """
    filtered_features = {}
    for (descriptor, data) in features.items():
        if len(data.shape) == 1:
            data.shape = (data.shape[0], 1)
        nCols = data.shape[1]
        nyqFreq = params['framerate'] / 2.
        mean = np.mean(data, axis=0)
        # std = np.std(data, axis=0)
        # if std == 0:
        #     std = 1.
        b = firwin(N, cutoff=cutoff / nyqFreq, window="hamming")
        # w, h = freqz(b)
        # plt.plot(w, 20 * np.log10(abs(h)), 'b')
        # plt.ylabel('Amplitude [dB]', color='b')
        # plt.xlabel('Frequency [rad/sample]')
        # plt.show()
        N2 = math.floor((N-1)/2)
        data = np.vstack((
            np.zeros((N2, nCols)),
            (data - mean),  # / std
            np.zeros((N2, nCols))
        ))
        filtered_data = lfilter(b, 1, data, axis=0)[N - 1:, :] + mean  # * std
        if nCols == 1:
            filtered_data = filtered_data[:, 0]
        filtered_features[descriptor] = filtered_data
    return filtered_features


def smoothFeatures(features, params, cutoff, N=201):
    data = features
    nCols = data.shape[1]
    nyqFreq = params['framerate'] / 2.
    mean = np.mean(data, axis=0)
   
    b = firwin(N, cutoff=cutoff / nyqFreq, window="hamming")
   
    N2 = math.floor((N-1)/2)
    data = np.vstack((
        np.zeros((N2, nCols)),
        (data - mean),  # / std
        np.zeros((N2, nCols))
    ))
    filtered_data = lfilter(b, 1, data, axis=0)[N - 1:, :] + mean  # * std
    
    return filtered_data


def extract_audio_jules(filename, params=DEF_PARAMS):    
    params['filename'] = filename
    audio_data = ess.MonoLoader(filename=params['filename'], sampleRate=params['audio_samplerate'])()

    features = {}
    features['raw'], params = extractAllFeatures(audio_data, params)
    features['filtered1'] = smoothPool(
        features['raw'], params, cutoff=1.
    )
    # features['filtered2'] = smoothPool(
        # features['raw'], params, cutoff=0.1
    # )
    
    return features, audio_data


def extract_audio_segment_jules(filename, start_sec, length, params=DEF_PARAMS, feat_type='filtered1'): 
    features, audio_data = extract_audio_jules(filename, params)
    
    v = features[feat_type]    
    file_len = v['sfx.oddtoevenharmonicenergyratio'].shape[0]

    sti = start_sec * params['framerate']
    eni = (start_sec + length)*params['framerate']

    if eni > file_len:
        import warnings
        warnings.warn("The requested length is larger than file length. Using all of the remaining frames in the file instead.")

    
    keys = [k for k in features[feat_type].keys()  if k != 'sfx.oddtoevenharmonicenergyratio' and k!= 'lowLevel.silence_rate_20dB']
    keys.sort() # this is necessary. otherwise each time the order of the keys might be different


    
    for k in keys:        
        if len(v[k].shape) < 2:
            v[k] = v[k][:,np.newaxis]

    # Concatenate all the features into one
    output = {}
    output['features'] = np.concatenate([v[k][sti:eni] for k in keys], axis=1)

    output['audio'] = audio_data[start_sec*params['audio_samplerate']:(start_sec + length)*params['audio_samplerate']]
    # output['audio'] = output['audio'].astype(np.int16)
    output['audio_samplerate'] = params['audio_samplerate']

    return output



### Omid's additions

def extract_all_features_o0(audio, params, smooth=False):
    """Extract all dynamic features using Essentia's standard extractor.

    Args:
        audio (essentia audio) - loaded audio file
        params (dict) - Dict of analysis parameters (req: framerate, hop_size,
                        window_size)

    Returns:
        features (dictionary) - Dict containing the results of the analysis
    """
    pool = ess.Extractor(
        sampleRate=params['audio_samplerate'],
        lowLevelFrameSize=params['window_size'],
        lowLevelHopSize=params['hop_size'],
        rhythm=True
    )(audio)
    num_frames = math.floor(len(audio) / params['hop_size'] + 1)
    
    print(num_frames)
    
    features = {}
    
    keys = [k for k in pool.descriptorNames()  if k != 'sfx.oddtoevenharmonicenergyratio' and k!= 'lowLevel.silence_rate_20dB']
    keys.sort() # this is necessary. otherwise each time the order of the keys might be different
    
    # print(num_frames)
    for k in keys:
#         print(k)
        val = pool[k]
        
        if type(val) is np.ndarray:
            #print(val.shape[0])
            if val.shape[0] == num_frames or val.shape[0] == num_frames + 1:
                features[k] = pool[k]
    
    if len(features.keys()) == 0:
    	print('no features extracted')
    	for k in keys:
#         print(k)
          val = pool[k]
        
          if type(val) is np.ndarray:
              print(val.shape)
    	
    if smooth:
        features = smoothPool(
            features, params, cutoff=3.
        )
    
    for k in features.keys(): 
        if len(features[k].shape) < 2:
            features[k] = features[k][:,np.newaxis]
    
        
    return features

def load_wav_seg(filename, start, length, sr=44100):
    audio_data = ess.MonoLoader(filename=filename, sampleRate=sr)()
    return audio_data[start*sr:(start+length)*sr]

def plot_feat_dict(feat):
    _=plt.figure(figsize=(16,8))
    for k in feat.keys():        
        plt.plot(feat[k])
        plt.ylim([-3,3])

def extract_wrapper_o0(filename, params, smooth=False):
    audio_data = ess.MonoLoader(filename=filename, sampleRate=params['audio_samplerate'])()
    feat_all = extract_all_features_o0(audio_data, params, smooth)
    features = {
        'all': feat_all,
        'lowLevel': {k: feat_all[k] for k in feat_all.keys() if k.startswith('lowLevel.')},
        'sfx': {k: feat_all[k] for k in feat_all.keys() if k.startswith('sfx.')},
        'tonal': {k: feat_all[k] for k in feat_all.keys() if k.startswith('tonal.')}
    }
    
    return features





def extract_audio_features_v2(raw_wave, params):
    selected_keys = [
     'sfx.tristimulus',
     'lowLevel.spectral_energy',
     'lowLevel.spectral_centroid',
     'lowLevel.spectral_complexity',
     'lowLevel.spectral_crest',
     'lowLevel.spectral_decrease',
     'lowLevel.spectral_energy',
     'lowLevel.spectral_energyband_high',
     'lowLevel.spectral_energyband_low',
     'lowLevel.spectral_energyband_middle_high',
     'lowLevel.spectral_energyband_middle_low',
     'lowLevel.spectral_flatness_db',
     'lowLevel.spectral_flux',
     'lowLevel.spectral_kurtosis',
     'lowLevel.spectral_rms',
     'lowLevel.spectral_rolloff',
     'lowLevel.spectral_skewness',
     'lowLevel.spectral_spread',
     'lowLevel.spectral_strongpeak']
    
#     raw_wave = ess.MonoLoader(filename=filename, sampleRate=params['audio_samplerate'])()
    
    pool = ess.Extractor(
        sampleRate=params['audio_samplerate'],
        lowLevelFrameSize=params['window_size'],
        lowLevelHopSize=params['hop_size'],
        rhythm=True
    )(raw_wave)
    
    num_frames = math.floor(len(raw_wave) / params['hop_size'] + 1)

    beat_pos = np.zeros(num_frames)
    beats_loudness = np.zeros(num_frames)
    bpm_intervals = np.zeros(num_frames)
    bpm_estimates = np.zeros(num_frames)
    beats_loudness_band_ratio = np.zeros((num_frames, 6))

    audio_features = {}
    audio_features['beat_pos'] = np.zeros(num_frames)
    audio_features['beat_signal'] = np.zeros(num_frames)
    audio_features['beats_bpm_intervals_signal'] = np.zeros(num_frames)
    audio_features['beats_loudness_signal'] = np.zeros(num_frames)
    audio_features['beats_loudness_band_ratio_signal'] = np.zeros((num_frames, 6))



    bb = []
    prev = 0


    for i in range(len(pool['rhythm.beats_position'])):
        bt = pool['rhythm.beats_position'][i]
        fr = int(bt * 30)
        bb.append((prev+bt)/2)
        bb.append(bt)
        prev = bt
        beat_pos[fr] = 1

        beats_loudness[fr] = pool['rhythm.beats_loudness'][i] if i < len(pool['rhythm.beats_loudness']) else 0
        bpm_intervals[fr] = pool['rhythm.bpm_intervals'][i] if i < len(pool['rhythm.bpm_intervals']) else 0
        bpm_estimates[fr] = pool['rhythm.bpm_estimates'][i] if i < len(pool['rhythm.bpm_estimates']) else 0
        beats_loudness_band_ratio[fr] = pool['rhythm.beats_loudness_band_ratio'][i] if i < len(pool['rhythm.beats_loudness_band_ratio']) else 0


    sig_pos = 0   
    d = 1
    for i in range(0,len(bb), 1):    
        fr = int(bb[i] * 30) #+ int(bb[i+1] * 30)


        if fr < len(beat_pos):
            for j in range(sig_pos, fr):    
                audio_features['beat_signal'][j] = d


        if i%2==1:                
            if fr < len(beat_pos):
                for j in range(sig_pos, fr):            
                    audio_features['beats_loudness_signal'][j] = beats_loudness[fr] * 10
                    audio_features['beats_loudness_band_ratio_signal'][j] = beats_loudness_band_ratio[fr] * 10
                    audio_features['beats_bpm_intervals_signal'][j] = bpm_intervals[fr] * 10

        sig_pos = fr

        d = d * -1
    
    ###
    
    audio_features['beat_pos'] = beat_pos
        
    for k in selected_keys:
        audio_features[k] = pool[k]
        
    for k in audio_features.keys(): 
        if len(audio_features[k].shape) < 2:
            audio_features[k] = audio_features[k][:,np.newaxis]
    
    return audio_features

###


def plot_feat_dict_rng(feat, rng=range(-1)):
    _=plt.figure(figsize=(16,8))
    for k in feat.keys():        
        plt.plot(feat[k][rng])
#         plt.ylim([-3,3])
