#!/usr/bin/env python3

import librosa
import numpy as np
from pathlib import Path
import pandas as pd


class WaveFile(object):
    filename = None
    duration = None

    def __init__(self, filename):
        
        y, sr = librosa.load(filename)
        self.duration = librosa.get_duration(y, sr)
        self.filename = str(filename)

        self.df = pd.DataFrame(np.abs(librosa.stft(y, n_fft=256)))
        self.df['bins'] = librosa.fft_frequencies(sr=sr, n_fft=256)
        self.df['avgamp'] = self.df.mean(axis=1)
        self.df = self.df[['bins', 'avgamp']]

        
    def almost_eq(self, other, threshold):        
        if abs(self.duration - other.duration) > 0.1:
            return False

        if max(abs(self.df['avgamp'] - other.df['avgamp'])) > threshold:
            return False

        return True
    

def main():
    print('loading ps1 waves')
    ps1waves = [WaveFile(path)
                for path in Path('stream1').rglob('*.wav')]


    print('loading PC waves')
    pcwaves = [WaveFile(path)
               for path in Path('extracted').rglob('*.wav')]


    print('comparing waves')
    for pc in pcwaves:
        candidates = []
        for threshold in [5,4,3,2,1,0.5,0.1]:
            candidates = [ps1.filename for ps1 in ps1waves if pc.almost_eq(ps1, threshold)]
            if len(candidates) == 1:
                print('%s;%s' % (pc.filename, candidates[0]))
                break

        else:
            print('%s;%s' % (pc.filename, candidates))

        
def test():
    Wario95d = WaveFile('samples/Wario95d.wav')
    Wario96c = WaveFile('samples/Wario96c.wav')
    wario96c = WaveFile('samples/wario96c.wav')

    assert Wario96c.almost_eq(wario96c)
    assert not Wario96c.almost_eq(Wario95d)



if __name__ == '__main__':
    main()