# YASARA BioTools
# Visit www.yasara.org for more...
# Copyright by Elmar Krieger

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA


import os,string,shutil,fnmatch,zipfile

#  ======================================================================
#                D I S C   F U N C T I O N   G R O U P  (D S C)
#  ======================================================================

# CALCULATE FILE CHECKSUM
# =======================
def checksum(filename):
  checksum=0
  file=open(filename,"rb")
  pos=1
  a=0
  b=0  
  while (1):
    data=file.read(1000000)
    datalen=len(data)
    i=0
    while (i<datalen):
      ch=data[i]
      a=(a+ch)&0xffff
      b=(b+pos*ch)&0xffff
      i+=1
      pos+=1
    if (datalen<1000000): break
  return(a|(b<<16))

# COPY FILE
# =========
def copy(srcfilename,dstfilename,timespreserved=1,mod=None):
  if (timespreserved):
    # copy2 CAN FAIL ON FAT PARTITIONS
    try: shutil.copy2(srcfilename,dstfilename)
    except: timespreserved=0
  if (not timespreserved):
    try: shutil.copy(srcfilename,dstfilename)
    except: shutil.copyfile(srcfilename,dstfilename)
  if (mod!=None): chmod(dstfilename,mod)

# COPY DIRECTORY
# ==============
# srcmatchname IS A WILDCARD MATCHING THE FILES TO COPY
# dstdirname IS THE NAME OF THE DESTINATION DIRECTORY
def copydir(srcmatchname,dstdirname,mod=None):
  makedirs(dstdirname)
  for srcfilename in dirlist(srcmatchname):
    dstfilename=os.path.join(dstdirname,os.path.basename(srcfilename))
    copy(srcfilename,dstfilename,1,mod)

# COPY ENTIRE DIRECTORY RECURSIVELY
# =================================
def copydirrec(srcdirname,dstdirname):
  rmdir(dstdirname)
  shutil.copytree(srcdirname,dstdirname,symlinks=True)

# RSYNC ENTIRE DIRECTORY TO REMOTE SERVER
# =======================================
# srcbasenamelist CAN BE EMPTY TO TAKE ALL FILES IN srcdirname
# FILES ON dstserver THAT ARE NEWER ARE SKIPPED (-u OPTION).
# THE IMPORTANT FEATURE IS THAT FILES THAT ARE TRANSFERRED GET THE LATEST MODIFICATION TIME.
# THIS ENSURES THAT THEY ARE REALLY DETECTED AS NEW BY E.G. THE SQUEEZER FOR INCREMENTAL UPDATES.
def rsyncto(srcdirname,srcbasenamelist,dstserver,dstdirname,deleted=1,recursive=0,excludelist=[]):
  if (recursive):
    if (srcbasenamelist!=None): error("Recursive rsync with file list is not supported")
    # RSYNC ENTIRE FOLDER RECURSIVELY
    command="rsync -v -u -p -r --delete --delete-before --exclude='.*' --exclude='.*/' "
    for exclude in excludelist: command+=" --exclude='%s'"%exclude
    command+=" '%s' '%s:%s'"%(srcdirname,dstserver,os.path.dirname(dstdirname))
    #print command
    if (os.system(command)): error("Failed to run rsync for entire folder %s"%srcdirname)
  else:
    if (srcbasenamelist==None):
      # GET A LIST OF ALL FILES IF NONE ARE PROVIDED
      srcbasenamelist=dirlist(os.path.join(srcdirname,"*"),inctype=recursive)
      for i in numbered(srcbasenamelist): srcbasenamelist[i]=os.path.basename(srcbasenamelist[i])
    """
    # SAVE THE LIST FOR RSYNC
    listfilename=tmpfilename("rsync_list")
    open(listfilename,"w").write('\n'.join(srcbasenamelist))
    # BUILD COMMAND
    command="rsync -v -u -p -r --files-from='%s'"%listfilename
    if (deleted): command+=" --delete --delete-after"
    command+=" '%s' '%s:%s'"%(srcdirname,dstserver,dstdirname)
    """
    # SAVE THE LIST FOR RSYNC
    listfilename=tmpfilename("rsync_list")
    listfile=open(listfilename,"w")
    for basename in srcbasenamelist: listfile.write("/%s\n"%basename)
    listfile.close()
    # BUILD COMMAND
    # https://stackoverflow.com/questions/1813907/rsync-delete-files-from-list-dest-does-not-delete-unwanted-files
    command="rsync -v -u -p -r --include-from='%s' --exclude=* "%listfilename
    if (deleted): command+=" --delete-excluded"
    command+=" '%s/' '%s:%s/'"%(srcdirname,dstserver,dstdirname)
    if (os.system(command)): error("Failed to run rsync for selected files")
    remove(listfilename)

# GET FREE DISC SPACE
# ===================
def free():
  stat=os.statvfs(os.getcwd())
  return(stat.f_bsize*stat.f_bavail)

# GET THE LAST N ELEMENTS IN A PATH
# =================================
def lastpathelements(path,elements):
  result=None
  for i in range(elements):
    if (result==None): result=os.path.basename(path)
    else: result=os.path.join(os.path.basename(path),result)
    path=os.path.dirname(path)
  return(result)

# REPLACE FILE
# ============
# dstfilename IS REPLACED WITH srcfilename, PERMISSIONS ARE SET TO mod.
def replace(srcfilename,dstfilename,mod=None):
  remove(dstfilename)
  # os.rename can't move files to a diffrent partition
  shutil.move(srcfilename,dstfilename)
  if (mod!=None): chmod(dstfilename,mod)

# GET SIZE OF FILE
# ================
def filesize(filename):
  if (not os.path.exists(filename)): return(0)
  return(os.stat(filename)[6])

# INCREASE NUMBER IN FILE NAME
# ============================
def incfilename(filename):
  i=len(filename)-1
  while (i>0 and filename[i] not in string.digits): i=i-1
  while (i>=0):
    num=ord(filename[i])
    if (num<48 or num>57): break
    num=num+1
    if (num==58): num=48
    filename=filename[:i]+chr(num)+filename[i+1:]
    if (num!=48): break
    i=i-1
  return(filename)

# MAKE DIRECTORY CHAIN
# ====================
def makedirs(path,permissions=0x1ed):
  if (not os.path.exists(path)): os.makedirs(path,permissions)
  # THE PERMISSIONS ARE NOT ALWAYS SET CORRECTLY, DO IT AGAIN
  chmod(path,permissions)

# GET MD5SUM
# ==========
def md5sum(filename):
  content=open(filename,"rb").read()
  md5sum=hashlib.md5(content)
  return(md5sum.hexdigest())

# GET THE MODIFICATION TIME OF A FILE
# ===================================
def modtime(filename):
  if (not os.path.exists(filename)): return(0)
  else: return(os.path.getmtime(filename))

# GET ALL MODIFICATION TIMES FOR A LIST OF FILES
# ==============================================
def modtimelist(filenamelist):
  timelist=[]
  for filename in filenamelist: timelist.append(modtime(filename))
  return(timelist)

# CHECK IF TWO FILES ARE THE SAME
# ===============================
def havesamecontent(filename1,filename2):
  size1=filesize(filename1)
  size2=filesize(filename2)
  if (size1==0 or size1!=size2): return(0)
  return(open(filename1,"rb").read()==open(filename2,"rb").read())

# BUILD TEMPORARY FILE NAME
# =========================
def tmpfilename(filename):
  dotpos=filename.rfind(".")
  slashpos=filename.rfind(os.sep)
  if (dotpos==-1 or dotpos<slashpos): dotpos=len(filename)
  filename=filename[:dotpos]+"_tmp%d"%os.getpid()+filename[dotpos:]
  return(filename)

# CHECK IF FILE(S) EXIST(S)
# =========================
# matchname IS A WILDCARD MATCHING THE FILENAME
def pathexists(matchname):
  # CREATE LIST OF ALL THE FILENAMES IN THE DIRECTORY SPECIFIED BY PATH
  (dirname,wildcard)=os.path.split(matchname)
  if (dirname==""): dirname="."
  try: basenamelist=os.listdir(dirname)
  except: return(0)
  # DELETE ALL LIST ENTRIES THAT DO NOT MATCH THE WILDCARD GIVEN IN PATH
  for basename in basenamelist:
    if (fnmatch.fnmatch(basename,wildcard)): return(1)
  return(0)

# CREATE DIRECTORY LISTING INCLUDING FULL PATH NAMES
# ==================================================
# matchname IS A WILDCARD MATCHING THE FILES TO LIST
# inctype==0: ONLY FILES WILL BE RETURNED
# inctype==1: ALL ENTRIES WILL BE RETURNED, INCLUDING DIRECTORIES
# inctype==-1: ONLY DIRECTORIES WILL BE RETURNED
def dirlist(matchname,inctype=0):
  # CREATE LIST OF ALL THE FILENAMES IN THE DIRECTORY SPECIFIED BY PATH
  (dirname,wildcard)=os.path.split(matchname)
  if (dirname==""): dirname="."
  elif (dirname.find("?")!=-1):
    # DIRECTORY NAME STILL CONTAINS A WILDCARD, E.G. /mnt/structure_data/raw/pdb/??/pdb????.ent.gz
    pathlist=[]
    for dirname in dirlist(dirname,-1): pathlist+=dirlist(os.path.join(dirname,wildcard),inctype)
    return(pathlist)
  try: pathlist=os.listdir(dirname)
  except: return([])
  # CYCLE THROUGH ALL FILES AND CHECK IF THEY MATCH WILDCARD
  i=0
  while (i<len(pathlist)):
    path=pathlist[i]
    if (path not in [".",".."] and fnmatch.fnmatch(path,wildcard)):
      pathlist[i]=os.path.join(dirname,path)
      if (inctype==1 or 
          (inctype==0 and os.path.isfile(pathlist[i])) or
          (inctype==-1 and os.path.isdir(pathlist[i]))):
        i+=1
        continue
    # NO MATCH
    del pathlist[i]
  pathlist.sort()
  return(pathlist)

# CREATE RECURSIVE DIRECTORY LISTING INCLUDING FULL PATH NAMES
# ============================================================
# matchname IS A WILDCARD MATCHING THE FILES TO LIST
def recursivedirlist(matchname):
  pathlist=[]
  # CREATE LIST OF ALL THE FILENAMES IN THE DIRECTORY SPECIFIED BY PATH
  (dirname,wildcard)=os.path.split(matchname)
  if (dirname==""): dirname="."
  try: basenamelist=os.listdir(dirname)
  except: return([])
  # CYCLE THROUGH ALL FILES AND CHECK IF THEY MATCH WILDCARD
  for basename in basenamelist:
    path=os.path.join(dirname,basename)
    if (os.path.isfile(path)):
      if (fnmatch.fnmatch(basename,wildcard)): pathlist.append(path)
    elif (basename not in [".",".."] and not os.path.islink(path)):
      # RECURSE DEEPER
      subdirlist=recursivedirlist(os.path.join(path,wildcard))
      if (subdirlist==[]): pathlist.append(path)
      else: pathlist+=subdirlist
  pathlist.sort()
  return(pathlist)

# DELETE A FILE OR LIST OF FILES
# ==============================
def remove(filename):
  if (type(filename)==type([])):
    filenamelist=filename
    for filename in filenamelist: remove(filename)
  elif (filename!=None and os.path.exists(filename)): os.remove(filename)

# RENAME A FILE
# =============
def rename(srcfilename,dstfilename):
  os.rename(srcfilename,dstfilename)

# DELETE ALL MATCHING FILES
# =========================
# matchname IS A WILDCARD MATCHING THE FILES TO DELETE
def removematch(matchname):
  if (type(matchname)==type([])):
    matchnamelist=matchname
    for matchname in matchnamelist: removematch(matchname)
  else:
    # CREATE LIST OF ALL THE FILENAMES IN THE DIRECTORY SPECIFIED
    (dirname,wildcard)=os.path.split(matchname)
    if (dirname==""): dirname="."
    try: basenamelist=os.listdir(dirname)
    except: return(0)
    # DELETE ALL LIST ENTRIES THAT DO NOT MATCH THE WILDCARD GIVEN IN PATH
    for basename in basenamelist:
      if (fnmatch.fnmatch(basename,wildcard)): remove(os.path.join(dirname,basename))
  return(0)

# SET MODIFICATION TIME
# =====================
def setmodtime(filename,modtime):
  os.utime(filename,(modtime,modtime))
  
# COPY MODIFICATION TIME
# ======================
def copymodtime(srcfilename,dstfilename):
  setmodtime(dstfilename,modtime(srcfilename))

# CHMOD A FILE
# ============
def chmod(path,mods):
  if (type(path)==type([])):
    pathlist=path
    for path in pathlist: chmod(path,mods)
  else:
    try: os.chmod(path,mods)
    except: print("Could not change permissions for ",path)

# DELETE AN ENTIRE DIRECTORY INCLUDING ALL THE FILES
# ==================================================
def rmdir(path):
  if (os.path.exists(path)): shutil.rmtree(path,1)

# REMOVE FILE EXTENSION
# =====================
def rmext(filename):
  dotpos=filename.rfind(".")
  if (dotpos!=-1): filename=filename[:dotpos]
  return(filename)

# GET FILE EXTENSION
# ==================
def ext(filename):
  dotpos=filename.rfind(".")
  if (dotpos!=-1): return(filename[dotpos+1:])
  return("")

# COUNT NUMBER OF FILES PRESENT
# =============================
def countpresent(filenamelist):
  files=0
  for filename in filenamelist:
    if (os.path.exists(filename)): files+=1
  return(files)

# UPDATE DIRECTORY
# ================
# - srcmatchname IS A WILDCARD MATCHING THE FILES TO UPDATE
# - dstdirname IS THE NAME OF THE DESTINATION DIRECTORY
# IF obsoleted, OBSOLETE FILES WILL BE DELETED.
# - excludelist CONTAINS EXCLUDED FILENAMES OR, .* TO EXCLUDE HIDDEN FILES
def updatedir(srcmatchname,dstdirname,contentchecked=1,recursive=0,obsoleted=0,mod=None,excludelist=[]):
  modfilelist=[]
  if (os.path.isdir(srcmatchname)): srcmatchname=os.path.join(srcmatchname,"*")
  srcbasenamelist=[]
  for srcfilename in dirlist(srcmatchname,recursive):
    srcbasename=os.path.basename(srcfilename)
    srcbasenamelist.append(srcbasename)
    if (srcbasename not in excludelist and (srcbasename[0]!='.' or ".*" not in excludelist)):
      dstfilename=os.path.join(dstdirname,srcbasename)
      if (os.path.isdir(srcfilename)):
        if (recursive):
          if (not os.path.exists(dstfilename)): makedirs(dstfilename)
          updatedir(srcfilename,dstfilename,contentchecked,1,obsoleted,mod,excludelist)
      else:
        if ((contentchecked and not havesamecontent(srcfilename,dstfilename)) or
            (not contentchecked and modtime(srcfilename)>modtime(dstfilename))):
          copy(srcfilename,dstfilename,0,mod)
          modfilelist.append(dstfilename) 
  if (obsoleted):
    # ALSO DELETE OBSOLETE FILES
    for dstfilename in dirlist(os.path.join(dstdirname,os.path.basename(srcmatchname))):
      dstbasename=os.path.basename(dstfilename)
      if (dstbasename not in srcbasenamelist): remove(dstfilename)
  return(modfilelist)

# ZIP A FILE
# ==========
def zip(filename):
  print("Deflating",filename)
  zipfilename=filename+".zip"
  zip=zipfile.ZipFile(zipfilename,"w",zipfile.ZIP_DEFLATED)
  zip.write(filename,os.path.basename(filename))
  zip.close()
  remove(filename)

# RETURN AN UNZIPPED FILE
# =======================
def unzipped(filename):
  zipfilename=filename+".zip"
  zip=zipfile.ZipFile(zipfilename,"r",zipfile.ZIP_DEFLATED)
  data=zip.read(os.path.basename(filename))
  zip.close()
  return(data)

# RETRUN ps -ef OUTPUT
# ====================
def psoutput():
  tmpfilename=tmpfilename("psoutput")
  os.system("ps -ef > "+tmpfilename)
  ps=open(tmpfilename).read()
  remove(tmpfilename)
  return(ps)