#!/usr/bin/python2
from __future__ import division
try:
    import cStringIO as StringIO
except:
    import StringIO

import gobject
gobject.threads_init()
import gst
from pylab import *
import time

samplfreq = 44100

def sample_to_nanotime(sample):
    return int(sample * gst.SECOND / 44100 +0.5)
def nanotime_to_sample(nanotime):
    return nanotime * 44100 / gst.SECOND

def test_data(N):
    largeN = ((N + (2 ** 15) - 1) // (2 ** 15)) * (2 ** 15)
    left = empty(largeN, dtype=int16)
    right = empty(largeN, dtype=int16)

    for i in xrange(0, largeN // (2 ** 15)):
        left[i * (2 ** 15):(i + 1) * (2 ** 15)] = i
        if i % 2 == 0:
            right[i * (2 ** 15):(i + 1) * (2 ** 15)] = arange(0, 2 ** 15, dtype=int16)
        else:
            right[i * (2 ** 15):(i + 1) * (2 ** 15)] = 2 ** 15 - 1 - arange(0, 2 ** 15, dtype=int16)

    stereo = empty(largeN * 2, dtype=int16)
    stereo[0::2] = left
    stereo[1::2] = right
    return stereo[:N * 2].tostring()

def generate_test_sound(filename, encoding_element_gstlaunch_string, samples):
    gstlaunch_string = 'appsrc name=source ! capsfilter caps=audio/x-raw-int,rate=44100,channels=2,depth=16,width=16,signed=true,endianness=1234 ! ' + encoding_element_gstlaunch_string + ' ! filesink location="' + filename + '"'
#    gstlaunch_string = 'audiotestsrc ! decodebin2 ! audioconvert dithering=0 !  ' + encoding_element_gstlaunch_string + ' ! filesink location="' + filename + '"'
#    gstlaunch_string = 'filesrc location=test.flac ! decodebin2 ! audioconvert  ! ' + encoding_element_gstlaunch_string + ' ! filesink location="' + filename + '"'
    pipeline = gst.parse_launch(gstlaunch_string)
    pipeline_src = pipeline.get_by_name('source')
    pipeline_bus = pipeline.get_bus()


    def watcher(bus, message):
        if message.type == gst.MESSAGE_EOS:
            loop.quit()
        return True
    pipeline_bus.add_watch(watcher)

    pipeline.set_state(gst.STATE_PLAYING)
    pipeline_src.emit('push-buffer', gst.Buffer(test_data(samples)))
    pipeline_src.emit('end-of-stream')

    loop = gobject.MainLoop()
    loop.run()

    pipeline.set_state(gst.STATE_NULL)


def build_answer(filename):
    result = StringIO.StringIO()
    pipeline = gst.parse_launch('filesrc location=' + filename + ' ! decodebin2  ! audioconvert dithering=0 ! appsink name=sink caps=audio/x-raw-int,rate=44100,channels=2,depth=16,width=16,signed=true,endianness=1234')
    pipeline_sink = pipeline.get_by_name('sink')
    pipeline_sink.set_sync(False)
    pipeline_sink.set_property('emit-signals', True)
    def buffer(sink):
        buffer = sink.emit('pull-buffer')
        result.write(buffer.data)
    pipeline_sink.connect('new_buffer', buffer)
    pipeline_bus = pipeline.get_bus()

    def watcher(bus, message):
        if message.type == gst.MESSAGE_EOS:
            loop.quit()
        return True
    pipeline_bus.add_watch(watcher)

    pipeline.set_state(gst.STATE_PLAYING)

    loop = gobject.MainLoop()
    loop.run()

    pipeline.set_state(gst.STATE_NULL)
    result.seek(0)
#    result_string = result.read()
#    result_int = fromstring(result_string, dtype=int16)
#    plot(result_int[1::2])
#    show()
    return result.read()


def test_seek_FORMAT_TIME_by_sample(filename, seek_positions):
    answer = build_answer(filename)
    pipeline = gst.parse_launch('filesrc location=' + filename + ' ! decodebin2  ! audioconvert dithering=0 ! appsink name=sink caps=audio/x-raw-int,rate=44100,channels=2,depth=16,width=16,signed=true,endianness=1234')
    pipeline_sink = pipeline.get_by_name('sink')
    pipeline_sink.set_sync(False)

#    pipeline_bus = pipeline.get_bus()
#    def watcher(bus, message):
#        print message
#        if message.type == gst.MESSAGE_EOS:
#            loop.quit()
#        return True
#    pipeline_bus.add_watch(watcher)


    pipeline.set_state(gst.STATE_PLAYING)

#    time.sleep(0.1) #TODO: wait for an event, for some reason, get_state() is enough
    print "{:>9}{:>30}{:>10}".format('requested', 'buffer_timestamp_to_sample', 'actual')
    pipeline.get_state()
    for seek_position in seek_positions:
        pipeline.seek_simple(gst.FORMAT_TIME, gst.SEEK_FLAG_FLUSH | gst.SEEK_FLAG_ACCURATE, sample_to_nanotime(seek_position))
        buffer = pipeline_sink.emit('pull-buffer')
        actual_position = answer.find(buffer.data) // 4
        buffer_timestamp_position = nanotime_to_sample(buffer.timestamp)
        print "{:>9}{:>30}{:>10}".format(seek_position,buffer_timestamp_position,actual_position)
#        print str(seek_position).rjust(9), str(buffer_timestamp_position).rjust(30),  str(actual_position).rjust(10)
    pipeline.set_state(gst.STATE_NULL)
#        plot_data(buffer.data)

if __name__ == '__main__':
    sec = 10
    generate_test_sound('test.wav', 'wavenc', 44100 * sec)
    generate_test_sound('test.flac', 'flacenc', 44100 * sec)
    generate_test_sound('test.ogg', 'audioconvert dithering=0 ! vorbisenc quality=1 ! oggmux', 44100 * sec)
    generate_test_sound('test.mp3', 'lamemp3enc bitrate=320', 44100 * sec)

    test_samples = range(44100, 88200, 5000)
    print 'wav'
    test_seek_FORMAT_TIME_by_sample('test.wav', test_samples)
    print
    print 'flac'
    test_seek_FORMAT_TIME_by_sample('test.flac', test_samples)
    print
    print 'ogg'
    test_seek_FORMAT_TIME_by_sample('test.ogg', test_samples)
    print
    print 'mp3'
    test_seek_FORMAT_TIME_by_sample('test.mp3', test_samples)
