#Script to adjust partial atom charges to approximate dielectric permeability of environment. Used within symmetrization script
set configFileName mscs-symm_conf.txt;	# ==>	Simmetrization Configuration file name
set workingFolder {/simulations/mscs-cry_symmetrization/adj-chrg_2_compress}
cd $workingFolder


set psf_filename {mscs-cry-cor-ini.psf};	# ==>	File name for the initial psf file
# set pdb_filename {mscs-cry-cor.pdb}
# set psf_filename {zz.psf}
# set pdb_filename {zz.pdb}
# set psf_filename_output "$psf_filename.out.psf"

# ==>	Define spherical shell parameters
# set shell_thickness 1.77
# set shell_thickness 3.54
set shell_thickness 6.5

# ==>	Mode of the nonuniform dielectric permeability:
#0 - uniform dielectric
#1 - nonuniform dielectric with fixed boundary radius profile criterion
#2 - nonuniform dielectric with pore/outside surface criterion
#3 - nonuniform dielectric with sequence criterion
set nonuniform_dielectric_permeability 1
set nonuniform_dielectric_permeability_filename {z_epsilon-outside_epsilon-inside_boundary-radius.txt}
set uniform_dielectric_permeability 80
set use_explicit_solvent 0

# ___ reading symmetrization status from the configuration and status files
set fileLook [open $configFileName r];	# Open configuration file
set parametersFound 0
while {([gets $fileLook line] >= 0)&($parametersFound<6)} {
	set lineRecognize [lindex $line 0]
	switch -regexp $lineRecognize {
		"^statusFileName$" {;	# Status file name
			set statusFileName [lindex $line 1]
			incr parametersFound
		}
		"^structure$" {;	# Protein structure file
			set psf_filename_output [lindex $line 1]
			incr parametersFound
		}
		"^coordinates$" {;	# Protein coordinates
			set pdb_filename [lindex $line 1]
			incr parametersFound
		}
		"^coordinatesFileType$" {;	# Coordinates File type
			set coordinatesFileType [lindex $line 1]
			incr parametersFound
		}
		"^symmetrizationStepsMax$" {;	# Maximal possible number of symmetrization steps, before it stops (if was not stopped earlier by other criterion).
			set symmetrizationStepsMax [lindex $line 1]
			incr parametersFound
		}
		"^logOutputFiles$" {;	# Option to keep log (numbered copies) of all output files - symmetrized structure, constants file, status file
			set logOutputFiles [lindex $line 1]
			incr parametersFound
		}
		default {}
	}
}
close $fileLook


# find out the last step that was performed
set fileLook [open $statusFileName r];	# open configuration file
set parametersFound 0
while {([gets $fileLook line] >= 0)&($parametersFound<1)} {;	# read configuration file and recognize data
	set lineRecognize [lindex $line 0]
	switch -regexp $lineRecognize {
		"^symmetrizationCyclesCompleted$" {;	# Number of minimization cycles completed up to the moment
			set symmetrizationCyclesCompleted [lindex $line 1]
			incr parametersFound
		}
		default {}
	}
}
close $fileLook


set start_time [clock seconds]

#Obtain from files (if necessary) and adjust values of dielectric permeability
#Since all the charges are adjusted, each of them should get square root of the value of the dielectric permeability
if {$nonuniform_dielectric_permeability == 0} {
	#0 - uniform dielectric
	set uniform_dielectric_permeability [expr {pow(double($uniform_dielectric_permeability),0.5)}]

} elseif {$nonuniform_dielectric_permeability == 1} {
	#1 - nonuniform dielectric with fixed boundary radius profile criterion
	set dielectric_file [open $nonuniform_dielectric_permeability_filename r]
	#It reads text file with 4 columns of data, separated by single spaces. File have 1st line with column headers.
	#Columns are:	Z;	Epsilon_outside;	Epsilon_inside;		Boundary_radius
	gets $dielectric_file line
	#Complete list of nonuniform dielectric properties
	set ZEER {}
	#Coordinates of non-uniform dielectric
	set Ze {}
	#Dielectric permeability of the media outside the channel (water/membrane system)
	set Eo {}
	#Dielectric permeability of the media inside the channel (water)
	set Ei {}
	#Radius profile of the boundary between the outer and inner media
	set Re {}
	while {[gets $dielectric_file line] > 0} {
	    lappend ZEER $line
	}
	close $dielectric_file
	set ZEER [lsort -real -increasing -index 0 $ZEER]
	foreach line $ZEER {
	    lappend Ze [lindex $line 0]
	    lappend Eo [expr {pow(double([lindex $line 1]),0.5)}] 
	    lappend Ei [expr {pow(double([lindex $line 2]),0.5)}]
	    lappend Re [lindex $line 3]
    }
    unset ZEER
    set Ze_min [lindex $Ze 0]
    set Ze_max [lindex $Ze [expr {[llength $Ze]-1}]]
    set Eo_min [lindex $Eo 0]
    set Eo_max [lindex $Eo [expr {[llength $Eo]-1}]]
	#Initialize the number of the bin from the previous nonuniform dielectric permeability search
	set last_Ze_bin 0
	#The number of the bin before the highest one 
	set Ze_bin_max [expr {[llength $Ze]-1}]


} elseif {$nonuniform_dielectric_permeability == 2} {
	#2 - nonuniform dielectric with pore/outside surface criterion
} elseif {$nonuniform_dielectric_permeability == 3} {
	#3 - nonuniform dielectric with sequence criterion
}




# open the psf file
set file_look [open $psf_filename r]


#Flag for the start of the charges section in the PSF file: before (-1), inside (0), and after (1) charges section
set atom_charges_section -1


puts "PSF reading started "
#Flag for continuing reading, until end of file will be reached and flag will be -1
set continue_reading 1
set full_psf_content {}
set charges_list {}
set line_counter -1
while {$continue_reading >= 0} {
	set continue_reading [gets $file_look line]
	set length [llength $line]
	if {$length > 0 | $continue_reading>=0} {
		lappend full_psf_content $line
		incr line_counter
	
		if {$atom_charges_section==-1} {
			#before atom_charges_section
			set atom_charges_flag [string compare [lindex $line 1] !NATOM]
			if {$atom_charges_flag==0} {
				set atom_charges_section 0
				set atom_charges_section_start [expr {$line_counter + 1}]
				set atom_charges_section_end [expr {$atom_charges_section_start + [lindex $line 0] - 1}]
			}
		} elseif {$atom_charges_section==0} {
			#atom_charges_section
			set atom_bonds_flag [string compare [lindex $line 1] !NBOND:]
			if {$atom_bonds_flag==0} {set atom_charges_section 1}
			#puts [lindex $line 6]
			if {[llength [lindex $line 6]] >0} {
				lappend charges_list [lindex $line 6]
			}
		} elseif {$atom_charges_section==1} {
			#after atom_charges_section
		}
	}
}
close $file_look
#close $fid_test

#puts $charges_list
puts "PSF reading completed "


#Load the molecule
mol new $psf_filename type psf
# mol addfile $pdb_filename type pdb
mol addfile $pdb_filename type namdbin

if {$use_explicit_solvent == 1} {
	puts "Water box creation started"
	#----------- use explicit solvent
	#Create the water box
	set minmax [measure minmax [atomselect top all]]
	foreach {min max} $minmax { break }
	foreach {xmin ymin zmin} $min { break }
	foreach {xmax ymax zmax} $max { break }
	molecule delete top
	eval "solvate $psf_filename $pdb_filename -minmax {{$xmin $ymin $zmin} {$xmax $ymax $zmax}} -o solvated_protein -t 15 -b 1.77"
	
	#Load water box
	mol new solvated_protein.psf type psf
	mol addfile solvated_protein.pdb type pdb
	puts "Water box creation completed"
	
	
	
	#assign NAMD VDW radii according to the values read from file with atom types and radii
	puts "VDW radii assignment started"
	#File with radii - first line - column names, ignored. Other lines - 2 values - atom type and radius - separated by tabs
	set filename "atom_type_vdw_radii.txt"
	# open the file
	set file_look [open $filename r]
	#read the data
	gets $file_look line
	set atom_type {}
	set atom_radius {}
	while {[gets $file_look line] > 0} {
		lappend atom_type [lindex $line 0]
		lappend atom_radius [lindex $line 1]
	}
	#Close the file
	close $file_look
	#Go through atom types and assign the radii. If atom type is unknown, radius is not changed
	foreach type $atom_type radius $atom_radius {
		set sel [atomselect top "type $type"]
		$sel set radius $radius
	}
	puts "VDW radii assignment completed"
	
	
	
	# set atom_radius 2.0
	# set shell_radius [expr {$atom_radius + $shell_thickness}]
	
	
	
	# Calculate new charges according to the solvent exposure
	puts "New charges calculation started"
	#Estimate total charge
	set total_charge_old [eval "vecadd $charges_list"]
	set number_of_charges [llength $charges_list]
	set new_charges_list {}
	#Calculate new charges
	set all [atomselect top all]
	$all set beta 0
	set water [atomselect top water]
	$water set beta -1
	for {set charge_number 0} {$charge_number < $number_of_charges} {incr charge_number} {
		set current_atom_selection [atomselect top "index $charge_number"]
		#atom charge from psf
		set atom_charge [lindex $charges_list $charge_number]
		set atom_radius [$current_atom_selection get radius]
		set shell_radius [expr {$atom_radius + $shell_thickness}]
		#Find the hydration shell
		set solvation_shell_selection [atomselect top "water and within $shell_radius of (index $charge_number)"]
		#Calculate relative density of the hydration shell
		set solvation_shell_density [expr {double([$solvation_shell_selection num])/(0.033177208*3*((pow($shell_radius,3)-pow($atom_radius,3))*3.1415926*4/3))}]
		#Remember relative density of the hydration shell
		$current_atom_selection set beta $solvation_shell_density
		#New adjusted charge
		if {$nonuniform_dielectric_permeability == 0} {
			set atom_charge [expr {$atom_charge/(1+($uniform_dielectric_permeability-1)*$solvation_shell_density)}]
		} elseif {$nonuniform_dielectric_permeability == 1} {
		} elseif {$nonuniform_dielectric_permeability == 2} {
		} elseif {$nonuniform_dielectric_permeability == 3} {
		}
		#Remeber new adjusted charge
		lappend  new_charges_list $atom_charge
		#Unset selections
		unset current_atom_selection
		unset atom_charge
		unset atom_radius
		unset solvation_shell_selection
		unset solvation_shell_density
		unset shell_radius
	}
	set charges_list $new_charges_list
	$all writepdb {solvated_protein_adjusted.pdb}
	molecule delete top
	puts "New charges calculation completed"
} else {
	#Dont use explicit solvent
		
	# Calculate new charges according to the solvent exposure
	puts "New charges calculation started"
	#Estimate total charge
	set total_charge_old [eval "vecadd $charges_list"]
	set number_of_charges [llength $charges_list]
	set new_charges_list {}
	#Calculate new charges
	set all [atomselect top all]
	$all set beta 0
	$all set radius 0
	set shell_radius $shell_thickness
	for {set charge_number 0} {$charge_number < $number_of_charges} {incr charge_number} { 
		set current_atom_selection [atomselect top "index $charge_number"]
		#atom charge from psf
		set atom_charge [lindex $charges_list $charge_number]
		#set atom_radius [$current_atom_selection get radius]
		#set shell_radius [expr {$atom_radius + $shell_thickness}]
		#Find the hydration shell
		set desolvation_shell_selection [atomselect top "within $shell_radius of (index $charge_number)"]
		#Calculate relative density of the hydration shell
		#set desolvation_shell_mass [expr {double([eval "vecadd [$desolvation_shell_selection get mass]"]-[$current_atom_selection get mass])}]
		#set solvation_shell_density [expr {1-$desolvation_shell_mass/(6*0.033177208*3*((pow($shell_radius,3)-pow($atom_radius,3))*3.1415926*4/3))}]
		#set solvation_shell_density [expr {1-double([$desolvation_shell_selection num]-1)/(0.033177208*3*((pow($shell_radius,3)-pow($atom_radius,3))*3.1415926*4/3))}]
		set solvation_shell_density [expr {1-double([$desolvation_shell_selection num])/(0.036*3*((pow($shell_radius,3))*3.1415926*4/3))}]
		#Limit out-of-range values of the water shell density
		if {$solvation_shell_density < 0} {
			set solvation_shell_density 0
		} elseif {$solvation_shell_density > 1} {
			set solvation_shell_density 1
		}
		#Remember relative density of the hydration shell
		$current_atom_selection set beta $solvation_shell_density
		#New adjusted charge
		if {$nonuniform_dielectric_permeability == 0} {
			#0 - uniform dielectric
			set atom_charge [expr {$atom_charge/(1+($uniform_dielectric_permeability-1)*$solvation_shell_density)}]
		} elseif {$nonuniform_dielectric_permeability == 1} {
			#1 - nonuniform dielectric with fixed boundary radius profile criterion
			set atom_z [$current_atom_selection get z]
			if {$atom_z < [lindex $Ze  $last_Ze_bin]} {
				#Moving down through the bins
				while {$last_Ze_bin >= 0}	{
					#we are above the lowest bin boundary
					if {$atom_z >= [lindex $Ze  $last_Ze_bin]} {
						#fit in the current bin
						break
					} else {
						#1 step down
						incr last_Ze_bin -1
					}
				}
			} else {
				#Moving up through the bins
				while {$last_Ze_bin < $Ze_bin_max}	{
					#we are below the topmost bin boundary
					if {$atom_z < [lindex $Ze  [expr {$last_Ze_bin+1}]]} {
						#fit in the current bin
						break
					} else {
						#1 step up
						incr last_Ze_bin
					}
				}
			}

			
			if {$last_Ze_bin < 0} {
				set local_dielectric_permeability $Eo_min
			} elseif {$last_Ze_bin >= $Ze_bin_max} {
				set local_dielectric_permeability $Eo_max
			} else {
				#Define radial position relative to the boundary between the outer and inner media
				set atom_R [expr {pow((pow([$current_atom_selection get x],2)+pow([$current_atom_selection get y],2)),0.5)}]
				if {$atom_R >= [lindex $Re $last_Ze_bin]} {
					#Atom in the outer dielectric
					set local_dielectric_permeability [expr {[lindex $Eo $last_Ze_bin] + ([lindex $Eo [expr {$last_Ze_bin+1}]]-[lindex $Eo $last_Ze_bin])*($atom_z-[lindex $Ze $last_Ze_bin])/([lindex $Ze [expr {$last_Ze_bin+1}]]-[lindex $Ze $last_Ze_bin])}]
				} else {
					#Atom in the inner dielectric
					set local_dielectric_permeability [expr {[lindex $Ei $last_Ze_bin] + ([lindex $Ei [expr {$last_Ze_bin+1}]]-[lindex $Ei $last_Ze_bin])*($atom_z-[lindex $Ze $last_Ze_bin])/([lindex $Ze [expr {$last_Ze_bin+1}]]-[lindex $Ze $last_Ze_bin])}]
				}
				unset atom_R
			}
			
			set atom_charge [expr {$atom_charge/(1+($local_dielectric_permeability-1)*$solvation_shell_density)}]
			#puts "$atom_z $last_Ze_bin      $atom_R $local_dielectric_permeability"
			unset local_dielectric_permeability
			unset atom_z
			
		} elseif {$nonuniform_dielectric_permeability == 2} {
			#2 - nonuniform dielectric with pore/outside surface criterion
		} elseif {$nonuniform_dielectric_permeability == 3} {
			#3 - nonuniform dielectric with sequence criterion
		}
		#Remeber new adjusted charge
		lappend  new_charges_list $atom_charge
		#Unset selections
# 		unset current_atom_selection
		$current_atom_selection delete
		unset atom_charge
		#unset atom_radius
		#unset desolvation_shell_selection
		$desolvation_shell_selection delete
		#unset desolvation_shell_mass
		unset solvation_shell_density
		#unset shell_radius
		#puts "$charge_number of $number_of_charges"
	}
	set charges_list $new_charges_list
	$all writepdb {solvated_protein_adjusted.pdb}
	molecule delete top
	puts "New charges calculation completed"

}

#Adjust total charge to keep it the same as before new charges calculation
puts "Total charge adjustment started"
set total_charge_new [eval "vecadd $charges_list"]
if {$total_charge_new != $total_charge_old} {
	#Charge adjustment is required
	set new_charges_list {}
	#Calculate shift charge - this small value will be added to all the atoms
	set shift_charge [expr {($total_charge_old - $total_charge_new)/$number_of_charges}]
	puts "Old total charge: $total_charge_old    New total charge: $total_charge_new    Correction charge per atom: $shift_charge"
	#Add shift charge to all the atoms
	for {set charge_number 0} {$charge_number < $number_of_charges} {incr charge_number} {
		lappend new_charges_list [expr {[lindex $charges_list $charge_number] + $shift_charge}]
	}
	set charges_list $new_charges_list
}
puts "Total charge adjustment completed"


#Insert modified charges into charges section of memorized PSF file
#1234567890123456789012345678901234567890123456789012345678901234567890
#       1 CER  1    CER  ABCD Habc -12.456789E+003 123.0080           0
#PSF charge columns: 35 to 49, or (without exponent) 35 to 44
puts "Charges replacement in PSF started"
set full_psf_content_insert {}
set charge_number 0
for {set line_number $atom_charges_section_start} {$line_number <= $atom_charges_section_end} {
	incr line_number
	incr charge_number
	} {
	set new_charge [format "%10.6f" [lindex $charges_list $charge_number]]
	lappend full_psf_content_insert [string replace [lindex $full_psf_content $line_number] 34 43 $new_charge]
}
set full_psf_content [eval "lreplace {$full_psf_content} $atom_charges_section_start $atom_charges_section_end $full_psf_content_insert"]

puts "Charges replacement in PSF completed"


#Record results to the new PSF file
puts "PSF writing started"
set fid [open $psf_filename_output w]
foreach line $full_psf_content {
	puts $fid $line
}
close $fid

puts "PSF writing completed"

# Make a copy for the log, if necessary
if {$logOutputFiles == 1} {
	file copy -force $psf_filename_output $psf_filename_output[format "_%0[string length $symmetrizationStepsMax]i" $symmetrizationCyclesCompleted]
	catch {zip -9oqm $psf_filename_output[format "_%0[string length $symmetrizationStepsMax]i" $symmetrizationCyclesCompleted].zip $psf_filename_output[format "_%0[string length $symmetrizationStepsMax]i" $symmetrizationCyclesCompleted]}
} 


puts "Finished !!!"


set spent_time [expr {[clock seconds]-$start_time}];
set spent_time_h [expr {int(floor($spent_time/3600))}];
set spent_time_m [expr {int(floor(($spent_time-$spent_time_h*3600)/60))}];
set spent_time_s [expr {int($spent_time-$spent_time_h*3600-$spent_time_m*60)}];

puts "time spent:  $spent_time_h hours  $spent_time_m min  $spent_time_s sec"

exit