# This script calculates the distance distribution for atoms within selections for all the frames of trajectory and records results as separate text files corresponding to the input files in filelist. Can be used for ion-'water oxygen' or water 'oxygen-hydrogen' etc. correlation function
# Andriy Anishkin (anishkin@icqmail.com) UMCP


#Open file and load part of DCD trajectory
#mol new mscs-cry-chnl-big-oct-cor_wtr_f.psf type psf
mol new mscs-cry-chnl-big-oct-cor_wtr_f_salt.psf type psf
#mol new mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt.psf type psf
mol off top
set first_frame 0
#set last_frame -1
set last_frame -1
#set filelist {sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_0.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_1.0.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_1.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_2.0.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_2.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_3.0.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_3.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_4.0.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_4.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_5.0.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_5.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_wrp_6.0.dcd}
#set filelist {sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_0.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_1.0.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_1.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_2.0.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_2.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_3.0.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_3.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_4.0.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_4.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_5.0.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_5.5.dcd sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_6.0.dcd}
#set filelist {sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_0.5.dcd sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_1.0.dcd sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_1.5.dcd sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_2.0.dcd sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_2.5.dcd sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_3.0.dcd sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_3.5.dcd sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_4.0.dcd sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_4.5.dcd sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_5.0.dcd sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_5.5.dcd sim_mscs-cry-chnl-big-oct-cor-l109s_wtr_f_salt_wrp_6.0.dcd}
#set filelist {sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_ionmove_e_up_wrp_a.dcd}
set filelist {sim_mscs-cry-chnl-big-oct-cor_wtr_f_salt_wrp_5.5-6.0_1to10.dcd}

#Set selection for atoms to estimate their hydration shell
set reference_atoms "name CLA and not ((within 10 of protein) or (within 10 of (resname OCT)))"
#Set selection for atoms of the shell
#set shell_atoms "name OH2"
set shell_atoms "name H1 H2"
#Set the maximal range for hydration shell analysis
set max_distance 10.025
#Set the minimal range for hydration shell analysis - usually the radius of ion and may be plus half-radius of water
#set min_distance [ expr {2.27+1.4/2 }]
set min_distance 1.725
#Set bins step for hydration shell analysis
set step_distance	0.05
set bins_medians {}
set bins_content {}

#Initialize the list of bins
for {set range [expr {$min_distance + $step_distance}]} {$range <= $max_distance} {} {
	lappend bins_medians [expr {$range - $step_distance/2}]
	lappend bins_content 0
	set range [expr {$range + $step_distance}]
}

#--------------- Block for counting ions distribution in hydration shell
#Cycle through the filelist
foreach crnt_file $filelist {
	#load file
	animate read dcd $crnt_file beg $first_frame end $last_frame waitfor all

	#Extract frames from file
	set num_steps [molinfo top get numframes]

	#Create selection of the molecules in the current layer of hydration shell
	set shell_layer_atoms_sel [atomselect top "(within $range of ($reference_atoms)) and $shell_atoms and beta 0" frame 0]
	set reference_atoms_sel [atomselect top "$reference_atoms"]
	set reference_atoms_total_number 0

	#Go through the frames and count atoms in the shell layers
	for {set frame 0} {$frame < $num_steps} {incr frame} {
	
		#clear all the marks of counted molecules in hydration shell
		set all [atomselect top all]
		$all set beta 0

		#Update the frame for reference_atoms and shell_layer_atoms
		set range $min_distance
		$reference_atoms_sel frame $frame
		$reference_atoms_sel update
# 		$shell_layer_atoms_sel frame $frame
# 		$shell_layer_atoms_sel update
		#Create selection of the molecules in the current layer of hydration shell
		set shell_layer_atoms_sel [atomselect top "(within $range of ($reference_atoms)) and $shell_atoms and beta 0" frame $frame]
		
		set reference_atoms_total_number [expr {	$reference_atoms_total_number + [$reference_atoms_sel num]}]
		
		#Mark atoms within the minimal distance
		$shell_layer_atoms_sel set beta 1
		$shell_layer_atoms_sel delete
		
		#Go through the shell layers and count atoms
		set i 0
		for {set range [expr {$min_distance + $step_distance}]} {$range <= $max_distance} {} {
			#count shell layer atoms and mark them as counted
			$reference_atoms_sel frame $frame
			$reference_atoms_sel update
# 			$shell_layer_atoms_sel frame $frame
# 			$shell_layer_atoms_sel update
			#Create selection of the molecules in the current layer of hydration shell
			set shell_layer_atoms_sel [atomselect top "(within $range of ($reference_atoms)) and $shell_atoms and beta 0" frame $frame]
			
			set shell_layer_atoms_number [$shell_layer_atoms_sel num]
			$shell_layer_atoms_sel set beta 1
			$shell_layer_atoms_sel delete
			#add value to the bins
			set bins_content [lreplace $bins_content $i $i [expr {[lindex $bins_content $i] + $shell_layer_atoms_number}]]
			incr i
			puts "Shell $range of $max_distance      $frame of [expr {($num_steps - 1)}]"
			set range [expr {$range + $step_distance}]
# 			puts -nonewline "."
		}
# 		puts ""
		puts "global frame [expr {$frame + $first_frame}]     file $crnt_file       frame $frame of [expr {($num_steps - 1)}] finished"
# 		puts "bins_content $bins_content"
	}
	
	#Open file for writing results
	set filename $crnt_file.shl_dstr.txt
	set fid [open $filename w]
	puts $fid "shell median, Å\tprobability, 1/Å^3]"
	#Go through shell layers and record results
	set layers_number [llength $bins_medians]
	for {set i 0} {$i < $layers_number} {incr i} {
		#normalize results by the total number of reference atoms
		set bins_content [lreplace $bins_content $i $i [expr {double([lindex $bins_content $i])/(double($reference_atoms_total_number)*((pow(([lindex $bins_medians $i]+$step_distance/2),3)-pow(([lindex $bins_medians $i]-$step_distance/2),3))*3.1415926*4/3))}]]
		#record to the file
		puts $fid "[lindex $bins_medians $i]\t[lindex $bins_content $i]"
	}
	#close written files
	close $fid

	set first_frame [expr {$first_frame + $num_steps}]
	animate delete all
}


bell
puts "Finished !!!"