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

class f2mtz(program):
  name="F2MTZ"
  binary="f2mtz"
  ccp4_parsing=True

  # warning: only fa's from shelxc and phs from shelxe supported as of now
  def TreatInput(self):
    inp_format = self.process.GetParam('inputformat')
    self.inp_cont = None
    if inp_format == 'hkl':
      self.inp_cont = self.inp.Get(filetype=inp_format,typ='fa',try_convert=False)
      if self.inp_cont:
        self.SetKey('format','\'(3I4,2F8.2,I4)\'')
    elif inp_format == 'phs':
      self.inp_cont = self.inp.Get(filetype=inp_format,try_convert=False)
      if self.inp_cont:
        self.SetKey('format','\'(3I4,2F9.2,F8.1,F9.2)\'')
    if self.inp_cont:
      self.SetArg('hklin', self.inp_cont.GetFileName())
    else:
      common.Error('No suitable file inputted to {0}'.format(self.name))
    mtzin = self.inp.Get(filetype='mtz',xname=self.inp_cont.GetCrystalName(),dname=self.inp_cont.GetDataName())
    if mtzin:
      self.SetKey('symm', mtzin.GetSpacegroup(self).replace(" ",""))
      self.SetKey('cell', ','.join(str(c) for c in mtzin.GetCell(self)))
    else:
      common.Error('No mtz file inputted to {0} (needed to obtain cell/spacegroup)'.format(self.name))

  def DefineOutput(self):
    self.outmtz=self.out.Add(self.inp_cont)
    self.out.AddFileToChild(self.outmtz, self.nick+'.mtz', 'mtz')
    if self.inp_cont.GetType()=='fa':
      self.outmtz.SetLabel(('f','sigf','alpha'))
    elif self.process.GetParam('inputformat')=='phs':
      self.outmtz.SetLabel( ['f','ph','fom'] )

  def TreatOutput(self):
    self.SetArg('hklout', self.out.Get().GetFileName('mtz'))
    if self.inp_cont.GetType()=='fa':
      self.AddToKey('labout', ('H','K','L',self.outmtz.GetLabel('f'), self.outmtz.GetLabel('sigf'), self.outmtz.GetLabel('alpha')))
      self.AddToKey('ctype', ('H','H','H','F','Q','P'))
    elif self.process.GetParam('inputformat')=='phs':
      self.AddToKey('labout', ('H','K','L',self.outmtz.GetLabel('f'), self.outmtz.GetLabel('fom'), self.outmtz.GetLabel('ph'), 'SIGF'))
      self.AddToKey('ctype', ('H','H','H','F','W','P','Q'))