import sys, os, time, math
from scitbx.array_family import flex
import iotbx.pdb
import scitbx.math
from cctbx import maptbx
import scitbx.math

pdb_str = """\
CRYST1   10.000   10.000   10.000  90.00  90.00  90.00 P 1
HETATM  120  CA  CA  A   1       5.000   5.000   5.000  1.00  0.01          CA
TER
END
"""

def convert_atom(xrs, scattering_type, i_seq):
  scatterer = xrs.scatterers()[i_seq]
  scatterer.scattering_type = scattering_type
  xrs.discard_scattering_type_registry()
  xrs.scattering_type_registry(
    table = "n_gaussian",
    d_min = 0.,
    types_without_a_scattering_contribution=["?"])
  return xrs  
    
def scale(o,c, i=None):
  o = flex.abs(o)
  c = flex.abs(c)
  if(i is not None):
    o = o[:i]
    c = c[:i]
  sc = flex.sum(o*c)/flex.sum(c*c)
  return sc  
      
def run(b_iso = 10.0):
  xrs = iotbx.pdb.input(source_info=None, lines=pdb_str).xray_structure_simple()
  xrs = xrs.set_b_iso(value=b_iso)
  xrs.scattering_type_registry(
    table = "n_gaussian",
    d_min = 0.,
    types_without_a_scattering_contribution=["?"])
  xrs.scattering_type_registry().show()
#  
  stol  = flex.double()
  ff_ca = flex.double()
  stol_ = 0
  while stol_<6:
     dss = (stol_ * 2)**2
     ff_ca.append(xrs._scattering_type_registry.unique_form_factors_at_d_star_sq(dss)[0])
     stol.append(stol_)
     stol_ +=0.01
#
  ff_zn = flex.double()
  xrs = convert_atom(xrs=xrs, scattering_type="Zn", i_seq=0)
  stol_ = 0
  while stol_<6:
     dss = (stol_ * 2)**2
     ff_zn.append(xrs._scattering_type_registry.unique_form_factors_at_d_star_sq(dss)[0])
     stol_ +=0.01
#  
  sc = scale(ff_ca,ff_zn)
  of = open("z","w")
  for xx,yy,zz in zip(stol,ff_ca,ff_zn):
    print >> of, "%12.6f %12.6f %12.6f %12.6f"%(xx,yy,zz, zz*sc)
  of.close()
#   
# If you plot density with different B (see run2.py) you will see at B~0...1
# there is a difference between shapes, while at higher B the only difference 
# between atom types is a single scale factor!


if (__name__ == "__main__"):
  run()