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

class substrdet(process):
  name="substructure determination"
  short_name=name
  supported_progs=["crunch2","shelxd","prasa"]
  supported_params = {}
  supported_params['target'] = par( desc='Experiment/refinement target', typ=str, cap=True )
  supported_params['threshold_stop'] = par( desc='Stop when this threshold is reached (CCweak for shelx, FOM for crunch)', typ=(float,bool) )
  supported_params['catch_output'] = par( desc='Catch program output continously (always True if threshold_stop)', typ=bool )
  supported_params['expected_num_atoms'] = par( desc='Expected number of substructure atoms', typ=int )
  supported_params['num_trials'] = par( desc='(Maximal) Number of substructure detection trials', typ=(int,bool) )
  supported_params['patterson_seed'] = par( desc='Seed from Patterson (crunch2 only)', typ=bool )
  supported_params['high_res_cutoff'] = par( desc='High resolution cutoff', typ=(float,bool) )

  def Init(self):
    self.min_score=1000.
    self.score=-1000.
    self.score_adj=-1000.
    self.fom_unopt=-1000.
    self.trial=0
    self.trial_adj=0
    self.trial_temp=0
    self.stop_message=""
    self.stop_sent=False
    self.stats_table=[]
    self.check_solution=(0,0)
    self.resel_cycle=400
    self.resel_init_half=self.resel_cycle/2
    self.thr_up, self.thr_down = 0., 0.

  def TreatInOutPar(self, set_all_par=False):
    # this should be modified after crunch plugin becomes available
    # perhaps shelx will be default if available and crunch otherwise (?)
    if not self.GetProg(supported=True):
      if program.from_name('shelxd',None).CheckBinary(silent=True):
        self.AddProg('shelxd')
      else:
        self.AddProg('crunch2')
    # catch output continously (eg to stop early if certain stat is reached) - shelxd only as of now.
    if self.GetParam('catch_output') is False:
      self.Info('Early stop on score (cc/fom) threshold disabled.')
    else:
      self.GetProg(supported=True).interact_output=True
      if self.GetParam('threshold_stop') in (None,True):
        if self.GetProg(supported=True).nick=='shelxd':
          self.SetParam('threshold_stop',30.0)
        if self.GetProg(supported=True).nick=='crunch2':
          self.SetParam('threshold_stop',1.0)
    # we are changing the shelx/crunch2 program defaults here!
    # the original defaults (if any) will be used when setting threshold_stop, num_trials, patterson_seed to False.
    if self.GetParam('num_trials') in (None,True): 
      if self.GetProg(supported=True).nick=='shelxd':
        self.SetParam('num_trials',500)
      if self.GetProg(supported=True).nick=='crunch2':
        self.SetParam('num_trials',20)
#      if self.GetProg(supported=True).nick=='prasa':
#        self.SetParam('num_trials',100)
    if self.GetParam('target')=='SAD' and self.GetParam('high_res_cutoff') in (None,True): 
      if self.GetProg(supported=True).nick=='crunch2':
        self.SetParam('high_res_cutoff', self.inp.Get(filetype='mtz').GetResolution(self)+0.5)
      elif self.GetProg(supported=True).nick=='prasa':
        self.SetParam('high_res_cutoff', self.EstimateCutOff())
    if self.GetParam('patterson_seed') is None and self.GetProg(supported=True).nick=='crunch2':
      self.SetParam('patterson_seed')
    process.TreatInOutPar(self,set_all_par)

  def EstimateCutOff(self):
    cutoff = self.inp.Get(filetype='mtz').GetResolution(self)+0.5
    # only do this if ecalc was used in the previous step - it does not work for E's from shelxc...
    cp = self.GetCrankParent()
    if cp and self in cp.processes and cp.processes.index(self)-1>=0 and \
       cp.processes[cp.processes.index(self)-1].nick=='faest' and \
       'ecalc' in [p.nick for p in cp.processes[cp.processes.index(self)-1].programs]:
      fa_obj = self.inp.Get('fsigf',typ='ea')
      # this can happen if in emulation mode (to get the default cutoff from GUI)
      # if too slow then can be disabled
      if not fa_obj:
        ecalc=self.AddProg('ecalc',propagate_out=False)
        try:
          ecalc.Run()
        except:
          sys.exc_clear()
        else:
          fa_obj=ecalc.out.Get('fsigf',typ='ea')
      if fa_obj:
        sft=self.AddProg('sftools')
        sft.runname=sft.name+'_compl'
        sft.SetKey('read',fa_obj.GetFileName('mtz'))
        sft.SetKey('select', ('col', fa_obj.GetLabel('f'), 'present'))
        sft.SetKey('select', ('col', fa_obj.GetLabel('sigf'), '>', '0'))
        sft.SetKey('purge\nY')
        sft.SetKey('checkhkl')
        sft.SetKey('complete', ('table', 'shells', '22', 'col', fa_obj.GetLabel('f')))
        sft.SetKey('quit\nY')
        sft.Run()
        filename=os.path.join(sft.rundir,'COMPLETE.TAB')
        filename2=os.path.join(sft.rundir,'COMPLETE2.TAB')
        shutil.move(filename, filename2)
        sft.ClearAllKeys()
        sft.runname=sft.name+'_sig'
        sft.SetKey('read',fa_obj.GetFileName('mtz'))
        sft.SetKey('select', ('col', fa_obj.GetLabel('f'), 'present'))
        sft.SetKey('select', ('col', fa_obj.GetLabel('sigf'), '>', '0'))
        sft.SetKey('purge\nY')
        sft.SetKey('calc', ('R','col','_delpm_','=','col','"'+fa_obj.GetLabel('f')+'"','abs'))
        sft.SetKey('calc', ('R','col','_delsigratio_','=','col','_delpm_','col','"'+fa_obj.GetLabel('sigf')+'"','/') )
#        sft.SetKey('calc', ('R','col','_sigpratio_','=','col','"'+fp_obj.GetLabel('f')+'"','col','"'+fp_obj.GetLabel('sigf')+'"','/'))
#        sft.SetKey('calc', ('R','col','_sigmratio_','=','col','"'+fm_obj.GetLabel('f')+'"','col','"'+fm_obj.GetLabel('sigf')+'"','/'))
#        sft.SetKey('absent', ('col', fa_obj.GetLabel('f'), 'if','col','_sigpratio_','<','4'))
#        sft.SetKey('absent', ('col', fa_obj.GetLabel('f'), 'if','col','_sigmratio_','<','4'))
#        sft.SetKey('plot', ('col', '_delsigratio_', 'resol'))
        sft.SetKey('absent', ('col', fa_obj.GetLabel('f'), 'if','col','_delsigratio_','<','2.75'))
        sft.SetKey('complete', ('table', 'shells', '22', 'col', fa_obj.GetLabel('f')))
        sft.SetKey('quit\nY')
        sft.Run()
        self.programs.remove(sft)
        if os.path.isfile(filename) and os.path.isfile(filename2):
          with open(filename) as f:
            with open(filename2) as g:
              resbin = []
              for line in f:
                resbin.append([float(v) for v in line.split()])
              for i,line in enumerate(g):
                if float(line.split()[-1])>0.01:
                  resbin[i][-1] /= (float(line.split()[-1])/100.0)
                else:
                  resbin[i][-1] = -1
          i=0
          for r in resbin[:]:
            if r[-1]<0:
              resbin.pop(i)
            else:
              i+=1
          # smooth
          s=[]
          for i,r in enumerate(resbin):
            print r[-1]
            smooth, tot_smooth = 0., 0.
            for j in range(max(0,i-2),min(i+3,len(resbin))):
              smooth += (1.0-0.45*abs(i-j))*resbin[j][-1]
              tot_smooth += 1.0-0.45*abs(i-j)
            s.append(smooth/tot_smooth)
          for i,r in enumerate(resbin):
            r[-1] = s[i]
          # 4.5A derived threshold
          thres,tot_w,part = 0., 0., 1.
          for r in resbin:
            if r[1]<4.5:
              part = (r[0]-4.5)/(r[0]-r[1])
            thres += part*r[2]*r[3]
            tot_w += part*r[2]
            if part<1.0:
              break
          # add bonus for very low res.
          for r in resbin:
            r[-1] *= max(1.0,2*(r[0]+r[1])-13.0)
            print r[-1]
          thres = thres/tot_w/6.
          print thres
          cutoff,cutoff0 = resbin[0][1],resbin[0][0]
          for r in resbin:
            # two subsequent bins above the threshold
            #if r[-1]>thres: 
            #  if abs(r[0]-cutoff0)<1e-5:
            #    cutoff = r[1]
            #  cutoff0 = r[1]
            # two subsequent bins below the threshold
            if r[-1]>thres:
              cutoff = r[1]
            else:
              if abs(r[0]-cutoff0)<1e-5:
                break
              cutoff0 = r[1]
          print cutoff
        else:
          common.Warning('Unexpected sftools problem when determining resolution cutoff.')
      else:
        common.Warning('Resolution cutoff could not be determined.')
    return cutoff

  def RunPreprocess(self,*args,**kwargs):
    process.RunPreprocess(self,*args,**kwargs)
    # run pmf for crunch2 (no special process defined as of now)
    if self.GetParam('patterson_seed'):
      pmf=self.AddProg('pmf',propagate_out=False)
      pmf.SetKey('TRY',(1,150))
      pmf.Run()
      self.GetProg(supported=True).inp.Add(pmf.out.Get('model'))
      self.programs.remove(pmf)

  def RunPostprocess(self,restore=True,*args,**kwargs):
    # for crunch2 we need to adjust the filename!
    if self.GetProg(supported=True).nick=='crunch2':
      if self.trial:
        outmod=self.GetProg(supported=True).out.Get('model')
        #outmod.SetFile( outmod.GetFileName().replace('trial1a','trial'+str(self.trial)+'a'), 'pdb' )
        outmod.SetFile( self.GetProg(supported=True).GetTrialPdb(self.trial), 'pdb' )
      else:
        # this should be removed and replaced by grepping the logfile to get the trial.
        common.Error('Could not pick model.')
    # taking the solution with largest adjusted score for prasa
    if self.GetProg(supported=True).nick=='prasa' and self.trial_adj:
      outmod=self.GetProg(supported=True).out.Get('model')
      outmod.SetFile( self.GetProg(supported=True).GetTrialPdb(self.trial_adj), 'pdb' )
      self.Info('Solution from trial {0} (adjusted score {1}) taken.'.format(self.trial_adj,self.score_adj))
    # "fix" the substructure pdb file
    fixsubstr=self.AddProcess('fixsubstrpdb')
    fixsubstr.inp.Set(self.GetProg(supported=True).out.model)
    fixsubstr.Run()
    # print the final loggraph
    self.PrintActualLogGraph(finished=True)
    process.RunPostprocess(self,restore,*args,**kwargs)

  def PrintActualLogGraph(self, finished=False):
    # print loggraph info for shelxd and crunch2
    prog=self.GetProg(supported=True)
    if prog.nick=='prasa':
      return
    if finished and not self.stop_message:
      self.stop_message='Maximum number of trials reached.'
    if self.stats_table:
      if prog.nick=='shelxd':
        sc,sc2=('CCweak','CC'),"correlation coefficients"
      elif prog.nick=='crunch2':
        sc,sc2=('FOM','FOM_unoptimized'),"figures of merit"
      if self.opened_loggraph:
        self.GetLogGraphHandle().seek(0,0)
        self.LGInfo(self.GetCCP4Header())
        self.LGInfo(self.GetLogGraphPrograms())
        self.LGInfo('\n$TABLE : {0} {1} versus trial:\n$SCATTER'.format(prog.name,sc[0]))
        self.LGInfo(' : {0} per trial :A:1,2,3:\n$$'.format(' and '.join(sc)))
        self.LGInfo('Trial       {0}$$ $$'.format('        '.join(sc)))
        for trial,ccw,cc,fom,fom_unopt in self.stats_table:
          if prog.nick=='shelxd':
            self.LGInfo('{0:5}  {1:10.2f} {2:10.2f}'.format(trial,ccw,cc))
          elif prog.nick=='crunch2':
            self.LGInfo('{0:5}  {1:10.3f} {2:10.3f}'.format(trial,fom,fom_unopt))
        self.LGInfo('$$\n\n')
        if finished:
          self.LGInfo('$TEXT:Result: $$ Final result $$')
          self.LGInfo(self.stop_message)
          self.LGInfo('Maximum {0}: {1} occurring at Trial {2}.\n$$'.format(sc[0], self.score, self.trial))
      if self.rv_report is not None:
        if not finished:
          trial,ccw,cc,fom,fom_unopt = self.stats_table[-1]
          # for shelxd we restrict the number of points as rvapi is slow or gets stuck dealing with many hundreds of points
          # we do this by re-selecting the points each resel_cycle cycles, thus keeping the max. displayed below roughly 500 (per line)
          tr_half, cc_ccw = trial/self.resel_init_half, cc*ccw
          if not hasattr(self,'rv_plot') or (trial==self.resel_cycle and prog.nick=='shelxd'):
            if trial==self.resel_cycle and prog.nick=='shelxd':
              self.rv_plot.parent.parent.Remove()
            self.rv_plot = self.rv_report.Plot( '{1} and {2} vs trial'.format(prog.name,sc[0],sc[1]), \
                                                "Trial", "{0} and {1}".format(sc[0],sc[1]), block="{0} {1}".format(prog.name,sc2) )
            self.rv_stat1 = self.rv_plot.PlotLine( ["x"], [sc[0]])
            self.rv_stat2 = self.rv_plot.PlotLine( ["x"], [sc[1]])
            if trial==self.resel_cycle and prog.nick=='shelxd':
              # new selection is made and plotted.
              stat_all = [s[1]*s[2] for s in self.stats_table]
              self.thr_up, self.thr_down = heapq.nlargest(20,stat_all)[-1], heapq.nsmallest(20,stat_all)[-1]
              stat_sel = [s for i,s in enumerate(self.stats_table) if i%tr_half==0 or s[1]*s[2]>=self.thr_up or s[1]*s[2]<self.thr_down]
              stat_sel = zip(*stat_sel)
              self.rv_stat1.Data(stat_sel[0],stat_sel[1]), self.rv_stat2.Data(stat_sel[0],stat_sel[2])
              self.resel_cycle*=2
          if prog.nick=='shelxd' and (trial<self.resel_init_half or trial%tr_half==0 or cc_ccw>=self.thr_up or cc_ccw<self.thr_down):
            self.rv_stat1.Data(trial,ccw), self.rv_stat2.Data(trial,cc)
          elif prog.nick=='crunch2':
            self.rv_stat1.Data(trial,fom), self.rv_stat2.Data(trial,fom_unopt)
          crvapi.Flush()
        else:  # if finished
          self.rv_report.Text("<i><small>Result:</small></i>") # title
          self.rv_report.Text('&emsp;'+self.stop_message)
          self.rv_report.Text('&emsp;Maximum {0}: {1} occurring at Trial {2}.'.format(sc[0], self.score, self.trial))
          self.rv_report.Text("<BR><i><small>Output:</small></i>")
          self.rv_report.DataFile(self.out.Get('model').GetFileName(), "xyz", "<small>Substructure</small>", flush=True)


  def UpdateInfo(self,line,popen=None,empty=False):
    # in shelx all comes in one line; in crunch not and we only have (optimized) fom outputted
    score, ccw, cc, fom_unopt = None, None, None, None
    prog=self.GetProg(supported=True)
    trial=prog.GetStatGrep('try', from_str=line)
    if prog.nick=='prasa':
      patfom=0.
    else:
      patfom=prog.GetStatGrep('patfom', from_str=line)
    # shelx part
    if prog.nick=='shelxd':
      ccw=prog.GetStatGrep('ccweak', from_str=line)
      cc=prog.GetStatGrep('cc', from_str=line)
      score=ccw
      if trial and score and score>self.score:
        stats_info=['corr.coef. {0}'.format(cc), 'corr.coef. weak {0}'.format(ccw), 'Patterson FOM {0}'.format(patfom)]
        stat_str='Stats'  if trial==1  else 'Best stats improved'
        self.Info('{0} in trial {1}: {2}'.format(stat_str,trial,', '.join(stats_info)))
    # prasa part
    if prog.nick=='prasa':
      cc=prog.GetStatGrep('cc', from_str=line)
      score=cc
      if trial and score and score>self.score:
        stats_info=['corr.coef. {0}'.format(cc),]
        stat_str='Stats'  if trial==1  else 'Best stats improved'
        self.Info('{0} in trial {1}: {2}'.format(stat_str,trial,', '.join(stats_info)))
    # crunch part
    if prog.nick=='crunch2':
      fom_unopt=prog.GetStatGrep('fom_unopt', from_str=line)
      if trial is not None:
        self.trial_temp=trial
        self.fom_unopt=-1000.
      elif self.trial_temp and patfom is not None:
        trial=self.trial_temp
        score=patfom
        self.Info('FOM after optimization in trial {0}: {1}'.format(trial,patfom))
      elif self.trial_temp and fom_unopt is not None and fom_unopt>self.fom_unopt:
        self.fom_unopt=fom_unopt
      fom_unopt=self.fom_unopt
    # save the best score and trial; save any stats to the table of stats; check thresholds; update loggraph
    if score and trial and not self.stop_message:
      if self.GetParam('threshold_stop') and score>=self.GetParam('threshold_stop') and \
         (prog.nick!='shelxd' or trial>=10):
        self.stop_message='Specified score threshold ({0}) reached.'.format(self.GetParam('threshold_stop'))
      elif self.GetParam('threshold_stop') and prog.nick=='crunch2':
        self.CheckExtraStopConditions(score,trial,prog,popen)
#      if prog.nick=='prasa' and 'candidate' in line:
      if prog.nick=='prasa' and score>self.score and score>15:
        self.check_solution=(trial,score)
      if score>self.score:
        self.score=score
        self.trial=trial
      if score<self.min_score:
        self.min_score=score
      self.stats_table.append( (trial,ccw,cc,patfom,fom_unopt) )
      self.PrintActualLogGraph()
    # let's stop early; with crunch2, we need to wait till the trial is finished
    if self.stop_message and not self.stop_sent and \
       (prog.nick!='crunch2' or prog.GetStatGrep('cyc_in_try', self.trial, from_str=line)):
      self.Info(' Stopping: '+self.stop_message)
      prog.Stop(popen)
      self.stop_sent = True

  def CheckExtraStopConditions(self,score,trial,prog,popen):
    # relative score threshold
    if score and score>0.3 and score>2.25*self.min_score:
      self.stop_message='Relative score threshold reached.'
    # running bp3 in check mode to preliminary stop on bp3 Luzzati threshold
    elif score and score>0.3 and self.trial_temp%3==1:
      fixsubstr=self.AddProcess('fixsubstrpdb',propagate_out=False)
      mod = fixsubstr.inp.AddCopy(prog.out.Get('model'))
      #mod.SetFile( self.out.Get('model').GetFileName().replace('1',str(self.trial_temp)), 'pdb' )
      mod.SetFile( prog.GetTrialPdb(self.trial_temp), 'pdb' )
      fixsubstr.Run()
      phas=self.AddProcess('phas',propagate_out=False)
      phas.inp.Add(fixsubstr.out.Get('model'))
      phas.SetParam('target', self.GetParam('target'))
      phas.AddProg('bp3').SetKey('check')
      try:
        phas.Run()
      except (phas.ProgramRunError, common.CrankError) as e:
        common.Warning('Could not test solution by quick phasing: {0}'.format(e.message))
      else:
        if phas.GetProg('bp3').GetStat('luzzati')>0.7:
          self.stop_message='BP3 Luzzati parameter threshold reached.'
      self.processes.remove(phas)
      self.processes.remove(fixsubstr)

  def CheckExtraStopConditionsPrasa(self,score,trial,prog,popen):
    self.Info('Fast testing of solution from trial {0} with CC of {1} follows.'.format(trial,score))
    trial_sol=self.GetProg(supported=True).out.GetCopy('model')
    trial_sol.SetFile( self.GetProg(supported=True).GetTrialPdb(trial), 'pdb' )
    if not os.path.isfile( trial_sol.GetFileName('pdb') ):
      self.Info("Skipping check of trial solution {0}.".format(trial) )
      return
    fixsubstr=self.AddProcess('fixsubstrpdb')
#    fixsubstr.inp.Set(self.GetProg(supported=True).out.Get('model'))
    fixsubstr.inp.Set(trial_sol)
    fixsubstr.Run()
#    if not fixsubstr.inp.Get('model',has_atomtypes=True):
#      fixsubstr.inp.Get('model').SetAtomTypes(atomtypes,atomtype1)
#    if fixsubstr.inp.Get('model').GetType()=='unknown':
#      fixsubstr.inp.Get('model').SetType('substr')
#    fixsubstr.Run()
#    if score<20.0:
#      return
#    afm=self.AddProcess('atomsfrommap')
#    afm.inp.Add(prog.out.Get('mapcoef',filetype='map'))
#    afm.SetParam('rms_threshold', 5.0)
#    if self.GetParam('expected_num_atoms'):
#      afm.SetParam('max_new_atoms', self.inp.GetParam('expected_num_atoms')+2)
#    elif self.inp.Get(has_num_atoms=True):
#      afm.SetParam('max_new_atoms', self.inp.Get(has_num_atoms=True).exp_num_atoms+2)
#    afm.Run()
    # this should be done somewhere in afm and also the mtz from shelxc should have x/dname!
#    afm.out.Get('model').xname=self.inp.Get('model',typ='substr').xname
    handdet=self.AddProcess('handdet',propagate_out=False)
    handdet.SetParam('target', self.GetParam('target'))
#    handdet.inp.Add(afm.out.Get('model'))
    handdet.inp.Add(fixsubstr.out.Get('model'))
    handdet.TreatInOutPar()
    handdet.phas.SetParam('cycles', 3)
    handdet.phas2.SetParam('cycles', 0)
    handdet.dmf[0].SetParam('dmcyc', 0)
    handdet.dmf[1].SetParam('dmcyc', 0)
    handdet.Run()
    score_adj=score*(handdet.phas.fom+handdet.phas2.fom)/2*max(0.017,abs(max(0.,handdet.cld[1])-max(0.,handdet.cld[0])))
    self.Info('Combined score of trial {0} is {1}'.format(trial,score_adj))
#    if (handdet.cld[0]*score_adj>0.2 or handdet.cld[1]*score_adj>0.2):
    if score_adj>0.3:
      self.stop_message='CLD*FOM*CC parameter threshold reached.'
      self.score_adj = score_adj
      self.trial_adj = trial
    elif score_adj>self.score_adj and not self.stop_message:
      self.score_adj = score_adj
      self.trial_adj = trial
#    self.out.Set(afm.out.Get('model'))
    self.out.Set(trial_sol)
    self.processes.remove(handdet)
#    self.processes.remove(afm)
    self.processes.remove(fixsubstr)