import math
import sys
import os

from iotbx import mtz
from cctbx import sgtbx
from cctbx import crystal
from cctbx import uctbx
from scitbx import matrix

def measure(hklin, spacegroup):
  '''Look at HKLIN, see how strong the absences (according to the given
  spacegroup) are... looking for IPR / SIGIPR.'''

  mtz_obj = mtz.object(hklin)

  sg = sgtbx.space_group(sgtbx.space_group_symbols(spacegroup).hall())
  mi = mtz_obj.extract_miller_indices()
  sg_m = mtz_obj.space_group()

  ipr_column = None
  sigipr_column = None

  for crystal in mtz_obj.crystals():
    for dataset in crystal.datasets():
      for column in dataset.columns():

        if column.label() == 'IPR':
          ipr_column = column
        elif column.label() == 'SIGIPR':
          sigipr_column = column

  assert(ipr_column != None)
  assert(sigipr_column != None)

  ipr_values = ipr_column.extract_values(not_a_number_substitute = 0.0)
  sigipr_values = sigipr_column.extract_values(not_a_number_substitute = 0.0)

  present = []
  absent = []

  for j in range(mi.size()):
    hkl = mi[j]

    if sg.is_sys_absent(hkl):
      absent.append(ipr_values[j] / sigipr_values[j])
    else:
      present.append(ipr_values[j] / sigipr_values[j])

  print 'A: %f' % (sum(absent) / len(absent))
  print 'P: %f' % (sum(present) / len(present))

if __name__ == '__main__':

  hklin = sys.argv[1]
  spacegroup = sys.argv[2]

  measure(hklin, spacegroup)