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

class parrot(program):
  name="Parrot"
  binary="cparrot"
  labelout_prefix="PARR_"
  modif='-'
  always_merge=True
  stat={}
  stat['ncs_operator'] = common.stats(name='NCS operator', multiple=True,
            regexp=r"-ncs-operator\s+(\S+,\S+,\S+,\S+,\S+,\S+,\S+,\S+,\S+)")
  stat['radius_auto'] = common.stats(name='Automatically determined filtering radius for solvent flattening',
            regexp=r"Suggested radius for solvent mask determination:\s*(\S+)")
  references = ( "Cowtan K (2000) General quadratic functions in real and reciprocal space and "+
                 "their applications to likelihood phasing. Acta Cryst. D56, 1612-1621.", )

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

  def TreatInput(self):
    # F,sigF
    self.fsf = self.inp.Get('fsigf',typ='average',col='f')
    if self.fsf:
      self.SetKey('mtzin-wrk', self.fsf.GetFileName('mtz'))
      self.SetKey('colin-wrk-fo', self.fsf.GetFullLabel('f','sigf'))
    else:
      common.Error('No data inputted for {0}'.format(self.name))
    # Parrot seems to always require phase distribution (even if not asked for phase combination and map is inputted...)
    for mapc in self.inp.GetAll('mapcoef',typ=('best','combined')):
      if mapc.GetLabel('hla'):
        self.SetKey('colin-hl', mapc.GetFullLabel('hla','hlb','hlc','hld'))
        break
      if mapc.GetLabel('ph') and mapc.GetLabel('fom'):
        self.SetKey('colin-phifom', mapc.GetFullLabel('ph','fom'))
        break
    else:
      common.Error('No phase distribution inputted for {0}'.format(self.name))
    # input map (optional - phase distr. and Fo used if not inputted)
    mapc=self.inp.Get('mapcoef',typ='weighted',col=('f','ph'))
    if mapc:
      self.SetKey('colin-wrk-fc', mapc.GetFullLabel('f','ph'))
    # solvent content - if supplied with an input container
    solv_obj = self.inp.Get(has_solvent_content=True)
    if solv_obj and not self.IsKey('solvent-content'):
      self.SetKey('solvent-content', solv_obj.solvent_content)
    # detect ncs
    if self.process.GetParam('ncs_det') is not False:
      heavy_model=self.inp.Get('model',typ='substr',filetype='pdb')
      if heavy_model and self.process.GetParam('ncs_det_ha') is not False:
        self.SetKey('pdbin-wrk-ha',heavy_model.GetFileName('pdb'))
      mr_model=self.inp.Get('model',typ='partial',custom='ncs',filetype='pdb')
      if mr_model and self.process.GetParam('ncs_det_mr') is not False:
        self.SetKey('pdbin-wrk-mr',mr_model.GetFileName('pdb'))
      if self.process.GetParam('ncs_det') and not heavy_model and not mr_model:
        common.Error('Asked for NCS detection by {0} but no appropriate model to detect from was supplied.'.format(self.name))

  def TreatParams(self):
    # real space DM
    if self.process.nick=='dm':
      self.SetKey('modify-map-and-terminate')
    if self.process.IsNonFalseVirtPar('cycles'):
      self.SetKey('cycles', self.process.GetVirtPar('cycles'))
    if self.process.IsNonFalseVirtPar('solvent_content'):
      self.SetKey('solvent-content', self.process.GetVirtPar('solvent_content'))
    program.TreatParams(self)

  def DefineOutput(self):
    out_mapc = self.out.AddNew( 'mapcoef', self.nick+'.mtz' )
    if self.process.nick=='dm':
      out_mapc.SetType('densmod')
      out_mapc.SetLabel( ['f','ph'] )
    else:
      out_mapc.SetType('combined')
      out_mapc.SetLabel( ['hla','hlb','hlc','hld'] )
    out_mapc.SetCrystalName(self.fsf.GetCrystalName())
    out_mapc.SetDataName(self.fsf.GetDataName())

  def TreatOutput(self):
    if self.process.nick=='dm':
      self.SetKey( 'colout-fc', self.out.mapcoef[-1].GetFullLabel('f','ph') )
    else:
      self.SetKey( 'colout-hl', self.out.mapcoef[-1].GetFullLabel('hla','hlb','hlc','hld') )
    self.SetKey( 'mtzout', self.out.mapcoef[-1].GetFileName('mtz') )