import re, os, os.path, glob, sys, string, md5
import naming

def ExtractResultsFromFile(fileName):
    myRegExp = re.compile(r"""
    \s*                           # Skip leading whitespace
    (?P<numFrames>[0-9]+)\s*      # capture num frames
    frames[ ]encoded:             # expected text
    \s*Y\s*(?P<snrY>[0-9.]+)\s*dB  # Y PSNR
    \s*U\s*(?P<snrU>[0-9.]+)\s*dB  # U PSNR
    \s*V\s*(?P<snrV>[0-9.]+)\s*dB  # V PSNR
    \s*$                           # skip leading whitespace to end of line
                                  # Next, capture the bit rate
    ^\s*                          # Skip leading whitespace
    average[ ]bit[ ]rate:\s*
    (?P<bitRate>[0-9.]+)\skbit/s  # capture bit rate
    
    """, re.VERBOSE | re.MULTILINE)
    match = myRegExp.search( open(fileName,'r').read() )
    if (match):
        return map(float, (match.group('bitRate'),match.group('snrY')))
    else:
        print '\n\nWARNING:  Failed to find results from file %s.\n' % fileName
        return None

def ExtractResultsFromDir(directoryName):
    taskPattern = re.compile('task_.*as_view_(?P<viewNum>[0-9]+).txt')
    results = {}
    for file in glob.glob(os.path.join(directoryName,'*task_*.txt')):
        match = taskPattern.search(file)
        if (match):
            viewNum = match.group('viewNum')
            while (results.has_key(viewNum)):
                print 'WARNING: Duplicate results for view number %i.'%viewNum
                viewNum = viewNum + 100                    
                print 'WARNING: Putting results for this view at view %i.' % (
                    viewNum)
            results[viewNum] = ExtractResultsFromFile(file)                
        else:
            print '\nWARNING: Could not extract view number from file %s.\n' %(
                file)
    return results

def ComputeMean( dataList ):
    return reduce( lambda a,b: a+b, dataList ) / float(len(dataList))

def AverageResults( resultDict ):
    averageRate = ComputeMean( map( lambda x: x[0], resultDict.values() ) )
    averageSNR =  ComputeMean( map( lambda x: x[1], resultDict.values() ) )
    return (averageRate, averageSNR)
        
def ExtractResultsIntoCSVFile(directoryName):
    results = ExtractResultsFromDir(directoryName)
    keys = results.keys()
    keys.sort()
    fd = open(os.path.join(directoryName,'results.csv'),'w')
    for i in keys:
        if (None != results[i]):
            fd.write('%s, %f, %f\n' % (`i`, results[i][0], results[i][1]))
    fd.close()

def CollectAllCSVResultsIntoFile(pathPattern, outputFile):
    fd = open(outputFile,'w')
    for file in glob.glob(os.path.join(pathPattern,'results.csv')):
        fd.write(file+'\n')
        fd.write( open(file,'r').read() + '\n\n' )
    fd.close()

def CollectFullResults(directory, outputDir):
    if (not os.path.exists(outputDir)):
        os.mkdir(outputDir)
    else:
        for item in glob.glob(os.path.join(outputDir,'*')):
            os.remove(item)

    dirDict = { directory : 1 }
        
    GetRelatedDirectories(dirDict, directory)

    reconFiles = {}
    outputFiles = {}
    taskFiles = {}
    configFiles = {}

    for dir in dirDict.keys():
        AddMatches(reconFiles, os.path.join(dir,'recon_*.yuv'))
        AddMatches(outputFiles, os.path.join(dir,'output_*.264'))
        AddMatches(configFiles, os.path.join(dir, '*.cfg'))
        AddMatches(taskFiles, os.path.join(dir, 'task_*.txt'))

    dups = []
    fd = open(os.path.join(outputDir,'make_results.py'),'w')
    fd.write('#!/usr/bin/python\n')
    fd.write('import os, os.path, sys\n')
    fd.write('reconFiles=%s\n' % ProcessMatchDict(reconFiles,dups,outputDir))
    fd.write('outputFiles=%s\n' % ProcessMatchDict(outputFiles,dups,outputDir))
    fd.write('taskFiles=%s\n' % ProcessMatchDict(taskFiles,dups,outputDir))
    fd.write('configFiles=%s\n' % ProcessMatchDict(configFiles,dups,outputDir))
    fd.write('\nduplicates=%s\n' % `dups`)
    fd.write("""
def MakeLinks(dictList):
    for dictItem in dictList:
        for item in dictItem.keys():
            if (len(dictItem[item]) > 1):
                print 'WARNING:  Element [%s] has duplicates.'%(item)
            for i in range(len(dictItem[item])):
                itemName = ('='*i) + item
                path =  dictItem[item][i]
                if (os.path.exists(itemName)):
                    assert os.path.islink(itemName)
                    os.remove(itemName)
                os.symlink(path, itemName)

MakeLinks([reconFiles,outputFiles,configFiles,taskFiles])

    """)
    fd.close()

def MakeLinks(dictList):
    for dictItem in dictList:
        for item in dictItem.keys():
            assert 1==len(dictItem[item]), 'Element [%s] has duplicates\n'%(
                item)
            path = dictItem[item][0]
            if (os.path.exists(path)):
                os.remove(path)
            os.symlink(path,item)

def ProcessMatchDict(matchDict, duplicates, outputDir):
    result = ['{\n']
    for k in matchDict.keys():
        result.append('\t%s : %s,\n' % (`k`, `matchDict[k]`))
        if (len(matchDict[k]) > 1):
            duplicates.append(k)
                
    result.append('\n}\n')
    return string.join(result,'')
    
def ComputeMD5(fileName,bufSize=16384):
    fd = open(fileName,'rb')
    hasher = md5.new()
    data = fd.read(bufSize)
    while ('' != data):
        hasher.update(data)
        data = fd.read(bufSize)
    return hasher.digest()
            

def AddMatches(matchDict,pattern):

    for file in glob.glob(os.path.join(pattern)):
        basename = os.path.basename(file)
        if (not matchDict.has_key(basename)):
            matchDict[basename] = []
        matchDict[basename].append(file)


def GetRelatedDirectories(dirDict, directory):

    sys.path.insert(0,directory)
    if (os.path.exists(os.path.join(directory,'configFile.py'))):
        config = __import__('configFile')
        if (config.paramsForMVC.has_key('alreadyEncodedViews')):
            for i in config.paramsForMVC['alreadyEncodedViews'].keys():
                directoryToTry=os.path.dirname(
                    config.paramsForMVC['alreadyEncodedViews'][i])
                if (not dirDict.has_key(directoryToTry)):
                    dirDict[directoryToTry] = 1
                    sys.path.remove(directory)
                    GetRelatedDirectories(dirDict, directoryToTry)
        else:
            sys.path.remove(directory)

if (__name__ == "__main__"):
    if (len(sys.argv) == 2):
        print 'Doing --> ExtractResultsIntoCSVFile(%s)\n' %(sys.argv[1])
        ExtractResultsIntoCSVFile(sys.argv[1])
    elif (len(sys.argv) == 3):
        print 'Doing --> CollectFullResults(%s,%s)\n'%(sys.argv[1],sys.argv[2])
        CollectFullResults(sys.argv[1],sys.argv[2])
    else:
        raise 'Wrong number of args; got %i args; expected either %i or %i.'%(
            len(sys.argv), 2,3)