
from __future__ import division
from libtbx import easy_run
import cStringIO
import os
import sys

def exercise (verbose=False) :
  expected_files = [
    "se_ca_frag.1.mtz", "se_ca_frag.1.hand.mtz",
    "se_ca_frag.1_dm.mtz", "se_ca_frag.1.hand_dm.mtz",
    "se_ca_frag.1.llgmaps.mtz", "se_ca_frag.1.hand.llgmaps.mtz",]
  for file_name in expected_files :
    if (os.path.exists(file_name)) :
      os.remove(file_name)
  write_pdb_file()
  fmodel_params = """
  high_resolution = 1.8
  low_resolution = None
  add_sigmas = True
  pdb_file = se_ca_frag.pdb
  output {
    label = F
    type = *real complex
    file_name = se_ca_frag_anom.mtz
  }
  anomalous_scatterers {
    group {
      selection = element SE
      f_prime = -8.016
      f_double_prime = 3.842
    }
    group {
      selection = element CA
      f_prime = 0.25
      f_double_prime = 0.50
    }
  }"""
  f = open("fmodel_se_ca_frag.eff", "w")
  f.write(fmodel_params)
  f.close()
  if (os.path.exists("se_ca_frag_anom.mtz")) :
    os.remove("se_ca_frag_anom.mtz")
  easy_run.fully_buffered("phenix.fmodel fmodel_se_ca_frag.eff"
    ).raise_if_errors()
  assert (os.path.isfile("se_ca_frag_anom.mtz"))
  sites = open("se_site.pdb", "w")
  sites.write("""\
CRYST1   21.937    4.866   23.477  90.00 107.08  90.00 P 1 21 1      2
ATOM     19 SE   SE  X   1      -4.069   3.906   7.12   0.90 20.00          SE
END
""")
  sites.close()
  from phaser.phenix_interface import driver
  f = open("phaser_se_ca.eff", "w")
  f.write("""
phaser {
  mode = *EP_AUTO
  crystal {
    xtal_id = "se_ca_frag"
    pdb_file = "se_site.pdb"
    dataset {
      hklin = "se_ca_frag_anom.mtz"
      wave_id = "semet"
      labin = "F(+),SIGF(+),F(-),SIGF(-)"
      wavelength = 0.9792
    }
  }
  composition {
    chain {
      sequence = GNMQQNY
      num = 1
    }
  }
  keywords {
    general.xyzout = False
    general.root = "se_ca_frag"
    hand = both
    llgcompletion {
      complete = True
      scatterer = "CA"
    }
    llgmaps = True
    macsad {
      protocol = *custom
      macrocycle {
        ref_k = False
        ref_b = False
        ref_sigma = False
        ref_pk = True
        ref_pb = False
        ref_xyz = False
        ref_occ = True
        ref_bfac = False
        ref_fdp = True
        ref_sa = False
        ref_sb = False
        ref_sp = False
        ref_sd = False
        ncycle = 50
        target = *not_anom_only
        minimizer = *bfgs
      }
      macrocycle {
        ref_k = False
        ref_b = False
        ref_sigma = False
        ref_pk = True
        ref_pb = True
        ref_xyz = True
        ref_occ = True
        ref_bfac = True
        ref_fdp = True
        ref_sa = False
        ref_sb = False
        ref_sp = True
        ref_sd = True
        ncycle = 50
        target = *not_anom_only
        minimizer = *bfgs
      }
    }
  }
}
density_modification {
  dm = True
}
""")
  f.close()
  if (verbose) :
    out = sys.stdout
  else :
    out = cStringIO.StringIO()
  driver.run(args=["phaser_se_ca.eff", "--test_mode"], out=out)
  for file_name in expected_files :
    assert os.path.isfile(file_name)
  from iotbx import file_reader
  pdb_in = file_reader.any_file("se_ca_frag.pdb")
  pdb_in.assert_file_type("pdb")
  hierarchy = pdb_in.file_object.construct_hierarchy()
  xrs = pdb_in.file_object.xray_structure_simple()
  fcalc = xrs.structure_factors(d_min=1.8).f_calc()
  fcalc_fft_map = fcalc.fft_map()
  fcalc_real_map = fcalc_fft_map.apply_sigma_scaling().real_map_unpadded()
  unit_cell = xrs.unit_cell()
  llg_map_in = file_reader.any_file("se_ca_frag.1.llgmaps.mtz")
  llg_fft_map = llg_map_in.file_server.miller_arrays[0].fft_map()
  llg_map = llg_fft_map.apply_sigma_scaling().real_map_unpadded()
  for atom in hierarchy.atoms() :
    if (atom.element == "SE" or atom.element == "CA") :
      site_frac = unit_cell.fractionalize(site_cart=atom.xyz)
      assert (llg_map.tricubic_interpolation(site_frac) < 1.0)
  phased_map_in = file_reader.any_file("se_ca_frag.1_dm.mtz")
  fft_map = phased_map_in.file_server.miller_arrays[0].fft_map()
  real_map = fft_map.apply_sigma_scaling().real_map_unpadded()
  from scitbx.array_family import flex
  cc = flex.linear_correlation(real_map.as_1d(),
    fcalc_real_map.as_1d()).coefficient()
  assert (cc > 0.65) # 0.726 on my Mac - what about Linux?

def write_pdb_file () :
  import iotbx.pdb.hierarchy
  pdb_in = iotbx.pdb.hierarchy.input(source_info=None, pdb_string="""\
REMARK   simple peptide for testing experimental phasing, derived from 1yjp
CRYST1   21.937    4.866   23.477  90.00 107.08  90.00 P 1 21 1      2
ORIGX1      1.000000  0.000000  0.000000        0.00000
ORIGX2      0.000000  1.000000  0.000000        0.00000
ORIGX3      0.000000  0.000000  1.000000        0.00000
SCALE1      0.045585  0.000000  0.014006        0.00000
SCALE2      0.000000  0.205508  0.000000        0.00000
SCALE3      0.000000  0.000000  0.044560        0.00000
ATOM      1  N   GLY A   1      -9.009   4.612   6.102  1.00 16.77           N
ATOM      2  CA  GLY A   1      -9.052   4.207   4.651  1.00 16.57           C
ATOM      3  C   GLY A   1      -8.015   3.140   4.419  1.00 16.16           C
ATOM      4  O   GLY A   1      -7.523   2.521   5.381  1.00 16.78           O
ATOM      5  N   ASN A   2      -7.656   2.923   3.155  1.00 15.02           N
ATOM      6  CA  ASN A   2      -6.522   2.038   2.831  1.00 14.10           C
ATOM      7  C   ASN A   2      -5.241   2.537   3.427  1.00 13.13           C
ATOM      8  O   ASN A   2      -4.978   3.742   3.426  1.00 11.91           O
ATOM      9  CB  ASN A   2      -6.346   1.881   1.341  1.00 15.38           C
ATOM     10  CG  ASN A   2      -7.584   1.342   0.692  1.00 14.08           C
ATOM     11  OD1 ASN A   2      -8.025   0.227   1.016  1.00 17.46           O
ATOM     12  ND2 ASN A   2      -8.204   2.155  -0.169  1.00 11.72           N
ATOM     13  N   MSE A   3      -4.438   1.590   3.905  1.00 12.26           N
ATOM     14  CA  MSE A   3      -3.193   1.904   4.589  1.00 11.74           C
ATOM     15  C   MSE A   3      -1.955   1.332   3.895  1.00 11.10           C
ATOM     16  O   MSE A   3      -1.872   0.119   3.648  1.00 10.42           O
ATOM     17  CB  MSE A   3      -3.256   1.362   6.016  1.00 20.00           C
ATOM     18  CG  MSE A   3      -4.304   2.043   6.882  1.00 20.00           C
ATOM     19 SE   MSE A   3      -4.066   3.936   7.075  0.90 17.00          SE
ATOM     20  CE  MSE A   3      -2.436   3.903   8.145  1.00 20.00           C
ATOM     21  N   GLN A   4      -1.005   2.228   3.598  1.00 10.29           N
ATOM     22  CA  GLN A   4       0.384   1.888   3.199  1.00 10.53           C
ATOM     23  C   GLN A   4       1.435   2.606   4.088  1.00 10.24           C
ATOM     24  O   GLN A   4       1.547   3.843   4.115  1.00  8.86           O
ATOM     25  CB  GLN A   4       0.656   2.148   1.711  1.00  9.80           C
ATOM     26  CG  GLN A   4       1.944   1.458   1.213  1.00 10.25           C
ATOM     27  CD  GLN A   4       2.504   2.044  -0.089  1.00 12.43           C
ATOM     28  OE1 GLN A   4       2.744   3.268  -0.190  1.00 14.62           O
ATOM     29  NE2 GLN A   4       2.750   1.161  -1.091  1.00  9.05           N
ATOM     30  N   GLN A   5       2.154   1.821   4.871  1.00 10.38           N
ATOM     31  CA  GLN A   5       3.270   2.361   5.640  1.00 11.39           C
ATOM     32  C   GLN A   5       4.594   1.768   5.172  1.00 11.52           C
ATOM     33  O   GLN A   5       4.768   0.546   5.054  1.00 12.05           O
ATOM     34  CB  GLN A   5       3.056   2.183   7.147  1.00 11.96           C
ATOM     35  CG  GLN A   5       1.829   2.950   7.647  1.00 10.81           C
ATOM     36  CD  GLN A   5       1.344   2.414   8.954  1.00 13.10           C
ATOM     37  OE1 GLN A   5       0.774   1.325   9.002  1.00 10.65           O
ATOM     38  NE2 GLN A   5       1.549   3.187  10.039  1.00 12.30           N
ATOM     39  N   ASN A   6       5.514   2.664   4.856  1.00 11.99           N
ATOM     40  CA  ASN A   6       6.831   2.310   4.318  1.00 12.30           C
ATOM     41  C   ASN A   6       7.854   2.761   5.324  1.00 13.40           C
ATOM     42  O   ASN A   6       8.219   3.943   5.374  1.00 13.92           O
ATOM     43  CB  ASN A   6       7.065   3.016   2.993  1.00 12.13           C
ATOM     44  CG  ASN A   6       5.961   2.735   2.003  1.00 12.77           C
ATOM     45  OD1 ASN A   6       5.798   1.604   1.551  1.00 14.27           O
ATOM     46  ND2 ASN A   6       5.195   3.747   1.679  1.00 10.07           N
ATOM     47  N   TYR A   7       8.292   1.817   6.147  1.00 14.70           N
ATOM     48  CA  TYR A   7       9.159   2.144   7.299  1.00 15.18           C
ATOM     49  C   TYR A   7      10.603   2.331   6.885  1.00 15.91           C
ATOM     50  O   TYR A   7      11.041   1.811   5.855  1.00 15.76           O
ATOM     51  CB  TYR A   7       9.061   1.065   8.369  1.00 15.35           C
ATOM     52  CG  TYR A   7       7.665   0.929   8.902  1.00 14.45           C
ATOM     53  CD1 TYR A   7       6.771   0.021   8.327  1.00 15.68           C
ATOM     54  CD2 TYR A   7       7.210   1.756   9.920  1.00 14.80           C
ATOM     55  CE1 TYR A   7       5.480  -0.094   8.796  1.00 13.46           C
ATOM     56  CE2 TYR A   7       5.904   1.649  10.416  1.00 14.33           C
ATOM     57  CZ  TYR A   7       5.047   0.729   9.831  1.00 15.09           C
ATOM     58  OH  TYR A   7       3.766   0.589  10.291  1.00 14.39           O
ATOM     59  OXT TYR A   7      11.358   2.999   7.612  1.00 17.49           O
TER      60      TYR A   7
HETATM   61  O   HOH A   8      -6.471   5.227   7.124  1.00 22.62           O
HETATM   62  O   HOH A   9      10.431   1.858   3.216  1.00 19.71           O
HETATM   63  O   HOH A  11      11.808   4.179   9.970  1.00 23.99           O
HETATM   64  O   HOH A  12      13.605   1.327   9.198  1.00 26.17           O
HETATM   65  O   HOH A  14      -1.500   0.682  10.967  1.00 43.49           O
HETATM   66 CA    CA B   1     -11.286   1.756  -1.468  1.00 10.00          CA
END
""")
  if (os.path.exists("se_ca_frag.pdb")) :
    os.remove("se_ca_frag.pdb")
  f = open("se_ca_frag.pdb", "w")
  f.write(pdb_in.hierarchy.as_pdb_string(pdb_in.input.crystal_symmetry()))
  f.close()

if (__name__ == "__main__") :
  exercise(verbose=("--verbose" in sys.argv))
  print "OK"
