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


class convert(process):
  name="data format conversion"
  short_name="data conversion"
  supported_progs=["fft", "invfft", "mtz2various", "sftools", "convert2mtz", "f2mtz"]
  supported_params = {}
  supported_params['outputformat'] = common.parameter( desc='Output format for mtz2various conversion', typ=str )
  supported_params['inputformat'] = common.parameter( desc='Input format for convert2mtz/f2mtz conversion', typ=str )
  no_reporting=True
  # setup for the conversion - filled automatically when creatin the convert process instance using RegisterConv
  setup = {}


  def RegisterConv(converts,func,fromlist,tolist):
    for f in fromlist:
      for t in tolist:
        converts[(f,t)] = func

  def ConvertMTZ2Various(self,oldtyp,newtyp):
    self.SetParam('outputformat',newtyp)
    mtz2var=self.AddProg('mtz2various')
    mtz2var.outfilename[newtyp] = self.inp.Get(filetype='mtz').GetFileName(trim_path=True).rsplit('.',1)[0]+'_'+mtz2var.name+'.'+newtyp
    # automatic scaling does not work correctly in mtz2various if max(I-)>max(I+)
    # as a workaround, find and set the scale here to prevent overflows with hkl.
    mtz_ip=self.inp.Get('fsigf',filetype='mtz',col='i',typ='plus')
    mtz_im=self.inp.Get('fsigf',filetype='mtz',col='i',typ='minus')
    if newtyp == 'hkl' and not self.GetProg().GetKey('SCAL') and mtz_ip and mtz_im:
      sftools=self.AddProg('sftools')
      sftools.SetKey("read", [mtz_ip.GetFileName('mtz'), 'col', "\"{0}\"".format(mtz_ip.GetLabel('i'))])
      sftools.SetKey("checkhkl")
      sftools.Run()
      maxp = sftools.GetStat('max', mtz_ip.GetLabel('i')[:12])
      sftools.SetKey("read", [mtz_im.GetFileName('mtz'), 'col', "\"{0}\"".format(mtz_im.GetLabel('i'))])
      sftools.SetKey("checkhkl")
      sftools.Run()
      self.GetProg().SetKey('SCAL', 99999./max(maxp,sftools.GetStat('max', mtz_im.GetLabel('i')[:12])))
      self.programs.remove(sftools)
  RegisterConv(setup, ConvertMTZ2Various, ['mtz'], ['hkl','sca','phs'])

  def ConvertMTZ2drear(self,oldtyp,newtyp):
    out_filename="convert.drear"
    sftools=self.AddProg('sftools')
    fa=self.inp.Get('fsigf',typ='fa',filetype='mtz')
    ea=self.inp.Get('fsigf',typ='ea',filetype='mtz')
    sftools.SetKey("read", fa.GetFileName('mtz'))
    sftools.SetKey("write", (out_filename, "format(3I5,2F9.2,2F9.3)", "col", \
      fa.GetLabel('f'), fa.GetLabel('sigf'), ea.GetLabel('f'), ea.GetLabel('sigf')))
    sftools.SetKey("quit")
    self.out.Add(fa)
    self.out.AddFileToChild(fa,out_filename,newtyp)
    self.out.Add(ea)
    self.out.AddFileToChild(ea,out_filename,newtyp)
  RegisterConv(setup, ConvertMTZ2drear, ['mtz'], ['drear'])

###!!!!! disabled as of now: from reasons yet unknown does not work for conversion of shelc's fa's
  def ConvertVarious2MTZ(self,oldtyp,newtyp):
    if self.inp.Get('fsigf',typ='fa'):
      self.parent_process.ConvertF2MTZ(oldtyp,newtyp)
    else:
      self.convert_proc.SetParam('inputformat',oldtyp)
      conv2mtz=self.AddProg('convert2mtz')
      conv2mtz.SetKey('no-complete')
      hklin=self.inp.Get('fsigf',filetype='hkl')
      if oldtyp=='hkl' and hklin:
        mtzin=self.inp.Get('fsigf',typ='average',filetype='mtz',try_convert=False)
        if mtzin:
          conv2mtz.SetKey('spacegroup', mtzin.GetSpacegroup(conv2mtz).replace(" ",""))
          conv2mtz.SetKey('cell', ','.join(str(c) for c in mtzin.GetCell(conv2mtz)))
        else:
          self.parent_process.no_conversion=True
    # automatic scaling does not work in convert2mtz, we will need to find and set the scale here 
    # to prevent overflows with hkl. this can be based on the code from mtz2various.
#  RegisterConv(setup, ConvertVarious2MTZ, ['hkl','sca','phs'], ['mtz'])

  def ConvertF2MTZ(self,oldtyp,newtyp):
    self.SetParam('inputformat',oldtyp)
    oo=self.inp.Get(filetype=oldtyp)
    mtzin=self.inp.Get('fsigf',typ='average',filetype='mtz',xname=oo.xname,dname=oo.dname,try_convert=False)
    if not mtzin:
      mtzin=self.parent_process.inp.Get('fsigf',typ='average',filetype='mtz',xname=oo.xname,dname=oo.dname,try_convert=False)
      if mtzin:
        self.inp.AddCopy(mtzin)
    # only supporting fa convertion (from shelx) as of now
    if mtzin and self.inp.Get('fsigf',typ='fa'):
      f2mtz=self.AddProg('f2mtz')
      f2mtz.Run()
      sftools=self.AddProg('sftools')
      sftools.inp.Add(f2mtz.out.Get(filetype='mtz'))
      sftools.SetKey("read", self.inp.Get('fsigf',typ='fa').GetFileName('hkl'))
      sftools.SetKey("reduce")
      sftools.SetKey("quit")
      self.programs.remove(f2mtz)
    elif oldtyp=='hkl':
#      self.ConvertVarious2MTZ(oldtyp,newtyp)
      if self.parent_process:
        self.parent_process.no_conversion=True
    else:
      f2mtz=self.AddProg('f2mtz')
  RegisterConv(setup, ConvertF2MTZ, ['hkl','phs'], ['mtz'])

  def ConvertMTZ2Map(self,oldtyp,newtyp):
    fft=self.AddProg('fft')
  RegisterConv(setup, ConvertMTZ2Map, ['mtz'], ['map'])

  def ConvertMap2MTZ(self,oldtyp,newtyp):
    invfft=self.AddProg('invfft')
    if not invfft.inp.Get('fsigf',filetype='mtz',col='f',try_convert=False):
      if self.parent_process:
        f = self.parent_process.inp.Get('fsigf', filetype='mtz', col='f', typ='average', \
          xname=self.inp.Get(filetype=oldtyp).GetCrystalName())
        if f:
          invfft.inp.Add(f)
  RegisterConv(setup, ConvertMap2MTZ, ['map'], ['mtz'])