import sys
import math
import os
from iotbx import mtz
from cctbx.eltbx import sasaki

def wavelength_energy(a):
  h = 6.6260693e-34
  c = 2.9979246e8
  e = 1.6021765e-19
  conv = 1.0e10 * h * c / e
  return conv / a

def guess_the_atom(hklin, nsites):
  '''Guess the atom which gives rise to the observed anomalous differences
  in intensities (i.e. I(+) and I(-)) though CCTBX code internally computes
  F(+) etc.'''

  mtz_obj = mtz.object(hklin)
  mi = mtz_obj.extract_miller_indices()

  sg = mtz_obj.space_group()

  for crystal in mtz_obj.crystals():
    if crystal.name() != 'HKL_base':
      uc = crystal.unit_cell()

  n_ops = len(sg.all_ops())
  v_asu = uc.volume() / n_ops
  mw = v_asu / 2.7

  atoms = ['Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Se', 'Br', 'S', 'P']

  tables = [sasaki.table(atom) for atom in atoms]

  for crystal in mtz_obj.crystals():
    if crystal.name() == 'HKL_base':
      continue
    assert(len(crystal.datasets()) == 1)
    for dataset in crystal.datasets():
      wavelength = dataset.wavelength()

  energy = wavelength_energy(wavelength)

  mas = mtz_obj.as_miller_arrays()

  best_atom = None
  best_diff = 100.0

  for ma in mas:
    columns = ma.info().label_string()
    if 'I(+)' in columns and 'I(-)' in columns:
      signal = ma.anomalous_signal()
      for j, atom in enumerate(atoms):
        for energy_offset in (-100.0, 0.0, 100.0):
          wavelength = wavelength_energy(energy + energy_offset)
          fdp = tables[j].at_angstrom(wavelength).fdp()
          p_signal = fdp * math.sqrt(nsites / mw)
          if math.fabs(p_signal - signal) < best_diff:
            best_diff = math.fabs(p_signal - signal)
            best_atom = atom

  return best_atom

if __name__ == '__main__':

  print guess_the_atom(sys.argv[1], int(sys.argv[2]))