pain in my ass

import os, sys, time, gzip, re, psutil, random
import numpy as np
import pandas as pd
from scipy.spatial import distance
import xgboost as xgb
from termcolor import colored


from sklearn.metrics import accuracy_score as acs
from sklearn.ensemble import RandomForestClassifier as rfc
from sklearn.svm import SVC as svc

acids = ['ALA', 'CYS', 'ASP', 'GLU', 'PHE', 'GLY', 'HIS', 'ILE', 'LYS', 'LEU', 'MET', 'ASN', 'PRO', 'GLN', 'ARG', 'SER', 'THR', 'VAL', 'TRP', 'TYR']

#open pdb if res is less than limit and if file exists
def openPDBFile(protein_name, resLim = 2.0):
    pdb = []
    res = None
    flag = False
    name = pathFinder(protein_name)
    if os.path.exists(name):
        with gzip.open(name,'rt') as f:
            for line in f:
                if flag == False:
                    res = checkResolution(line)
                    flag = True
                    if res > resLim:
                        return None
                pdb.append(line)

        return pdb
    else:
        return None

#return resolution of prot
def checkResolution(line):
    res = 0
    if ('RESOLUTION.' in line) and ('ANGSTROMS' in line):
        res = [float(s) for s in re.findall(r'-?\d+\.?\d*', line)][1]
    return res

#dir where the pdb files are saved
def pathFinder(protein_name):
    path = '/home/artur/a/pdb/' + protein_name[1].lower() + protein_name[2].lower() + '/pdb' + protein_name.lower() + '.ent.gz'
    return path


def getXYZ(line):
    x = float(line[30:38])
    y = float(line[38:46])
    z = float(line[46:54])

    return [x, y, z]


def getID(line):
    i = int(line[6:11])
    return i

#Get atom's name
def getAtomID(line):
    i = line[13:16]
    return i

#Get amino acids name
def getResName(line):
    i = line[17:20]
    return i


def getProteinAndWater(pdb):
    protein = []
    water = []
    proteinCoords = []
    for line in pdb:
        if len(line) < 78:
            continue
        if (line[17:20] == "HOH" and line[77] == "O"):
            water.append(getXYZ(line))
        if (line[16] != "A" and line[16] != " ") or (line[27] != "A" and line[27] != " "):
            continue
        if (line[17:20] != "HOH" and line[0:4] == "ATOM"):
            protein.append(line)
    return protein, water

def toPdb(xyz,B= 0, atom_id=None, res_id="  0", element= "U"):
    if (B>999):
        B = 999.0
    x = "{:8.3f}".format(xyz[0])
    y = "{:8.3f}".format(xyz[1])
    z = "{:8.3f}".format(xyz[2])
    B = "{:6.2f}".format(float(B))

    if res_id != "  0":
        if res_id>999:
            res_id-=int(res_id/1000)*1000
        res_id = str(res_id)
        res_id = " ".join([""]*(3-len(res_id)))+res_id
    line = "ATOM         "+element+"   "+element * 3+"    "+res_id+"     42.931 -14.533  18.887        0.00           "+element
    line_ = line[:30]+x+y+z+line[54:60]+B+line[66:]+"\n"
    return line_


def getDistanceBetween(chosenMolecule, molecules):
    return distance.cdist(chosenMolecule, molecules)

#get distances between each of water molecules
def getWatersArray(water):
    cd = getDistanceBetween(water, water)
    id_cd = np.where((cd < 4.0) & (cd > 0.0))
    cd_ = cd[id_cd]
    cd_.sort()
    return cd_

#get all the carbon, oxygen, nitro atoms from protein
def getCON(protein): #carbon oxygen nitro
    C = []
    O = []
    N = []
    id_C = id_O = id_N = []

    for i in protein:
        if i[13] == 'O':
            O.append(getXYZ(i))
            id_O.append(getID(i))
        if i[13] == 'C':
            C.append(getXYZ(i))
            id_C.append(getID(i))
        if i[13] == 'N':
            N.append(getXYZ(i))
            id_N.append(getID(i))

    return [C, O, N, id_C, id_O, id_N]

#get nearly located C O N atoms from each water molecule
def getAtomsArrays(protein, water):
    global acids
    [C, O, N,  id_C, id_O, id_N] = getCON(protein)

    cd_C = getDistanceBetween(water, C)
    cd_O = getDistanceBetween(water, O)
    cd_N = getDistanceBetween(water, N)

    id_cd_C = np.where(cd_C < 4.0)
    id_cd_O = np.where(cd_O < 4.0)
    id_cd_N = np.where(cd_N < 4.0)

    cd_C_ = np.array(cd_C[id_cd_C])
    cd_O_ = np.array(cd_O[id_cd_O])
    cd_N_ = np.array(cd_N[id_cd_N])

    cd_C_.sort()
    cd_O_.sort()
    cd_N_.sort()

    return cd_C_, cd_N_, cd_O_

#get fake water molecules
def getFakeWater(water):
    cd = getDistanceBetween(water, water)
    fwater = []
    id_cd = np.where((cd < 4.0) & (cd > 0))

    idx = id_cd[0].tolist()
    jdx = id_cd[1].tolist()

    for i in range(len(idx)):
                if (jdx[i] > idx[i]):
                    fwater.append(((np.array(water[idx[i]]) + np.array(water[jdx[i]]) )/ 2))
    return fwater

#return the (number = lim) number of randomized prots from different clusters with the short names alike 3y3q_ai
def filterProts(lim):
    file = open('bc-90.out', 'r')
    global forTest
    fprots = []

    count = 0
    for line in file:
        count += 1
        prots = line.split(" ")
        for i in prots:
            if len(i) > 6:
                prots.remove(i) #delete proteins like 3y3q_ai

        if lim >= len(prots):
            fprots.extend(prots)
        else:
            random = np.random.randint(len(prots), size= lim)
            for i in random:
                fprots.append(prots[i])

    return fprots


#shuffle and return number of water molecules
def filterWater(water, lim = 100):
    water_ = []
    if len(water) <= lim:
        water_ = water
    else:
        random = np.random.randint(len(water), size= lim)
        for i in random:
            water_.append(water[i])
    return water_


def prepareData(lim = 1, protsNum = 25):

    global protsUsed

    global testProts
    global testCON
    global testFWater

    prots = filterProts(lim)
    random.shuffle(prots)

    count = 0
    forTest = 2
    #len(prots) < ~10000
    maxp = protsNum
    posData = []
    negData = []

    for i in prots:
        name = i[0:4]
        pdb = openPDBFile(name)
        if pdb == None:
            continue
        else:
            try:

                [protein, water] = getProteinAndWater(pdb)
                water_ = filterWater(water)
                fwater = getFakeWater(water_)
                count += 1
                if count == maxp:
                    break
                [C, O, N] = getAtomsArrays(protein, water_)
                posData.append([C, O, N])
                [C1, O1, N1] = getAtomsArrays(protein, fwater)
                negData.append([C, O, N])
                protsUsed.append(name)
            except ValueError:
                continue

    for i in range(forTest):
        testProts.append(protsUsed.pop())
        testCON.append(posData.pop())
        testFWater.append(negData.pop())

    return posData, negData




def balanceX(x):
    maxLen = max([len(i) for i in x])
    for i in x:
        if len(i) < maxLen:
            i = np.concatenate((np.zeros(maxLen - len(i)),i))
    print x.shape
    return x



def createXY(posData, negData, maxlen):
    print colored(len(negData), 'green')
    x = posData + negData
    x_= []
    foo = []
    for i in range(len(x)):
        foo = np.zeros(maxlen - len(x[i][0])).tolist() + x[i][0].tolist() + np.zeros(maxlen - len(x[i][1])).tolist() + x[i][1].tolist() + np.zeros(maxlen - len(x[i][2])).tolist() + x[i][2].tolist()
        x_.append(np.array(foo))
    y = np.zeros(len(x))
    for i in range(len(posData)):
        y[i] = 1
    return np.array(x_), y


def getMax(posData, negData, testCON, testFWater):
    x = posData + negData + testCON + testFWater
    maxlen = 0
    for i in x:
        maxx = len(max(i, key = len))
        if maxx > maxlen:
            maxlen = maxx
    return maxlen




def createXYHist(posData, negData, bins = 20):
    x = posData + negData
    x_ = []
    y = np.zeros(len(x))
    for i in range(len(posData)):
        y[i] = 1
    for i in range(len(x)):
        histC = np.histogram(x[i][0], bins= bins)
        histO = np.histogram(x[i][1], bins= bins)
        histN = np.histogram(x[i][2], bins= bins)
        foo = np.concatenate((histC[0], histO[0], histN[0]))
        x_.append(foo)
    return np.array(x_), y




def printXY(x, y, protsUsed, testProts):
    f = open('log.out', 'w')
    print x.shape
    print y.shape
    f.write('x shape ' + str(x.shape[0]) + ' ' + str(x.shape[1]) + '\n')
    f.write('y shape ' + str(y.shape[0]) + '\n')
    f.write('prots used \n')

    for i in protsUsed:
        f.write(i + ' ')
    f.write('\ntest prots \n')
    for i in testProts:
        f.write(i+ ' ')

    f.close()
    np.savetxt('testX.out', x)
    np.savetxt('testY.out', y)


def train(posData, negData):
    global testProts
    global testCON #test posData
    global testFWater # test negData
    global protsUsed
    maxlen = getMax(posData, negData, testCON, testFWater)
    [x0, y0] = createXY(posData, negData, maxlen)
    [xt0, yt0] = createXY(testCON, testFWater, maxlen)
    [x1, y1] = createXYHist(posData, negData)
    [xt1, yt1] = createXYHist(testCON, testFWater)
    printXY(x0, y0, protsUsed, testProts)
    """#xgboost
    model0 = xgb.XGBClassifier()
    model0.fit(x0, y0)

    yp0 = model0.predict(xt0)
    predictions0 = [round(value) for value in yp0]
    # evaluate predictions
    accuracy0 = acs(yt0, predictions0)
    print colored('accuracy of prediction xgb ' + str(accuracy0), 'red')

    #xgboost hist
    model1 = xgb.XGBClassifier()
    model1.fit(x1, y1)

    yp1 = model1.predict(xt1)
    predictions1 = [round(value) for value in yp1]
    # evaluate predictions
    accuracy1 = acs(yt1, predictions1)
    print colored('accuracy of prediction xgb hist ' + str(accuracy1), 'red')

    #svc
    model2 = svc(gamma= 'auto')
    model2.fit(x0, y0)
    yp2 = model2.predict(xt0)
    predictions2 = [round(value) for value in yp2]
    accuracy2 = acs(yt0, predictions2)
    print colored('accuracy of pred svm svc ' + str(accuracy2), 'red')

    #svc with hist
    model3 = svc(gamma= 'auto')
    model3.fit(x1, y1)
    yp3 = model3.predict(xt1)
    predictions3 = [round(value) for value in yp3]
    accuracy3 = acs(yt1, predictions3)
    print colored('accuracy of pred svm svc with hist ' + str(accuracy3), 'red')

    #random forest
    model4 = rfc()
    model4.fit(x0, y0)
    yp4 = model4.predict(xt0)
    predictions4 = [round(value) for value in yp4]
    accuracy4 = acs(yt0, predictions4)
    print colored('accuracy of pred rfc ' + str(accuracy4), 'red')

    #random forest hist
    model5 = rfc()
    model5.fit(x1, y1)
    yp5 = model5.predict(xt1)
    predictions5 = [round(value) for value in yp5]
    accuracy5 = acs(yt1, predictions5)
    print colored('accuracy of pred rfc with hist ' + str(accuracy5), 'red')

    print 'train complete'"""



#x vectors : c , o, n
#y vector : fake water
timer = time.time()

protsUsed = []
testProts = []
testCON = []
testFWater = []

[posData, negData] = prepareData()

train(posData, negData)

#print colored(testProts, 'yellow')
#print "These", len(protsUsed), "proteins are used for preparing of data"
#print colored(protsUsed, 'cyan')


print "My program took", time.time() - timer, "to run and memory is used", psutil.virtual_memory()
#fwater = getFakeWater(protein, water)