from __future__ import division
from libtbx.program_utils.result import program_result
from libtbx.utils import Sorry, Usage, null_out, multi_out
from libtbx import runtime_utils
from libtbx import easy_pickle
import libtbx.load_env
import libtbx.phil
import os
import sys

keywords_phil_str = """
  nma
    .short_caption = NMA input options
    .expert_level = 1
    .style = phaser:mode:NMAYXZ,SCEDS auto_align box
  {
    mode = None
      .type = floats
      .short_caption = Modes
    combinations = None
      .type=float
  }

  enm
    .short_caption = NMA method options
    .expert_level = 2
    .style = phaser:mode:NMAXYZ,SCEDS auto_align box
  {
    oscillators = *rtb ca all
      .type=choice
      .short_caption = Atoms used for NMA
      .caption = Rotation-translation_block C-alpha_atoms_only All_atoms
    rtb_maxblocks = None
      .type=int
      .short_caption = RTB max. blocks
    rtb_nresidues = None
      .type=int
      .short_caption = RTB number of residues
    radius = None
      .type=float
      .short_caption = Elastic Network Model interaction radius
    force = None
      .type=float
      .short_caption = Elastic Network Model force constant
  }

  perturb
    .short_caption = Normal mode calculation options
    .expert_level = 1
    .style = phaser:mode:NMAYXZ,SCEDS auto_align box
  {
    increment = *rms dq
      .type=choice
    rms {
      step = None
        .type=float
      stretch = None
        .type=float
      clash = None
        .type=float
      maximum = None
        .type=float
      direction = *forward backward tofro
        .type=choice
    }
    dq = None
      .type = float
      .multiple=True
      .optional=True
      .short_caption = DQ factor
  }

"""

keywords_phil_nmad_str = keywords_phil_str + """
  sceds
    .short_caption = NMA domain finding
    .expert_level = 1
    .style = phaser:mode:NMAYXZ,SCEDS auto_align box
  {
    enable = None
      .type = bool
    ndom = None
      .type = int
      .short_caption = Number of domains to find
    weight {
      equality = None
        .type = float
        .short_caption = Weight for equality test
      sphericity = None
        .type = float
        .short_caption = Weight for sphericity test
      density = None
        .type = float
        .short_caption = Weight for density test
      continuity = None
        .type = float
        .short_caption = Weight for continuity test
    }
  }

  ddm
    .short_caption = NMA domain options
    .expert_level = 1
    .style = phaser:mode:NMAYXZ,SCEDS auto_align box
  {
    slider = None
      .type = int
    distance
      .style = box auto_align
    {
      step = None
        .type = int
      min = None
        .type = int
      max = None
        .type = int
    }
    separation
      .style = box auto_align
    {
      min = None
        .type = int
       max = None
        .type = int
    }
    sequence
      .style = box auto_align
    {
      min = None
        .type = int
      max = None
        .type = int
    }
    join
      .style = box auto_align
    {
      min = None
        .type = int
      max = None
        .type = int
    }
    cluster = *max first last
      .type = choice
  }
"""

def master_params (include_domains=False) :
  keywords_phil_str_ = keywords_phil_str
  other_phil = ""
  if (include_domains) :
    keywords_phil_str_ = keywords_phil_nmad_str
    other_phil = """
validate_domains = True
  .type = bool
  .short_caption = Validate domain selections
"""
  master_phil_str = """
file_name = None
  .type = path
  .short_caption = PDB file
  .style = file_type:pdb bold input_file
root = None
  .type = str
  .short_caption = Output prefix
  .input_size = 200
gui_output_dir = None
  .type = path
  .short_caption = Output directory
  .style = bold output_dir default_cwd directory
include scope libtbx.phil.interface.tracking_params
%s
keywords {
  %s
}""" % (other_phil, keywords_phil_str_)
  defaults_file = libtbx.env.under_build("include/phaser_nma_defaults.phil")
  if (not os.path.exists(defaults_file)) :
    raise RuntimeError("Missing default parameters - re-run libtbx.configure phaser")
  # FIXME if include_domains=True, set modes to None
  phil1 = libtbx.phil.parse(master_phil_str, process_includes=True)
  sources = [ libtbx.phil.parse(file_name=defaults_file) ]
  if (include_domains) :
    sources.append(libtbx.phil.parse("""keywords.nmapdb.mode=None"""))
  return phil1.fetch(sources=sources)

def setup_logging (log_file, out) :
  f = open(log_file, "w")
  out2 = multi_out()
  out2.register("logfile", f)
  out2.register("stdout", out)
  return out2

def run (args, out=None) :
  if (out is None) : out = sys.stdout
  master_phil = master_params()
  if (len(args) == 0) :
    raise Usage("""phaser.nma_perturb model.pdb [options]

Full keywords:
%s""" % master_phil.as_str())
  import iotbx.phil
  cmdline = iotbx.phil.process_command_line_with_files(
    args=args,
    master_phil=master_phil,
    pdb_file_def="file_name")
  params = cmdline.work.extract()
  validate_params_basic(params)
  if (params.root is None) :
    params.root = os.path.basename(os.path.splitext(params.file_name)[0]) + \
      "_NMA"
  log_file = params.root + ".log"
  out2 = setup_logging(log_file, out)
  phaser_result = run_phaser(
    pdb_file=params.file_name,
    root=params.root,
    params=params.keywords,
    domain_finding=False,
    log=out2)
  out2.close()
  pdb_files = phaser_result.getPdbFiles()
  pdb_files = [ os.path.basename(f) for f in pdb_files ]
  return nma_perturb_result(
    program_name="phaser.nma_perturb",
    job_title=params.job_title,
    directory=os.getcwd(),
    log_file=log_file,
    pdb_files=pdb_files)

def run_phaser (
    pdb_file,
    root,
    params=None,
    log=None,
    quiet=False,
    domain_finding=False,
    n_domains=None) :
  if (log is None) : log = null_out()
  if (params is None) :
    params = master_params().extract().keywords
  if (n_domains is not None) :
    assert (domain_finding)
    params.sced.number = n_domains
  master_phil = master_params(include_domains=domain_finding)
  keyword_phil = master_phil.get("keywords").objects[0].format(
    python_object=params)
  from phaser.phenix_interface.driver import convert_keywords
  from phaser.phenix_interface import callbacks
  import phaser
  input = phaser.InputNMA()
  input.addENSE_PDB_ID("model", pdb_file, 1.0)
  input.setROOT(root)
  if (domain_finding) :
    input.setEIGE_WRIT(False)
  if (quiet) :
    input.setMUTE(True)
  convert_keywords(
    input=input,
    keyword_phil=keyword_phil,
    mode="NMA",
    script_out=null_out())
  output = phaser.Output()
  output.setPackagePhenix(log)
  #if (not quiet) :
  #  callback = callbacks.manager()
  #  output.setPhenixCallback(callback)
  if (domain_finding) :
    run_nma = phaser.runSCEDS(input, output)
  else :
    run_nma = phaser.runNMAXYZ(input, output)
  if (run_nma.Failed()) :
    raise Sorry(str(run_nma.ErrorMessage()))
  if (not quiet) :
    easy_pickle.dump(".phaser_result.pkl", run_nma)
  return run_nma

def validate_params_basic (params) :
  if (params.file_name is None) :
    raise Sorry("PDB file not specified.")
  elif (not os.path.isfile(params.file_name)) :
    raise Sorry("%s is not a valid file." % params.file_name)
  return True

#-----------------------------------------------------------------------
# GUI stuff below here
def validate_params (params) :
  return validate_params_basic(params)

class nma_perturb_result (program_result) :
  def finish_job (self) :
    output_files = [ (f, "Perturbed model") for f in self.get_pdb_files() ]
    return output_files, {}

class launcher (runtime_utils.target_with_save_result) :
  def run (self) :
    os.makedirs(self.output_dir)
    os.chdir(self.output_dir)
    return run(self.args, out=sys.stdout)

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