#!/usr/bin/python
import os,sys,re
from process import process
import common


class sepsubstrprot(process):
  name="separation of substructure and protein into 2 PDB files"
  short_name="PDB models separation"
  # assumes substr+prot pdb and substructure atom types on input 
  # output is a substructure pdb and a protein pdb (with empty intersection between them)
  # at the moment, just performed in a simplistic way, based on substructure atomtypes
  # this may be not indended eg for S+methionine/cysteine - perhaps (some previous) substr pdb 
  # should be inputted too and distance criterion from previous substr atoms added?

  def RunBody(self,*args,**kwargs):
    if not self.GetProg('pdbcur'):
      self.pdbcur=self.AddProg('pdbcur')
    else:
      self.pdbcur=self.GetProg('pdbcur')
    if not self.inp.Get('model',typ='partial+substr'):
      common.Error('No protein+substructure model inputted - separation cannot be done')
    if not self.inp.Get('model',typ=('substr','partial+substr'),has_atomtypes=True):
      common.Error('Substructure atom types not specified - separation of substructure and protein cannot be done')
    self.pdbcur.inp.Set(self.inp.Get('model',typ='partial+substr'))
    # if the prot+substr input model does not contain ha types then include them from substr
    if not self.pdbcur.inp.Get('model',has_atomtypes=True):
      substr=self.inp.Get('model',typ='substr',has_atomtypes=True)
      self.pdbcur.inp.Get('model').SetAtomTypes( substr.GetAtomTypes(), atomtype1=substr.GetAtomType() )
    # first get heavy atoms PDB
    self.pdbcur.ClearAnyParams()
    self.pdbcur.runname=self.pdbcur.name+'_heavy'
    at_str=','.join(self.inp.Get('model',typ=('substr','partial+substr'),has_atomtypes=True).GetAtomTypes())
    # at the moment, an exception is needed for TA as its name is TA1 in the only residue existing
    # we could only check the atom type in the future but for the moment this is safer
    if at_str=='TA':
      self.pdbcur.SetKey('lvatom', '"{0}1[{0}]:*"'.format(at_str))
    else:
      self.pdbcur.SetKey('lvatom', '"{0}[{0}]:*"'.format(at_str))
    self.pdbcur.outfilename['pdb']='heavy.pdb'
    self.pdbcur.Run()
    self.pdbcur.out.Get('model').SetType('substr')
    # manually remove all cards before CRYST1 (otherwise wrong cards eg non-existing bonds may remain 
    # and cause trouble later - pdbcur should remove them but does not)
    with open(self.pdbcur.out.Get('model').GetFileName('pdb')) as f:
      pdb=f.read()
      re_res=re.search('CRYST1.+(\n.*)+',pdb)
      if not re_res:
        common.Error('CRYST1 card missing in PDB file {0}'.format(self.pdbcur.out.Get('model').GetFileName('pdb')))
    with open(self.pdbcur.out.Get('model').GetFileName('pdb'),'w') as g:
      g.write(re_res.group(0))
    # now get partial model PDB
    self.pdbcur.ClearAnyParams()
    self.pdbcur.runname=self.pdbcur.name+'_part'
    for at in self.inp.Get('model',typ=('substr','partial+substr'),has_atomtypes=True).GetAtomTypes():
      self.pdbcur.SetKey('delatom', '"{0}[{0}]:*"'.format(at))
      if at=='TA':
        self.pdbcur.SetKey('delatom', '"{0}1[{0}]:*"'.format(at))
    self.pdbcur.outfilename['pdb']='part.pdb'
    self.pdbcur.Run()
    self.pdbcur.out.Get('model').SetType('partial')