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

import sys
import getopt
import platform
import re

from vector3d import *
from pdb import *
from system import *
from sysmgr import *

# Order of options is critical in this program.
# Thus, I cannot use optparse/argparse, since the order is ignored in them.

def usage():
  print "usage: builder.py [options] > output.xml"

def showHelp( opts ):
  if len( opts ) != 0 and opts[0].lower() in ( "select", "selection" ):
    print >> sys.stderr, ">> SELECTION"
    print >> sys.stderr, ">>"
    print >> sys.stderr, ">> There are three types of way to select atoms."
    print >> sys.stderr, ">>"
    print >> sys.stderr, ">> 1. selection by a single number"
    print >> sys.stderr, ">>"
    print >> sys.stderr, ">>   You can select an/a atom/residue/molecule by a single number."
    print >> sys.stderr, ">>   For example,"
    print >> sys.stderr, ">>     atom 4   - select atom(s) having index(2nd column of pdb) of 4"
    print >> sys.stderr, ">>     resid 10 - select atoms in 10th residue"
    print >> sys.stderr, ">>     mol 1    - select atoms in second molecule"
    print >> sys.stderr, ">>                (molecules are separated by \"TER\" in pdb,"
    print >> sys.stderr, ">>                 and the first molecule is 0th molecule, not 1st)"
    print >> sys.stderr, ">> Regular expressions such as 10? are NOT ALLOWED in this method."
    print >> sys.stderr, ">> Please use range selection 100-109 instead if you need something like that."
    print >> sys.stderr, ">>"
    print >> sys.stderr, ">> 2. selection by range; with two numbers"
    print >> sys.stderr, ">>"
    print >> sys.stderr, ">>    You can select atoms within ranges."
    print >> sys.stderr, ">>    A valid range is described by two integer numbers separated by \"-\","
    print >> sys.stderr, ">>    where right value may be larger than left one."
    print >> sys.stderr, ">>    For example, 1-10 is a valid range."
    print >> sys.stderr, ">>    A range 10-1 is also valid, but no atoms will be selected."
    print >> sys.stderr, ">>"
    print >> sys.stderr, ">>    Atom index (2nd column in pdb) is used for atom selection."
    print >> sys.stderr, ">>    Residue index (5th column in pdb) is used for residue selection."
    print >> sys.stderr, ">>    Molecular index (order in pdb file) is used for molecular selection."
    print >> sys.stderr, ">>"
    print >> sys.stderr, ">> 3. selection by a string"
    print >> sys.stderr, ">>"
    print >> sys.stderr, ">>    You can select atoms by their names, of course."
    print >> sys.stderr, ">>    But this can be used only for atoms and residues, not molecules."
    print >> sys.stderr, ">>    Regular expressions such as GL?, AS* are allowed in this method."
    print >> sys.stderr, ">>    String in 3rd column of pdb is used for atom names,"
    print >> sys.stderr, ">>    that in 4th column of pdb is used for residue names,"
    print >> sys.stderr, ">>"
    print >> sys.stderr, ">> All the selections above are dealed in the manner of \"or\"."
    print >> sys.stderr, ">> If you did the selection like this:"
    print >> sys.stderr, ">>   resid MET"
    print >> sys.stderr, ">>   resid 10-20"
    print >> sys.stderr, ">>   select"
    print >> sys.stderr, ">> All the atoms in MET residue \"AND\" atoms in residue 10-20 are selected,"
    print >> sys.stderr, ">> not the MET residue(s) in 10-20."
    print >> sys.stderr, ">> If you want MET residue in 10-20, try the following:"
    print >> sys.stderr, ">>   resid MET"
    print >> sys.stderr, ">>   select"
    print >> sys.stderr, ">>   resid 10-20"
    print >> sys.stderr, ">>   select"
    print >> sys.stderr, ">>"
  else:
    print >> sys.stderr, ">> LIST of COMMANDs"
    print >> sys.stderr, ">>"
    print >> sys.stderr, ">>   help          : show this message (abbreviations: h )"
    print >> sys.stderr, ">>   load [FILE]   : load pdb file ( l, p, pdb )"
    print >> sys.stderr, ">>   clear         : clear cuurent selection ( c )"
    print >> sys.stderr, ">>   allclear      : clear pdb data and selection ( ac )"
    print >> sys.stderr, ">>   atom [str]    : register atom query string: \"str\""
    print >> sys.stderr, ">>                 ( a, selectatom )"
    print >> sys.stderr, ">>   resid [str]   : register residue query string: \"str\""
    print >> sys.stderr, ">>                 ( r, res, residue, selectres, selectresidue )"
    print >> sys.stderr, ">>   mol [str]     : register molecules query string \"str\""
    print >> sys.stderr, ">>                 ( m, molecule, selectmol, selectmolecule )"
    print >> sys.stderr, ">>   select        : select atoms which match registered queries ( s, sel )"
    print >> sys.stderr, ">>   genatoms      : generate xml data of ATOMs ( ga, genatom )"
    print >> sys.stderr, ">>   genbonds      : generate xml data of BONDs ( gb, genbond )"
    print >> sys.stderr, ">>   genchains     : generate xml data of CHAINs ( gc, genchain )"
    print >> sys.stderr, ">>   gencoils      : generate xml data of CHAINs ( gco, gencoil )"
    print >> sys.stderr, ">>   write         : write selected atom information to stdout in pdb format"
    print >> sys.stderr, ">>   showatoms     : show registered atoms (sa)"
    print >> sys.stderr, ">>   clearatoms    : clear all registered atoms (ca)"
    print >> sys.stderr, ">>   ratom [str]   : register an atom (ra,registeratom)"
    print >> sys.stderr, ">>   hcol [str]    : modify helix color (hc,helixcolor)"
    print >> sys.stderr, ">>   scol [str]    : modify strand color (sc,scolor,strandcolor)"
    print >> sys.stderr, ">>   ccol [str]    : modify coil color (cc,ccolor,coilcolor)"
    print >> sys.stderr, ">>   hrad [float]  : modify helix radius (hr,helixradius)"
    print >> sys.stderr, ">>   srad [float]  : modify strand radius (sr,strandradius)"
    print >> sys.stderr, ">>   crad [float]  : modify coil radius (cr,coilradius)"
    print >> sys.stderr, ">>   showrad       : show radius of helix, strand, coil (showradius)"
    print >> sys.stderr, ">>   showcol       : show color of helix, strand, coil (showcolor)"
    print >> sys.stderr, ">>   wmxml         : add <WMXML> and </WMXML> in XML output"
    print >> sys.stderr, ">>   suffix [char] : add [char] to element name"
    print >> sys.stderr, ">>"
    print >> sys.stderr, ">>   exit, quit    : exit program. same as Ctrl+D"
    return

def interactive():
  print >> sys.stderr, ">> entering interactive mode."
  print >> sys.stderr, ">> type \"exit\" or press Ctrl+D to exit."
  print >> sys.stderr, ">> type \"help\" to show list of commands."

  s = SystemMgr()

  #print "<WMXML>"
  for line in iter( sys.stdin.readline, "" ):
    line = line.strip()
    if len( line ) == 0:
      continue
    sline = line.split()
    # number of argument in a command is 0 or 1
    arg = " ".join( sline[1:] )
    if sline[0].lower() in ( "h", "help" ):
      showHelp( sline[1:] )
    elif sline[0].lower() in ( "l", "p", "pdb", "load" ):
      try:
        fh = open( arg, "r" )
      except IOError, ( errno, msg ):
        if len( arg ) == 0:
          print >> sys.stderr, ">> error: filename missing."
        else:
          print >> sys.stderr, ">> error: cannot open %s for reading" % arg
        print >> sys.stderr, ">> error: errno: [%d] msg: [%s]" % ( errno, msg )
        continue
      s.readPDB( fh )
      print >> sys.stderr, ">> read %s atoms from the pdb file" % s.totatom()
    elif sline[0].lower() in ( "c", "clear" ):
      s.clear()
      print >> sys.stderr, ">> clear current selection."
    elif sline[0].lower() in ( "ac", "allclear" ):
      s.allclear()
      print >> sys.stderr, ">> clear all the data."
    elif sline[0].lower() in ( "a", "atom", "selectatom" ):
      s.addAtomQuery( arg )
      print >> sys.stderr, ">> add atom query to the selection."
    elif sline[0].lower() in ( "r", "res", "resid", "residue", "selectres", "selectresid", "selectresidue" ):
      s.addResQuery( arg )
      print >> sys.stderr, ">> add residue query to the selection."
    elif sline[0].lower() in ( "m", "mol", "molecule", "selectmol", "selectmolecule" ):
      s.addMolQuery( arg )
      print >> sys.stderr, ">> add molecule query to the selection."
    elif sline[0].lower() in ( "s", "sel", "select" ):
      s.select()
      print >> sys.stderr, ">> select %s atoms" % s.natom()
    elif sline[0].lower() in ( "ga", "genatom", "genatoms" ):
      nelem = s.genAtoms( arg )
      if nelem > 0:
        print >> sys.stderr, ">> finished to generate atoms. (" + str(nelem) + " XML elements)"
      else:
        print >> sys.stderr, ">> no XML elements added (did you select?)"
    elif sline[0].lower() in ( "gb", "genbond", "genbonds" ):
      nelem = s.genBonds( arg )
      if nelem > 0:
        print >> sys.stderr, ">> finished to generate bonds. (" + str(nelem) + " XML elements)"
      else:
        print >> sys.stderr, ">> no XML elements added (did you select?)"
    elif sline[0].lower() in ( "gc", "genchain", "genchains" ):
      nelem = s.genChains()
      if nelem > 0:
        print >> sys.stderr, ">> finished to generate chains. (" + str(nelem) + " XML elements)"
      else:
        print >> sys.stderr, ">> no XML elements added (did you select?)"
    elif sline[0].lower() in ( "gco", "gencoil", "gencoils" ):
      nelem = s.genChains( True )
      if nelem > 0:
        print >> sys.stderr, ">> finished to generate chains. (" + str(nelem) + " XML elements)"
      else:
        print >> sys.stderr, ">> no XML elements added (did you select?)"
    elif sline[0].lower() in ( "q", "exit", "quit" ):
      print >> sys.stderr, ">> bye."
      break
    elif sline[0].lower() in ( "w", "write", "writepdb" ):
      s.write( sys.stdout )
    elif sline[0].lower() in ( "sa", "showatoms" ):
      s.showAtoms()
    elif sline[0].lower() in ( "ca", "clearatoms" ):
      s.clearAtoms()
      print >> sys.stderr, ">> forget all registered atoms."
    elif sline[0].lower() in ( "ra", "ratom", "registeratom" ):
      if s.registerAtom( arg ):
        print >> sys.stderr, ">> atom registered"
      else:
        print >> sys.stderr, ">> failed to register atom"
    elif sline[0].lower() in ( "ac", "acol", "acolor", "allcolor" ):
      System.helixcolor = arg
      System.strandcolor = arg
      System.coilcolor = arg
      print >> sys.stderr, ">> color of helix, strand, coil has changed"
    elif sline[0].lower() in ( "hc", "hcol", "hcolor", "helixcolor" ):
      System.helixcolor = arg
      print >> sys.stderr, ">> color of helix has changed"
    elif sline[0].lower() in ( "sc", "scol", "scolor", "strandcolor" ):
      System.strandcolor = arg
      print >> sys.stderr, ">> color of strand has changed"
    elif sline[0].lower() in ( "cc", "ccol", "ccolor", "coilcolor" ):
      System.coilcolor = arg
      print >> sys.stderr, ">> color of coil has changed"
    elif sline[0].lower() in ( "hr", "hrad", "helixradius" ):
      System.helixrad = float(arg)
      print >> sys.stderr, ">> radius of helix has changed"
    elif sline[0].lower() in ( "sr", "srad", "strandradius" ):
      System.strandrad = float(arg)
      print >> sys.stderr, ">> radius of strand has changed"
    elif sline[0].lower() in ( "cr", "crad", "coilradius" ):
      System.coilrad = float(arg)
      print >> sys.stderr, ">> radius of coil has changed"
    elif sline[0].lower() in ( "showrad", "showradius" ):
      print >> sys.stderr, ">> helix:  " + str( s.helixrad )
      print >> sys.stderr, ">> strand: " + str( s.strandrad )
      print >> sys.stderr, ">> coil:   " + str( s.coilrad )
    elif sline[0].lower() in ( "showcol", "showcolor" ):
      print >> sys.stderr, ">> helix:  " + str( s.helixcolor )
      print >> sys.stderr, ">> strand: " + str( s.strandcolor )
      print >> sys.stderr, ">> coil:   " + str( s.coilcolor )
    elif sline[0].lower() in ( "wmxml" ):
      s.addwmxml = True
      print >> sys.stderr, ">> WMXML start and end tags will be added."
    elif sline[0].lower() in ( "suffix" ):
      System.setNewSuffix( arg )
      print >> sys.stderr, ">> add suffix " + arg + " to the element name."
    else:
      # invalid command
      print >> sys.stderr, ">> unknown command:", sline[0]
      print >> sys.stderr, ">> type \"help\" to see available command list."
  #print "</WMXML>"
  s.purgeData()

if __name__ == "__main__":
  if len( sys.argv ) <= 1:
    interactive()
    sys.exit(0)

  try:
    opts, args = getopt.getopt( sys.argv[1:],
                                "l:p:cCa:r:m:swxyzWX:Y:Z",
                                [ "pdb=", "load=", "clear", "allclear",
                                  "atom=", "residue=", "molecule=",
                                  "select", "gatom", "gbond", "gribbon",
                                  "gcoil", "write", "clearatom",
                                  "hcolor=", "scolor=", "ccolor=",
                                  "hrad=", "srad=", "crad=",
                                  "regatom=", "wmxml", "suffix=" ] )
  except getopt.GetoptError:
    usage()
    sys.exit(2)

  s = SystemMgr()
  #print "<WMXML>"
  optflag = "-"
  if platform.system() == "Windows":
    optflag = "/"
  regopt = re.compile( "^" + optflag + "+" )
  for o, a in opts:
    myo = regopt.sub( "", o )
    if myo in ( "l", "load", "p", "pdb" ):
      fh = open( a, "r" )
      s.readPDB( fh )
    elif myo in ( "c", "clear" ):
      s.clear()
    elif myo in ( "C", "allclear" ):
      s.allclear()
    elif myo in ( "a", "atom" ):
      s.addAtomQuery( a )
    elif myo in ( "r", "residue" ):
      s.addResQuery( a )
    elif myo in ( "m", "molecule" ):
      s.addMolQuery( a )
    elif myo in ( "s", "select" ):
      s.select()
    elif myo in ( "w", "write" ):
      s.write( sys.stdout )
    elif myo in ( "x", "gatom" ):
      s.genAtoms()
    elif myo in ( "X" ):
      s.genAtoms( a )
    elif myo in ( "y", "gbond" ):
      s.genBonds()
    elif myo in ( "Y" ):
      s.genBonds( a )
    elif myo in ( "z", "gribbon" ):
      s.genChains()
    elif myo in ( "Z", "gcoil" ):
      s.genChains( True )
    elif myo in ( "clearatom" ):
      s.clearAtoms()
    elif myo in ( "hcolor" ):
      System.helixcolor = a
    elif myo in ( "scolor" ):
      System.strandcolor = a
    elif myo in ( "ccolor" ):
      System.coilcolor = a
    elif myo in ( "hrad" ):
      System.helixrad = float(a)
    elif myo in ( "srad" ):
      System.strandrad = float(a)
    elif myo in ( "crad" ):
      System.coilrad = float(a)
    elif myo in ( "regatom" ):
      System.registerAtom( a )
    elif myo in ( "W", "wmxml" ):
      s.addwmxml = True
    elif myo in ( "suffix" ):
      System.setNewSuffix( a )
  #print "</WMXML>"
  s.purgeData()
