# Script to reposition all the atoms in specified selections, rotating them around z axis so that they all align along the y axis, i.e. from {x y z} to {0 r z}, where r - radial distance from the atom to the z axis. Then a <smoothed> average is calculated and plotted as semitransparent spheres.
# Andriy Anishkin (anishkin@icqmail.com) UMCP

#Selection criterion for atoms which will be repositioned
#set atoms_selection_criterion "resname POPC"
#set atoms_selection_criterion "name P1"
#set atoms_selection_criterion "name P1 and x >= 0 and y > 0 and z > 0"
# set atoms_selection_criterion "name P1 and z > 0"
set atoms_selection_criterion "name C218 C316"
# set atoms_selection_criterion "resname POPC and name N and z > 0"
# set atoms_selection_criterion "resname POPC and name C21 and z > 0"

#Option to average z position among the neighbors within avrg_dist radial distance from the atom - running average
set averaging 1
set avrg_width 5
set avrg_step 1
set neighborSmoothingTimes 3
set avrg_color "yellow"
set avrg_radius "1"

# set selectionsColorsList {};	# An alternative to atoms_selection_criterion and avrg_color. If this is not empty, selection-color pairs will be used from this list. 
set selectionsColorsList {
	{resname POPC and name N and z > 0} {blue}
	{resname POPC and name N and z < 0} {blue}
	{resname POPC and name P1 and z > 0} {tan}
	{resname POPC and name P1 and z < 0} {tan}
	{resname POPC and name O22 O21 O32 O31 and z > 0} {red}
	{resname POPC and name O22 O21 O32 O31 and z < 0} {red}
	{resname POPC and name C29 C210 and z > 0} {cyan}
	{resname POPC and name C29 C210 and z < 0} {cyan}
	{resname POPC and name C218 C316} {yellow}
};	# An alternative to atoms_selection_criterion and avrg_color. If this is not empty, selection-color pairs will be used from this list



#For averaging - option to handle atoms depending on z coordinate: <0> - handle all z together, as is, <1> use absolute value of z and reflect negative z up, <2> - calculate separately positive and negative z
set z_mode 0


if {$selectionsColorsList==""} {set selectionsColorsList [list $atoms_selection_criterion $avrg_color]}

for {set s 0} {$s < [llength $selectionsColorsList]/2} {incr s} {; # Cycle through the selection-color list
	draw color [lindex $selectionsColorsList [expr {$s*2+1}]]
	draw material Transparent
	
	set atoms_selection [atomselect top [lindex $selectionsColorsList [expr {$s*2}]]]
	set atoms_index [$atoms_selection get index]
	set numatoms [$atoms_selection num]
	
	
	#Cycle through the atoms to calculate radial distance and reposition"
	set counter 1
	set rs {}
	set xs [$atoms_selection get x]
	set ys [$atoms_selection get y]
	# set fileWrite [open [molinfo top get name]_radii.txt]
	set fileWrite [open radii_sel[set s]_[clock format [clock seconds] -format "%Y.%m.%d_%Hh%Mm%Ss"].txt w]
	puts $fileWrite "Radius\tz\t[lrange $selectionsColorsList [expr {$s*2}] [expr {$s*2+1}]]"
	
	foreach x $xs y $ys {
		lappend rs [expr {sqrt(pow($x,2) + pow($y,2))}]
# 		puts "atom $counter of $numatoms"
		incr counter
	}
	
	$atoms_selection set x 0
	$atoms_selection set y $rs
	
	#Running average of z among the neighbors
	if {$averaging == 1} {
		#Select z handling mode
		puts "Averaging for selection $s started"
		if {$z_mode == 0} {
			#handle all z together, as is
# 			puts "$z_mode == 0 - handle all z together, as is"
			
			#Get the list of radial positions (in y) and z coordinates
			set yzs [$atoms_selection get {y z}]
					
			#Sort the list of y and z by y
			set yzs [lsort -real -index 0 $yzs]
			set y_min [lindex [lindex $yzs 0] 0]
			set y_max [lindex [lindex $yzs [expr {[llength $yzs]-1}]] 0]
# 			puts "bins from $y_min to $y_max"
			
			
			#go through the bins and average
			set current_atom 0
			set atoms_max $numatoms
			set y_newList {}
			set z_newList {}
			for {set b $y_min} {$b < ($y_max-$avrg_width/2)} {set b [expr {$b + $avrg_step}]} {
				set ub [expr {$b + $avrg_width}]
# 				puts "Lower bin boundary at $b   between $y_min and $y_max"
				set y_count 0
				set y_new 0.0
				set z_new 0.0
				
				#Go through atoms, select atoms fitting into the bin and average them
				for {set a $current_atom} {$a < $atoms_max} {set ttt 0} {
					set y [lindex [lindex $yzs $a] 0]
					if {$y < $ub} {
						#y coordinate is smaller than upper boundary, include it and z
						set z [lindex [lindex $yzs $a] 1]
						set y_new [expr {$y_new + $y}]
						set z_new [expr {$z_new + $z}]
						incr y_count
						incr a
					} else {
						#y coordinate is larger than upper boundary, stop scanning and exit the current bin
						set current_atom $a
						set a $atoms_max
					}
				}
				if {$y_count > 0} {
					#There were some atoms in the bin
					#Average values between them
					lappend y_newList [expr {$y_new / $y_count}]
					lappend z_newList [expr {$z_new / $y_count}]
					
				}
			}
			
			for {set nn 0} {$nn < $neighborSmoothingTimes} {incr nn} {
				set y_newListSmooth {}
				set z_newListSmooth {}
				for {set a 1} {$a < [llength $y_newList]-1} {incr a} {
					lappend y_newListSmooth [expr {([lindex $y_newList [expr {$a-1}]]+[lindex $y_newList $a]+[lindex $y_newList [expr {$a+1}]])/3}]
					lappend z_newListSmooth [expr {([lindex $z_newList [expr {$a-1}]]+[lindex $z_newList $a]+[lindex $z_newList [expr {$a+1}]])/3}]
				}
				set y_newList [eval "lreplace \$y_newList 1 end-1 $y_newListSmooth"]
				set z_newList [eval "lreplace \$z_newList 1 end-1 $z_newListSmooth"]
			}
			foreach y_new $y_newList z_new $z_newList {
				#Draw sphere at the averaged position of the atom
				set sphere_coord "0 $y_new $z_new"
				draw sphere $sphere_coord radius $avrg_radius
				puts $fileWrite "$y_new $z_new"
			}
		} elseif {$z==1} {
			#use absolute value of z and reflect negative z up
		} elseif {$z==2} {
			#calculate separately positive and negative z
		}
	}
	
	close $fileWrite
}; # Cycle through the selection-color list
puts "Finished !!!"