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

Jacob Bradford, Dimitri Perrin. 2019.

This script takes the normalised (with scores, generated by normalise-Xu2015.py)
data and plots it to visualise the recall and accuracy of each tool.

Run:
    python plotXu2015.py
    
Input:
    - normalised data (with scores) in directory specified in TOOLS_NORMALISED_DATA_DIR 
    - (process the raw data using normalise-Xu2015.py)
    
Output:
    - PNG file
    - see last line of file for details
'''

import matplotlib.pyplot as plt
import argparse, csv, numpy, time, os, re, string
from matplotlib.pyplot import cm, colors
import matplotlib.patches as mpatches
import matplotlib as mpl
import math

''' Config '''
XU_2014_DATA = r"Xu-2015_Is-Efficient.csv"
TOOLS_NORMALISED_DATA_DIR = r"normalised-bauer-lab-updates-exp-with-scores"
PLOT_OUTPUT_DPI = 300
NUMBER_OF_ROWS_IN_PLOT = 2
WIDTH = 20
HEIGHT = 15

PRETTY_NAMES = {
    'chopchop' : 'CHOPCHOP', 
    'gtscan' : 'GT-Scan', 
    'crispor' : 'CRISPOR', 
    'crisprera' : 'CRISPR-ERA', 
    'ctfinder' : 'CT-Finder', 
    'casdesigner' : 'Cas-Designer', 
    'mm10crisprdatabase' : 'mm10db', 
    'sgrnascorer20' : 'sgRNAScorer2', 
    'guidescan' : 'GuideScan', 
    'casfinder' : 'CasFinder',
    'wucrispr' : 'WU-CRISPR',
    'phytocrispex' : 'PhytoCRISP-Ex',
    'flashfry' : 'FlashFry',
    'sgrnacas9' : 'sgRNAcas9',
    'ssc' : 'SSC',
    'cctop' : 'CCTop',
    'crisprdo' : 'CRISPR-DO',
    'tuscan' : 'TUSCAN'
}

PLOT_ORDER = ['chopchop', 'mm10crisprdatabase', 'tuscan', 'sgrnascorer20','crisprdo',  'wucrispr', 'flashfry', 'ssc']

''' Do not modify '''
X_COL_INDEX = 5
Y_COL_INDEX = 6
SEQ_INDEX = 4
IS_EFFICIENT_COL_INDEX = 7
NORMALISED_TOOL_NAME_INDEX = 0
NORMALISED_GUIDE_INDEX = 1
NORMALISED_SCORE_INDEX = 5

GUIDE_START_POS = 10
GUIDE_LEN = 23 #- 3 # ignore the PAM for now

''' Data handlers '''
X_XU_NO_MATCH_DATA = []
Y_XU_NO_MATCH_DATA = []
X_XU_MATCH_DATA = []
Y_XU_MATCH_DATA = []
SEQ_XU_DATA = []
TOOL_DATA = {} # key: tool name, value: array of guides it reported
TOOL_SCORES_DATA = {} # key: tool name, value: array of guides it reported

''' Copied from mm10db target_identification_viaC.py '''
def rc(dna):
    complements = string.maketrans('acgtrymkbdhvACGTRYMKBDHV', 'tgcayrkmvhdbTGCAYRKMVHDB')
    rcseq = dna.translate(complements)[::-1]
    return rcseq

''' Load in the normalised data '''
for root, dirs, files in os.walk(TOOLS_NORMALISED_DATA_DIR):
    for file in files:
        if '.normalised' not in file:
            continue
            
        fileName = os.path.join(TOOLS_NORMALISED_DATA_DIR, file)
        print 'Reading: %s' % file
        with open(fileName, 'r') as fOpen:
            
            # expecting: crisprera,TACAGCCAGACATAACACCT...,1826,1848,+
            # split the file by row, then split by comma
            for line in [x.split(',') for x in fOpen.read().split('\n')[:-1]]:
                # we want every guide from all tools except chopchop.
                # in the case of chopchop, we only want those that are reported 
                # to be efficient. same for tuscan.
                if (line[NORMALISED_TOOL_NAME_INDEX] not in ['chopchop', 'tuscan']) or (line[NORMALISED_TOOL_NAME_INDEX] == 'chopchop' and line[NORMALISED_SCORE_INDEX] == '1.00') or (line[NORMALISED_TOOL_NAME_INDEX] == 'tuscan' and line[NORMALISED_SCORE_INDEX] == '1.0'):
                    if line[NORMALISED_TOOL_NAME_INDEX] not in TOOL_DATA:
                        TOOL_DATA[line[NORMALISED_TOOL_NAME_INDEX]] = []
                        TOOL_SCORES_DATA[line[NORMALISED_TOOL_NAME_INDEX]] = []
                        
                    TOOL_DATA[line[NORMALISED_TOOL_NAME_INDEX]].append(line[NORMALISED_GUIDE_INDEX])#[:-3])
                    
                    if line[NORMALISED_SCORE_INDEX] in ['None', 'N/A']:
                        TOOL_SCORES_DATA[line[NORMALISED_TOOL_NAME_INDEX]].append(None)
                    else:
                        TOOL_SCORES_DATA[line[NORMALISED_TOOL_NAME_INDEX]].append(float(line[NORMALISED_SCORE_INDEX]))

''' Repeat for each tool ''' 
plotCount = 1 # one index
print 'tool,len(X_XU_MATCH_DATA),len(X_XU_NO_MATCH_DATA),MATCH_XU_EFFICIENT_COUNT,MATCH_XU_INEFFICIENT_COUNT'
for tool in PLOT_ORDER:
    #effGuideFile = open('%s.eff' % tool, 'w')
    #if tool != 'mm10crisprdatabase':
    #    continue
    count = 0
    X_XU_NO_MATCH_DATA = []
    Y_XU_NO_MATCH_DATA = []
    X_XU_MATCH_DATA = []
    Y_XU_MATCH_DATA = []
    MATCH_XU_EFFICIENT_COUNT = 0
    MATCH_XU_INEFFICIENT_COUNT = 0
    XU_GUIDE_SEQS_MATCHES_ONLY = []
    SCORES = []
    XU_IS_EFFICIENT = []
    
    # Load in the Xu 2014 data

    with open(XU_2014_DATA, 'r') as fDataRead:

        # split the file by row, then split by comma
        for line in [x.split(',') for x in fDataRead.read().split('\n')[1:-1]]:
            SEQ_XU_DATA.append(line[SEQ_INDEX])
            x = float(line[X_COL_INDEX])
            y = float(line[Y_COL_INDEX])

            indexInNormalisedData = None

            XuGuide = line[SEQ_INDEX][GUIDE_START_POS:(GUIDE_START_POS + GUIDE_LEN)][:-3]

            forwardMatch = XuGuide in [i[:-3] for i in TOOL_DATA[tool]]

            if forwardMatch:
                indexInNormalisedData = [i[:-3] for i in TOOL_DATA[tool]].index(XuGuide)
            
            reverseMatch = False
            if XuGuide[:2] == 'CC':
                reverseMatch = rc(XuGuide) in [i[:-3] for i in TOOL_DATA[tool]]
           
            if indexInNormalisedData is not None:
                score = TOOL_SCORES_DATA[tool][indexInNormalisedData]
                if score not in [None, 'N/A']:
                    SCORES.append(score)
           
            if forwardMatch or reverseMatch:
                count = count + 1 
                
                X_XU_MATCH_DATA.append(x)
                Y_XU_MATCH_DATA.append(y)
                XU_GUIDE_SEQS_MATCHES_ONLY.append(XuGuide)
                if line[IS_EFFICIENT_COL_INDEX] == '1':
                    XU_IS_EFFICIENT.append(1)
                    MATCH_XU_EFFICIENT_COUNT += 1
                else:
                    XU_IS_EFFICIENT.append(0)
                    MATCH_XU_INEFFICIENT_COUNT += 1
            else:
                X_XU_NO_MATCH_DATA.append(x)
                Y_XU_NO_MATCH_DATA.append(y)

    if len(SCORES) != len(X_XU_MATCH_DATA):
        SCORES=[255.0] * len(X_XU_MATCH_DATA)
   
    ax = plt.subplot(NUMBER_OF_ROWS_IN_PLOT, math.ceil(len(PLOT_ORDER) / NUMBER_OF_ROWS_IN_PLOT), plotCount)
    if plotCount in [1,5]:
        plt.ylabel('log2 fold-change, KBM-7')
    plt.xlabel('log2 fold-change, HL-60')

    plt.title(PRETTY_NAMES[tool])
    plotCount += 1
    
    if tool == 'sgrnascorer20': 
        # we want to know how many guides were deemed efficient vs inefficient 
        # according to Xu in each percentile
        ninetyAbove = numpy.percentile(SCORES, 90)
        seventyFive = numpy.percentile(SCORES, 75)
        fifty = numpy.percentile(SCORES, 50)
        
        # colours
        colNinetyAbove  = '#d7191c'
        colSeventyFive  = '#fdae61'
        colFifty        = '#abd9e9'
        colLessFifty    = '0.4'    
        
        colNinetyAbove  = '#00b8ff'
        colSeventyFive  = '#6aff8d'
        colFifty        = '#ff7e00'
        colLessFifty    = '#b20000'
        
        counts = [0, 0, 0, 0] # each index represents a percentile. see below
        isEfficientCounts = [[0,0], [0,0], [0,0], [0,0]]  # [[inefficient, efficient], ..]

        for i in xrange(len(SCORES)):

            score = SCORES[i]
            
            # this is so hacky its ridic
            if score >= ninetyAbove:
                SCORES[i] = colNinetyAbove
                counts[0] += 1
                
                # doing this will allow us to know how many guides were deemed by
                # Xu as being efficient or inefficient in this particular percentile
                isEfficientCounts[0][XU_IS_EFFICIENT[i]] += 1
            if score >= seventyFive and score < ninetyAbove:
                SCORES[i] = colSeventyFive 
                counts[1] += 1
                isEfficientCounts[1][XU_IS_EFFICIENT[i]] += 1
            if score >= fifty and score < seventyFive:
                SCORES[i] = colFifty 
                counts[2] += 1
                isEfficientCounts[2][XU_IS_EFFICIENT[i]] += 1
            if score < fifty:
                SCORES[i] = colLessFifty
                counts[3] += 1
                isEfficientCounts[3][XU_IS_EFFICIENT[i]] += 1

        plt.scatter(X_XU_NO_MATCH_DATA, Y_XU_NO_MATCH_DATA, s=15, marker='.', label='No Match', c='0.85')
        plt.scatter(X_XU_MATCH_DATA, Y_XU_MATCH_DATA, s=15, marker='.', label='Match', c=SCORES)
    
        plt.legend(handles=[
            mpatches.Patch(color=colNinetyAbove, label='Above 90%'),
            mpatches.Patch(color=colSeventyFive, label='75-89%'),
            mpatches.Patch(color=colFifty, label='50-74%'),
            mpatches.Patch(color=colLessFifty, label='Below 50%'),
            mpatches.Patch(color='0.85', label='Not Reported')
            ], loc='upper left')
            
        print 'sgrnascorer20: %i %i %i %i %i %i' % (counts[0], counts[1], counts[2], counts[3], MATCH_XU_EFFICIENT_COUNT, MATCH_XU_INEFFICIENT_COUNT)
        
    elif tool in ['flashfry', 'casdesigner', 'wucrispr', 'crispor', 'crisprdo', 'ssc']:   
        minn = min(SCORES) - min(SCORES) * 0.1
        maxx = max(SCORES) + max(SCORES) * 0.1
        norm = mpl.colors.Normalize(vmin=minn, vmax=maxx)
        plt.scatter(X_XU_NO_MATCH_DATA, Y_XU_NO_MATCH_DATA, s=15, marker='.', label='No Match', color='0.85')  
        plt.scatter(X_XU_MATCH_DATA, Y_XU_MATCH_DATA, s=15, marker='.', label='Match', c=SCORES, cmap=cm.jet_r, norm=norm)
        plt.colorbar()
        
        if tool in ['wucrispr', 'crisprdo']:
            plt.legend(handles=[
            mpatches.Patch(color='0.85', label='Not Reported')
            ], loc='upper left')
        
    else:
        
        plt.scatter(X_XU_NO_MATCH_DATA, Y_XU_NO_MATCH_DATA, s=15, marker='.', label='No Match', color='0.85')
        SCORES = ['#1f78b4' * len(SCORES)]
        plt.scatter(X_XU_MATCH_DATA, Y_XU_MATCH_DATA, s=15, marker='.', label='Match', c=['#1f78b4' * len(SCORES)], cmap=cm.copper)
        plt.legend(handles=[
            mpatches.Patch(color='#1f78b4', label='Accepted'),
            mpatches.Patch(color='0.85', label='Rejected')
            ], loc='upper left')
       
       
    print '%s\t%i\t%i\t%i\t%i' % (tool, len(X_XU_MATCH_DATA), len(X_XU_NO_MATCH_DATA), MATCH_XU_EFFICIENT_COUNT, MATCH_XU_INEFFICIENT_COUNT)    

    plt.axhline(0, linestyle='-', color='k', linewidth=0.5)
    plt.axvline(0, linestyle='-', color='k', linewidth=0.5)
    
    '''    		(calculated in excel)			
                                efficient		                inefficient	
                                min	            max	            min	            max
    x	log2 fold change, HL60	-5.305412391	-1.001239638	-0.497285084	2.519349582
    y	log2 fold change, KBM7	-5.498241555	-1.022733387	-0.49447052	    2.97226217
    '''
    # efficient (q3)
    ax.hlines(-1.022733387, -5.75, -1.001239638, linestyle='-', color='0.65', linewidth=0.5)
    ax.vlines(-1.001239638, -5.75, -1.022733387, linestyle='-', color='0.65', linewidth=0.5)
    
    # inefficient (q1)
    ax.hlines(-0.49447052, -0.497285084, 2.75, linestyle='-', color='0.65', linewidth=0.5)
    ax.vlines(-0.497285084, -0.49447052, 3, linestyle='-', color='0.65', linewidth=0.5)

    plt.text(-5.6, -1.9, 'Efficient', color='0.1', fontsize='small', fontstyle='italic')
    plt.text(0.4, 2.4, 'Inefficient', color='0.1', fontsize='small', fontstyle='italic')
    
    effMatchPercent = float(MATCH_XU_EFFICIENT_COUNT) / 731.0 * 100
    ineffMatchPercent = float(MATCH_XU_INEFFICIENT_COUNT) / (1169 - 731.0) * 100
    plt.text(-5.65, -2.2, '(%3.1f%%)' % (effMatchPercent), color='0.1', fontsize='small', fontstyle='italic')
    plt.text(0.5, 2.1, '(%3.1f%%)' % (ineffMatchPercent), color='0.1', fontsize='small', fontstyle='italic')
    
    plt.rcParams.update({'font.size': 13})
    
fig = plt.gcf()

fig.set_size_inches(WIDTH, HEIGHT)

fig.savefig('plotXu2015.png', dpi=PLOT_OUTPUT_DPI)
#fig.savefig('plotXu2015.eps', format='eps', dpi=PLOT_OUTPUT_DPI)
