import csv
import itertools
import numpy as np
import getopt
from os import path
from os import system
from sys import argv

try:
    opts, args = getopt.getopt(argv[1:], 'f:')
except getopt.GetoptError:
    print 'GopBitRate.py -f <frame rate> <stats csv>'
print opts
print args

frameRate = 30
for opt, arg in opts:
    if opt == '-f':
        frameRate = int(arg)
    else:
       print 'unknow opt: ', opt
print 'frame rate: %d' % frameRate

esFile = args[0]
stats = path.splitext(path.basename(esFile))[0] + '.csv'
system('ffprobe -of csv -show_frames ' + esFile + '>' + stats)

frameType = ''
count = 0
totalFrames = 0
totalSize = 0
gopFrames = []
gopSize = []
size = 0

with open(stats, 'r') as f:
    reader = csv.reader(f)
    for row in itertools.islice(reader, 0, None):
        try:
            count = int(row[-5])
            frameType = row[-6]
            totalFrames += 1
            size = int(row[13])
            totalSize += size
            if frameType == 'I':
                gopFrames.append(0)
                gopSize.append(0)
            gopFrames[-1] += 1
            gopSize[-1] += size
        except Exception, e:
            print str(e) + ': %s' % row

totalSize *= 8
print 'Total size (%d frames): %d bits' % (totalFrames, totalSize)
print 'Average bitrate: %f Mbps' % (totalSize * frameRate / totalFrames / 1000000.0)

normGopSize = map(lambda x,y:x/y*frameRate*8/1000000.0, gopSize, gopFrames)
maxGopSize = max(normGopSize)
maxGop = normGopSize.index(maxGopSize)
print 'Maximum bitrate(Gop#%d): %f Mbps' % (maxGop, normGopSize[maxGop])

for i in np.argsort(normGopSize)[::-1][:10]:
    print 'GOP#%3d(%4d): %8d bytes, %6d frames -> %f Mbps' % (i, sum(gopFrames[0:i+1]) - gopFrames[i], gopSize[i], gopFrames[i], normGopSize[i])
