'''
A Benchmark of Computational CRISPR-Cas9 Guide Design Methods

Jacob Bradford, Dimitri Perrin. 2019.


This scripts compares the guides generated by one tool, with the guides generated
by every other tool. 
Example:
    ATCGATCGATCGAGG,2000,2023,+,[guidescan, mm10db, ctfinder]
    Meaning:
        The guide was found by guidescan, mm10db and ctfinder

Guides are only compared if their on the same strand. 
The data used in this script must be normalised via normaliser.py first.

The output is matrix (in CSV format), where the toolname column describes
the current tool being compared with each tool along the toolname row.

eg: 
(casdesigner desigend a total of 6411 guides. mm10db had 1150 of those.)
(casfinder designed a total of 740932 guides. guidescan had 43587 of those.)
,casdesigner,casfinder,chopchop,crispor,crisprera,ctfinder,gtscan,guidescan,mm10crisprdatabase,total
casdesigner,6411,6379,6098,6411,6411,3576,6411,3882,1150,6411
casfinder,6379,74092,71481,74092,74092,39705,74092,43587,1250,74092
...


'''

import os

FOR_EXON_GUIDES_ONLY = True
EXON_ONLY_DIR = 'exon-only'

# directories to scan (these directories must contain the normalised output file)
if not FOR_EXON_GUIDES_ONLY:
    scannableDirs = {
        '500k' :    r'normalised\500k',
        #'1m' :     r'normalised\1m',
        #'5m' :     r'normalised\5m',
        #'full' :   r'normalised\full',
    }
else:
    scannableDirs = {
        '500k' :    r'normalised\exon-only\500k',
        #'1m' :     r'normalised\exon-only\1m',
        #'5m' :     r'normalised\exon-only\5m',
        #'full' :   r'normalised\exon-only\full',
    }

humanToIntSizes = {
    '500k' : 500000,
    '1m' : 1000000,
    '5m' : 5000000,
    'full' : 61431566,
}


OUTPUT_DIR = 'compare-one-with-all-others-output'

toolsToConsider = [
    'casdesigner',
    'casfinder',
    'cctop',
    'chopchop',
    'crispor',
    'crisprdo',
    'crisprera',
    'ctfinder',
    'flashfry',
    'gtscan1.3',
    'guidescan',
    'mm10db',
    'phytocrispex',
    'sgrnacas9',
    'sgrnascorer',
    'ssc',
    'tuscan',
    'wucrispr'
]

def doesFileNameMatchToolsToConsider(fileName):
    matches = False
    for toolToConsider in toolsToConsider:
        if toolToConsider.lower() in fileName.replace('-', '').lower():
            matches = True
    return matches
    
# for each test size (separated by directory)
for scannableDir in scannableDirs:
    print '\nProcessing: %s' % scannableDir
    
    hashtable = {'+' : {}, '-' : {}}
    toolNames = []


    directoryPath = os.path.join(scannableDirs[scannableDir])
    
    # walk over the directory
    for (dirpath, dirnames, filenames) in os.walk(directoryPath):
                        
        # for file within the directory
        for currentFilename in filenames:
            if doesFileNameMatchToolsToConsider(currentFilename) == False:
                continue
            
            currentFullpath = os.path.join(scannableDirs[scannableDir], currentFilename)
                
            with open(currentFullpath, 'r') as fCurrentRead:
                #print 'Reading: %s' % currentFullpath
                guideCount = 0
                for currentLine in fCurrentRead:
                    guideCount = guideCount + 1
                    currentLine = currentLine.strip('\n').split(',')
                    currToolName = currentLine[0]
                    currStartPos = int(currentLine[2])
                    currEndPos = int(currentLine[3])
                    currSeq = currentLine[1]
                    currStrand = currentLine[4]
                    
                    if currStrand == '+':
                        hashKey = currEndPos
                    else:
                        hashKey = currStartPos
                
                    if currToolName not in toolNames:
                        toolNames.append(currToolName)

                    if hashKey in hashtable[currStrand]:
                        #print 'test'
                        #exit()
                        hashtable[currStrand][hashKey][-1].append(currToolName)
                    else:
                        hashtable[currStrand][hashKey] = [currSeq, currStrand, currStartPos, currEndPos, [currToolName]]

                        
                        
    # here, we write the raw data to file and also calculate the similarity
    # between tools. the output of this similarity report is done further down.
    if FOR_EXON_GUIDES_ONLY:
        rawDataFileName = '%s/%s_exons_raw_data.csv' % (OUTPUT_DIR, scannableDir)
    else:
        rawDataFileName = '%s/%s_raw_data.csv' % (OUTPUT_DIR, scannableDir)
    print 'Writing to file: %s' % rawDataFileName
    
    #setup the matrix
    commonGuideCount = {}
    for toolA in toolNames:
        commonGuideCount[toolA] = {'total' : 0}
        for toolB in toolNames:
            commonGuideCount[toolA][toolB] = 0
    
    with open(rawDataFileName, 'w+') as fRaw:
        # for each strand
        for strand in hashtable:
        
            # for each guide on the strand
            for guide in hashtable[strand]:
            
                # write it to file
                fRaw.write('%s\n' % hashtable[strand][guide])

                # for each tool that found this guide
                for toolA in hashtable[strand][guide][-1]:
                
                    # for each tool that found this guide, again
                    for toolB in hashtable[strand][guide][-1]:
                        # increase the commonGuideCount by 1
                        commonGuideCount[toolA][toolB] = commonGuideCount[toolA][toolB] + 1

                    # we also want to know how many guides this tool found in total
                    commonGuideCount[toolA]['total'] = commonGuideCount[toolA]['total'] + 1

                    
                    
                    
                    
                    
                    
                    
    # output a nice CSV form of the hashtable   
    if FOR_EXON_GUIDES_ONLY:    
        csvDataFileName = '%s/%s_exons_csv_data.csv' % (OUTPUT_DIR, scannableDir)
    else:
        csvDataFileName = '%s/%s_csv_data.csv' % (OUTPUT_DIR, scannableDir)
        
    print 'Writing to file: %s' % csvDataFileName
    with open(csvDataFileName, 'w+') as fRawCsv:
        fRawCsv.write('sequence,strand,start,end,is_%s\n' % (',is_'.join(toolNames))  )
        for strand in hashtable:
            # for each guide on the strand
            for hashItem in hashtable[strand]:
                guide = hashtable[strand][hashItem]
                strOut = '%s,%s,%s,%s' % (guide[0], guide[1], guide[2], guide[3])
                for tool in toolNames:
                    strOut = '%s,%s' % (strOut, tool in guide[-1])
                        
                fRawCsv.write('%s\n' % strOut)
                
                
    # output a nice CSV matrix. The actual data is calculated within the rawDataFileName
    # context manager.
    if FOR_EXON_GUIDES_ONLY:   
        resultsFileName = '%s/%s_exons_results.csv' % (OUTPUT_DIR, scannableDir)
    else:
        resultsFileName = '%s/%s_results.csv' % (OUTPUT_DIR, scannableDir)
        
    print 'Writing to file: %s' % (resultsFileName)
    with open(resultsFileName, 'w+') as fStats:
        fStats.write('%s\n' % (','.join([''] + toolNames + ['total'])))
        for row in toolNames: # (ie: row in the CSV matrix)
            strOut = ""
            for tool in (toolNames + ['total']):
                if strOut == "":
                    strOut = '%s' % (row)
                strOut = '%s,%s' % (strOut, commonGuideCount[row][tool])
                
            fStats.write('%s\n' % (strOut))