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

class shelxc(program):
  name="SHELXC"
  binary="shelxc"
  ccp4_parsing=True
  stat={}
  stat['resol'] = common.stats(name='resol',multiple=True,
                               regexp=[r" Resl.\s+Inf."+r'(?:\s+(\d+\.\d+))?'*15] ) 
  stat['dano'] = common.stats(name='dano',multiple=True,
                              regexp=[r" <d\"/sig>"+r'(?:\s+(\d+\.\d+))?'*15] ) 

  references = ( "Sheldrick GM (2008) A short history of SHELX. " +
                 "Acta Cryst. A64, 112-122.", )

  def Init(self):
    self.SetArg('shelx')

  def TreatInput(self):
    exper=self.process.GetVirtPar('target')
    # get required data for the target; try I's first and then F's
    fi_opt=''
    obj,N=self.inp.GetTargetObjects(exper, inten=True, no_model=True, no_warn=True)
    if obj is None:
      obj,N=self.inp.GetTargetObjects(exper, inten=False, no_model=True)
      fi_opt='-f'
    if obj is None:
      common.Error('The input objects for target {0} of {1} could not be retrieved'.format(exper,self.name))
    self.obj=obj
    # set cell and spacegroup
    if not self.IsKey('cell') or not self.IsKey('spag'):
      # get cell/spgr (we might want to implement this directly to GetCell/Spacegroup?)
      self.cell=obj['f+'].cell
      self.spgr=obj['f+'].spgr
      if not self.cell or not self.spgr:
        if obj['f+'].GetFile('mtz'):
          obj_mtz=obj['f+']
        else:
          obj_mtz=self.inp.Get(filetype='mtz',typ='plus',xname=obj['f+'].GetCrystalName(),dname=obj['f+'].GetDataName())
          if obj_mtz is None:
            self.Error("Cell and spacegroup not known for crystal {0}, data {1}.".format(obj['f+'].GetCrystalName(),obj['f+'].GetDataName()) )
        self.cell=obj_mtz.GetCell(self)
        self.spgr=obj_mtz.GetSpacegroup(self)
      if not self.IsKey('cell'):
        self.SetKey('cell', ' '.join((str(c) for c in self.cell)))
      if not self.IsKey('spag'):
        self.SetKey('spag', self.spgr)
    # set atomtype
    if not self.IsKey('sfac') and 'mod' in obj and obj['mod'].atomtypes:
      self.SetKey('sfac', obj['mod'].GetAtomType())
    # expected number of atoms from substr. object - maybe this should be moved to shelxd?
    if self.inp.Get(has_num_atoms=True) and not self.GetKey('FIND'):
      self.SetKey('FIND', self.inp.Get(has_num_atoms=True).exp_num_atoms)
    # SAD, MAD and SIRAS input
    # shelxc does not deal with long path strings so we use relative path at input
    if exper=='SIRAS' or exper=='SAD':
      self.SetKey(exper, (fi_opt, obj['f+'].GetFileName(('hkl','sca'),conv_parent=self,inp_cont=obj['f-'])))
      if 'fn' in obj:
        self.SetKey('NAT', (fi_opt, obj['fn'].GetFileName(('hkl','sca'),conv_parent=self)))
    elif exper=='MAD':
      if N>4:
        common.Warning('{0} datasets supplied, only 4 will be used by {1}'.format(N,self.name))
      keys=['PEAK','INFL','HREM','LREM']
      mad_assign={}
      # try to see whether the assignment is defined in dataname or custom attributes
      for f in ['f'+str(n)+'+' for n in range(1,N)]:
        for key in keys:
          if key in [ c for c in obj[f].custom ] or key in obj[f].dname:
            mad_assign[key] = f
      # if not then just define by order
      for f in ['f'+str(n)+'+' for n in range(1,N)]:
        if len(mad_assign)==len(keys):
          break
        if f not in mad_assign.itervalues():
          key=[k for k in keys if k not in mad_assign][0]
          mad_assign[key] = f
      for key,val in mad_assign:
        self.SetKey(key, (fi_opt, obj[val].GetFileName(('hkl','sca'),conv_parent=self,inp_cont=obj[val[:-1]+'-'])))

  def DefineOutput(self):
    # the shelxc output filenames are hardcoded...
    project=self.arg_list[0]
    # SAD only output yet. SIRAS/MAD will need to be added later.
    self.out.AddNew('fsigf', project+'.hkl', filetype='hkl', typ='average', xname=self.obj['f+'].xname, dname=self.obj['f+'].dname)
    fa=self.out.AddNew('fsigf', project+'_fa.hkl', filetype='hkl', typ='fa', xname=self.obj['f+'].xname, dname=self.obj['f+'].dname)
    ins=self.out.AddNew('datafile', project+'_fa.ins', typ='ins')