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


class mergemtz(process):
  name="checking mtz files and merging them"
  short_name="merging mtz files"
  supported_progs=["sftools","cad"]
  no_reporting=True


  def RunBody(self,*args,**kwargs):
    # switch between CAD and SFtools; SFtools does not support xname/dname setting!
    use_cad=True
    consistent=kwargs.pop('consistent',None)
    consistency_check_sufficient=kwargs.pop('cons_check_suff',None)
    keep_nodata=kwargs.pop('keep_nodata',None)
    prog=kwargs.pop('prog',None)
    filelist=program.GetFileList(prog,self.inp.GetAll())
    if consistent is None:
      consistent = self.CheckMTZ(filelist)
    if consistent and consistency_check_sufficient:
      self.out.AddNew('datafile', filelist[0], 'mtz')
      return
    # write the output mtz
    outmtz=os.path.join(prog.rundir,prog.nick+'_inpmerge.mtz')
    if not use_cad:
      self.lbl,self.lbl_uniq=program.GetUniqueLabels(prog,objects,filelist)
      sft = self.AddProg('sftools')
      sft.runname=sft.name+'_merge'
      filelist.reverse()
      for f in filelist:
        for l in self.lbl[f]:
          sft.SetKey( 'read', (f, 'col', "\"{0}\"".format(l)) )
        for l_ren in [l  for l in zip(self.lbl[f],self.lbl_uniq[f])  if l[0]!=l[1]]:
          sft.SetKey('set label col',l_ren[0])
          sft.SetKey(l_ren[1])
      if not keep_nodata:
        sft.SetKey('purge', 'nodata')
      sft.SetKey('write',outmtz)
      sft.SetKey('Y')
      sft.SetKey('quit\nY')
      sft.Run()
      # assuming outmtz was passed as full path here
      sft.out.AddNew('datafile', outmtz, 'mtz')
    else:
      cad = self.AddProg('cad')
      cad.runname=cad.name+'_merge'
      cad.outfilename['mtz']=outmtz
      if not keep_nodata:
        cad.SetKey('VALM', ('NaN', 'NOOUTPUT'))
      cad.Run()
      self.lbl,self.lbl_uniq=cad.lbl,cad.lbl_uniq


  def CheckMTZ(self,filelist):
    """Check the consistency between mtz dname/xname and object dname/xname
       MTZ d/x-names are set to object d/x-names unless object d/x-name is default_unknown
    """
    consistent=True
    sft = self.AddProg('sftools')
    sft_run=None
    for i,fname in enumerate(filelist):
      if not consistent: break
      for o in self.inp.GetAll():
        if not consistent: break
        if o.GetFileName('mtz')==fname and (o.xname!=o.default_unknown or o.dname!=o.default_unknown):
          if sft_run!=i:
            sft.ClearAnyParams()
            sft.SetKey('read',fname)
            sft.SetKey('list id')
            sft.SetKey('quit\nY')
            sft.runname=sft.name+'_check'+str(i)
            sft.Run()
            sft_run=i
          for l in [o.GetLabel(c) for c in o.col_list  if o.GetLabel(c)]:
            xname=sft.GetStat('xname',l)
            dname=sft.GetStat('dname',l)
            if xname!=o.xname or dname!=o.dname:
              consistent=False
              break
    self.programs.remove(sft)
    return consistent