#!/usr/bin/env python
#
# grab audio from bag
#

import rosbag, rospy, numpy as np
import struct
import sys, os, cv2, glob
from itertools import izip, repeat
from subprocess import call
import subprocess
import argparse

depth_to_string = {16 : 'int16', 32: 'int32', None: 'int32'}


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Extract audio from bag files.')
    parser.add_argument('--start', '-s', action='store', default=rospy.Time(0), type=rospy.Time,
                        help='Rostime representing where to start in the bag.')
    parser.add_argument('--end', '-e', action='store', default=rospy.Time(sys.maxint), type=rospy.Time,
                        help='Rostime representing where to stop in the bag.')
    parser.add_argument('--depth', '-d', action='store', default=None, type=int,
                        help='force depth in bits (16 or 32) of audio samples.')
    parser.add_argument('topic')
    parser.add_argument('bagfile')

    args = parser.parse_args()

    out_file_name  = "audio.raw"
    wav_file_name  = "audio.wav"
    plot_file_name = "audio.dat"

    sample_list     = []
    tot_num_samples = 0
    num_samples     = 0 # does not include first packet
    num_packets     = 0 # does not include first packet
    t1              = 0
    t               = 0
    t_start         = 0 # start time of bag
    t_last          = 0
    drop_packets    = 5
    packet_cnt      = 0
    outfile = open(out_file_name, 'w')

    print "reading bag file..."
    depth = depth_to_string[args.depth]
    channels = 2
    rate = 48000
    for bagfile in glob.glob(args.bagfile):
        bag = rosbag.Bag(bagfile, 'r')
        t_start = bag.get_start_time()
        iterator = bag.read_messages(topics=args.topic, start_time=args.start, end_time=args.end)
        for (topic, msg, time) in iterator:
            if msg._type == 'audio_common_msgs/AudioDataStamped':
                if packet_cnt < drop_packets:
                    packet_cnt += 1
                    continue
                data = np.asarray(msg.data)
                outfile.write(data)
                channels = msg.format.channels
                rate = msg.format.rate
                if not args.depth:
                    depth = depth_to_string[msg.format.depth]
                arr = np.reshape(np.fromstring(msg.data, dtype=depth),(-1, channels))
                sample_list.append(arr)
                nsamples = arr.shape[0]
                tot_num_samples += nsamples
                t_record  = time.to_sec()
                t  = msg.header.stamp.to_sec()
                if t1 == 0:
                    # not counting the first packet!
                    t1 = t
                else:
                    num_samples += nsamples
                    num_packets += 1
                    t_last = t
    outfile.close()

    print "channels: %d, depth: %s" % (channels, depth)
    dt = (t - t1) / num_samples
    print "num samples: %d estimated rate: %f, recorded rate: %f" % (num_samples, (1.0/dt), rate)
    samples = np.zeros((tot_num_samples,channels), dtype=int)
    print samples.shape
    idx = 0
    for s in sample_list:
        samples[idx:(idx+s.shape[0]),:] = s
        idx = idx + s.shape[0]

    avg_packet_size = num_samples / num_packets
    t0 = t1 - avg_packet_size * dt  # shift to account for first packet lag

    print "writing to file  ", plot_file_name
    outfile = open(plot_file_name, 'wb')
    outfile.write(struct.pack('i', 0))  # file format version #
    outfile.write(depth) 
    outfile.write(struct.pack('i', tot_num_samples))
    outfile.write(struct.pack('i', channels))  #args.channels))
    # ros sample time ts = t0 + sample_number * dt
    # relative time = ts - t_start
    outfile.write(struct.pack('f'*5, *(rate, t_start, t0, t1, dt)))

    #perm = [3,6,7,5,4,0,1,2]   # XXX permutation rearrange mic's in right order
    #s_perm = samples[:,perm]
    s_perm = samples
    s_perm.astype(depth).tofile(outfile)

    outfile.close()


    print "wrote raw data to %s" % out_file_name
    enc = 's16le' if depth == 'int16' else 's32le'
    bashCommand = "ffmpeg -y -f %s -ar %f -ac %d -i %s %s" % (enc, float(rate), channels, out_file_name, wav_file_name)
    print bashCommand
    process = subprocess.Popen(bashCommand.split(), stdout=subprocess.PIPE)
    output, error = process.communicate()

    
