#!/usr/bin/python
# -- coding: utf-8 --
#source activate my-rdkit-env
#source deactivate
from __future__ import print_function
from sys import argv
from rdkit import Chem
import numpy as np
import re
from rdkit.Chem import rdMolDescriptors
from os import popen as po
def logP(listfile):
    logP={}
    f=open(listfile)
    l=f.read().split("\n")
    f.close()
    for line in l:
      if line=="":
        continue
      aa=line.split(",")
      logP[aa[0]]={}
      for a in aa[1:]:
        atom,LOGP=a.split(":")
        logP[aa[0]][atom.rstrip().upper()]=float(LOGP)
    return logP


def read_mol(file):
    ro=Chem.MolFromMolFile(file)
    n=ro.GetNumAtoms()
    element=[i.GetSymbol() for i in ro.GetAtoms()]
    crippen = rdMolDescriptors._CalcCrippenContribs(ro)
    atom=[]
    for i in range(n):
        pos=(ro.GetConformer().GetAtomPosition(i))
        atom.append([element[i],pos.x,pos.y,pos.z,crippen[i][0]])
    with open(file,"r")as f:
        s=f.read().split("\n")
        s.pop(-1)
    return [atom,ro,s]

def read_HETATM(file):
    atom=[]
    line_num=[]
    with open(file,"r")as f:
        s=f.read().split("\n")
        s.pop(-1)
    n=0
    for i in s:
        if i[:6]=='HETATM':
            a=i[75:].strip().upper()
            atom.append(
                [re.split('[-+0-9]',a)[0],
                float(i[30:38]),
                float(i[38:46]),
                float(i[46:54])]
                )
            line_num.append(n)
        elif i[:4]=='ATOM':
            atom.append(
                [i[13],
                float(i[30:38]),
                float(i[38:46]),
                float(i[46:54])]
                )
            line_num.append(n)
        n+=1
    return [atom,Chem.MolFromPDBFile(file),s,line_num]

def read_ATOM(file):
    with open(file,"r")as f:
        s=f.read().split("\n")
        s.pop(-1)
    atom={}
    residue=[]
    for i in s:
        if i[:4]=='ATOM':
            res=i[17:26]
            if not res in atom:
                atom[res]=[]
                residue.append(res)
            atom[res].append(
                [i[13],
                i[13:16].rstrip().upper(),
                float(i[30:38]),
                float(i[38:46]),
                float(i[46:54])]
                )
    return [residue,atom,s]

def get_pos_mol(mol):
    n=mol.GetNumAtoms()
    mol_pos=[]
    for i in range(n):
        pos=(mol.GetConformer().GetAtomPosition(i))
        mol_pos.append([pos.x,pos.y,pos.z])
    # print("mol_pos=")
    # for i in mol_pos:
    #     prin/t(i)
    mol_pos=np.array(mol_pos)
    return mol_pos

def reduce_pdb(pdb,mol,dist):
    residue=[]
    for res in pdb[0]:
        flag=False
        for c in pdb[1][res]:
            for j in range(len(mol)):
                for k in range(len(mol[j][0])):
                    if (c[2]-mol[j][0][k][1])**2+(c[3]-mol[j][0][k][2])**2+(c[4]-mol[j][0][k][3])**2<=dist**2:
                        residue.append(res)
                        flag=True
                        break
                if flag:
                    break
            if flag:
                break
    pdb[0]=residue
    atom={}
    for i in residue:
        atom[i]=pdb[1][i]
    pdb[1]=atom
    s=[]
    for i in pdb[2]:
        if not i[:4]=="ATOM":
            s.append(i)
        elif i[21:26] in residue:
            s.append(i)
    pdb[2]=s

def wait_memory(rt,ti):
  free=list(po("free"))
  s=free[1].split()
  total=int(s[1])
  used=int(s[2])
  p=used*100/total
  while p>=rt:
    print("memory used=",used,"total=",total,"p=",'{:.1f}'.format(p))
    t=os.system("sleep "+str(ti))
    free=list(po("free"))
    s=free[1].split()
    used=int(s[2])
    p=used*100/total