#!/usr/bin/python

import pyaudio
import wave
import opus
from opus import encoder, decoder
import time

CHUNK = 320
FORMAT = pyaudio.paInt16
SAMPLE_WIDTH = 2
CHANNELS = 1
RATE = 16000
RECORD_SECONDS = 5
WAVE_OUTPUT_FILENAME = "recorder_output.wav"

RECORD_CHUNKS = RECORD_SECONDS * 1000 / 20

print "Will Record %d chunks" % RECORD_CHUNKS

pcmdata = []

record_chunk_count = 0

def record_callback(in_data, frame_count, time_info, status):
    global pcmdata, RECORD_CHUNKS, record_chunk_count

    pcmdata.append(in_data)

    record_chunk_count += 1
    if record_chunk_count % 50 == 0:
        #print "recorded %ds, chunk bytes: %d" % ((i / 50 + 1),len(in_data))
        print "recorded %d bytes, chunk: %d" % (len(in_data),record_chunk_count)

    flag = pyaudio.paContinue
    if record_chunk_count == RECORD_CHUNKS:
        flag = pyaudio.paComplete
    return None, flag


p = pyaudio.PyAudio()
stream = p.open(format=FORMAT,
                channels=CHANNELS,
                rate=RATE,
                input=True,
                frames_per_buffer=CHUNK,
                start=False,
                stream_callback=record_callback)

stream_output = p.open(format=FORMAT,
                       channels=CHANNELS,
                       rate=RATE,
                       output=True,
                       frames_per_buffer=CHUNK)

print("* recording")
stream.start_stream()

# Wait the stream end
while stream.is_active():
    time.sleep(0.1)


# for i in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
#     data = stream.read(CHUNK)
#     if i % 50 == 0:
#         print "recorded %ds, chunk bytes: %d" % ((i / 50 + 1),len(data))
#     pcmdata.append(data)

pcmbytes = b''.join(pcmdata);
expected_record_pcmdata_len = CHANNELS * SAMPLE_WIDTH * RECORD_SECONDS * RATE
print("* done recording, expected %d bytes, recorded %d bytes" % (expected_record_pcmdata_len, len(pcmbytes)))
stream.stop_stream()
stream.close()

#codec setup
opus_chunk = 20 # 20ms
frames_per_chunk = RATE * opus_chunk / 1000 # 320 frames per 20ms chunk
bytes_per_chunk = frames_per_chunk * SAMPLE_WIDTH # 640 bytes per 20ms chunk
enc = encoder.Encoder(RATE ,CHANNELS,'voip')
enc.complexity = 8
enc.vbr = 1
enc.bitrate = 27800
enc.signal = 3001 # signal voice
dec = decoder.Decoder(RATE ,CHANNELS)

#encoding
print "Encode PCM frames ..."
start = time.time()
encdata = []
for i in range(0, len(pcmbytes) , bytes_per_chunk):
    x = pcmbytes[i: i + bytes_per_chunk]
    y = enc.encode(x, frames_per_chunk)
    encdata.append(y) #encode each frame into array
    #if (i / frames_per_chunk) % 10 == 0:
    #    print "chunk %d (%d/%d bytes)/ %d" % (i / frames_per_chunk, len(y), len(x), f.getnframes() / frames_per_chunk)
    #streamout.write(x)
elapsed = time.time() - start
print "Encode completed, elapsed %f s" % elapsed


#decoding
print "Decode OPUS frames ..."
start = time.time()
decdata = ''
for x in encdata:
    decdata += dec.decode(x,frames_per_chunk)
elapsed = time.time() - start
print "Decode completed, elapsed %f s" % elapsed

#save to wave
print "Saving PCM frames ..."
wf = wave.open(WAVE_OUTPUT_FILENAME, 'wb')
wf.setnchannels(CHANNELS)
wf.setsampwidth(p.get_sample_size(FORMAT))
wf.setframerate(RATE)
wf.writeframes(decdata)
wf.close()

# play the dec data
print "Playing decoded PCM frames ..."
stream_output.write(decdata)
stream_output.stop_stream()
stream_output.close()
p.terminate()


