import IMP
import IMP.atom
import IMP.core
import IMP.saxs
import os

#! read experimental profile
exp_profile = IMP.saxs.Profile(IMP.saxs.get_example_path('saxs.dat'))

print 'min_q = ' + str(exp_profile.get_min_q())
print 'max_q = ' + str(exp_profile.get_max_q())
print 'delta_q = ' + str(exp_profile.get_delta_q())
saxs_score = IMP.saxs.Score(exp_profile)


m = IMP.Model()
sum_profile = IMP.saxs.Profile()

#! read PDBs
for struc in range(2):
  mp= IMP.atom.read_pdb(IMP.saxs.get_example_path('%i.fit'%struc), m,
                      IMP.atom.NonWaterNonHydrogenPDBSelector(), True, True)

  #! select particles from the model
  particles = IMP.atom.get_by_type(mp, IMP.atom.ATOM_TYPE)

  #! add radius for water layer computation
  #ft = IMP.saxs.default_form_factor_table()
  ft = IMP.saxs.FormFactorTable()
  for i in range(0, len(particles)):
    radius = ft.get_radius(particles[i])
    IMP.core.XYZR(particles[i]).set_radius(radius)
  # compute surface accessibility
  s = IMP.saxs.SolventAccessibleSurface()
  surface_area = s.get_solvent_accessibility(IMP.core.XYZRs(particles))

  #! calculate SAXS profile
  model_profile = IMP.saxs.Profile()
  model_profile.calculate_profile_partial(particles, surface_area)
# sum_profile = model_profile.add_partial_profiles(sum_profile)
# sum_profile = sum_profile.add_partial_profiles(model_profile)
  sum_profile.add_partial_profiles(model_profile)
  
  #! print and write out individual fit
  chi = saxs_score.compute_chi_score(model_profile)
  print ' ... Chi without parameter fitting = ' + str(chi)
  model_profile.write_SAXS_file('indiv_%i.dat'%struc)
  print " ... Fitting structure %i"%struc

#! 1. fit profile, calculate chi score and output the fit 
chi = saxs_score.fit_profile(sum_profile, False, 'fit_out.dat')
print 'Chi after adjustment of excluded volume and water layer parameters = ' + str(chi)
# 2. output profile that fits best to the experimental one
sum_profile.write_SAXS_file('out.dat')