#!/usr/bin/env python
#
# read and plot audio file
#

import numpy as np
import matplotlib.pyplot as plt;
import struct
import sys, os
import argparse
import scipy.signal
import ipdb
from scipy.io import wavfile

def plot_wave(t, s):
    if s.ndim > 1:
        for i in range(0, s.shape[1]):
            plt.subplot(s.shape[1], 1, i + 1)
            plt.plot(t, s[:,i])
    else:
        plt.plot(t, s)
    plt.show()

def fft_correlate(a, b):
    af = np.fft.fft(a, axis=0)
    bf = np.fft.fft(b, axis=0)
    ft = np.multiply(np.conj(af), bf)
    cc = np.real(np.fft.ifft(ft))
    return cc

def read_file(f):
    ffmt = struct.unpack('i', f.read(struct.calcsize('i')))
    depth = f.read(5)
    sz   = struct.unpack('2i', f.read(struct.calcsize('2i')))
    # t_start: time at which the rosbag started
    # t1: time at which first audio packet was received
    # t0: time at which first sample must have been generated
    # dt: delta time between samples, matched to recording
    # t = t0 + idx * dt   = best rostime estimate
    # t - t_start = best estimate for time since start of bag recording
    rate, t_start, t0, t1, dt = struct.unpack('5f', f.read(struct.calcsize('5f')))
    print "bag start: %f, rate: %f, depth: %s, number of samples: %d, number of channels: %d" % (t_start, rate, depth, sz[0], sz[1])
    ts      = t0 + np.arange(0,sz[0]) * dt
    dtype   = np.int16 if depth == 'int16' else np.int32
    smpl    = np.fromfile(f,dtype=dtype)
    samples = np.reshape(smpl, sz)
    t       = t0 + np.arange(0, sz[0]) * dt - t_start;  # gives best estimate
    return t,  rate, samples

    
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='find audio onset.')
    parser.add_argument('--start', '-s', action='store', default=0.0, type=float,
                        help='start time')
    parser.add_argument('--end', '-e', action='store', default=1e30, type=float,
                        help='end time')
    parser.add_argument('--song', action='store', default='song.wav', help = 'wav file with pattern to search for')
    parser.add_argument('--file', action='store', required=True, help = 'audio.dat file generated from rosbag')
    parser.add_argument('--songstart', action='store', default=0, type=float, help = 'start time of call within song file')

    args = parser.parse_args()

    f = open(args.file)
    t, rate_recording, recording = read_file(f)
    rate_song, song = wavfile.read(args.song)
    #song = np.pad(song, (0,int(rate_recording*20)), 'constant', constant_values=(0,0)) # XXX
    
    # average channels of recording
    recording    = np.sum(recording, axis=1)  # add across channels to get average
    # resample song to recording frequency
    song = scipy.signal.resample(song, int(song.shape[0] * rate_recording / rate_song))
    # pad whichever is shorter with zeros
    num_pad = abs(recording.shape[0] - song.shape[0])
    if (song.shape[0] < recording.shape[0]):
        print "padding %d zero samples to song" % num_pad
        song = np.expand_dims(np.pad(song, (0,num_pad), 'constant', constant_values=(0,0)), -1)
    if (recording.shape[0] < song.shape[0]):
        print "padding %d zero samples to recording" % num_pad
        recording = np.pad(recording, (0,num_pad), 'constant', constant_values=(0,0))
    # expand dimensions so we can stack it
    #song = np.expand_dims(song, axis=-1)
    print 'recording rate: %d, song rate: %d' % (rate_recording, rate_song)
    recording_exp = np.expand_dims(recording, axis=-1)
    #cc = fft_correlate(song[:,0], samples[:,0])  # use this for individual channels
    cc = fft_correlate(song[:,0], recording)
    amax  = np.argmax(cc, axis=0)  # find maximum correlation
    song_rolled = np.roll(song, amax)
    idx     = np.where((t >= args.start) & (t <= args.end))
    t_ss = args.songstart + amax / rate_recording
    if t_ss > song.shape[0] / rate_recording:
        t_ss = t_ss - song.shape[0] / rate_recording
    print 'max correlation roll: %d, song start: %f' % (amax, t_ss)
    # stacked = np.hstack((song_rolled[idx[0],:], song[idx[0],:], cc[idx[0],:]))
    # stacked = np.hstack((song_rolled[idx[0],:], recording_exp[idx[0],:]))
    stacked = np.hstack((song[idx[0],:], recording_exp[idx[0],:], np.expand_dims(cc[idx[0]],-1)))
    plot_wave(t[idx[0]], stacked)
