#!/usr/bin/python
import os,sys,shutil,time
from program import program
import common

class refmac(program):
  name="REFMAC5"
  binary="refmac5"
  labelout_prefix="REFM_"
  ccp4_parsing=True
  stat={}
  stat['fom'] = common.stats(name='FOM', regexp=r"Overall figure of merit\s+=\s+(\S+)")
  stat['correl'] = common.stats(name='correlation coefficient', regexp=r"Overall correlation coefficient\s+=\s+(\S+)")
  stat['rfact'] = common.stats(name='R factor', regexp=r"Overall R factor\s+=\s+(\S+)")
  stat['rfree'] = common.stats(name='R-free factor', regexp=r"Free R factor\s+=\s+(\S+)")
  stat['version'] = common.stats(name='version', regexp=r"refmac5;\s+version\s+(\S+)")
  stat['lg_fom'] = common.stats(name='loggraph FOM info', 
                   regexp=r"\$TABLE: Cycle\s+\S+\s+Fom\(<cos\(DelPhi\)>")
  references = ( "Murshudov GN, Skubak P, Lebedev AA, Pannu NS, Steiner RA, Nicholls RA, " +
                 "Winn MD, Long F & Vagin AA (2011) REFMAC5 for the refinement of " +
                 "macromolecular crystal structures. Acta Cryst. D67, 355-367.", )

  def Init(self):
    self.outfilename = { 'pdb': self.name+'.pdb', \
                         'pdbnat': self.name+'_nat.pdb', \
                         'mtz': self.name+'.mtz' \
                       }

  def CheckBinary(self):
    program.CheckBinary(self)
    self.SetRunDir('checkbinary')
    self.runname='{0}_version'.format(self.name)
    self.ExternalRun( [self.binary, '-i'], [] )
    ver=self.GetStat('version')
    if ver in ('5.7.0028','5.7.0029','5.7.0030','5.7.0031'):
      common.Warning('Version {0} of {1} is known to contain bugs leading to significantly worse results.'.format(ver,self.name))
      common.Warning('It is STRONGLY RECOMMENDED TO STOP NOW and update {0} to 5.7.0032 or higher.'.format(self.name))
      time.sleep(3)
    if not ver.startswith('5.'):
      common.Warning('Version {0} of {1} might not be supported. It is recommended to use the current CCP4 version.'.format(ver,self.name))
      time.sleep(3)

  def Interact_output(self, line, popen, empty):
    #if not empty:
    self.process.UpdateInfo(line,popen)

  def TreatInput(self):
    exper=self.process.GetParam('target',capitalized=True)
    obj,N=self.inp.GetTargetObjects(exper,modtyp='pdb',fsftyp='mtz')
    if obj is None:
      common.Error('The input objects for target {0} of {1} could not be retrieved'.format(exper,self.name))
    # SAD and SIRAS input
    if exper=='SIRAS':
      self.fn=obj['fn']
      self.SetKey( 'labin', ('FN='+obj['fn'].GetLabel('f'),'SIGFN='+obj['fn'].GetLabel('sigf')) )
    if exper=='SAD' or exper=='SIRAS': 
      self.SetArg( 'hklin', obj['f+'].GetFileName('mtz') )
      self.AddToKey('labin', ('F+='+obj['f+'].GetLabel('f'),'SIGF+='+obj['f+'].GetLabel('sigf')) )
      self.AddToKey('labin', ('F-='+obj['f-'].GetLabel('f'),'SIGF-='+obj['f-'].GetLabel('sigf')) )
      for at,f in obj['mod'].GetAtomTypes().iteritems():
        if f[0] is not None and f[1] is not None:
          self.SetKey('ANOM FORM', (at,f[0],f[1]) )
    # MLHL and RICE input
    elif exper=='MLHL':
      self.SetArg( 'hklin', obj['f'].GetFileName('mtz') )
      self.SetKey(  'labin', ('FP='+obj['f'].GetLabel('f'),'SIGFP='+obj['f'].GetLabel('sigf')) )
      self.AddToKey('labin', ('HLA='+obj['hl'].GetLabel('hla'),'HLB='+obj['hl'].GetLabel('hlb')))
      self.AddToKey('labin', ('HLC='+obj['hl'].GetLabel('hlc'),'HLD='+obj['hl'].GetLabel('hld')))
    elif exper=='RICE':
      self.SetArg( 'hklin', obj['f'].GetFileName('mtz') )
      self.SetKey( 'labin', ('FP='+obj['f'].GetLabel('f'),'SIGFP='+obj['f'].GetLabel('sigf')) )
    # the map we want to combine with (at the moment, phase comb. key is DM in refmac but any map can be used)
    if self.GetKey('DM') or self.GetKey('DM2'):
      self.SetKey('scpa','1')
      comb_typ=self.process.GetVirtPar('comb_with_typ')
      ph = self.inp.Get('mapcoef',typ=comb_typ,col='ph',filetype='mtz')
      if ph:
        self.AddToKey('labin', ('FPART1='+ph.GetLabel('f'),'PHIP1='+ph.GetLabel('ph')))
      else:
        common.Error("No phase information to combine with for {0}".format(self.name))
    # model(s) to be refined
    if not obj['mod']:
      common.Error('No model to be refined by {0}'.format(self.name))
    # just saving so that output can reach the objects easily
    for f in [f  for f in ('f+','fn','f') if f in obj]:
      self.f=obj[f]
    self.mdl,self.mdlnat = obj['mod'],None
    # xyzin
    if 'modn' in obj:
      self.mdlnat=obj['modn']
      self.SetArg('xyzin',obj['modn'].GetFileName('pdb'))
      self.SetArg('xyzin2',obj['mod'].GetFileName('pdb'))
    else:
      self.SetArg('xyzin',obj['mod'].GetFileName('pdb'))
    # free assignment
    free=self.inp.Get('exclude',typ='freeR',col='free')
    if free:
      self.AddToKey( 'labin', "FREE={0}".format(free.GetLabel('free')) )
    # read in luzzatti D parameters
    if self.inp.Get('datafile',typ='dluz') and self.rundir:
      if not self.inp.Get('datafile',typ='dluz').GetFileName().startswith(self.rundir):
        shutil.copy(self.inp.Get('datafile',typ='dluz').GetFileName(), self.rundir)
      self.SetKey('REFI','RLUZZ')

  def TreatParams(self):
    if self.process.IsNonFalseVirtPar('cycles') and not self.IsKey('ncyc'):
      self.SetKey('ncyc', self.process.GetVirtPar('cycles'))
    if self.process.IsNonFalseVirtPar('beta') and not self.IsKey('dmbl'):
      self.SetKey('dmbl', self.process.GetVirtPar('beta'))
    if self.process.nick=='phcomb' or (self.process.nick=='ref' and self.process.GetVirtPar('phcomb')):
      if self.process.GetVirtPar('target')=='SIRAS' or self.process.nick=='ref':
        self.SetKey('DM2')
      else:
        self.SetKey('DM')
      if not self.IsKey('ncyc'):
        self.SetKey('ncyc',0)
    if self.process.nick!='ref':
      self.SetKey('solv','no')
    if (self.process.nick!='ref' or self.process.GetVirtPar('phcomb')): 
      if not any('SUBS' in v for v in self.GetKey('refi',allval=True,capitalized=True)):
        self.SetKey('refi',('subs','yes'))
    self.SetKey('REFI','WLUZZ')
    program.TreatParams(self)

  def DefineOutput(self):
    besttyp='combined'
    if self.process.nick=='phas':
      besttyp='best'
    self.outmapcc = self.out.AddNew( 'mapcoef', self.outfilename['mtz'], typ=besttyp, xname=self.f.xname, dname=self.f.dname)
    if self.process.nick!='ref':
      self.outmapcc.SetLabel( ['f','ph','fom','hla','hlb','hlc','hld'] )
    else:
      self.outmapcc.SetLabel( ['ph','fom','hla','hlb','hlc','hld'] )
    self.outmapcw = self.out.AddNew( 'mapcoef', self.outfilename['mtz'], typ='weighted', xname=self.f.xname, dname=self.f.dname)
    self.outmapcw.SetLabel( ['f','ph'] )
    self.outmapcad = None
    if self.process.GetVirtPar('target') in ('SAD','SIRAS'):
      self.outmapcad = self.out.AddNew( 'mapcoef', self.outfilename['mtz'], typ='anom-diff', xname=self.f.xname, dname=self.f.dname )
      self.outmapcad.SetLabel( ['f','ph'] )
    self.outmodel = self.out.AddCopy(self.mdl)
    self.out.SetFileToChild(self.outmodel, self.outfilename['pdb'], 'pdb')
    #self.outmodel = self.out.AddNew( 'model', self.outfilename['pdb'], typ=self.mdl.GetType(), atomtypes=self.mdl.GetAtomTypes() )
    if self.mdlnat:
      self.outmodelnat = self.out.AddCopy(self.mdlnat)
      self.out.SetFileToChild(self.outmodelnat, self.outfilename['pdbnat'], 'pdb')
      #self.outmodelnat = self.out.AddNew( 'model', self.outfilename['pdbnat'], typ=self.mdlnat.GetType(), atomtypes=self.mdlnat.GetAtomTypes() )
    if any(val for val in self.GetKey('REFI',allval=True,capitalized=True) if 'WLUZ' in val):
      self.out.AddNew( 'datafile', typ='dluz', filename='luzzd' )

  def TreatOutput(self):
    self.SetKey('phout')
    if self.process.nick!='ref':
      self.SetKey( 'labout', ('FB='+self.outmapcc.GetLabel('f'),'PHIB='+self.outmapcc.GetLabel('ph')) )
    else:
      self.SetKey( 'labout', 'PHCOMB='+self.outmapcc.GetLabel('ph') )
    self.AddToKey( 'labout', 'FOM='+self.outmapcc.GetLabel('fom') )
    self.AddToKey( 'labout', ('HLACOMB='+self.outmapcc.GetLabel('hla'),'HLBCOMB='+self.outmapcc.GetLabel('hlb')) )
    self.AddToKey( 'labout', ('HLCCOMB='+self.outmapcc.GetLabel('hlc'),'HLDCOMB='+self.outmapcc.GetLabel('hld')) )
    self.AddToKey( 'labout', ('FWT='+self.outmapcw.GetLabel('f'),'PHWT='+self.outmapcw.GetLabel('ph')) )
    if self.outmapcad:
      self.AddToKey( 'labout', ('DELFAN='+self.outmapcad.GetLabel('f'),'PHDELAN='+self.outmapcad.GetLabel('ph')) )
    self.SetArg( 'hklout', self.outmapcc.GetFileName('mtz') )
    if self.mdlnat:
      self.SetArg( 'xyzout', self.outmodelnat.GetFileName('pdb') )
      self.SetArg( 'xyzout2', self.outmodel.GetFileName('pdb') )
    else:
      self.SetArg( 'xyzout', self.outmodel.GetFileName('pdb') )