#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys, os, re

if __name__ == "__main__":
  # Make sure we can import stuff from this file's directory
  sys.path.append(os.path.join(os.path.abspath(os.path.dirname(sys.argv[0])), "..", "lib-python"))
  sys.path.append(os.path.abspath(os.path.dirname(sys.argv[0])))


from prosci.common import basename_noext
from prosci.util.tmalign import superimpose
from prosci.util.pdb import Pdb


if __name__ == "__main__":
  
  from prosci.shell import Params
  
  params = Params(withargument=("suffix", "options", "out1", "out2", "minscore", "fasta"), allowed=("suffix", "options", "out1", "out2", "minscore", "fasta", "no1", "no2"))

  if not len(params.args) == 2:
    print "USAGE:"
    print "    %s [OPTIONS] structure.pdb target.pdb" % (params.scriptname)
    print
    print "The ATOMs of structure.pdb will be superimposed onto target.pdb,"
    print "without modifying the latter's coordinates."
    print "By default, two files will be created, containing the new ATOM coordinates, after"
    print "the superimposition (by default with the file extension '.sup')."
    print "Creation of the second PDB file can be suppressed with the option '--no2'."
    print
    print "OPTIONS:"
    
    print "    --minscore NUM     Only write output files if TM-score >= NUM."
    print "    --suffix '.xxx'    Suffix for structure output filenames (default: '.sup')."
    print "    --out1 FILENAME    Output filename of the transformed structure.pdb. '--suffix' will be ignored."
    print "    --out2 FILENAME    Output filename of the transformed target.pdb. '--suffix' will be ignored."
    print "    --options XXXX     Pass XXXX as options to TMalign. For available options, type 'TMalign'."
    print "    --no1              Do not write structure 1 (filtered & transformed structure.pdb)."
    print "    --no2              Do not write structure 2 (filtered target.pdb)."
    print "    --fasta FILENAME   Write sequence alignment into FASTA file."
    print
    sys.exit(1);
 
  # filenames of PDB files
  pdb1_filename = params.args[0]
  pdb2_filename = params.args[1]
  
  if not "options" in params.opts:
    params.opts["options"] = ""
  if not "suffix" in params.opts:
    params.opts["suffix"] = ".sup"
  if not "out1" in params.opts:
    params.opts["out1"] = pdb1_filename + params.opts["suffix"]
  if not "out2" in params.opts:
    params.opts["out2"] = pdb2_filename + params.opts["suffix"]
  if not "minscore" in params.opts:
    params.opts["minscore"] = -1.0
  else:
    params.opts["minscore"] = float(params.opts["minscore"])
  
  
  struc1_allchains = Pdb(pdb1_filename)
  struc2_allchains = Pdb(pdb2_filename)
  
  seq1, seq2, alignment_info = superimpose(struc1_allchains, struc2_allchains, \
#                                                            fname1=pdb1_filename, \
#                                                            fname2=pdb2_filename, \
#                                                            align_atoms=["N", "CA", "C", "O"], \
                                                            options=params.opts["options"], \
                                                            modify_structures=("no1" not in params.opts))
  
  for k in sorted(alignment_info):
    print "%s=%s" % (k, str(alignment_info[k]))
  
  tmscore = alignment_info["TM-score"]
  
  if tmscore >= params.opts["minscore"]:
    if "no1" not in params.opts:
      f = open(params.opts["out1"], "w")
      f.write(str(struc1_allchains))
      f.close()
    if "no2" not in params.opts:
      f = open(params.opts["out2"], "w")
      f.write(str(struc2_allchains))
      f.close()
  else:
    sys.stderr.write("NOTICE: TM-score < cut-off (%.5f). Not writing output file(s).\n" % (params.opts["minscore"]))
  
  if params.isOpt("fasta"):
    f=open(params.getOpt("fasta"), "w")
    f.write(">%s\n%s\n>%s\n%s\n" % (struc1_allchains.code, seq1, struc2_allchains.code, seq2))
    f.close()
  
