# -*- coding: utf-8 -*-

# Copyright (c) 2018: Elodie Laine
# Copyright (c) 2022-2023: Mustafa Tekpinar
# This code is part of the gemme package and governed by its license.
# Please see the LICENSE.txt file included as part of this package.

import sys
import os
import time
import argparse
import re
import subprocess
import math
import numpy as np
import matplotlib.pylab as plt

from prody import *
from scipy.stats import rankdata

from gemmeAnal import *

import pandas as pd

def extractQuerySeq(filename):
	"""
		# Extract the query sequence from the input alignment
	"""

	fIN = open(filename,"r")
	lines = fIN.readlines()
	fIN.close()

	if lines[0][0]!=">":
		raise Exception('bad FASTA format')
	else:
		prot = re.compile("[^A-Z0-9a-z]").split(lines[0][1:])[0]
		fOUT = open(prot+".fasta","w")
		fOUT.write(">"+prot+"\n")
		seq=""
		i = 1
		while lines[i][0]!=">":
			seq = seq + lines[i].strip().strip(".").strip("-")
			fOUT.write(lines[i])
			i = i + 1
	fOUT.close()
	return prot,seq,i-1

def rankSortProteinData(dataArray, inverted=True):
    """
        This function ranksorts protein data and converts it
        to values [0,1.0].
    
    Parameters:
    ----------
    dataArray: numpy array of arrays
               data read by numpy.loadtxt

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

#    np.savetxt(outfile, np.array([dataArray, normalizedRankedDataArray]).T)
    if(inverted==True):
        return (1.0 - normalizedRankedDataArray)
    else:
        return (normalizedRankedDataArray)

def calcPercentDFI(pdbfile, outfile, nmodes=None, attenuate="true", ranksorted="true"):
    """
        Calculate rank normalized dynamic flexibility index. 
    """
    
    #confProDy(auto_show=False)

    structure = parsePDB(pdbfile)
    calphas = structure.select('name CA')

    anm = ANM('ANM')
    anm.buildHessian(calphas)
    if(nmodes==None):
        anm.calcModes(n_modes='all')
    else:
        anm.calcModes(n_modes=int(nmodes))

    dfi = calcDynamicFlexibilityIndex(anm, calphas, select='ca', norm=True)
    percentDfi = rankSortProteinData(dfi)

    if(attenuate.lower()=='true'):
        percentDfi = attenuateEndPoints(percentDfi)

    # if(outfile != None):
    #     np.savetxt(outfile, percentDfi)
    # else:
    #     np.savetxt(outfile, dfi)

    if(ranksorted.lower()=='true'):
        return percentDfi
    else:
        return dfi

def calcBfactors(pdbfile, outfile, nmodes=None, attenuate="true", 
                ranksorted="true", inverted=False):
    """
        Calculate rank normalized mean squared fluctuations,
        which correlate well with Bfactors. 
    """
    
    #confProDy(auto_show=False)

    structure = parsePDB(pdbfile)
    calphas = structure.select('name CA')

    anm = ANM('ANM')
    anm.buildHessian(calphas)
    if(nmodes==None):
        anm.calcModes(n_modes='all')
    else:
        anm.calcModes(n_modes=int(nmodes))

    bfactors = calcSqFlucts(anm)

    percentBfactors = rankSortProteinData(bfactors, inverted=False)

    if(attenuate.lower()=='true'):
        percentBfactors = attenuateEndPoints(percentBfactors)

    # if(outfile != None):
    #     np.savetxt(outfile, percentBfactors)
    # else:
    #     np.savetxt(outfile, dfi)

    if(ranksorted.lower()=='true'):
        if(inverted==True):
            return (1.0 - percentBfactors)
        else:
            return percentBfactors
    else:
        if(inverted==True):
            return (np.max(bfactors) - bfactors)
        else:
            return bfactors

def getBfactors(pdbfile, outfile, attenuate="true", \
                ranksorted="true",  inverted=True):
    """
        Calculate rank normalized dynamic flexibility index. 
    """
    
    #confProDy(auto_show=False)

    structure = parsePDB(pdbfile)
    calphas = structure.select('name CA')

    bfactors = calphas.getBetas()

    percentBfactors = rankSortProteinData(bfactors, inverted=False)

    if(attenuate.lower()=='true'):
        percentBfactors = attenuateEndPoints(percentBfactors)

    # if(outfile != None):
    #     np.savetxt(outfile, percentBfactors)
    # else:
    #     np.savetxt(outfile, dfi)

    if(ranksorted.lower()=='true'):
        if(inverted==True):
            return (1.0 - percentBfactors)
        else:
            return percentBfactors
    else:
        if(inverted==True):
            return (np.max(bfactors) - bfactors)
        else:
            return bfactors
    
def attenuateEndPoints(dfi):
    """
        Reduce DFI on N and C terminals with a sigmoid function.
    """

    debug = 0

    n = len(dfi)
    if(debug):
        #Check array lengths
        print(n/2)
        print(len(dfi))

    if(n%2 == 0):
        x1 = np.linspace(-0, 10, int(n/2))
        x2 = np.linspace(-10, 0, int(n/2))
    else:
        x1 = np.linspace(-0, 10, int((n-1)/2))
        x2 = np.linspace(-10, 0, int((n+1)/2))
    
    z1 = 1/(1 + np.exp(-2.50*x1))
    z2 = 1.0 - (1/(1 + np.exp(-2.50*x2)))


    z = np.append(z1, z2)

    if(debug):
        #Check array lengths
        print(len(z1))
        print(len(z2))
        print(len(z))
        #Plot the arrays
        plt.plot(z)
        plt.plot(dfi)
        plt.plot(dfi*z)
        plt.xlabel("x")
        plt.ylabel("Sigmoid(X)")
        plt.show()


    if(len(dfi) == len(z)):
        return (dfi*z)
    else:
        print("ERROR: Can not attenuate N and C terminal dfi signals!")
        sys.exit(-1)


alphabeticalAminoAcidsList = ['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L',
                              'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y']
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()
        for item in data:
            if item[0] == "\"":
                aaColumn.append(item)
            elif (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

def plotGEMMEmatrix(scanningMatrix, outFile, beg, end, \
                    colorMap = 'coolwarm', offSet=0, pixelType='square'):
    """
        A function to plot deep mutational scanning matrices. 
  
    Parameters
    ----------
    scanningMatrix: numpy array of arrays
        Data matrix to plot

    outFile: string
        Name of the output png image
    
    colorMap: matplotlib cmap
        Any colormap existing in matplotlib.
        Default is coolwarm. 

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

    pixelType: string
        It can only have 'square' or 'rectangle' values.
        It is a matter of taste but I added it as an option.
        Default is 'square'

    Returns
    -------
    Nothing

    """

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

    # print("Beginning: "+str(beg))
    # print("End      : "+str(end))
    
    # print(len(scanningMatrix[0]))
    subMatrix = scanningMatrix[:, (beg-1):end]
    #print(subMatrix)

    ##########################################################################
    # Set plotting parameters
    nres_shown = len(subMatrix[0])
    fig_height=8
    # figure proportions
    fig_width = fig_height/2  # inches
    fig_width *= nres_shown/20
    
    fig, ax = plt.subplots(figsize=(fig_width, fig_height))

    if (nres_shown >=200):
        majorTics = 50
    else:
        majorTics = 25


    major_nums_x = np.arange(majorTics, len(subMatrix[0]), majorTics, dtype=int)
    major_nums_x = major_nums_x -1
    major_nums_x = np.insert(major_nums_x, 0, 0)
    # print(major_nums_x)

    minor_nums_x = np.arange(10, len(subMatrix[0]), 10, dtype=int)
    minor_nums_x = minor_nums_x - 1
    minor_nums_x = np.insert(minor_nums_x, 0, 0)
    # print(minor_nums_x)

    major_labels_x = major_nums_x + 1 + offSet

    major_nums_y = np.arange(0, 20, 1, dtype=int)
    major_labels_y = alphabeticalAminoAcidsList

    plt.xticks(major_nums_x, major_labels_x, size=28)
    ax.set_xticks(minor_nums_x, minor=True)
    
    plt.yticks(major_nums_y, major_labels_y, size=16)
    ax.set_yticklabels(major_labels_y, ha='left')
    ax.tick_params(axis='y', which='major', pad=30)

    #############################################################################
    if(pixelType=='square'):
        #For plotting square pixels
        plt.imshow(subMatrix, cmap=colorMap)
    elif(pixelType=='rectangle'):
        #For plotting rectangular pixels
        plt.imshow(subMatrix, cmap=colorMap, aspect=3.0)
    else:
        print("\nERROR: Unknown pixelType specified!\n")
        sys.exit(-1)

    #plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=1.0)
    plt.tight_layout()
    plt.savefig(outFile)
    #plt.show()
    

    #plt.imsave('output.png', subMatrix)
    plt.close()


def check_argument_groups(parser, arg_dict, group, argument):
    """
    Check for use of arguments.
    Raise an error if the parser uses an argument that belongs to an argument
    group (i.e. mode) but the group flag is not used or if the argument is
    required in the argument group and it's not used.
    Notes:
    - argument and group should be strings and should start with - or --
    - group should be a flag with action 'store_true'
    - argument shouldn't be a flag, and should be None if it isn't used.
    >>> import argparse
    >>> parser = argparse.ArgumentParser()
    >>> parser.add_argument('--phylo', action='store_true')
    _StoreTrueAction(...)
    >>> parser.add_argument('--inseq')
    _StoreAction(...)
    >>> args = parser.parse_args(['--phylo', '--inseq', 'not_none'])
    >>> check_argument_groups(parser, vars(args), '--phylo', '--inseq', True)
    """
    c = 0
    for myArg in argument:
    	arg_name = myArg.replace("-", "")
    	if arg_dict[arg_name]!='':
    		c = c+1
    group_name = group.replace("-", "")
    if arg_dict[group_name]=="input":
    	if c!=1:
    		parser.error("gemme requires " + str(argument) +
                         " if " + group + " is set to input.")
    else:
        if c>0:
            parser.error("gemme requires " + group +
                         " to be set to input if " + str(argument) + " is used.")
    return None

def parse_command_line():
    """
    Parse command line.
    It uses argparse to parse phylosofs' command line arguments and returns the
    argparse parser.
    """
    parser = argparse.ArgumentParser(
        prog="gemme",
        description="""
        GEMME (Global Epistasis Model for predicting Mutational Effects) 
        is a tool to predict mutational outcomes based on sequence analysis
        """,
        epilog="""
        It has been developed at LCQB (Laboratory of Computational and
        Quantitative Biology), UMR 7238 CNRS, Sorbonne Université.
        If you use it, please cite:
        Laine E, Karami Y, Carbone A. Predicting Protein Mutational Landscapes
        using Evolutionary Conservation and Global Epistasis
        """,
    )

    parser.add_argument(
    	'input', metavar='aliFile', type=str,
        help='input alignment file in FASTA format (first sequence = ungapped query sequence)',
    )

    parser.add_argument(
        '-m', '--mutations',
        help='text file containing the mutations of interest. Each line of the file should contain a mutation (e.g. D136R) or combination of mutations separated by commas and ordered according to their positions in the sequence (e.g. D136R,V271A)',
        default=''
    )

    retMet_args = parser.add_argument_group('Conservation levels calculation', """
        Arguments used for computing the conservation levels.
        """)
    
    retMet_args.add_argument(
        '-n', '--nIter',
        help='number of iterations to compute the conservation levels',
        default='1'
    )

    retMet_args.add_argument(
        '-N', '--NSeqs',
        help='maximum number of sequences to compute the conservation levels',
        default='40000'
    )
        
    retMet_args.add_argument(
        '-r', '--retrievingMethod',
        help='mode to retrieve related sequences for conservation levels calculation (input, local or server)',
        default='local'
    )

    retMet_args.add_argument(
        '-b', '--blastFile',
        help='psiblast file containing related sequences',
        default=''
    )

    retMet_args.add_argument(
        '-f', '--fastaFile',
        help='fasta file containing related sequences',
        default=''
    )

    retMet_args.add_argument('--jetfile', dest='jetfile', type=str, \
        help="If a jet file is provided, it will skip JET2 calculation and use the precalculated JET2 data in that file. Default is None",
        required=False, default=None)

    retMet_args.add_argument('--alphabet', dest='alphabet', type=str, \
        help="Which alphabet to use. Default is lw-i.7",
        required=False, default="lw-i.7")

    retMet_args.add_argument('-p', '--pdbfile', dest='pdbfile', type=str, \
        help="If a pdb file is provided, it will skip fake pdb file production step and use that file. Default is None",
        required=False, default=None)
    
    retMet_args.add_argument('--normweightmode', dest='normweightmode', type=str, \
        help="It can be one of these: 'trace', 'tracemovingaverage', 'cv', 'pc', 'dfi', 'maxtracesc',"+\
             "'maxtracepc', 'maxtracecv', 'maxtracepccv', 'maxtracepcsc', 'maxtracepchalfpccv' or 'halfcvpc'. Default is 'trace'.",
        required=False, default="trace")

    args = parser.parse_args()

    arg_dict = vars(args)

    check_argument_groups(parser, arg_dict, '--retrievingMethod', ['--blastFile','--fastaFile'])
 
    # Check flag arguments
    if args.input is None:
    	parser.error("GEMME requires an input alignment.")

    return args


def doit(inAli,mutFile,retMet,bFile,fFile,n,N, jetfile, pdbfile, normWeightMode, alphabet):
    """
    Perfect explanation for a function: typing the function call exactly!
    doit is basically the main function in disguise!
    doit(args.input,args.mutations,args.retrievingMethod,args.blastFile,args.fastaFile, args.jetfile)
    """
    if(mutFile != ''):
        simple = False
    else:
        simple = True

    
    prot,seq,nl=extractQuerySeq(inAli)

    if(pdbfile == None):
        createPDB(prot,seq)
    else:
        if(os.path.exists(prot+".pdb")):
            shutil.move(prot+".pdb", prot+"_original.pdb")
            shutil.copy2(pdbfile, prot+".pdb")
        else:
            shutil.copy2(pdbfile, prot+".pdb")

    print("query protein: "+prot)
    
    if( (normWeightMode != 'trace') and \
        (normWeightMode != 'cv') and \
        (normWeightMode != 'pc') and \
        (normWeightMode != 'dfi') and \
        (normWeightMode != 'bfactor') and \
        (normWeightMode != 'maxtracepc') and \
        (normWeightMode != 'maxpctrace') and \
        (normWeightMode != 'maxpccv') and \
        (normWeightMode != 'maxcvpc') and \
        (normWeightMode != 'maxtracecv') and \
        (normWeightMode != 'maxtracesc') and \
        (normWeightMode != 'maxcvtrace') and \
        (normWeightMode != 'maxtracedfi') and \
        (normWeightMode != 'maxtracebfactor') and \
        (normWeightMode != 'maxdfitrace') and \
        (normWeightMode != 'maxbfactortrace') and \
        (normWeightMode != 'maxtracepccv') and \
        (normWeightMode != 'maxtracepcsc') and \
        (normWeightMode != 'maxtracepchalfpccv') and \
        (normWeightMode != 'maxtracecvpc') and \
        (normWeightMode != 'halfpccv') and \
        (normWeightMode != 'halfcvpc') and \
        (normWeightMode != 'halftracepc') and \
        (normWeightMode != 'combinedv1') and \
        (normWeightMode != 'combinedv2') and \
        (normWeightMode != 'combinedv3') and \
        (normWeightMode != 'combinedv4') and \
        (normWeightMode != 'maxtracehalfcvpc') and \
        (normWeightMode != 'maxtracehalftracecvhalfcvpc') and \
        (normWeightMode != 'maxhalftracepchalftracecvhalfcvpc') and \
        (normWeightMode != 'tracemovingaverage') and \
        (normWeightMode != 'maxtracehalfpccv')):
        print("ERROR: normWeightMode can only be 'trace', 'tracemovingaverage', 'cv', 'pc', 'dfi', bfactor,\n"+\
              "       'maxtracepc', 'maxtracedfi', 'maxpccv', 'maxtracesc',"+\
              "       'maxtracebfactor', 'maxtracepcsc', 'maxhalftracepchalftracecvhalfcvpc', "+\
              "       'maxtracecv', 'maxtracepccv', maxtracepchalfpccv"+\
              "       'combinedv1', 'combinedv2', 'combinedv3', 'combinedv4'"+\
              "       'halfcvpc', 'halftracepc', 'maxtracehalfcvpc' or maxtracehalftracecvhalfcvpc!")
        sys.exit(-1)

    structure = parsePDB(prot+".pdb")
    calphas = structure.select('name CA')

    chains = list(set(calphas.getChids()))
    
    #print(chains)

    if((jetfile) == None):
        #I intend to run JET2 completely externally!!
        #It is too much buggy and it has too many dependencies. 
        #Using it with a Docker or Singularity may be the best solution!
        print("computing conservation levels...")
        launchJET(prot,retMet,bFile,fFile, pdbfile, chains, n,N,nl)
        print("done")
    else:
        print("using previously calculated JET2 data from "+jetfile+"...")
        
        #Doing this copy operation just due to a peculiar behaviour of JET2.
        #According to JET2, all input files must have the same name obtained 
        #from the fasta file right after the > symbol.
        #Weird but true!
        if(jetfile!=prot+"_jet.res"):
            shutil.copy2(jetfile, prot+"_jet.res")
        print("done")

    #If a real pdb file is given, calculate dfi for the residues. 
    if(((normWeightMode=='dfi') or (normWeightMode=='maxtracedfi') or (normWeightMode=='maxdfitrace'))):
        if (pdbfile == None):
            print("ERROR: There is not any pdb file to calculate DFI.")
            sys.exit(-1)
        else:
            print("Calculating %DFI data per residue!")
            dfi = calcPercentDFI(pdbfile, outfile=None, nmodes=None, attenuate="true", ranksorted="true")
            dfi = dfi
            df = pd.read_table(prot+"_jet.res")
            print(df)

            df['dfi'] = dfi.round(4)
            df.to_csv(prot+"_jet.res", header=True, index=None, sep='\t', mode='w')

            print(df)

    #If a real pdb file is given, calculate dfi for the residues. 
    if(((normWeightMode=='tracemovingaverage'))):
        print("Calculating trace moving average per residue!")
        # tma = tracemovingaverage
        df = pd.read_csv(prot+"_jet.res", delimiter=r"\s+") 
        print(df.columns)
        # Get the window of series
        # of observations of specified window size
        window_size = 3
        windows = df['trace'].rolling(window_size)
        
        # Create a series of moving
        # averages of each window
        moving_averages = windows.mean()
        tma = moving_averages

        print(df)

        df['tracemovingaverage'] = tma.round(4)
        #The first two elements are just copied bc they don't exist due to the moving average. 
        df['tracemovingaverage'].iloc[0] = df['trace'].iloc[0].round(4)
        df['tracemovingaverage'].iloc[1] = df['trace'].iloc[1].round(4)
        df.to_csv(prot+"_jet.res", header=True, index=None, sep='\t', mode='w')

        print(df)
    #If a real pdb file is given, calculate or get Bfactors for the residues. 
    if(((normWeightMode=='bfactor') or (normWeightMode=='maxtracebfactor') or (normWeightMode=='maxbfactortrace'))):
        isCalc = True
        if (pdbfile == None):
            print("ERROR: There is not any pdb file to calculate or get Bfactors!")
            sys.exit(-1)
        else:        
            if (isCalc):
                print("Computing Bfactors from the user-provided pdb file using ANM with all modes.")
                #Original
                # bfactors = calcBfactors(pdbfile, outfile=None, nmodes=None, \
                #           attenuate="true", ranksorted="true", inverted=False)
                bfactors = calcBfactors(pdbfile, outfile=None, nmodes=10, \
                            attenuate="false", ranksorted="true", inverted=True)
            else:
                print("Getting Bfactors from the user-provided pdb file.")
                bfactors = getBfactors(pdbfile, outfile=None,\
                            attenuate="true", ranksorted="true", inverted=False)

            bfactors = bfactors
            df = pd.read_table(prot+"_jet.res")
            print(df)

            df['bfactor'] = bfactors.round(4)
            df.to_csv(prot+"_jet.res", header=True, index=None, sep='\t', mode='w')

            print(df)
        
    launchPred(prot,inAli,mutFile, normWeightMode, alphabet)

    #Do Python plotting here
    #TODO: Eventually, I will do the map plotting with a completely independent 
    #      module and call the module here! demust module was created!
    #TODO: Mark the original (wildtype) residue locations with a dot or something
    #      special to show the original amino acid.
    #TODO: You can even put letters on the top line like in EVmutation output. 
    

    if(simple):
        print("generating the plots...")
        #Check if the normalized data files were created. 
        if(os.path.exists(prot+"_normPred_evolEpi.txt")):
            gemmeData = parseGEMMEoutput(prot+"_normPred_evolEpi.txt", verbose=False)
            plotGEMMEmatrix(gemmeData, prot+"_normPred_evolEpi.png", 1, None,\
                            colorMap='Blues_r', offSet=0, pixelType='square')
        else:
            print("ERROR: "+prot+"_normPred_evolEpi.txt file does not exist!")
            print("       Can not generate "+prot+"_normPred_evolEpi.png file!")
        
        #Check if the normalized data files were created. 
        if(os.path.exists(prot+"_normPred_evolInd.txt")):
            gemmeData = parseGEMMEoutput(prot+"_normPred_evolInd.txt", verbose=False)
            plotGEMMEmatrix(gemmeData, prot+"_normPred_evolInd.png", 1, None,\
                            colorMap='Greens_r', offSet=0, pixelType='square')
        else:
            print("ERROR: "+prot+"_normPred_evolInd.txt file does not exist!")
            print("       Can not generate "+prot+"_normPred_evolInd.png file!")

        #Check if the normalized data files were created. 
        if(os.path.exists(prot+"_normPred_evolCombi.txt")):        
            gemmeData = parseGEMMEoutput(prot+"_normPred_evolCombi.txt", verbose=False)
            plotGEMMEmatrix(gemmeData, prot+"_normPred_evolCombi.png", 1, None,\
                            colorMap='Oranges_r', offSet=0, pixelType='square')
        else:
            print("ERROR: "+prot+"_normPred_evolCombi.txt file does not exist!")
            print("       Can not generate "+prot+"_normPred_evolCombi.png file!")
        print("done")

    cleanTheMess(prot,bFile,fFile, chainID=chains[0])


def main():
    tic = time.perf_counter()

    args = parse_command_line()
    doit(args.input,args.mutations,args.retrievingMethod,args.blastFile,\
         args.fastaFile,args.nIter,args.NSeqs, args.jetfile, args.pdbfile,\
         args.normweightmode, args.alphabet)

    toc = time.perf_counter()
    
    print(f"S-GEMME computation finished in {toc - tic:0.4f} seconds!")

if (__name__ == '__main__'):
    main()

