# XXX: this code is required by the GUI for configuration - everything else
# is in driver.py to keep GUI startup time to a minimum

# TODO expand finish_job callback, include output files

from __future__ import division
import iotbx.phil
from libtbx import object_oriented_patterns as oop
from libtbx.phil import interface as phil_interface
from libtbx.utils import Sorry
import libtbx.load_env
import os
import re
import sys

master_phil_ = master_params_ = None
def master_phil (validate=False) :
  global master_phil_, master_params_
  if (master_phil_ is None) :
    master_phil_ = iotbx.phil.read_default(__file__)
    master_params_ = master_phil
    defaults_file = libtbx.env.under_build("include/phaser_defaults.params")
    if os.path.isfile(defaults_file) :
      defaults_phil = iotbx.phil.parse(file_name=defaults_file)
      master_phil_, unused = master_phil_.fetch(source=defaults_phil,
        track_unused_definitions=True)
      if (validate) and (len(unused) > 0) :
        for param in unused :
          sys.stderr.write("  WARNING: unknown parameter %s\n" % param.path)
      master_params_ = master_phil
  return master_phil_

def master_params () :
  return master_phil()

class _index_phaser (oop.injector, phil_interface.index) :
  def get_phaser_modes (self, scope_name) :
    if not hasattr(self, "_styles") :
      self._styles = {}
    if not scope_name in self._styles :
      phil_scope = self.get_scope_by_name(scope_name)
      if isinstance(phil_scope, list) :
        phil_scope = phil_scope[0]
      self._styles[scope_name] = phaser_style(phil_scope)
    style = self._styles[scope_name]
    if style.mode is None :
      return []
    else :
      return style.mode

  def get_mode_keywords (self, mode) :
    if not hasattr(self, "_keywords_by_scope") :
      self._keywords_by_mode = {}
    if not mode in self._keywords_by_mode :
      current_keywords = []
      keyword_phil = self.get_scope_by_name("phaser.keywords")
      for object in keyword_phil.objects :
        if object.style is not None :
          object_styles = object.style.split()
          if "phaser:noindex" in object_styles :
            continue
        scope_name = object.full_path()
        keyword_modes = self.get_phaser_modes(scope_name)
        if match_modes(mode, keyword_modes) :
          current_keywords.append(scope_name)
      self._keywords_by_mode[mode] = current_keywords
    return self._keywords_by_mode[mode]

  def get_ensemble_ids (self) :
    ensemble_phil = self.get_scope_by_name("phaser.ensemble")
    if isinstance(ensemble_phil, list) :
      ensembles = [ ep.extract() for ep in ensemble_phil ]
    else :
      ensembles = [ ensemble_phil.extract() ]
    model_ids = []
    for ense in ensembles :
      if ense.model_id is not None :
        model_ids.append(ense.model_id)
    return model_ids

class phaser_style (object) :
  def __init__ (self, phil_object) :
    if phil_object.style is not None :
      for style_word in phil_object.style.split() :
        if style_word.startswith("phaser:") :
          words = style_word.split(":")
          if len(words) == 1 :
            continue
          elif len(words) == 2 :
            setattr(self, words[1], True)
          else :
            setattr(self, words[1], words[2].split(","))

  def __getattr__ (self, name) :
    return None

def match_modes (current_mode, keyword_modes, inherit_mode=False) :
  # XXX: if no modes defined, inherit from parent
  if keyword_modes is None :
    return inherit_mode
  elif current_mode in keyword_modes :
    return True
  else :
    for valid_mode in keyword_modes :
      if valid_mode == '*' :
        return True
      elif valid_mode.endswith('*') :
        if current_mode.startswith(valid_mode[:-1]) :
          return True
  return False

def check_for_model_identity_remark (pdb_in) :
  from phaser import tbx_utils
  found_remarks = 0
  for r in pdb_in.remark_section() :
    try:
      (mid,identity) = tbx_utils.parse_phaser_model_identity_remark(text=r)
    except ValueError:
      continue
    else :
      found_remarks += 1
  return found_remarks

def validate_params (params, index=None) :
  if index is None :
    working_phil = master_phil.fetch()
    index = phil_interface.index(master_phil, working_phil,
      parse=iotbx.phil.parse)
  mode = params.phaser.mode
  not_alnum = re.compile("[^a-zA-Z0-9_\-]{1,}")
  if params.phaser.keywords.general.root is None :
    raise Sorry("Please define a root filename for this run.")
  elif not_alnum.search(params.phaser.keywords.general.root) is not None :
    raise Sorry("The root filename must be alphanumeric (underscores OK).")
  if match_modes(mode, index.get_phaser_modes("phaser.hklin")) : # XXX: MR_*
    if params.phaser.hklin is None :
      raise Sorry("You must specify a reflections file in mode %s." % mode)
    elif params.phaser.labin is None :
      # MR modes
      raise Sorry("You must specify reflection data labels in mode %s."%mode)
  elif match_modes(mode, index.get_phaser_modes("phaser.crystal")) :
    # TODO: eventually, allow multiples of everything
    if len(params.phaser.crystal) == 0 :
      raise Sorry("Crystal data required for mode %s." % mode)
    elif len(params.phaser.crystal) > 1 :
      raise Sorry("Only one crystal supported at this time.")
    xtal = params.phaser.crystal[0]
    if len(xtal.dataset) > 1 :
      raise Sorry("Only one wavelength supported at this time.")
    # XXX: for now, still have one global hklin
    for xtal in params.phaser.crystal :
      xtalid = xtal.xtal_id
      if not None in [xtal.pdb_file, xtal.ha_file] :
        raise Sorry("For each crystal, you may supply heavy-atom sites as "+
          "EITHER a PDB file OR 'ha' file, but not both.")
      for dataset in xtal.dataset :
        waveid = dataset.wave_id
        labels = dataset.labin.split(",")
        if len(labels) != 4 :
          pass
          # XXX don't do this yet!
          #raise Sorry(("The data labels %s do not appear to correspond to "+
          # "the appropriate data type.  For experimental phasing, data must "+
          #  "be merged, anomalous amplitudes in MTZ format.  The graphical "+
          #  "reflection file editor (phenix.reflection_file_editor) can be "+
          #  "used to convert the data as neeeded.") % dataset.labin)
        if (dataset.wavelength is not None) :
          if (dataset.energy is not None) :
            raise Sorry("Both the energy and the wavelength for this dataset "+
              "have been provided - please enter only one of these values.")
          elif (dataset.wavelength < 0.5) or (dataset.energy > 4) :
            raise Sorry("The specified data collection wavelength (%g) is not "+
              "suitable for "+
              "SAD phasing; please make sure you have entered the correct "+
              "value in Angstroms, or provide the energy instead.")
        elif (dataset.energy is not None) :
          if (dataset.energy < 1000) or (dataset.energy > 20000) :
            raise Sorry("The specified energy is not a suitable value for "+
              "SAD phasing; please make sure you have entered the correct "+
              "value in eV (not KeV), or provide the wavelength instead.")
        else :
          raise Sorry("Please specify either the wavelength (in Angstroms) "+
            "or the energy (in eV) for this dataset.")
  have_composition = False
  if match_modes(mode, index.get_phaser_modes("phaser.ensemble")) :
    have_models = []
    if (len(params.phaser.model) > 0) :
      if ((params.phaser.model_rmsd is None) and
          (params.phaser.model_identity is None)) :
        raise Sorry("You must specify either the RMSD or the identity for "+
          "the single-model input.")
      elif ((params.phaser.model_rmsd is not None) and
            (params.phaser.model_identity is not None)) :
        raise Sorry("Either the model RMSDs or the sequence identities "+
          "must be defined, but not both.")
      elif (params.phaser.model_rmsd is not None) :
        if (len(params.phaser.model_rmsd) != len(params.phaser.model)) :
          raise Sorry("You must specify an RMSD for each model.")
      else :
        if (len(params.phaser.model_identity) != len(params.phaser.model)) :
          raise Sorry("You must specify a sequence identity for each model.")
      have_models.append("ense_1")
      if (len(params.phaser.ensemble) > 0) :
        raise Sorry("You may not specify the single-model input search "+
          "parameter (phaser.model) when the ensemble scope is used.")
    for i, ensemble in enumerate(params.phaser.ensemble) :
      model_id = ensemble.model_id
      if model_id is None :
        raise Sorry("Missing model ID for ensemble #%d." % (i+1))
      if (model_id in have_models) :
        raise Sorry("Duplicate model ID '%s'." % model_id)
      have_models.append(model_id)
      if len(ensemble.coordinates) > 0 :
        if ensemble.map_hklin is not None :
          # XXX: does Phaser already check this?
          raise Sorry(("Please specify either PDB files or an MTZ file "+
            "for ensemble %s, but not both.") % model_id)
        for j, coords in enumerate(ensemble.coordinates) :
          if coords.pdb is None :
            raise Sorry(("Missing PDB file for ensemble %s coordinate set "+
              "%d.") % (model_id, j+1))
          if (not coords.read_variance_from_pdb_remarks) :
            if [coords.rmsd, coords.identity].count(None) != 1 :
              raise Sorry(("Please specify EITHER an RMSD OR fractional "+
                "identity for PDB file %s.") % coords.pdb)
          elif (len(ensemble.coordinates) > 1) :
            raise Sorry("Reading of model variances from PDB REMARK lines "+
              "is only allowed when a single multi-model PDB file is used "+
              "as input.  (Hint: you can use the Ensembler tool to prepare "+
              "a multi-model ensemble.)")
      elif ensemble.map_hklin is not None :
        ens = ensemble
        if ens.map_labels is None :
          raise Sorry("You must specify labels for map data for %s." %
                      model_id)
        if ((ens.map_extent is None or len(ens.map_extent) != 3) or
              (ens.map_centre is None or len(ens.map_centre) != 3)) :
          # TODO: confirm this
          raise Sorry(("You must specify the extent and center of the map "+
            "used for ensemble %s; this must both take the form of three "+
            "decimal numbers between zero and one.") % model_id)
        if (None in [ens.map_protein_mw, ens.map_nucleic_mw] or
            (ens.map_protein_mw == 0 and ens.map_nucleic_mw == 0)) :
          raise Sorry(("Please specify protein and/or nucleic acid masses "+
            "for ensemble %s.") % model_id)
      if ensemble.fractional_comp is not None :
        have_composition = True
  if match_modes(mode, index.get_phaser_modes("phaser.search")) :
    if ((len(params.phaser.model) > 0) and
        (params.phaser.search_copies >= 1)) :
      pass
    elif len(params.phaser.search) == 0 :
      raise Sorry("No search procedure defined.")
    for i, component in enumerate(params.phaser.search) :
      if component.ensembles is None :
        raise Sorry("No model IDs supplied for search component %d." % (i+1))
      elif (component.copies is None) :
        raise Sorry(("Please specify a number of copies for search component "+
          "%d.") % (i+1))
  if match_modes(mode, index.get_phaser_modes("phaser.composition")) :
    if have_composition :
      pass
    elif (params.phaser.seq_file is not None) :
      have_composition = True
    elif params.phaser.composition.solvent is not None :
      if len(params.phaser.composition.chain) > 0 :
        raise Sorry("Please specify EITHER fractional solvent content OR "+
          "individual components, but not both.")
      # XXX: is there currently no way to set this?
    elif len(params.phaser.composition.chain) > 0 :
      for i, chain in enumerate(params.phaser.composition.chain) :
        mass_attr = ["sequence_file","sequence","nres","mw"]
        masses = [ getattr(chain, name) for name in mass_attr]
        if masses.count(None) != 3 :
          raise Sorry(("Please specify ONE source of mass information for "+
            "chain #%d: sequence file, sequence string, # of residues, or "+
            "molecular weight (Da).") % (i+1))
        if chain.num is None :
          raise Sorry("Please specify a number of copies for component #%d."%
            (i+1))
    for i, atom in enumerate(params.phaser.composition.atom) :
      if atom.element is None or atom.num is None :
        raise Sorry(("Please specify an element and number of atoms "+
          "for heavy atom %d.") % (i+1))
      from cctbx.eltbx import tiny_pse
      try :
        table = tiny_pse.table(atom.element)
      except RuntimeError, e :
        raise Sorry("Invalid element symbol '%s'!" % atom.element)
  if match_modes(mode,index.get_phaser_modes("phaser.keywords.sgalternative")):
    if ((params.phaser.keywords.sgalternative.select == "list") and
        (len(params.phaser.keywords.sgalternative.test) == 0)) :
      raise Sorry("You have set Phaser to try listed alternative space "+
        "groups, but no additional space groups were defined.  Please "+
        "enter additional space group symbols (click the 'More' button next "+
        "to the control on the \"Search tab\"), or change the space group "+
        "search settings.")
  if (mode in ["MR_FTF", "MR_BTF"]) :
    if (params.phaser.solution is None) :
      raise Sorry("You must specify a previous solution to start from when "+
        "running a translation search.")
  if (params.density_modification.dm) :
    pass
  return True

# PDB files for individual ensembles
def get_ensemble_pdb_output (dir_name, root) :
  ensemble_files = []
  for file_name in os.listdir(dir_name) :
    if (file_name.startswith(root)) :
      # ensemble output will have name root.X.Y.pdb, where Y is component num.
      match = re.search("\.([0-9])\.([0-9])\.pdb", file_name)
      if (match is not None) :
        ensemble_files.append((os.path.join(dir_name, file_name),
          "Solution %s, Component %s ensemble" % (match.group(1),
            match.group(2))))
  return ensemble_files

class phaser_result (object) :
  def __init__ (self, params, phaser_object, run_dir) :
    self.phaser_object = phaser_object
    self.mode = params.phaser.mode
    self.root = params.phaser.keywords.general.root
    self.hand = params.phaser.keywords.hand
    self.run_dir = run_dir

  def get_phaser_object (self) :
    return self.phaser_object

  def get_output_files (self) :
    output_files = []
    if (self.mode.startswith("MR")) :
      mtz_file = os.path.join(self.run_dir, self.root + ".1.mtz")
      pdb_file = os.path.join(self.run_dir, self.root + ".1.pdb")
      if (os.path.isfile(mtz_file)) and (os.path.getsize(mtz_file) > 0) :
        output_files.append((mtz_file, "Phases and maps"))
      if (os.path.isfile(pdb_file)) and (os.path.getsize(pdb_file) > 0) :
        output_files.append((pdb_file, "Output model"))
      dm_file = os.path.join(self.run_dir, self.root + ".1_dm.mtz")
      if (os.path.isfile(dm_file)) and (os.path.getsize(dm_file) > 0) :
        output_files.append((dm_file, "Density-modified map"))
      output_files.extend(get_ensemble_pdb_output(self.run_dir, self.root))
    elif (self.mode.startswith("EP")) :
      hand = self.hand
      use_original_hand = (hand in ["both", "off"])
      use_inverse_hand = (hand in ["both", "on"])
      if (use_original_hand) :
        mtz_file = os.path.join(self.run_dir, self.root + ".1.mtz")
        pdb_file = os.path.join(self.run_dir, self.root + ".1.pdb")
        if (os.path.isfile(mtz_file)) and (os.path.getsize(mtz_file) > 0) :
          output_files.append((mtz_file, "Phases and maps"))
        if (os.path.isfile(pdb_file)) and (os.path.getsize(pdb_file) > 0) :
          output_files.append((pdb_file, "Heavy-atom sites"))
        dm_file = os.path.join(self.run_dir, self.root + ".1_dm.mtz")
        if (os.path.isfile(dm_file)) and (os.path.getsize(dm_file) > 0) :
          output_files.append((dm_file, "Density-modified map"))
        if (use_inverse_hand) :
          mtz_file2 = os.path.join(self.run_dir, self.root + ".1.hand.mtz")
          pdb_file2 = os.path.join(self.run_dir, self.root + ".1.hand.pdb")
          if (os.path.isfile(mtz_file2)) and (os.path.getsize(mtz_file2) > 0) :
            output_files.append((mtz_file2, "Phases and maps (inverse hand)"))
          if (os.path.isfile(pdb_file2)) and (os.path.getsize(pdb_file2) > 0) :
            output_files.append((pdb_file2, "Heavy-atom sites (inverse hand)"))
          dm_file2 = os.path.join(self.run_dir, self.root + ".1_dm.mtz")
          if (os.path.isfile(dm_file2)) and (os.path.getsize(dm_file2) > 0) :
            output_files.append((dm_file2, "Density-modified map"))
      elif (use_inverse_hand) :
        mtz_file = os.path.join(self.run_dir, self.root + ".1.mtz")
        pdb_file = os.path.join(self.run_dir, self.root + ".1.pdb")
        if (os.path.isfile(mtz_file)) and (os.path.getsize(mtz_file) > 0) :
          output_files.append((mtz_file, "Phases and maps (inverse hand)"))
        if (os.path.isfile(pdb_file)) and (os.path.getsize(pdb_file) > 0) :
          output_files.append((pdb_file, "Heavy-atom sites (inverse hand)"))
        dm_file = os.path.join(self.run_dir, self.root + ".1_dm.mtz")
        if (os.path.isfile(dm_file)) and (os.path.getsize(dm_file) > 0) :
          output_files.append((dm_file, "Density-modified map (inverse hand)"))
    data_file = os.path.join(self.run_dir, self.root + "_data.mtz")
    if (os.path.isfile(data_file)) :
      output_files.append((data_file, "Processed input data"))
    return output_files

  def finish_job (self) :
    output_files = self.get_output_files()
    stats = []
    if hasattr(self.phaser_object, "getTopLLG") :
      llg = self.phaser_object.getTopLLG()
      stats.append(("Log-likelihood gain", "%.2f" % llg))
    if hasattr(self.phaser_object, "getTopTFZeq") :
      tfz = self.phaser_object.getTopTFZeq()
      stats.append(("TFZ", "%.1f" % tfz))
    elif hasattr(self.phaser_object, "getTopTFZ") :
      tfz = self.phaser_object.getTopTFZ()
      stats.append(("TFZ", "%.1f" % tfz))
    return (output_files, stats)