import os,sys
import pkgutil,shutil,copy
import common,processes,ccp4i2crank,parse
from process import process, crvapi
from program import program

class crank(process):
  name='CRANK2'
  ccp4i2=None
  # parameters
  supported_params = {}
  supported_params['target'] = common.parameter( desc='Experiment/refinement target', typ=str, cap=True )
  supported_params['replace_met_mse'] = common.parameter( desc='Replace methionines by selenomethionines in input model(s) with SE substructure', typ=bool, cap=True )
  # all processes
  supported_procs = [ n for _,n,_ in pkgutil.iter_modules( [os.path.dirname(processes.__file__)] ) ]
  # filehandle of the loggraphfile - LGInfo() will send output to this filehandle if not None
  loggraphfilehandle=None
  loggraph_actfilehandle=None
  # ccp4 version is determined at the beginning at saved here (for loggraph printing)
  ccp4_version=None
  # references
  references = ( 'Skubak P, Pannu NS (2013) Automatic protein structure solution from weak X-ray data.' + 
                 'Nature Communications 4, 2777.',
                 'Pannu NS, Waterreus WJ, Skubak P, Sikharulidze I, Abrahams JP & ' +
                 'de Graaff RAG (2011) Recent advances in the CRANK software suite ' +
                 'for experimental phasing. Acta Cryst. D67, 331-337.', )

  def TreatInOutPar(self, set_all_par=False):
    if self.GetParam('replace_met_mse') is None:
      for mod in self.inp.GetAll('model',has_atomtypes=True):
        if 'SE' in mod.GetAtomTypes():
          self.SetParam('replace_met_mse', True)
    process.TreatInOutPar(self,set_all_par)

  def RunPreprocess(self,emulate=False,*args,**kwargs):
    if hasattr(self,'logout') and not hasattr(self,'stdout_save'):
      self.Info('Output will be redirected to {0}'.format(self.logout))
      self.stdout_save = sys.stdout
      sys.stdout = open(self.logout,'w',1)
    global crvapi
    if hasattr(self,'disable_rvapi') and crvapi:
      self.Info('Disabling rvapi output.')
      crvapi = False
    if hasattr(self,'rvapi_uri_prefix') and crvapi:
      crvapi.uri_prefix = self.rvapi_uri_prefix
    process.RunPreprocess(self,emulate=emulate,crvapi=crvapi,*args,**kwargs)
    self.DetermineUnknownModel(emulate=emulate)
    if not emulate:
      self.FixInpModels()

  def RunBody(self, emulate=False, set_all_par=False, *args, **kwargs):
    """Runs CRANK2 job."""
    self.Initialize(self)
    self.run.TreatInOutPar(set_all_par=set_all_par)
    # the emulation goes through the TreatInOutPar of all steps without actually running them
    # if set_all_par is True then for all the steps all the parameters are set and printed
    for ip,p in enumerate(self.processes):
      #  at this point we are dealing with "dummy" preparation process instances
      #  (this is in order to write desired XML files by conversion from dummy instances)
      # first pass output from the previous step to input of the current step
      #if not emulate:
      self.PassPrevStepsOutput(p,ip,emulate)
      # set target for the process if needed
      self.SetTargetDefault(p)
      # pass previous phasing/dm jobs to the handdet process
      if p.nick=='handdet':
        self.PassToHandDet(p,ip)
      # write XML of process p and create its run instance by reading it
      p_run = self.SubProcess_WriteXML_CreateRun(p,ip,emulate=emulate)
      # any other general process initialization if needed
      self.Initialize(p_run)
      if not emulate:
        # run the p_run process
        self.RunSubProcess(p_run)
        # pass the chosen hand mtz's/pdb's to all the subsequent steps
        if p.nick=='handdet':
          self.PassChosenHand(p_run,ip)
      else:
        p_run.TreatInOutPar(set_all_par)
    # emulation cleanup
    if emulate:
      os.chdir(self.rundir)
      for p_run in self.run.processes:
        if os.path.isdir(p_run.rundir) and os.path.realpath(os.path.realpath(p_run.rundir)).startswith(self.rundir):
          shutil.rmtree(p_run.rundir)

  def CheckBinaries(self):
    # checks whether the binaries needed exist (and their versions if such a check is defined for program)
    # first let us construct a copy of the crank used for the run and run TreatInOutPar() 
    # for each step so that the default programs are assigned
    self.emulate = self.GetProcessCopy()
    #crank = self.AddProcessCopy(self)
    #self.processes.remove(crank)
    self.emulate.RunBody(emulate=True)
    checked=[]
    for p in self.emulate.run.GetFlatSubTree() + [program.from_name('sftools',None)]:
      if hasattr(p,'binary'):
        program.TreatParams(p)
        # make sure each program is included only once
        if p.binary not in checked:
          p.CheckBinary()
          checked.append(p.binary)

  def GetCCP4Version(self):
    # check/get ccp4 library version
    if not self.ccp4_version:
      bp3=program.from_name('bp3',None)
      bp3.SetRunDir('checkbinary')
      bp3.runname='{0}_ccp4_version'.format(self.name)
      bp3.ExternalRun( [bp3.binary, '-i'], [] )
      self.ccp4_version = bp3.GetStat('ccp4_version')
    return self.ccp4_version

  def DetermineUnknownModel(self,emulate=False):
    for p in self.GetFlatSubTree():
      # try to determine type of 'unknown' pdb from user input
      model_to_fix=p.inp.Get('model',typ='unknown',filetype='pdb')
      if model_to_fix:
        at_obj=p.inp.Get('model',has_atomtypes=True)
        if not at_obj: at_obj=self.inp.Get('model',has_atomtypes=True)
        if at_obj:
          pdbcur=self.AddProg('pdbcur',propagate_out=False,propagate_inp=False)
          pdbcur.inp.Add(model_to_fix)
          for at in at_obj.GetAtomTypes():
            pdbcur.SetKey('delatom', "[{0}]".format(at))
          if self.GetParam('replace_met_mse') and 'S' not in at_obj.GetAtomTypes():
            pdbcur.SetKey('delatom', "(MET)/SD[S]".format(at))
          pdbcur.Run()
          anom_atoms_present = pdbcur.GetStat('atoms_deleted')
          pdbcur.inp.Set(pdbcur.out.model)
          pdbcur.SetKey('delatom', "[C]", keep_previous=False)
          pdbcur.Run()
          prot_atoms_present =  pdbcur.GetStat('atoms_deleted')
          if anom_atoms_present:
            model_to_fix.SetType('substr')
            if prot_atoms_present:
              model_to_fix.SetType('partial+substr')
          elif prot_atoms_present:
            model_to_fix.SetType('partial')
          else:
            common.Error('Neither substructure nor partial model present in file {0}.'.format(model_to_fix.GetFileName('pdb')))
          self.Info('Type of {0} model set to {1}.'.format(model_to_fix.GetFileName('pdb'), model_to_fix.GetType()))
          if emulate:
            shutil.rmtree(pdbcur.rundir)
          self.programs.remove(pdbcur)
        else:
          common.Error('Could not classify {0} as substructure atom types not specified.'.format(model_to_fix.GetFileName('pdb')))

  def FixInpModels(self):
    for p in self.GetFlatSubTree():
      # try fixing substructure pdb from user input
      model_to_fix=p.inp.Get('model',typ='substr',filetype='pdb')
      if model_to_fix:
        fixsub=self.AddProcess('fixsubstrpdb',propagate_out=False,propagate_inp=False)
        fixsub.inp.Add(model_to_fix)
        fixsub.AddProg('pdbcur')
        fixsub.GetProg().outfilename['pdb'] = "fixed_"+p.nick+".pdb"
        fixsub.Run()
        p.inp.model.remove(model_to_fix)
        p.inp.Add(fixsub.out.Get('model'), propagate=False)
        self.processes.remove(fixsub)
      if self.GetParam('replace_met_mse'):
        for mod in p.inp.GetAll('model',typ=('partial','partial+substr'),has_atomtypes=True):
          if 'SE' in mod.GetAtomTypes():
            pdbcur=self.AddProg('pdbcur',propagate_out=False,propagate_inp=False)
            pdbcur.inp.Add(mod)
            pdbcur.SetKey('renresidue', "(MET) 'MSE'")
            pdbcur.SetKey('renatom', "SD[S] 'SE'")
            pdbcur.SetKey('renelement', "SE[S] 'SE'")
            pdbcur.outfilename['pdb'] = "mse_repl_"+p.nick+".pdb"
            pdbcur.Run()
            p.inp.model.remove(mod)
            p.inp.Add(pdbcur.out.Get('model'), propagate=False)
            self.programs.remove(pdbcur)


  def SetTargetDefault(self, proc):
    # exclude processes that never need/support targets first
    if proc.nick not in ['phas','faest','substrdet','handdet','dmfull','phcomb','ref','comb_phdmmb','mbref','refatompick']:
      return
    if proc.nick=='dmfull' and self.GetProg(supported=True):
      return
    # if overall target/experiment is set
    if self.IsParam('target') and not proc.IsParam('target'):
      # now set the default target for the current process
      if proc.nick in ['phas','faest','substrdet'] or self.GetParam('target') in ('SAD','RICE') or \
         (self.GetParam('target')=='SIRAS' and (proc.nick=='dmfull' or proc.nick=='phcomb')):
        proc.SetParam('target', self.GetParam('target'))
      else:
        proc.SetParam('target', 'MLHL')
    # if target/experiment is not set
    if not self.IsParam('target') and not proc.IsParam('target'):
      if proc.nick in ['ref','phcomb','mbref']:
        proc.SetParam('target', 'RICE')
      else:
        common.Error('No target specified for process of {0}'.format(proc.name))


  def PassToHandDet(self, handdet, ip):
    tobepassed=("phas","dmfull")
    for tbp in tobepassed:
      for pp in [ i for i in range(ip) if self.processes[i].nick==tbp ]:
        handdet.processes.append(self.processes[i])
        for o in self.run.processes[i].out.GetAll(stored_order=True):
          handdet.processes[-1].out.AddCopy(o)
      #for pp in [ self.run.processes[i] for i in range(ip) if self.run.processes[i].nick==tbp ]:
      #  handdet.processes.append(pp)

  def PassChosenHand(self, handdet_run, ip):
    for pp in [ self.processes[i] for i in range(ip+1,len(self.processes)) ]:
      # we invert any previously inputted models to subsequent steps
      # I would need to create createotherhand process before writing this.
      #for mod in pp.inp.GetAll('model'):
      #  handdet.CreateHand2Model(mod, pp.inp.Get('fsigf',xname=mod.xname).GetSpacegroup(handdet))
      #pp.inp.SetCopy()
      # we simply replace all mtz's - if something else is needed it should be added/adjusted here
      pp.inp.SetCopy(handdet_run.out.fsigf,propagate=False)

  def SubProcess_WriteXML_CreateRun(self,p,ip,emulate=False):
    # write the process input xml
    p_xml=os.path.join(self.rundir, p.nick+'.xml')
    crank_p=copy.copy(self)
    crank_p.processes=[p,]
#    if emulate:
#      p_xml = crank_p.Data2XML()
#    else:
#      common.WriteXML(crank_p, p_xml)
    common.WriteXML(crank_p, p_xml)
    # "undummify" by re-creating the process from the written xml
    crank_p=self.from_xml(p_xml, rundir=self.rundir)
    p_run = crank_p.processes[0]
    self.run.processes[ip] = p_run
    p_run.parent_process=self.run
    # initialize the 'actual' crank (meta)attributes not saved in xml
    if ip==0:
      self.run.prep=self
      for at in ('logfilehandle','loggraphfilehandle','opened_log','opened_loggraph','ccp4_version','rv_report','emulate') + \
                parse.parse(dummy=True).share_with_crank:
        if hasattr(self,at):
          setattr(self.run, at, getattr(self,at))
    p_run.SetRunDir(os.path.join(self.rundir,str(ip)+'-'+p.nick))
    return p_run

  def Initialize(self,p_run):
    # any other general process initialization
    #convert I's to F's in case no F was supplied
    for o in p_run.inp.GetAll('fsigf',col='i',typ=('plus','average')):
      if not p_run.inp.Get('fsigf',col='f',typ=o.GetType(),xname=o.GetCrystalName(),dname=o.GetDataName()):
#    if not p_run.inp.Get('fsigf',col='f') and p_run.inp.Get('fsigf',col='i'):
        orig_i = [copy.deepcopy(o),]
        if o.GetType()=='plus':
          orig_i.append(copy.deepcopy(p_run.inp.Get('fsigf',typ='minus',col='i',xname=o.GetCrystalName(),dname=o.GetDataName())))
        trun=p_run.AddProg('truncate',propagate_out=False)
        trun.inp.Set(orig_i)
        o_base = o.GetFileName(trim_path=True).rsplit('.',1)[0]
        trun.outfilename['mtz'] = o_base+'_'+trun.name+'.mtz'
        try:
          trun.Run(check_bin=True)
        except trun.ProgramRunError as e:
          sys.exc_clear()
          trun.out.ClearAll()
          ctrun=p_run.AddProg('ctruncate',propagate_out=False)
          ctrun.inp.Set(orig_i)
          ctrun.outfilename['mtz'] = o_base+'_'+ctrun.name+'.mtz'
          self.Info('{0} failed with error message: {1}'.format(trun.name,e))
          self.Info('Will try {0} instead.'.format(ctrun.name))
          ctrun.Run(check_bin=True)
          for o in ctrun.out.GetAll('fsigf'):
            p_run.inp.Add(o)
          p_run.programs.remove(ctrun)
        else:
          for o in trun.out.GetAll('fsigf'):
            p_run.inp.Add(o)
        p_run.programs.remove(trun)
    #convert F+/F- to F average (+FD) in case no F average was not supplied
    for o in p_run.inp.GetAll('fsigf',col='f',typ='plus'):
      if not p_run.inp.Get('fsigf',typ='average',col='f',xname=o.GetCrystalName(),dname=o.GetDataName()):
        mtzmadmod=p_run.AddProg('mtzMADmod',propagate_out=False)
        o_base = o.GetFileName(trim_path=True).rsplit('.',1)[0]
        mtzmadmod.outfilename['mtz'] = o_base+'_'+mtzmadmod.name+'.mtz'
        mtzmadmod.Run(check_bin=True)
        p_run.inp.Add( mtzmadmod.out.Get(filetype='mtz',typ='average'), ind=p_run.inp.fsigf.index(p_run.inp.Get('fsigf',typ='plus')), propagate=False )
        p_run.programs.remove(mtzmadmod)
    # setting solvent content - this is only done in this procedure so that the gui1 easily gets the value...
    # this could be removed in the future as gui1 is obsolete and a general check done in gui2
    if not self.inp.Get(has_solvent_content=True) and self.inp.Get('sequence') and self.inp.Get('fsigf',filetype='mtz'):
      matthews=p_run.AddProcess('matthews', propagate_out=False)
      matthews.no_reporting = True
      matthews.Run()
      p_run.processes.remove(matthews)

  def RunSubProcess(self,p_run):
    # register process with ccp4i2 if run from ccp4i2 and run the process
    if self.run.ccp4i2:
      ccp4i2crank.RegisterProcessToCCP4i2(self.run.ccp4i2, p_run)
      # run the current process
      error=False
      try:
        p_run.Run(from_ccp4i2=True)
      except:
        error=True
      # register process output with ccp4i2
      ccp4i2crank.RegisterOutputToCCP4i2(p_run,error)
    else:
      p_run.Run()

  def PassPrevStepsOutput(self,p,ip,emulate=False):
    # passes output containers from previous steps to input of this process p
      if ip>0 and not self.run.processes[ip-1].IsNonFalseVirtPar('no_output_to_next_step'):
        # pass all except the specified exceptions from the previous step
        for o in self.run.processes[ip-1].out.GetAll(stored_order=True):
          # we exclude any luzd output as that is generally not meant to be used by next steps
          # similarly, we use best rather than weighted mapcoef - any exceptions should be specified here
          if o.typ!='dluz' and o.typ!='weighted':
            p.inp.AddCopy(o,propagate=False)
      # pass shelxc output to shelxe - not needed anymore (for the shelxc-d-e pipeline at least)
      #if ip>1 and not self.run.processes[ip-2].IsNonFalseVirtPar('no_output_to_next_step') and \
      #   'shelxe' in [r.nick for r in p.programs] and 'shelxc' in [r.nick for r in self.run.processes[ip-2].programs]:
      #  for o in self.run.processes[ip-2].out.GetAll(stored_order=True):
      #    p.inp.AddCopy(o,propagate=False)
      # pass what was specified by the input
      if not emulate:
        p.inp.DelayedInit(ip,self)


  def RunPostprocess(self,restore=True,emulate=False,*args,**kwargs):
    if emulate:
      os.chdir(self.previous_cwd)
      try:
        os.remove(self.GetLogFileName())
        os.rmdir(self.rundir)
      except (OSError,):
        sys.exc_clear()
    else:
      outmtz=self.run.out.Get(filetype='mtz')
      if outmtz and hasattr(self,'hklout') and outmtz.GetFileName()!=self.hklout:
        shutil.copy(outmtz.GetFileName(), self.hklout)
      outpdb=self.run.out.Get(filetype='pdb')
      if outpdb and hasattr(self,'xyzout') and outpdb.GetFileName()!=self.xyzout:
        shutil.copy(outpdb.GetFileName(), self.xyzout)
      process.RunPostprocess(self,restore)
    if hasattr(self,'logout'):
      sys.stdout = self.stdout_save