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

class phdmmb(process):
  name="phasing, density modification and model building"
  short_name="model tracing with DM"
  supported_progs = ['shelxe']
  supported_params={}
  supported_params['catch_output'] = par( desc='Catch program output continously (True by default)', typ=bool )
  supported_params['bigcyc'] = par( desc='Number of model building cycles', typ=int )
  supported_params['dmcyc'] = par( desc='Number of density modification cycles', typ=int )
  supported_params['solvent_content'] = par( desc='(Expected) solvent fraction of crystal', typ=float, share=True )
  supported_params['handdet'] = par( desc='Run with both hands (True by default)', typ=bool )

  def Init(self):
    self.stats_table={}
    self.stats_table['1'], self.stats_table['2'] = {}, {}
    self.stat_nicks = ('cc','contrast','res_built')
    self.stat_names = ("Correlation coefficient", "Contrast", "Residues")

  def TreatInOutPar(self, set_all_par=False):
    shelxe=self.GetOrAddProg('shelxe')
    if self.GetParam('catch_output') is False:
      self.Info('Early stop and continuous reporting are disabled.')
    else:
      self.GetProg(supported=True).interact_output=True
    if self.GetVirtPar('handdet') is None:
      self.SetVirtPar('handdet', True)

  def RunHand(self,hand_str,lock=None):
    if lock:
      lock.acquire()
    prog = getattr(self,'prog'+hand_str)
    prog.hand=hand_str
    self.Info(" Running shelxe with hand {0} ".format(hand_str))
    prog.Run( rundir=os.path.join(self.rundir,'hand'+hand_str), lock=lock )
    #self.stats_table.append( [prog.GetStat("cc"),prog.GetStat("contrast"),prog.GetStat("res_built")] )
    #self.PrintActualLogGraph()

  def RunBody(self,*args,**kwargs):
    if self.GetProg('shelxe'):
      self.prog1 = self.GetProg('shelxe')
      if self.GetVirtPar('handdet'):
        self.prog2 = self.AddProgCopy(self.prog1, deeper_copy=True)
        self.prog2.SetArg('i')
      if self.GetVirtPar('handdet'):
        lock=threading.Lock()
        run1 = threading.Thread(target=self.RunHand, args=('1',lock))
        run1.start()
        self.RunHand('2',lock)
        run1.join()
      else:
        self.RunHand('1')

  def UpdateInfo(self,line,prog):
    for stat_nick in self.stat_nicks:
      sv=prog.GetStatGrep(stat_nick, from_str=line)
      if sv:
        if not stat_nick in self.stats_table[prog.hand]:
          self.stats_table[prog.hand][stat_nick]=[]
        self.stats_table[prog.hand][stat_nick].append(sv)
        if stat_nick=='cc':
          self.Info('CC after {0}. model building with hand {1} is {2}'.format(
                      len(self.stats_table[prog.hand][stat_nick]),prog.hand,sv) )
        self.PrintActualLogGraph(prog.hand, stat_nick)

  def PrintActualLogGraph(self, hand='1', stat_nick=None, finished=False):
    prog=self.GetProg(supported=True)
    if finished:
      #result1="Maximum contrast for hand 1 is {0} and for the other hand is {1}".format(max(self.stats_table[0][0]),max(self.stats_table[1][0]))
      #result2="Maximum number of residues built for hand 1 is {0} and for the other hand {1}".format(max(self.stats_table[0][2]),max(self.stats_table[1][2]))
      if not self.prog1.GetStat('untraceable'):
        result1="Maximum correlation coefficient for hand 1 is {0};".format(self.prog1.GetStat('best_cc'))
      else:
        result1="Unable to trace map for hand 1;"
      if self.GetVirtPar('handdet'):
        if not self.prog2.GetStat('untraceable'):
          result1+=" max. correlation coef. for the other hand is {0}. Hand {1} is chosen.".format(self.prog2.GetStat('best_cc'),self.chosen_hand)
        else:
          result1+=" unable to trace map for hand 2."
          if self.chosen_hand:  result1+=" Hand 1 is chosen."
      if not self.chosen_hand:
        result2="Tracing unsuccessful."
      else:
        result2="Max. number of residues built (for the chosen hand) is {0}".format(max(self.stats_table[self.chosen_hand]['res_built']))
    if self.opened_loggraph:
      self.GetLogGraphHandle().seek(0,0)
      self.LGInfo(self.GetCCP4Header())
      self.LGInfo(self.GetLogGraphPrograms())
      if self.stats_table:
        for si,stat in enumerate(self.stat_nicks):
          self.LGInfo('\n $TABLE : {0} of each hand per cycle:'.format(self.stat_names[si]))
          if stat in self.stats_table['1'] and stat in self.stats_table['2']:
            self.LGInfo('$GRAPHS    : {0} of each hand per cycle :A:1,2,3:\n$$'.format(self.stat_names[si]))
            self.LGInfo('Cycle   Hand1   Hand2 $$ $$')
            for i in range(min(len(self.stats_table['1'][stat]),len(self.stats_table['2'][stat]))):
              self.LGInfo('{0:3}  {1:8.3f} {2:8.3f}'.format(i+1, self.stats_table['1'][stat][i], 
                                                                 self.stats_table['2'][stat][i]))
          elif self.stats_table['1'] and stat in self.stats_table['1']:
            self.LGInfo('$GRAPHS    : {0} of each hand per cycle :A:1,2:\n$$'.format(self.stat_names[si]))
            self.LGInfo('Cycle   Hand1 $$ $$')
            for i,sh1 in enumerate(self.stats_table['1'][stat]):
              self.LGInfo('{0:3}  {1:8.3f}'.format(i+1, sh1))
          self.LGInfo('$$\n')
        if finished:
          self.LGInfo('\n$TEXT:Result: $$ Correlation coefficient results $$')
          self.LGInfo(result1+'\n$$')
          self.LGInfo('\n$TEXT:Result: $$ Model building results $$')
          self.LGInfo(result2+'\n$$')
    if self.rv_report is not None: 
      if 'cc' in self.stats_table[hand]:
        if not hasattr(self,'rv_plot_cc'):
          self.rv_plot_cc = self.rv_report.Plot( 'Correlation coef. vs building cycle', "Cycle", "Correlation coefficient", block="CC vs cycle")
          self.rv_plot_res_built = self.rv_plot_cc.parent.parent.Plot( 'Residues vs building cycle', "Cycle", "Backbone residues built", block="Residues vs cycle" )
          self.rv_cc, self.rv_res_built = {}, {}
        if stat_nick and hasattr(self, 'rv_'+stat_nick):
          rv_line_stat, rv_plot_stat = getattr(self, 'rv_'+stat_nick), getattr(self, 'rv_plot_'+stat_nick)
          if not rv_line_stat:
            rv_line_stat['1'] = rv_plot_stat.PlotLine( ["x"], ["hand1"] )
            if self.GetVirtPar('handdet'):
              rv_line_stat['2'] = rv_plot_stat.PlotLine( ["x"], ["hand2"] )
          rv_line_stat[hand].Data( len(self.stats_table[hand][stat_nick]),
                                       self.stats_table[hand][stat_nick][-1], x_type=int, flush=True )
      if finished:
        self.rv_report.Text("<i><small>Result:</small></i>")
        self.rv_report.Text('&emsp;'+result1)
        self.rv_report.Text('&emsp;'+result2, flush=True)
        if self.chosen_hand:
          self.rv_report.Text("<BR><i><small>Output:</small></i>")
          self.rv_files1 = self.rv_report.DataFile(self.out.Get('model').GetFileName(), "xyz", "<small>Built model (chosen hand)</small>")
          if self.GetVirtPar('handdet'):
            self.rv_files2 = self.rv_report.DataFile(getattr(self,'prog'+str(1+abs(int(self.chosen_hand)-2))).out.Get('model',typ='partial').GetFileName(), "xyz", "<small>Built model (other hand)</small>")
        #self.rv_files.DataFile(self.out.Get('mapcoef',typ='combined').GetFileName(), "hkl:map", flush=True)


  def RunPostprocess(self,restore=True,*args,**kwargs):
    self.chosen_hand = None  if self.prog1.GetStat('untraceable')  else '1'
    if self.GetVirtPar('handdet'):
      if not self.prog2.GetStat('untraceable') and (self.prog1.GetStat('untraceable') or  \
                              self.prog1.GetStat('best_cc')<self.prog2.GetStat('best_cc') ):
        self.chosen_hand = '2'
      if self.chosen_hand:
        self.out.Set(getattr(self,'prog'+self.chosen_hand).out.Get('model',typ='substr'))
        self.out.Set(getattr(self,'prog'+self.chosen_hand).out.Get('mapcoef',filetype='phs',typ='best'))
        if getattr(self,'prog'+self.chosen_hand).out.Get('model',typ='partial'):
          self.out.Add(getattr(self,'prog'+self.chosen_hand).out.Get('model',typ='partial'))
    self.PrintActualLogGraph( finished=True )
    if not self.chosen_hand:
      common.Warning('Unable to determine the hand and trace the model. Stopping.')
      # not returning an error at the moment due to gui's - might be changed later
      sys.exit(0)
    # shelxe sometimes leaves lines with only zeros and asterixes in the phs file - removing these here
    if self.GetProg('shelxe'):
      phs=self.GetProg('shelxe').out.Get('mapcoef', filetype='phs')
      if phs:
        lines, removed = [], False
        with open(phs.GetFileName('phs')) as f:
          for line in f:
            if '***' in line:
              removed=True
            else:
              lines.append(line)
        if removed:
          newfile = phs.GetFileName('phs')+'_fixed.phs'
          with open(newfile,'w') as g:
            for line in lines:
              g.write(line)
          phs_new=self.out.AddCopy(phs)
          self.out.SetFileToChild(phs_new,newfile,filetype='phs')
    process.RunPostprocess(self,restore,*args,**kwargs)