# TODO tests

from __future__ import division
from iotbx import crystal_symmetry_from_any
import iotbx.bioinformatics
import iotbx.phil
from cctbx import crystal
from libtbx.utils import Sorry
import sys

master_phil_str = """
data = None
  .type = path
labels = None
  .type = str
sequence = None
  .type = path
model = None
  .type = path
space_group = None
  .type = space_group
unit_cell = None
  .type = unit_cell
n_residues = None
  .type = int(value_min=1)
  .optional = True
n_bases = None
  .type = int(value_min=1)
  .optional = True
"""

def run (args, out=sys.stdout) :
  cmdline = iotbx.phil.process_command_line_with_files(
    args=args,
    master_phil_string=master_phil_str,
    pdb_file_def="model",
    reflection_file_def="data",
    seq_file_def="sequence",
    space_group_def="space_group",
    unit_cell_def="unit_cell",
    integer_def="n_residues",
    usage_string="""\
phenix.matthews [data.hkl] [space_group] [unit_cel] [sequence] [n_residues] ...

Calculate the expected Matthews coefficient given the crystal symmetry and
crystallized molecule(s).
""")
  params = cmdline.work.extract()
  if (params.space_group is None) or (params.unit_cell is None) :
    if (params.data is None) :
      raise Sorry("You must supply both a space group and a unit cell (or "+
        "a data file containing this information).")
    else :
      symm = crystal_symmetry_from_any.extract_from(file_name=params.data)
      space_group_from_file = symm.space_group()
      if (params.space_group is None) :
        if (space_group_from_file is not None) :
          params.space_group = symm.space_group()
      elif (space_group_from_file is not None) :
        if (space_group_from_file != params.space_group) :
          print >> out, "WARNING: space group mismatch between command line "+\
            "and file:"
          print >> out, "  %s (cmdline), %s (file)" % (params.space_group,
            space_group_from_file)
      if (params.unit_cell is None) :
        params.unit_cell = symm.unit_cell()
  validate_params(params, check_symmetry=True)
  if (params.sequence is not None) :
    assert (params.n_residues == params.n_bases == None)
    seq_comp = iotbx.bioinformatics.composition_from_sequence_file(
      file_name=params.sequence,
      log=out)
    if (seq_comp is not None) :
      params.n_residues = seq_comp.n_residues
      params.n_bases = seq_comp.n_bases
    else :
      raise Sorry("No composition information could be obtained from the "+
        "sequence file.")
  elif (params.model is not None) :
    assert (params.n_residues == params.n_bases == None)
    from iotbx.file_reader import any_file
    params.n_residues = 0
    params.n_bases = 0
    pdb_in = any_file(params.model)
    hierarchy = pdb_in.file_object.hierarchy
    for chain in hierarchy.models()[0].chains() :
      if chain.is_protein() :
        params.n_residues += chain.residue_groups_size()
      elif chain.is_na() :
        params.n_bases += chain.residue_groups_size()
  print >> out, "Space group: %s" % params.space_group
  print >> out, "Unit cell: %s" % params.unit_cell
  if (params.n_residues > 0) :
    print >> out, "Number of residues: %d" % params.n_residues
  if (params.n_bases > 0) :
    print >> out, "Number of bases: %d" % params.n_bases
  symm = crystal.symmetry(
    space_group_info=params.space_group,
    unit_cell=params.unit_cell)
  from mmtbx.scaling import matthews
  result = matthews.matthews_rupp(
    crystal_symmetry=symm,
    n_residues=params.n_residues,
    n_bases=params.n_bases)
  result.show(out=out)
  return result

def validate_params (params, check_symmetry=False) :
  if (params.sequence is None) and (params.model is None) :
    if (params.n_residues is None) and (params.n_bases is None) :
      pass
      #raise Sorry("You must specify the composition of the crystallized "+
      #  "entity - either a sequence, a partial model, or the number of "+
      #  "protein residues or nucleic acid bases.")
  else :
    if (params.n_residues is not None) or (params.n_bases is not None) :
      raise Sorry("You may only specify a sequence file OR a partial model "+
        "OR the numbers of residues and/or bases.")
  if (check_symmetry) :
    if (params.space_group is None) or (params.unit_cell is None) :
      raise Sorry("You must supply both a space group and a unit cell (or "+
        "a data file containing this information).")
  return True

if (__name__ == "__main__") :
  run(sys.argv[1:])