from __future__ import division
from phaser.phenix_interface import driver
from phaser import phenix_interface
import libtbx.load_env
from libtbx import easy_pickle
import libtbx.utils
from cStringIO import StringIO
import shutil
import os
import sys

libtbx.utils.disable_tracebacklimit = True

class test_cb (object) :
  def __init__ (self) :
    self.messages = []

  def __call__ (self, message, data, accumulate=True, cached=True) :
    self.messages.append(message)

def copy_file_to_cwd (file_name) :
  full_path = libtbx.env.find_in_repositories(
    relative_path="phaser/test/%s" % file_name,
    test=os.path.isfile)
  if (full_path is None) :
    full_path = libtbx.env.find_in_repositories(
      relative_path="phaser/tutorial/%s" % file_name,
      test=os.path.isfile)
  assert (full_path is not None)
  dest_path = os.path.join(os.getcwd(), file_name)
  shutil.copy(full_path, dest_path)

def exercise_mr (args) :
  input_files = ["beta.pdb", "blip.pdb", "beta_blip_P3221.mtz"]
  for file_name in input_files :
    copy_file_to_cwd(file_name)
  eff_file = libtbx.env.find_in_repositories(
    relative_path="phaser/test/phenix_interface/beta_blip.eff",
    test=os.path.isfile)
  out = StringIO()
  driver.run(args=["--show_defaults"], out=out)
  out2 = StringIO()
  phenix_interface.master_phil().show(out=out2)
  assert out.getvalue() == out2.getvalue()
  if ("--verbose" in args) :
    out = sys.stdout
  else :
    out = StringIO()
  cb = test_cb()
  libtbx.call_back.register_handler(cb)
  driver.run(args=[eff_file,"--show_script","--test_mode"],
    out=out)
  if (not "--no-run" in args) :
    if ("--verbose" in args) :
      out = sys.stdout
    else :
      out = StringIO()
    result = driver.run(args=[eff_file,], out=out)
    result_ = result.get_phaser_object()
#   assert len(cb.messages) > 1000
    assert result_.getTopLLG() > 200 # ~256.6
    assert os.path.isfile(result_.getTopPdbFile())
    easy_pickle.dump("mr.pkl", result)
    result2 = easy_pickle.load("mr.pkl")
    assert result2.get_phaser_object().getTopLLG() > 200
    # run 1: beta only
    eff_file = libtbx.env.find_in_repositories(
      relative_path="phaser/test/phenix_interface/beta_blip1.eff",
      test=os.path.isfile)
    result = driver.run(args=[eff_file,], out=out)
    easy_pickle.dump("beta_solution.pkl", result)
    # run 2: start from partial solution (as pickle file)
    eff_file = libtbx.env.find_in_repositories(
      relative_path="phaser/test/phenix_interface/beta_blip2.eff",
      test=os.path.isfile)
    result = driver.run(args=[eff_file,], out=out)
    result_ = result.get_phaser_object()
#   assert len(cb.messages) > 1000
    assert result_.getTopLLG() > 200 # ~256.6
    assert os.path.isfile(result_.getTopPdbFile())
    easy_pickle.dump("mr.pkl", result)
    result2 = easy_pickle.load("mr.pkl")
    assert result2.get_phaser_object().getTopLLG() > 200

def exercise_sad (args) :
  input_files = ["S-insulin_hyss.pdb", "S-insulin.mtz", "S-insulin.seq"]
  for file_name in input_files :
    copy_file_to_cwd(file_name)
  eff_file = libtbx.env.find_in_repositories(
    relative_path="phaser/test/phenix_interface/insulin.eff",
    test=os.path.isfile)
  assert (eff_file is not None)
  for file_name in ["insulin.mtz", "insulin_data.mtz"] :
    if os.path.exists(file_name) :
      os.remove(file_name)
  if ("--verbose" in args) :
    out = sys.stdout
  else :
    out = StringIO()
  out = StringIO()
  driver.run(args=[eff_file,"--show_script","--test_mode"],
    out=out)
  if (not "--no-run" in args) :
    if ("--verbose" in args) :
      out = sys.stdout
    else :
      out = StringIO()
    result = driver.run(args=[eff_file, "--test_mode"], out=out)
    easy_pickle.dump("sad.pkl", result)
    assert os.path.isfile("insulin.1.mtz")

def exercise_sad_simple (args) :
  from phaser import tst_phenix_interface_ep
  tst_phenix_interface_ep.exercise(verbose=("--verbose" in args))

def exercise_mr_simple (args) :
  pdb_file = libtbx.env.find_in_repositories(
    relative_path="phenix_regression/wizards/coords.pdb",
    test=os.path.isfile)
  sca_file = libtbx.env.find_in_repositories(
    relative_path="phenix_regression/wizards/native.sca",
    test=os.path.isfile)
  seq_file = libtbx.env.find_in_repositories(
    relative_path="phenix_regression/wizards/seq.dat",
    test=os.path.isfile)
  if (pdb_file is None) :
    print "phenix_regression not found, skipping exercise_mr_simple()"
    return
  args = [pdb_file, sca_file, seq_file, "model_rmsd=0.85", "root=simple",
          "--test_mode"]
  if ("--verbose" in args) :
    out = sys.stdout
  else :
    out = StringIO()
  result = driver.run(args=args,
    out=out)

def exercise (args) :
  print "validating PHIL defaults..."
  from phaser import phenix_interface
  master_phil = phenix_interface.master_phil(validate=True)
  print "exercising simple MR run..."
  exercise_mr_simple(args)
  print "exercising simple EP run..."
  exercise_sad_simple(args)
  if (not "--quick" in args) :
    print "exercising MR mode..."
    exercise_mr(args)
    print "exercising EP mode..."
    exercise_sad(args)
  print "OK"

if __name__ == "__main__" :
  if (not os.path.isdir("tst_phenix_interface")) :
    os.mkdir("tst_phenix_interface")
  os.chdir("tst_phenix_interface")
  exercise(sys.argv[1:])