from __future__ import division
from libtbx import slots_getstate_setstate
from libtbx.str_utils import format_value
import sys

class entity (slots_getstate_setstate) :
  """
  Base class for all validation results.  This includes a boolean outlier flag,
  the information used to zoom in the Phenix GUI (optional, but strongly
  recommended), and some kind of numerical score (also optional, but strongly
  recommended - although some analyses may require multiple distinct scores).
  """
  __slots__ = [
    "atom_selection",
    "xyz",
    "outlier",
    "score",
  ]
  score_format = "%s"
  molprobity_table_labels = []

  def __init__ (self, **kwds) :
    for name in kwds.keys() :
      assert (name in self.__slots__), name
    for name in self.__slots__ :
      setattr(self, name, kwds.get(name, None))

  @staticmethod
  def header () :
    """
    Format for header in result listings.
    """
    raise NotImplementedError()

  def format_score (self, replace_none_with="None") :
    return format_value(self.score_format, self.score,
      replace_none_with=replace_none_with)

  def is_outlier (self) :
    return self.outlier

  def as_string (self, prefix="") :
    raise NotImplementedError()

  def __str__ (self) :
    return self.as_string()

  def id_str (self, ignore_altloc=None) :
    """
    Returns a formatted (probably fixed-width) string describing the molecular
    entity being validation, independent of the analysis type.
    """
    raise NotImplementedError()

  def __hash__ (self) :
    return self.id_str().__hash__()

  def as_list (self) :
    """
    Optional; returns old format used by some tools in mmtbx.validation.
    """
    raise NotImplementedError()

  def as_table_row_molprobity (self) :
    """
    Returns a list of formatted table cells for display by MolProbity.
    """
    raise NotImplementedError()

  def as_table_row_phenix (self) :
    """
    Returns a list of formatted table cells for display by Phenix.
    """
    raise NotImplementedError()

  def as_kinemage (self) :
    """
    Returns a kinemage string for displaying an outlier.
    """
    raise NotImplementedError()

  def format_old (self) :
    raise NotImplementedError()

  def __eq__ (self, other) :
    """
    Compare two validation results to determine whether they correspond to the
    same molecular entity and analysis type.  This is intended to be used for
    analysis of a structure before-and-after refinement (etc.).
    """
    raise NotImplementedError()

  def __cmp__ (self, other) :
    return cmp(self.score, other.score)

  def is_single_residue_object (self) :
    raise NotImplementedError()

  def as_selection_string (self) :
    """
    Returns PDB atom selection string for the atom(s) involved.
    """
    return None

  def zoom_info (self) :
    """
    Returns data needed to zoom/recenter the graphics programs from the Phenix
    GUI.
    """
    return [ self.as_selection_string(), self.xyz ]

__residue_attr__ = [
  "chain_id",
  "resseq",
  "icode",
  "resname",
  "altloc",
  "segid",
]

class residue (entity) :
  """
  Base class for validation information about a single residue, which depending
  on context could mean either any one of the residue_group, atom_group, or
  residue objects from the PDB hierarchy.
  """
  __slots__ = entity.__slots__ + __residue_attr__ + ["occupancy"]
  def _copy_constructor (self, other) :
    for attr in __residue_attr__ :
      setattr(self, attr, getattr(other, attr))

  def assert_all_attributes_defined (self) :
    for name in self.__residue_attr__ :
      assert (getattr(self, name) is not None) or (name == "segid")

  def id_str (self, ignore_altloc=False) :
    base = "%2s%4s%1s" % (self.chain_id, self.resseq, self.icode)
    if (not ignore_altloc) :
      base += "%1s" % self.altloc
    else :
      base += " "
    base += "%3s" % self.resname
    if (self.segid is not None) :
      base += " segid='%4s'" % self.segid
    return base

  def resseq_as_int (self) :
    from iotbx.pdb import hybrid_36
    return hybrid_36.hy36decode(len(self.resseq), self.resseq)

  @property
  def resid (self) :
    return "%4s%1s" % (self.resseq, self.icode)

  def residue_id (self, ignore_altloc=False) :
    return self.id_str(ignore_altloc=ignore_altloc)

  def simple_id (self) :
    return ("%s%s%s" % (self.chain_id, self.resseq, self.icode)).strip()

  # XXX probably needs to be flexible about altloc...
  def is_same_residue (self, other, ignore_altloc=False) :
    if hasattr(other, "residue_id") :
      return (self.residue_id(ignore_altloc=ignore_altloc) ==
              other.residue_id(ignore_altloc=ignore_altloc))
    return (self.id_str(ignore_altloc=ignore_altloc) ==
            other.id_str(ignore_altloc=ignore_altloc))

  def is_same_residue_group (self, other) :
    return ((self.chain_id == other.chain_id) and
            (self.resseq == other.resseq) and
            (self.icode == other.icode) and
            (self.segid == other.segid))

  def atom_group_id_str (self) :
    return "%1s%3s%2s%4s%1s" % (self.altloc, self.resname, self.chain_id,
      self.resseq, self.icode)

  def residue_group_id_str (self) :
    return "%2s%4s%1s" % (self.chain_id, self.resseq, self.icode)

  def __eq__ (self, other) :
    assert type(self) == type(other), type(other)
    return self.id_str() == other.id_str()

  def is_single_residue_object (self) :
    return True

  def atom_selection_string (self) :
    return "(chain '%s' and resid '%s' and resname %s and altloc '%s')" % \
      (self.chain_id, self.resid, self.resname, self.altloc)

  def set_coordinates_from_hierarchy (self, pdb_hierarchy,
      atom_selection_cache=None) :
    if (atom_selection_cache is None) :
      atom_selection_cache = pdb_hierarchy.atom_selection_cache()
    sel = atom_selection_cache.selection(self.atom_selection_string())
    assert (len(sel) > 0)
    self.xyz = pdb_hierarchy.atoms().select(sel).extract_xyz().mean()

class atoms (entity) :
  """
  Base class for validation results involving a specific set of atoms, such
  as covalent geometry restraints, clashes, etc.
  """
  __atoms_attr__ = [
    "atoms_info",
  ]
  __slots__ = entity.__slots__ + __atoms_attr__

  def n_atoms (self) :
    return len(self.atoms_info)

  def __eq__ (self, other) :
    assert type(self) == type(other), type(other)
    return sorted(self.atoms_info) == sorted(other.atoms_info)

  def is_single_residue_object (self) :
    return False

  def get_altloc (self) :
    consensus_altloc = ''
    for atom in self.atoms_info :
      if (atom.altloc.strip() != '') :
        if (consensus_altloc == '') :
          consensus_altloc = atom.altloc
        else :
          assert (atom.altloc == consensus_altloc)
    return consensus_altloc

  def sites_cart (self) :
    return [ a.xyz for a in self.atoms_info ]

  def is_in_chain (self, chain_id) :
    if (chain_id == None) :
      return True
    for a in self.atoms_info :
      if (a.chain_id == chain_id) :
        return True
    return False

class atom_base (slots_getstate_setstate) :
  """
  Container for metadata for a single atom, in the context of validation
  results involving multiple atoms.  Intended to be used as-is inside various
  atoms classes.
  """
  __atom_attr__ = [
    "name",
    "element", # XXX is this necessary?
    "xyz",
    "symop",
    "occupancy",
    "b_iso",
  ]
  __atom_slots__ = __residue_attr__ + __atom_attr__
  # XXX __slots__ should be left empty here

  def __init__ (self, **kwds) :
    pdb_atom = kwds.get("pdb_atom", None)
    if (pdb_atom is not None) :
      del kwds['pdb_atom']
    for name in kwds.keys() :
      assert (name in self.__slots__), name
    for name in self.__slots__ :
      setattr(self, name, kwds.get(name, None))
    if (pdb_atom is not None) :
      labels = pdb_atom.fetch_labels()
      self.chain_id = labels.chain_id
      self.resseq = labels.resseq
      self.icode = labels.icode
      self.resname = labels.resname
      self.altloc = labels.altloc
      self.segid = labels.segid
      self.name = pdb_atom.name
      self.xyz = pdb_atom.xyz
      self.occupancy = pdb_atom.occ
      self.b_iso = pdb_atom.b
      self.element = pdb_atom.element

  def __cmp__ (self, other) :
    return cmp(self.id_str(), other.id_str())

  def __eq__ (self, other) :
    assert isinstance(other, atom_base), type(other)
    return self.id_str() == other.id_str()

  def id_str (self, ignore_altloc=False) :
    base = "%2s%4s%1s" % (self.chain_id, self.resseq, self.icode)
    if (not ignore_altloc) :
      base += "%1s" % self.altloc
    else :
      base += " "
    base += "%3s %4s" % (self.resname, self.name)
    if (self.segid is not None) :
      base += " segid='%4s'" % self.segid
    return base

  def residue_group_id_str (self) :
    return "%2s%4s%1s" % (self.chain_id, self.resseq, self.icode)

  def atom_group_id_str (self) :
    return "%1s%3s%2s%4s%1s" % (self.altloc, self.resname, self.chain_id,
      self.resseq, self.icode)

class atom_info (atom_base) :
  """
  Container for metadata for a single atom, in the context of validation
  results involving multiple atoms.  Intended to be used as-is inside various
  atoms classes.
  """
  __slots__ = atom_base.__atom_slots__ + ["symop"]

def get_atoms_info (pdb_atoms, iselection,
      use_segids_in_place_of_chainids=False) :
  proxy_atoms = []
  for n, i_seq in enumerate(iselection):
    atom = pdb_atoms[i_seq]
    labels = atom.fetch_labels()
    if use_segids_in_place_of_chainids:
      chain_id = atom.segid
    else:
      chain_id = labels.chain_id
    info = atom_info(
      name=atom.name,
      element=atom.element,
      chain_id=chain_id,
      resseq=labels.resseq,
      icode=labels.icode,
      resname=labels.resname,
      altloc=labels.altloc,
      occupancy=atom.occ,
      #segid=atom.segid,
      xyz=atom.xyz)
    proxy_atoms.append(info)
  return proxy_atoms

class atom (atom_base, entity) :
  """
  Base class for validation results for a single atom.  This is distinct from
  the atom_info class above, which is used to track individual atoms within
  a multi-atom validation result.
  """
  __slots__ = entity.__slots__ + atom_base.__atom_slots__

  def is_single_residue_object (self) :
    return True

#-----------------------------------------------------------------------
class validation (slots_getstate_setstate) :
  """
  Container for a set of results from a single analysis (rotamers, clashes,
  etc.).  This is responsible for the console display of these results and
  associated statistics.  Individual modules will subclass this and override
  the unimplemented methods.
  """
  __slots__ = ["n_outliers", "n_total", "results", "_cache"]
  program_description = None
  output_header = None
  gui_list_headers = [] # for Phenix GUI ListCtrl widgets
  gui_formats = []      # for Phenix GUI ListCtrl widgets
  def __init__ (self) :
    self.n_outliers = 0
    self.n_total = 0
    self.results = []
    self._cache = None

  def get_outliers_count_and_fraction(self):
    if (self.n_total != 0):
      fraction = float(self.n_outliers) / self.n_total
      assert fraction <= 1.0
      return self.n_outliers, fraction
    return 0, 0.

  @property
  def percent_outliers (self) :
    n_outliers, frac_outliers = self.get_outliers_count_and_fraction()
    return frac_outliers * 100.

  def outlier_selection (self) :
    """
    Return a flex.size_t object containing the i_seqs of atoms flagged as
    outliers (either individually or as part of an atom_group).  This needs
    to be implemented in the underlying classes unless they include a
    pre-built _outlier_i_seqs attribute.
    """
    if hasattr(self, "_outlier_i_seqs") :
      return self._outlier_i_seqs
    raise NotImplementedError()

  def get_outliers_goal (self) :
    raise NotImplementedError()

  def get_result_class (self) :
    raise NotImplementedError()

  def coot_todo (self) :
    return ""

  def show_old_output (self, out=sys.stdout, verbose=False) :
    """
    For backwards compatibility with output formats of older utilities
    (phenix.ramalyze et al.).
    """
    if (verbose) :
      assert (self.output_header is not None)
      print >> out, self.output_header
    for result in self.results :
      print >> out, result.format_old()
    if (verbose) :
      self.show_summary(out)

  def show_summary (self, out=sys.stdout, prefix="") :
    raise NotImplementedError()

  def show (self, out=sys.stdout, prefix="  ", outliers_only=True,
      verbose=True) :
    if (len(self.results) > 0) :
      print >> out, prefix + self.get_result_class().header()
      for result in self.iter_results(outliers_only) :
        print >> out, prefix + str(result)
    self.show_summary(out=out, prefix=prefix)

  def iter_results (self, outliers_only=True) :
    for result in self.results :
      if (not outliers_only) or (result.is_outlier()) :
        yield result

  def as_kinemage (self) :
    return None

  def as_coot_data (self) :
    """
    Return results in a format suitable for unpickling in Coot.
    """
    raise NotImplementedError()

  def as_gui_table_data (self, outliers_only=True, include_zoom=False) :
    """
    Format results for display in the Phenix GUI.
    """
    table = []
    for result in self.iter_results(outliers_only) :
      extra = []
      if (include_zoom) :
        extra = result.zoom_info()
      row = result.as_table_row_phenix()
      assert (len(row) == len(self.gui_list_headers) == len(self.gui_formats))
      table.append(row + extra)
    return table

  def find_residue (self, other=None, residue_id_str=None) :
    assert ([other, residue_id_str].count(None) == 1)
    if (other is not None) :
      if hasattr(other, "residue_group_id_str") :
        residue_id_str = other.residue_group_id_str()
      elif hasattr(other, "id_str") :
        residue_id_str = other.id_str()
      else :
        residue_id_str = str(other)
    if (self._cache is None) :
      self._cache = {}
      for i_res, result in enumerate(self.results) :
        result_id_str = result.residue_group_id_str()
        self._cache[result_id_str] = i_res
    i_res = self._cache.get(residue_id_str, None)
    if (i_res is not None) :
      return self.results[i_res]
    return None

  def find_atom_group (self, other=None, atom_group_id_str=None) :
    """
    Attempt to locate a result corresponding to a given atom_group object.
    """
    assert ([other, atom_group_id_str].count(None) == 1)
    if (other is not None) :
      if hasattr(other, "atom_group_group_id_str") :
        atom_group_id_str = other.atom_group_group_id_str()
      elif hasattr(other, "id_str") :
        atom_group_id_str = other.id_str()
      else :
        atom_group_id_str = str(other)
    if (self._cache is None) :
      self._cache = {}
      for i_res, result in enumerate(self.results) :
        result_id_str = result.atom_group_group_id_str()
        self._cache[result_id_str] = i_res
    i_res = self._cache.get(atom_group_id_str, None)
    if (i_res is not None) :
      return self.results[i_res]
    return None

class dummy_validation (object) :
  """
  Placeholder for cases where values may be undefined because of molecule type
  (e.g. all-RNA structures) but we want to substitute None automatically.
  """
  def __getattr__ (self, name) :
    return None

  def __nonzero__ (self) :
    return False

molprobity_cmdline_phil_str = """
  model = None
      .type = path
      .help = "Model file (PDB or mmCIF)"

  outliers_only = False
    .type = bool
    .help = "Only display outliers"

  verbose = True
    .type = bool
    .help = '''Verbose'''
"""