#! /bin/awk -f
#
#   Constrain an mlphare script to have sites with
#   the same B-factor and overall occupancy,
#   regardless of wavelength
#
#   That is, impose the constraint that atoms with the same "ATOM#"
#   card should correspond to the same "site", and have the same xyz 
#   coordinate, B-factor, and "true" occupancy, regardless of wavelength.
#   The mlphare occupancies, however are modified by an fp and fpp 
#   value that is the same for every atom in a particular DERIV block (wavelength).
#
#   occ  = fp*tocc(site)
#   aocc = fpp*tocc(site)
#   Bfac = Bfac(site)
#
#   This script assumes a single value for fp and fpp for each wavelength, 
#   and, therefore, should not be used for double-edge MAD experiments! 
#
BEGIN{
    #  same_occ==1 means force all sites to have same (effective) occupancy
    if(same_occ=="") same_occ = 0
    #  same_B==1 means force all sites to have same B-factor
    if(same_B=="")   same_B   = 0
}

$1 ~ /^LABIN/ || labinline{
    labinline= ( $NF == "-" )
    for(i=1;i<=NF;++i){
	if($i ~ /^FP=/)
	{
	    FP = substr($i,index($i,"=")+1)
	}
	if($i ~ /^FPH/)
	{
	    wave = substr($i,4)+0
	    FPH[wave] = substr($i,index($i,"=")+1)
	}
	if($i ~ /^DPH/)
	{
	    wave = substr($i,4)+0
	    DPH[wave] = substr($i,index($i,"=")+1)
	}
    }
}

$1 ~ /^DERIV/{
    ++waves
}

$1 ~ /^ATOM/{
    # use atom number as "site" identifier
    atomnum = substr($1, 5)+0
    Bfac[atomnum] += $NF;
    X[atomnum] += $3; Y[atomnum] += $4; Z[atomnum] += $5;
    ++count[atomnum];
    
     fp[waves] += $6;
    fpp[waves] += $7;
    ++atoms[waves];
    
    # remember each atom's occupancy
     Occ[waves " " atomnum] = $6;
    aOcc[waves " " atomnum] = $7;
    ++occs[num];

    ++tatoms;
    B += $NF
}

{line[NR] = $0}

END{
    if(! waves) exit
    if(! tatoms) exit
    # overall, average B-factor
    B /= tatoms

    # compute mean f' and f" for each wavelength
    for(wave in fp)
    {
	# normalize anomalous and dispersive differences
	if(atoms[wave]) fp[wave]  =  fp[wave]/atoms[wave];
	if(atoms[wave]) fpp[wave] = fpp[wave]/atoms[wave];
    }
    
    # add up combined occupancy for each atom (averaged over all wavelengths)
    for(atomnum in Bfac)
    {
	# compute mean B-factor for each site
	if(count[atomnum]) Bfac[atomnum] = Bfac[atomnum]/count[atomnum];
	
	# compute mean XYZ position for each site
	if(count[atomnum])
	{
	    X[atomnum] = X[atomnum] / count[atomnum];
	    Y[atomnum] = Y[atomnum] / count[atomnum];
	    Z[atomnum] = Z[atomnum] / count[atomnum];
	}
	
	# compute mean occ,aocc for each wavelength
	occ_count = 0
	for(wave in fp)
	{
	    if( fp[wave])
	    {
		tOcc[atomnum] += Occ[wave " " atomnum]/fp[wave]
		++occ_count
	    }
	    if( fpp[wave])
	    {
		tOcc[atomnum] += aOcc[wave " " atomnum]/fpp[wave]
		++occ_count
	    }
	}
	# divide by total number of occupancies contributing to this "total"
	if(occ_count) tOcc[atomnum] = tOcc[atomnum]/occ_count
	
	# optionally set all occs to same value
	if(same_occ) tOcc[atomnum] = 1;
	if(same_B)   Bfac[atomnum] = B;
    }
    
    
    # now re-write the lines in the script
    wave = 0
    for(i=1;i<=NR;++i)
    {
	split(line[i],w);
	if(w[1] ~ /^DERIV/) ++wave
	
	if(w[1] ~ /^ATOM/)
	{
	    atomnum = substr(w[1], 5)+0
	    
	    line[i] = sprintf(" ATOM%-3d %3s  %6.3f %6.3f %6.3f %6.3f %6.3f BFAC %8.3f",\
	    atomnum, w[2], X[atomnum], Y[atomnum], Z[atomnum], \
	    fp[wave]*tOcc[atomnum], fpp[wave]*tOcc[atomnum], Bfac[atomnum])
	}
	
	print line[i];
    }
    
    print "Summary:"
    for(wave=1;wave<=waves;++wave)
    {
	printf "%s - %s = %6.3f \n", FPH[wave], FP, fp[wave];
    }
    for(wave=1;wave<=waves;++wave)
    {
	printf "%s = %6.3f \n", DPH[wave], fpp[wave];
    }
    for(wave=1;wave<=waves;++wave)
    {
#	printf "DERIV %d %6.3f %6.3f\n", wave, fp[wave], fpp[wave];
    }
}

