#!/usr/bin/python
import os,sys
from process import process
import common
par=common.parameter


class refatompick(process):
  name="iterative refinement and (substructure) atom picking"
  short_name="iterative atom picking"
  supported_procs=["ref", "phas", "atomsfrommap"]
  supported_params = {}
  supported_params['target'] = par( desc='Experiment/refinement target', typ=str, cap=True, share=True )
  supported_params['refcyc'] = par( desc='Number of refinement cycles in each iteration', typ=int )
  supported_params['num_iter'] = par( desc='Number of refinement+picking iterations', typ=int )
  supported_params['rms_threshold'] = par( desc='RMS threshold for atom picking from map', typ=float, share=True )
  supported_params['skip_init_ref'] = par( desc='Skip the initial refinement', typ=bool )
  supported_params['skip_last_pick'] = par( desc='Skip atom picking in the last iteration', typ=bool )
  supported_params['max_new_atoms'] = common.parameter( desc='Maximal number of atoms added in one iteration', typ=(int,bool), share=True )
  supported_params['bfactor'] = par( desc='Set B factor of the new atoms to this value', typ=(float,bool), share=True )

  def TreatInOutPar(self, set_all_par=False):
    # input
    self.ph = self.GetProcess('phas')
    if not self.ph:
      self.ph = self.GetProcess('ref')
    if not self.ph:
      if self.inp.Get('model',typ='substr',filetype='pdb'):
        self.ph = self.AddProcess('phas')
      else:
        self.ph = self.AddProcess('ref')
    if not self.ph.GetProg():
      self.ph.AddProg('refmac')
    self.ph.TreatInOutPar()
    self.afm = self.GetOrAddProcess('atomsfrommap')
    self.afm.TreatInOutPar()
    # parameters
    if not self.IsParam('num_iter'):
      self.SetParam('num_iter', 3)
    if not self.GetParam('rms_threshold'):
      self.SetParam('rms_threshold', 4.5)
    if self.GetVirtPar('refcyc') is None and self.ph.nick=='ref': 
      self.SetVirtPar('refcyc', 10)
      if self.ph.GetParam('low_res_par'):
        self.SetVirtPar('refcyc', 20)
    if self.IsVirtPar('refcyc') and not self.ph.IsVirtPar('cycles'):
      self.ph.SetParam('cycles', self.GetVirtPar('refcyc'))
    process.TreatInOutPar(self,set_all_par)

  def RunBody(self,*args,**kwargs):
    for i in range(self.GetParam('num_iter')):
      self.LGInfo('\n$TEXT: Iteration {0}: $$ Beta estimation result $$'.format(i+1))
      if i>0 or not self.GetParam('skip_init_ref'):
        self.ph.GetProg().runname=self.ph.GetProg().name+'_cyc'+str(i)
        self.ph.Run()
        self.afm.inp.AddCopy(self.ph.out.Get('model'))
        self.afm.inp.Set(self.ph.out.GetAll('mapcoef'))
      if i<self.GetParam('num_iter')-1 or not self.GetParam('skip_last_pick'):
        self.afm.GetProg().runname=self.afm.GetProg().name+'_cyc'+str(i)
        self.afm.Run()
        self.ph.inp.AddCopy(self.afm.out.Get('model'))
      self.LGInfo('\n$$')