
from __future__ import division
from libtbx import group_args
import libtbx.callbacks # import dependency
import re

phaser_messages = ["phaser_start",
                   "phaser_end",
                   "MODE",
                   "summary",
                   "current_best_solution",
                   "loggraph",
                   "start_progress_bar",
                   "increment_progress_bar",
                   "end_progress_bar",
                   "warn"]
cached_messages = ["phaser_start", "phaser_end", "MODE", "summary",
  "current_best_solution", "loggraph", "warn"]

class manager (object) :
  def __init__ (self) :
    self._summary_buffer = ""

  def startProgressBar (self, label, size) :
    #assert size > 0
    libtbx.call_back(message="start_progress_bar",
      data=group_args(label=label, size=size),
      cached=False)

  def incrementProgressBar (self) :
    # TODO: buffering?
    libtbx.call_back(message="increment_progress_bar",
      data=group_args(chunk=1),
      cached=False)

  def endProgressBar (self) :
    libtbx.call_back(message="end_progress_bar", data=None, cached=False)

  def warn (self, message) :
    libtbx.call_back(message="warn", data=message, accumulate=True)

  def loggraph (self, title, data) :
    if data is None :
      return
    from iotbx import data_plots
    table_data = data_plots.table_data(title=title)
    try :
      table_data.import_loggraph(data)
    except TypeError : # FIXME loggraph with no rows causes TypeError
      return
    else :
      final_data = group_args(title=title,graph=table_data)
      libtbx.call_back(message="loggraph",
        data=final_data,
        accumulate=True)

  def call_back (self, message, data) :
    if message == "summary" :
      if data.startswith("**********") :
        self.flush_summary()
      self._summary_buffer += data
    elif message == "current best solution" :
      try :
        scores = extract_annotation_scores(data)
      except ValueError :
        pass
      else :
        libtbx.call_back(message="current_best_solution", data=scores,
          accumulate=True)
    else :
      libtbx.call_back(message=message, data=data)

  def flush_summary (self) :
    if self._summary_buffer != "" :
      libtbx.call_back(message="summary", data=self._summary_buffer,
        accumulate=True)
    self._summary_buffer = ""

def extract_annotation_scores (annotation) :
  lines = annotation.splitlines()
  solution = group_args(rfz=[], tfz=[], pak=[], llg=[], rlg=[])
  for line in lines :
    if line.startswith("SOLU SET") :
      fields = re.sub("SOLU SET ", "", line).split()
      i = 0
      last_field = None
      while i < len(fields) :
        stat, value = fields[i].split("=")
        if stat == "llg" and last_field == "llg" :
          stat = "rlg"
        getattr(solution, stat.lower()).append(float(value))
        i += 1
        last_field = stat
  return solution
