# Copyright (c) 2022-2023: Mustafa Tekpinar
# This code is part of the prescott package and governed by its license.
# Please see the LICENSE.txt file included as part of this package.
import os
import sys
import matplotlib.pyplot as plt
# import seaborn as sns
import pandas as pd
from sklearn import mixture
from sklearn import metrics
import numpy as np
import argparse
from Bio import SeqIO

one_letter ={'VAL':'V', 'ILE':'I', 'LEU':'L', 'GLU':'E', 'GLN':'Q', \
            'ASP':'D', 'ASN':'N', 'HIS':'H', 'TRP':'W', 'PHE':'F', 'TYR':'Y',    \
            'ARG':'R', 'LYS':'K', 'SER':'S', 'THR':'T', 'MET':'M', 'ALA':'A',    \
            'GLY':'G', 'PRO':'P', 'CYS':'C'}

def plotROCandAUCV2(labels, pred):
    fpr, tpr, thresholds = metrics.roc_curve(labels, pred, pos_label=1)
    roc_auc = metrics.auc(fpr, tpr)

    return fpr, tpr, roc_auc

def getGnomADOverallFrequency(infile, usePopMax="true"):
    # df = pd.read_csv("P53_gnomAD_v3.1.2_ENSG00000141510_2023_07_04_16_10_48.csv")
    df = pd.read_csv(infile)
    # print(df.columns)

    # #Select only columns containing missense variants!
    # df = df.loc[df[]=='']
    #print(df1.columns)
    dfMissense = df.loc[df['VEP Annotation']=='missense_variant',
                    ['HGVS Consequence', 'Protein Consequence', 'Transcript Consequence',
                    'VEP Annotation', 'ClinVar Clinical Significance', 'ClinVar Variation ID', 'Flags', 
                    'Allele Count', 'Allele Number', 'Allele Frequency', 
                    'Homozygote Count', 'Hemizygote Count',
                    'Allele Count African/African American',
                    'Allele Number African/African American',
                    'Homozygote Count African/African American',
                    'Hemizygote Count African/African American',
                    'Allele Count Latino/Admixed American',
                    'Allele Number Latino/Admixed American',
                    'Homozygote Count Latino/Admixed American',
                    'Hemizygote Count Latino/Admixed American',
                    'Allele Count Ashkenazi Jewish', 
                    'Allele Number Ashkenazi Jewish',
                    'Homozygote Count Ashkenazi Jewish',
                    'Hemizygote Count Ashkenazi Jewish', 
                    'Allele Count East Asian',
                    'Allele Number East Asian', 
                    'Homozygote Count East Asian',
                    'Hemizygote Count East Asian', 
                    'Allele Count European (Finnish)',
                    'Allele Number European (Finnish)',
                    'Homozygote Count European (Finnish)',
                    'Hemizygote Count European (Finnish)',
                    'Allele Count European (non-Finnish)',
                    'Allele Number European (non-Finnish)',
                    'Homozygote Count European (non-Finnish)',
                    'Hemizygote Count European (non-Finnish)', 
                    'Allele Count Other',
                    'Allele Number Other', 
                    'Homozygote Count Other',
                    'Hemizygote Count Other', 
                    'Allele Count South Asian',
                    'Allele Number South Asian', 
                    'Homozygote Count South Asian',
                    'Hemizygote Count South Asian']]
    

    dfMissense = dfMissense.reset_index()

    dfMissense['Allele Frequency']
    dfMissense['Allele Frequency Log'] = ""
    if(usePopMax.lower()=="true"):
        alleleCountList = ['Allele Count African/African American',
                        'Allele Count Latino/Admixed American',
                        'Allele Count Ashkenazi Jewish', 
                        'Allele Count East Asian',
                        'Allele Count European (Finnish)',
                        'Allele Count European (non-Finnish)',
                        'Allele Count Other',
                        'Allele Count South Asian']
        alleleNumberList = ['Allele Number African/African American',
                        'Allele Number Latino/Admixed American',
                        'Allele Number Ashkenazi Jewish',
                        'Allele Number East Asian', 
                        'Allele Number European (Finnish)',
                        'Allele Number European (non-Finnish)',
                        'Allele Number Other', 
                        'Allele Number South Asian']
        for index, row in dfMissense.iterrows():
            maxFreq = 0.0
            tempIndex = 0
            for i in range(len(alleleCountList)):
                if(row[alleleNumberList[i]]!=0):
                    
                    tempValue = (row[alleleCountList[i]]/row[alleleNumberList[i]])
                    if(tempValue>maxFreq):
                        maxFreq=tempValue
                        tempIndex = i
            # Avoid zero frequency error by setting it to a very low number such as 10**-10
            if(maxFreq==0.0):
                maxFreq = 10**-10 # Which means 1 in 10 billion, which is the estimated population in 2050. 
            # print(maxFreq)
            dfMissense.at[index,'Selected Population'] = alleleCountList[tempIndex]
            dfMissense.at[index,'Allele Frequency Log'] = np.log10(maxFreq)
            
            #print(np.log10(maxFreq))
    else:
        # Avoid zero frequency error by setting it to a very low number such as 10**-10
        for index, row in dfMissense.iterrows():
            if(row['Allele Frequency']==0.0):
                dfMissense.at[index,'Allele Frequency Log'] =  np.log10(10**-10)
            else:
                dfMissense.at[index,'Allele Frequency Log'] =  np.log10(row['Allele Frequency'])
        
    # sys.exit(-1)
    #print(dfMissense['Allele Frequency Log'])
    #print(dfMissense[['Allele Frequency', 'ClinVar Clinical Significance']])

    # dfMissense.dropna(subset = ['ClinVar Clinical Significance'], inplace=True)
    dfMissense = dfMissense.reset_index()
    #print(dfMissense[['Allele Frequency', 'ClinVar Clinical Significance']])



    # plt.figure()
    # plt.hist(dfMissense['Allele Frequency Log'], density=False, color='red', label='pathogenic')
    # plt.show()

    mutantList = []
    proteinNameList = []
    for index, row in dfMissense.iterrows():
        source = one_letter[row['Protein Consequence'][2:5].upper()]
        position = (row['Protein Consequence'][5:-3])
        target = one_letter[(row['Protein Consequence'][-3:]).upper()]
        mutant = source+position+target
        mutantList.append(mutant)
        # print(mutant, row['Protein Consequence'])
        # print(infile.split('_')[2])
        proteinNameList.append(os.path.basename(infile).split('_')[2])
        
    dfMissense['mutant'] = mutantList
    dfMissense['protein'] = proteinNameList
    return (dfMissense)

def getGnomADOverallFrequencyV2(infile, usePopMax="true"):
    """
        In V2 of this function, I'll assing (practically) zero frequency (10**-10)
        to the mutations that are observed only in one community out of 8. 
    """
    # df = pd.read_csv("P53_gnomAD_v3.1.2_ENSG00000141510_2023_07_04_16_10_48.csv")
    df = pd.read_csv(infile)
    # print(df.columns)

    # #Select only columns containing missense variants!
    # df = df.loc[df[]=='']
    #print(df1.columns)
    dfMissense = df.loc[df['VEP Annotation']=='missense_variant',
                    ['HGVS Consequence', 'Protein Consequence', 'Transcript Consequence',
                    'VEP Annotation', 'ClinVar Clinical Significance', 'ClinVar Variation ID', 'Flags', 
                    'Allele Count', 'Allele Number', 'Allele Frequency', 
                    'Homozygote Count', 'Hemizygote Count',
                    'Allele Count African/African American',
                    'Allele Number African/African American',
                    'Homozygote Count African/African American',
                    'Hemizygote Count African/African American',
                    'Allele Count Latino/Admixed American',
                    'Allele Number Latino/Admixed American',
                    'Homozygote Count Latino/Admixed American',
                    'Hemizygote Count Latino/Admixed American',
                    'Allele Count Ashkenazi Jewish', 
                    'Allele Number Ashkenazi Jewish',
                    'Homozygote Count Ashkenazi Jewish',
                    'Hemizygote Count Ashkenazi Jewish', 
                    'Allele Count East Asian',
                    'Allele Number East Asian', 
                    'Homozygote Count East Asian',
                    'Hemizygote Count East Asian', 
                    'Allele Count European (Finnish)',
                    'Allele Number European (Finnish)',
                    'Homozygote Count European (Finnish)',
                    'Hemizygote Count European (Finnish)',
                    'Allele Count European (non-Finnish)',
                    'Allele Number European (non-Finnish)',
                    'Homozygote Count European (non-Finnish)',
                    'Hemizygote Count European (non-Finnish)', 
                    'Allele Count Other',
                    'Allele Number Other', 
                    'Homozygote Count Other',
                    'Hemizygote Count Other', 
                    'Allele Count South Asian',
                    'Allele Number South Asian', 
                    'Homozygote Count South Asian',
                    'Hemizygote Count South Asian']]
    

    dfMissense = dfMissense.reset_index()

    dfMissense['Allele Frequency']
    dfMissense['Allele Frequency Log'] = ""
    if(usePopMax.lower()=="true"):
        alleleCountList = ['Allele Count African/African American',
                        'Allele Count Latino/Admixed American',
                        'Allele Count Ashkenazi Jewish', 
                        'Allele Count East Asian',
                        'Allele Count European (Finnish)',
                        'Allele Count European (non-Finnish)',
                        'Allele Count Other',
                        'Allele Count South Asian']
        alleleNumberList = ['Allele Number African/African American',
                        'Allele Number Latino/Admixed American',
                        'Allele Number Ashkenazi Jewish',
                        'Allele Number East Asian', 
                        'Allele Number European (Finnish)',
                        'Allele Number European (non-Finnish)',
                        'Allele Number Other', 
                        'Allele Number South Asian']
        for index, row in dfMissense.iterrows():
            maxFreq = 0.0
            tempIndex = 0
            
            # Get counts for 8 populations.
            alleleCountsAccrossPopulations = []
            for i in range(len(alleleCountList)):
                alleleCountsAccrossPopulations.append(row[alleleCountList[i]])

            zeroOccuranceCount = alleleCountsAccrossPopulations.count(0)
            if (zeroOccuranceCount>1):
                for i in range(len(alleleCountList)):
                    if(row[alleleNumberList[i]]!=0):
                        
                        tempValue = (row[alleleCountList[i]]/row[alleleNumberList[i]])
                        if(tempValue>maxFreq):
                            maxFreq=tempValue
                            tempIndex = i
                # Avoid zero frequency error by setting it to a very low number such as 10**-10
                if(maxFreq==0.0):
                    maxFreq = 10**-10 # Which means 1 in 10 billion, which is the estimated population in 2050. 
                # print(maxFreq)
                dfMissense.at[index,'Selected Population'] = alleleCountList[tempIndex]
                dfMissense.at[index,'Allele Frequency Log'] = np.log10(maxFreq)

            else:
                dfMissense.at[index,'Selected Population'] = None
                dfMissense.at[index,'Allele Frequency Log'] = np.log10(10**-10)
            
            #print(np.log10(maxFreq))
    else:
        # Avoid zero frequency error by setting it to a very low number such as 10**-10
        for index, row in dfMissense.iterrows():
            if(row['Allele Frequency']==0.0):
                dfMissense.at[index,'Allele Frequency Log'] =  np.log10(10**-10)
            else:
                dfMissense.at[index,'Allele Frequency Log'] =  np.log10(row['Allele Frequency'])
        
    # sys.exit(-1)
    #print(dfMissense['Allele Frequency Log'])
    #print(dfMissense[['Allele Frequency', 'ClinVar Clinical Significance']])

    # dfMissense.dropna(subset = ['ClinVar Clinical Significance'], inplace=True)
    dfMissense = dfMissense.reset_index()
    #print(dfMissense[['Allele Frequency', 'ClinVar Clinical Significance']])



    # plt.figure()
    # plt.hist(dfMissense['Allele Frequency Log'], density=False, color='red', label='pathogenic')
    # plt.show()

    mutantList = []
    proteinNameList = []
    for index, row in dfMissense.iterrows():
        source = one_letter[row['Protein Consequence'][2:5].upper()]
        position = (row['Protein Consequence'][5:-3])
        target = one_letter[(row['Protein Consequence'][-3:]).upper()]
        mutant = source+position+target
        mutantList.append(mutant)
        # print(mutant, row['Protein Consequence'])
        # print(infile.split('_')[2])
        proteinNameList.append(os.path.basename(infile).split('_')[2])
        
    dfMissense['mutant'] = mutantList
    dfMissense['protein'] = proteinNameList
    return (dfMissense)

alphabeticalAminoAcidsList = ['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L',
                              'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y']
def writeSinglelineFormat(scanningMatrix, outFile, residueList, 
                    beg, end, \
                    aaOrder = alphabeticalAminoAcidsList, \
                    offSet=0):
    """
        A function to convert GEMME matrix format data to experimental 
        deep mutational scanning format. 

    The format has only two columns. In the first column, one-letter wild
    type amino acid, residue number and mutated amino acid are combined without any
    gap or any other character, such as A19C. The second column is typically a float
    value of the predicted effect.
    
    Residue numbering starts with 1.
    "A1A 0.211
    "A1C 0.124
    .
    .
    .
    "T262C 0.434

    Parameters
    ----------
    scanningMatrix: numpy array of arrays
        Data matrix to plot

    outFile: string
        Name of the output png image
    
    residueList: Python list
        A list of the residues in the protein such as ['X1', 'X2', ..., 'XN']. 
        X is one-letter amino acid code and N is the total number of residues. 
        Default is None.
        If not provided, the program will write them starting from 1 till the end. 

    beg: int
        The first residue in the fasta file may not necessarily be the first residue
        in the protein. This beg parameter gives us a chance to adjust the first residue
        other than one. It only effects the DMS format text output. Default value is 0.
    
    end: int
        The last residue to use. It is used to select a subrange 
        of amino acids.

    aaOrder: Python list
        A list of 20 amino acid characters such as ['A', 'C',...., 'Y'].  
        Default is alphabetical list.

    offSet: int
        It is used to match the residue IDs in your PDB
        file with 0 based indices read from scanningMatrix matrix

    Returns
    -------
    Nothing

    """
    debug = 0

    #We subtract 1 from beg bc matrix indices starts from 0
    if(end == None):
        end = len(scanningMatrix[0])

    if(debug):
        print("Beginning: "+str(beg))
        print("End      : "+str(end))
    
        print(len(scanningMatrix[0]))

    # I don't want to write the submatrices for now!
    # Let's keep it simple for now! 
    # subMatrix = scanningMatrix[:, (beg-1):end]
    realResidueList = []
    if(residueList == None):
        for i in range(0,len(scanningMatrix[0])):
            tempString = 'X'+str(i+1+beg)
            realResidueList.append(tempString)
    else:
        for i in range(0,len(scanningMatrix[0])):
            tempString = residueList[i].upper()+str(i+1+beg)
            realResidueList.append(tempString)

    with open(outFile, 'w') as file:
        for res in range (len(scanningMatrix[0])):
            for i in range(20):
                file.write(realResidueList[res]+aaOrder[i]+" ")
                file.write(str(scanningMatrix[alphabeticalAminoAcidsList.index(aaOrder[i])][res])+"\n")

def parseGEMMEoutput(inputFile, verbose):
    """
        Parse normalized (I don't know how?!) GEMME output files: 
        Independent, Epistatic and Combined
        Return: Data File as a Numpy array, where each col contains 
        deleteriousness score for an amino acid.
        verbose is a boolean value
    """
    gemmeDataFile =open(inputFile, 'r')
    allLines = gemmeDataFile.readlines()
    gemmeDataFile.close()

    headerLine = allLines[0]
    if(verbose):
        print(headerLine)
    matrixData = []
    aaColumn = []
    for line in allLines[1:]:
        tempList = []
        data = line.split()
        aaColumn.append(data[0])
        for item in data[1:]:
            if (item == 'NA'):
                tempList.append(0.0000000)
            else:
                tempList.append(float(item))
        matrixData.append(tempList)
    mutationsData = np.array(matrixData)
    # mutatedTransposed = mutationsData.T
    # return mutatedTransposed
    return mutationsData

from scipy.stats import rankdata
def rankSortData(dataArray):
    """
        This function ranksorts protein data and converts it
        to values between [0,1.0].
    
    Parameters:
    ----------
    dataArray: numpy array of arrays
               data read by numpy.loadtxt

    Returns:
    -------
    normalizedRankedDataArray: numpy array
    """
    normalizedRankedDataArray = rankdata(dataArray)/float(len(dataArray))

    return (normalizedRankedDataArray)

def main():
    # Adding the main parser
    main_parser = argparse.ArgumentParser(description=\
        "PRESCOTT is used to incorporate GnomAD frequency information into missense mutation scores of ESCOTT.")
    
    main_parser.add_argument('-e', '--escottfile', dest='escottfile', type=str, \
        help='Main input for prescott is an ESCOTT single point mutation file.', \
        required=True, default=None)
    
    main_parser.add_argument('-s', '--sequencefile', dest='sequencefile', type=str, \
        help='Sequence file of only the query protein sequence in fasta format.', \
        required=True, default=None)

    main_parser.add_argument('-g', '--gnomadfile', dest='gnomadfile', type=str, \
        help='The second input for prescott is a gnomad frequency file for the protein\n.'+
             'You can obtain it from https://gnomad.broadinstitute.org/ by searching your favorite protein (such as P53).\n'+
             'After you found your protein, go to the \'Export variants to CSV\' link and click it do download the csv file you needed!', \
        required=True, default=None)

    main_parser.add_argument('-o', '--outputfile', dest='outputfile', type=str, \
        help='Name of the output file without file extension. Default extension is txt.', \
        required=False, default='prescott-scores')

    main_parser.add_argument('--usepopmax', dest='usepopmax', type=str, \
        help='A true or false value to use population max. frequency of one of the eight populations available in GnomAD (Default is true).',
        required=False, default="true")
    
    main_parser.add_argument('--equation', dest='equation', type=int, \
        help='An integer value. Default is 2, which means it only reduces scores if it is too frequent!',
        required=False, default=2)

    main_parser.add_argument('-c', '--coefficient', dest='coefficient', type=float, \
        help='Scaling coefficient to adjust contribution of frequency information to the scores. \n It is denoted with c in the model equation. Default is 1.0', \
        required=False, default=1.0)
    main_parser.add_argument('-f', '--frequencycutoff', dest='frequencycutoff', type=float, \
        help='Frequency cutoff value. \n It is denoted with Fc in the model equation. Default is -4.0.', \
        required=False, default=-4.0)
    
    main_parser.add_argument('--usefrequencies', dest='usefrequencies', type=str, \
        help='Do not touch this if you don\'t know what you are doing! Default is true',
        required=False, default='true')
    
    main_parser.add_argument('--ranksorted', dest='ranksorted', type=str, \
        help='If your data is already ranksorted, change this argument to true. Default is false',
        required=False, default='false')
    
    main_parser.add_argument('--escottformat', dest='escottformat', type=str, \
        help='Main format of escott file. There are two possibilities: gemme or singleline. \n'+\
            'gemme: a horizontal format of 20 rows and N columns.\n'+\
            'singleline: each line contains a mutation and its value separated by a space.\n'+\
            'M1A 0.378\n', required=False, default='gemme')
    # main_parser.add_argument('--colormap', dest='colormap', type=str, \
    #     help='A colormap as defined in matplotlib',
    #     required=False, default='coolwarm_r')

    # main_parser.add_argument('--field', dest='field', type=int, \
    #     help='An integer value starting from 0 for reading the rhapsody output',
    #     required=False, default=10)
    
    # main_parser.add_argument('-b', '--beginning', dest='beginning', type=int, \
    #     help='An integer to indicate the first residue index.',
    #     required=False, default=1)

    # main_parser.add_argument('-e', '--end', dest='end', type=int, \
    #     help='An integer to indicate the final residue index.',
    #     required=False, default=None)



    args = main_parser.parse_args()
    print("\n\n@> Running PRESCOTT with the following parameters:\n\n")
    print("@> ESCOTT file                      : {}".format(args.escottfile))
    print("@> GNOMAD frequency file            : {}".format(args.gnomadfile))
    print("@> Use population max. freq         : {}".format(str(args.usepopmax).lower()))
    print("@> Which equation to use (Default=2): {}".format(str(args.equation)))
    print("@> Scaling coefficient (Default=1.0): {}".format(args.coefficient))
    print("@> Frequency cutoff (Default=-4.0)  : {}".format(args.frequencycutoff))
    print("@> Name of the output file          : {}".format(args.outputfile))
    # End of argument parsing!

    protein = os.path.splitext(os.path.basename(args.escottfile))[0]
    outfile = os.path.splitext(args.outputfile)[0]

    # Check if file exists
    usePopMaxOrNot = args.usepopmax.lower()
    version = args.equation
    if (os.path.exists(args.escottfile)):

        if(args.escottformat=='gemme'):

            #Parse the file containing raw ESCOTT scores. 
            scanningMatrix = parseGEMMEoutput(args.escottfile, verbose=False)
            
            #Convert the matrix format to singleline format
            localResidueList = None
            if(args.sequencefile != None):
                referenceSeq = SeqIO.read(args.sequencefile, 'fasta')
                localResidueList = list(referenceSeq.seq)
            aaOrderList = list('ACDEFGHIKLMNPQRSTVWY')
            writeSinglelineFormat(scanningMatrix, protein+'_singleline.txt', residueList = localResidueList,\
                            beg=0, end=None, aaOrder = aaOrderList, \
                            offSet=0)
            
            #Mostyl, I am using normPred_Combi_singleline as input file and it doesn't have a header.
            df = pd.read_table(protein+'_singleline.txt', sep="\s+", header=None)
            
        elif(args.escottformat=='singleline'):
            
            df = pd.read_table(args.escottfile, sep="\s+", header=None)
        else:
            print('@> ERROR: Unknown escott format. It should be gemme or singleline!')
            sys.exit(-1)

        if(args.ranksorted == 'false'):
            #data = np.genfromtxt(args.input,dtype=None)
            data = df.to_numpy()
            rawData = data.T[1]
            processedData = 1.0 - rankSortData(rawData)
            with open(protein+'_singleline_1-ranksort.txt', 'w') as f:
                #f.write("#Resid Value\n")
                for i in range (len(processedData)):
                    f.write("{:} {:6.2f}\n".format(data.T[0][i], processedData[i]))
            
            dfESCOTT = pd.read_table(protein+'_singleline_1-ranksort.txt', sep='\s+', header=None)
        else:
            dfESCOTT = df

        dfESCOTT.columns = ['mutant', 'ESCOTT']
        dfESCOTT['protein']=protein
    else:
        print("ERROR: ESCOTT input file does not exist!")
        sys.exit(-1)

    myBigMergedDF = pd.DataFrame()
    myBigMergedDF = pd.concat([myBigMergedDF, dfESCOTT], ignore_index=True)

    gnomadDF = getGnomADOverallFrequency(args.gnomadfile, usePopMax=usePopMaxOrNot)
    
    # Assign labels to pathogenic/benign mutations for performance evaluation
    gnomadDF['labels'] = ""
    for index, row in gnomadDF.iterrows():
        if ((row['ClinVar Clinical Significance']=='Benign/Likely benign') or \
            (row['ClinVar Clinical Significance']=='Benign') or \
            (row['ClinVar Clinical Significance']=='Likely benign')):
            gnomadDF.at[index,'labels'] = 0

        if((row['ClinVar Clinical Significance']=='Pathogenic/Likely pathogenic') or \
            (row['ClinVar Clinical Significance']=='Pathogenic') or \
            (row['ClinVar Clinical Significance']=='Likely pathogenic')):
            gnomadDF.at[index,'labels'] = 1
    if (len(gnomadDF.loc[(gnomadDF['labels']==0) | (gnomadDF['labels']==1)]) > 0):   
        print(gnomadDF.loc[(gnomadDF['labels']==0) | (gnomadDF['labels']==1)])
    # print(gnomadDF['ClinVar Clinical Significance'])
    # Add frequency column and a dummy frequency to each row in myBigMergedDF
    myBigMergedDF['frequency'] = 999.0
    myBigMergedDF['labels'] = np.nan
    myBigMergedDF['position'] = ""

    useFrequencies = args.usefrequencies
    selectedPositionsList = []
    selectedValuesList = []
    selectedMutantsList = []

    # Assign ESCOTT scores to PRESCOTT scores.
    # Then, we will modify them according to different conditions. 
    myBigMergedDF['PRESCOTT'] = myBigMergedDF['ESCOTT']
    labelsList = []
    if (useFrequencies.lower() == 'true'):

        # print(myBigMergedDF)
        for index, row in myBigMergedDF.iterrows():
            myBigMergedDF.at[index,'position'] = row['mutant'][1:-1]
            # print(row['mutant'], row['ESCOTT'])
            # print(row['mutant'][1:-1])
            temp = (gnomadDF.loc[gnomadDF['mutant'] == row['mutant'], 'Allele Frequency Log'].values)
            #print(temp)
            if (len(temp) > 0):
                myBigMergedDF.at[index,'frequency'] = temp[0]
                myBigMergedDF.at[index,'labels'] = gnomadDF.loc[gnomadDF['mutant'] == row['mutant'], 'labels'].values[0]

        # print(myBigMergedDF)
        scalingCoeff = args.coefficient
        freqCutoff = args.frequencycutoff
        for index, row in myBigMergedDF.iterrows():
            if(row['frequency'] != 999.0):
                
                # print(row['frequency'])
                # freq = np.log10(row['frequency'])
                freq = row['frequency']
                # print(freq)
                temp1 = row['PRESCOTT']
                label = row['labels']
                if(version==1):
                    if(freq>freqCutoff):
                        temp2 = temp1 - freq*scalingCoeff/freqCutoff
                        if(temp2<0.0):
                            myBigMergedDF.at[index,'PRESCOTT'] = 0.0
                            selectedValuesList.append(0.0)
                        else:
                            myBigMergedDF.at[index,'PRESCOTT'] = temp2
                            selectedValuesList.append(temp2)
                        selectedPositionsList.append(row['position'])
                        selectedMutantsList.append(row['mutant'])
                
                if(version==2):
                    if(freq>freqCutoff):
                        temp2 = temp1 - scalingCoeff*(freqCutoff - freq)/freqCutoff
                        if(temp2<0.0):
                            temp2 = 0.0
                        myBigMergedDF.at[index,'PRESCOTT'] = temp2
                        if(label==0 or label==1):
                            selectedValuesList.append(temp2)
                            selectedPositionsList.append(row['position'])
                            selectedMutantsList.append(row['mutant'])
                if(version==3):
                    temp2 = temp1 - scalingCoeff*(freqCutoff - freq)/freqCutoff
                    if(freq>freqCutoff):
                        if(temp2<0.0):
                            temp2 = 0.0
                        myBigMergedDF.at[index,'PRESCOTT'] = temp2
                        if(label==0 or label==1):
                            selectedValuesList.append(temp2)
                            selectedPositionsList.append(row['position'])
                            selectedMutantsList.append(row['mutant'])
                    else:
                        if(temp2>1.0):
                            temp2 = 1.0

                        myBigMergedDF.at[index,'PRESCOTT'] = temp2
                        if(label==0 or label==1):
                            selectedValuesList.append(temp2)
                            selectedPositionsList.append(row['position'])
                            selectedMutantsList.append(row['mutant'])
                if(version==4):
                    temp2 = temp1 - scalingCoeff*(freqCutoff - freq)/freqCutoff
                    if(freq>freqCutoff):
                        if(temp2<0.0):
                            temp2 = 0.0
                        myBigMergedDF.at[index,'PRESCOTT'] = temp2
                        if(label==0 or label==1):
                            selectedValuesList.append(temp2)
                            selectedPositionsList.append(row['position'])
                            selectedMutantsList.append(row['mutant'])
                    else:
                        print(myBigMergedDF.loc[index, 'Selected Population'])
                        sys.exit(-1)
                        if (myBigMergedDF.iloc[index,'Selected Population'].values==None):                
                            if(temp2>1.0):
                                temp2 = 1.0

                            myBigMergedDF.at[index,'PRESCOTT'] = temp2
                            if(label==0 or label==1):
                                selectedValuesList.append(temp2)
                                selectedPositionsList.append(row['position'])
                                selectedMutantsList.append(row['mutant'])                        

    # myBigMergedDF.dropna(subset = ['labels'], inplace=True)

    clinvarLabeledDF = myBigMergedDF.loc[(myBigMergedDF['labels']==0) | (myBigMergedDF['labels']==1)]
    clinvarLabeledDF['labels'] = clinvarLabeledDF['labels'].astype('int64')
    if(len(clinvarLabeledDF)>0):
        print("\nMutations with ClinVar labels according to the gnomAD file:\n")
        print(clinvarLabeledDF)
    #print(myBigMergedDF.loc[(myBigMergedDF['labels']=='0') | (myBigMergedDF['labels']=='1'), 'labels'])
    # print(clinvarLabeledDF['labels'].values)
    # print(clinvarLabeledDF['ESCOTT'].values)
    numPathogenic = len(myBigMergedDF.loc[(myBigMergedDF['labels']==1)])
    numBenign = len(myBigMergedDF.loc[(myBigMergedDF['labels']==0)])

    if ((numPathogenic >=1) and (numBenign>=1)):
        fprESCOTT, tprESCOTT, AUC_ESCOTT = plotROCandAUCV2(clinvarLabeledDF['labels'], \
                                                        clinvarLabeledDF['ESCOTT'])
        
        fprPRESCOTT, tprPRESCOTT, AUC_PRESCOTT = plotROCandAUCV2(clinvarLabeledDF['labels'], \
                                                        clinvarLabeledDF['PRESCOTT'])
        # fprPRESCOTT, tprPRESCOTT, AUC_PRESCOTT = plotROCandAUCV2(myBigMergedDF.loc[(myBigMergedDF['labels']==0) | (myBigMergedDF['labels']==1), 'labels'], \
        #                                                   myBigMergedDF.loc[(myBigMergedDF['labels']==0) | (myBigMergedDF['labels']==1), 'PRESCOTT'])

        fig = plt.figure(figsize=(12,6))
        # plt.rcParams.update({'font.size': 18})
        plt.grid(linestyle='--')
        # plt.title(protName + " - "+method+" AUC={:.2f}".format(AUC_ESCOTT))
        plt.title("AUC={:.2f} -> AUC={:.2f}".format(AUC_ESCOTT, AUC_PRESCOTT))
        plt.ylim([0.0, 1.0])
        #plt.xlim([1000, 1863])
        plt.scatter(myBigMergedDF.loc[myBigMergedDF['labels'] == 1, 'position'], myBigMergedDF.loc[myBigMergedDF['labels'] == 1, 'ESCOTT'], marker='o', color='red', label='pathogenic')
        plt.scatter(myBigMergedDF.loc[myBigMergedDF['labels'] == 0, 'position'], myBigMergedDF.loc[myBigMergedDF['labels'] == 0, 'ESCOTT'], marker='o', color='blue', label='benign')
        if (useFrequencies.lower() == 'true'):
            #print(selectedPositionsList)
            #print(selectedValuesList)
            plt.scatter(selectedPositionsList, selectedValuesList, marker='o', color='olive', label='PRESCOTT')

            # Add vertical lines connecting old and new values
            for i in range(len(selectedPositionsList)):
                plt.annotate("", xy=(selectedPositionsList[i], selectedValuesList[i]), xycoords='data', \
                            xytext=(selectedPositionsList[i], myBigMergedDF.loc[myBigMergedDF['mutant'] == selectedMutantsList[i], 'ESCOTT'].values[0]), textcoords='data',
                arrowprops=dict(arrowstyle="->", connectionstyle="arc3"))

        plt.xticks(rotation = 90)
        plt.ylabel("Ranksorted Score")
        plt.xlabel("Position")
        plt.legend(loc='upper right')
        plt.tight_layout()
        plt.savefig("clinvar-vs-position.png")
        plt.close()
        print("@> AUC= {:.3f} {:.3f}".format( AUC_ESCOTT, AUC_PRESCOTT))

    myBigMergedDF.to_csv(outfile+'-details.csv', index=None)
    myBigMergedDF.to_csv(outfile+'.txt', columns=['mutant', 'PRESCOTT'], index=False, header=None, sep=' ')
    os.remove(protein+'_singleline.txt')

if __name__ == "__main__":
    main()
