#order parameter calculator 
#created May 2007 by Richard Swenson



#*************************************************************************#
# normal returns the normal of the least squares fit of a plane to the    #
# phosphorus atoms within the radius of point and on the same leaflet     #
#       point    index of the center phosphorus                           #
#       radius   radius in angstroms					  #
#	frame    current trajectory frame                                 #
#*************************************************************************#
proc normal {point radius frame} {
	

        #dimmensions of the periodic boundaries	
        set box_x [molinfo top get a frame $frame]
        set box_y [molinfo top get b frame $frame]
        set box_z [molinfo top get c frame $frame]
        #puts "x: $box_x y: $box_y z: $box_z"


        set p_point [atomselect top "(same fragment as index $point) and name P" frame $frame]
        set reference "O2"
	set r_point [atomselect top "(same fragment as index $point) and name $reference" frame $frame]
        set leaf 1
        if {[$p_point get z] < [$r_point get z]} {
                set leaf -1
        }
	#puts "[$p_point get z] < [$r_point get z] leaf: $leaf"	


        #select atoms within the radius of selected atom
        set close [atomselect top "name P and (within $radius of index $point)" frame $frame]
        set closelist [$close get index]
        $close delete

        #select atoms within the radius of selected atom's projections due to periodicity
        set original_x [$p_point get x]
        set original_y [$p_point get y]
        set pos_x [expr $original_x + $box_x]
        set neg_x [expr $original_x - $box_x]
        set pos_y [expr $original_y + $box_y]
        set neg_y [expr $original_y - $box_y]

        $p_point set x $pos_x
        set imagepx [atomselect top "name P and (within $radius of index $point)" frame $frame]
        set imagepxlist [$imagepx get index]
        $imagepx delete
        $p_point set x $neg_x
        set imagenx [atomselect top "name P and (within $radius of index $point)" frame $frame]
        set imagenxlist [$imagenx get index]
        $imagenx delete
        $p_point set x $original_x
        $p_point set y $pos_y
	set imagepy [atomselect top "name P and (within $radius of index $point)" frame $frame]
        set imagepylist [$imagepy get index]
        $imagepy delete
        $p_point set y $neg_y
        set imageny [atomselect top "name P and (within $radius of index $point)" frame $frame]
        set imagenylist [$imageny get index]
        $imageny delete
        $p_point set y $original_y

        #combine list and remove duplicate atoms
        set templist [concat $closelist $imagepxlist $imagenxlist $imagepylist $imagenylist]
        set nearlist [lsort -unique $templist]
        #puts "$templist -> $nearlist"
	#puts "$templist"
	#puts "$nearlist"

        #average Z for sigma=(1/N)sum((zi-aveZ)^2)
        set N 0
        set aveZ 0

        foreach p_index $nearlist {
                set p [atomselect top "index $p_index" frame $frame]
                set r [atomselect top "(same fragment as index $p_index) and name $reference" frame $frame]
                set pz [$p get z]
                set rz [$r get z]
                #puts "pz: $pz rz: $rz p_index: $p_index"
                set diff [expr ($pz-$rz)*$leaf]
                #puts "pz: $pz rz: $rz diff: $diff"

                if {$diff > 0} {
                        set N [expr $N + 1]
                        set aveZ [expr $aveZ + $pz]
                }
		#puts "aveZ: $aveZ"
                $p delete
                $r delete
        }
	#puts "averaging"
	if {$N == 0} {
		puts "no atoms near [$p_point get index] in calculating normal"
		return [list 0 0 0]
	}
        set aveZ [expr $aveZ/$N]
	#puts "aveZ: $aveZ"
	#calculate least squares fit to a plain
        #  F(x,y)=ax+by+c  X^2=sum(((zi-axi-byi-ci)/sigma)^2)
        #  e=sum(xi/((sigma)^2))
        #  f=sum(yi/((sigma)^2))
        #  g=sum(zi/((sigma)^2))
        #  h=sum((xi*xi)/((sigma)^2))
        #  i=sum((yi*yi)/((sigma)^2))
        #  j=sum((zi*zi)/((sigma)^2))
        #  k=sum((xi*yi)/((sigma)^2))
        #  l=sum((xi*zi)/((sigma)^2))
        #  m=sum((yi*zi)/((sigma)^2))

        set e 0
        set f 0
        set g 0
        set h 0
        set i 0
        set j 0
        set k 0
        set l 0
        set m 0

        foreach p_index $nearlist {
                set p [atomselect top "index $p_index" frame $frame]
                set r [atomselect top "(same fragment as index $p_index) and name $reference" frame $frame]
                set xi [$p get x]
                set yi [$p get y]
                set zi [$p get z]
                set zr [$r get z]
                #puts "pz: $pz rz: $rz p_index: $p_index"
                set diff [expr ($zi-$zr)*$leaf]
                #puts "zi: $zi zr: $zr diff: $diff"
		if {$diff > 0} {
                        set sigma [expr (($pz-$aveZ)*($pz-$aveZ))]
			#puts "pz: $pz aveZ:$aveZ sigma: $sigma"
			if {$sigma == 0} {
				set sigma .000000000001
			}
			set sigma 1
                        set e [expr $e+($xi/(($sigma)*($sigma)))]
                        set f [expr $f+($yi/(($sigma)*($sigma)))]
                        set g [expr $g+($zi/(($sigma)*($sigma)))]
                        set h [expr $h+(($xi*$xi)/(($sigma)*($sigma)))]
                        set i [expr $i+(($yi*$yi)/(($sigma)*($sigma)))]
                        set j [expr $j+(($zi*$zi)/(($sigma)*($sigma)))]
                        set k [expr $k+(($xi*$yi)/(($sigma)*($sigma)))]
                        set l [expr $l+(($xi*$zi)/(($sigma)*($sigma)))]
                        set m [expr $m+(($yi*$zi)/(($sigma)*($sigma)))]
                }

                $p delete
                $r delete
        }
	set oi [expr ($e*$f-$N*$k)*($e*$f-$N*$k)-($e*$e-$N*$h)*($f*$f-$N*$i)]
	set pi [expr $oi*($e*$e-$N*$h)]
        set si [expr $e*($e*$f-$N*$k)+$f*($N*$h-$e*$e)]
        set ti [expr $N*($e*$e-$N*$h)]
	set qi [expr $oi*$e+$si*($N*$k-$e*$f)]
	set ri [expr $ti*($N*$k-$e*$f)]
	set di [expr -$N*($e*$f-$N*$k)]
	set aai [expr -$oi*$N-$di*($e*$f-$N*$k)]
	set ui [expr (1-$e*$qi/$pi-$f*$si/$oi)/$N]
	set vi [expr (-$e*$aai/$pi-$f*$di/$oi)/$N]
	set wi [expr -$oi*$N-$di*($e*$f-$N*$k)]
	
	#puts "calculating a,b, and c"
        set a [expr $g*$ui+$vi*$l+$m*$wi]
	#puts "z: $a"
        set b [expr $g*$qi/$pi+$l*$aai/+$m*$N/$pi]
	#puts "x: $b"
        set c [expr $g*$si/$oi+$l*$di/$oi+$m*$ti/$oi]
	#puts "y: $c"

        #normal (ni nj nk)
        set ni $b
        set nj $c
        set nk $a
	#puts "N: $N"

        $p_point delete
        $r_point delete
	return [list $ni $nj $nk]
}



#*************************************************************************#
# order_parameter_zpos calculates the order of a carbon in a fatty-acid chain. 
# The bilayer normal is assumed to be in the z direction.                 #
# 	c 	 index of the carbon atom                                 #
# 	h1       index of the first hydrogen bound to c                   #
# 	h2       index of the second hydrogen bound to c                  #
# 	firstf   first trajectory frame                                   #
# 	lastf    last trajectory frame                                    #
# 	step     frequency for reading frames                             #
#*************************************************************************#
proc order_parameter_zpos {c h1 h2 firstf lastf step output} {
	set count 0
	set z 0
	set order 0
	for {set i $firstf} {$i <= $lastf} {incr i $step} {
		set phos [atomselect top "name P and same fragment as index $c" frame $i]
		set carb [atomselect top "index $c" frame $i]
		set hyd1 [atomselect top "index $h1" frame $i]
		set hyd2 [atomselect top "index $h2" frame $i]	
		set px [$phos get x]
                set py [$phos get y]
                set pz [$phos get z]
		set cx [$carb get x]
		set cy [$carb get y]
		set cz [$carb get z]
		set hax [$hyd1 get x]
                set hay [$hyd1 get y]
                set haz [$hyd1 get z]
		set hbx [$hyd2 get x]
                set hby [$hyd2 get y]
                set hbz [$hyd2 get z]
		set dix [expr ([$carb get x]-([$hyd1 get x]+[$hyd2 get x])/2)]
		set diy [expr ([$carb get y]-([$hyd1 get y]+[$hyd2 get y])/2)]
		set diz [expr ([$carb get z]-([$hyd1 get z]+[$hyd2 get z])/2)]
		#puts "$cx $cy $cz $hax $hay $haz $hbx $hby $hbz"
		
		set ztemp [expr (3*($diz*$diz/($diz*$diz+$diy*$diy+$dix*$dix))-1)/-2]
		set z [expr $z+$ztemp]
		set count [expr $count + 1]
		#puts "z_pos: $px $py $pz frame$i order: $ztemp"
		$carb delete
		$hyd1 delete
		$hyd2 delete
		$phos delete
		unset carb
                unset hyd1
                unset hyd2
                unset dix
		unset diy
		unset diz
                unset ztemp
	}
	#puts "computing"
	set order [expr $z/$count]
	puts "z_pos order parameter from frame $firstf to frame $lastf: $order"
	unset count
	unset z
        unset c
        unset h1
        unset h2
        unset firstf
        unset lastf
        unset step
	return $order
}

#*************************************************************************#
# order_parameter_angles calculates the order of a carbon in a fatty-acid chain. 
# The bilayer normal is defined by norm_x, normy, and norm_z.             #
#       c        index of the carbon atom                                 #
#       pc       index of the first carbon bound to c                     #
#       nc       index of the second carbon bound to c                    #
#       firstf   first trajectory frame                                   #
#       lastf    last trajectory frame                                    #
#       step     frequency for reading frames                             #
#*************************************************************************#
proc order_parameter_angles {c pc nc firstf lastf step output} {
	#puts "order_parameter_angle"
	set radius 25
        set count 0
	set scd_sum 0
        set order 0
	set leaf 1
	set sumPx 0
	set sumPy 0
	set sumPz 0
	
        for {set i $firstf} {$i <= $lastf} {incr i $step} {
		#puts "selecting atoms"
                set phos [atomselect top "name P and same fragment as index $c" frame $i]
                set carb [atomselect top "index $c" frame $i]
		set prev [atomselect top "index $pc" frame $i]
		set next [atomselect top "index $nc" frame $i]
		#puts "setting vector componants"
                set sumPx [expr ($sumPx + [$phos get x])]
                set sumPy [expr ($sumPy + [$phos get y])]
                set sumPz [expr ($sumPz + [$phos get z])]
                set cx [$carb get x]
                set cy [$carb get y]
                set cz [$carb get z]
                set px [$prev get x]
                set py [$prev get y]
                set pz [$prev get z]
		set nx [$next get x]
                set ny [$next get y]
                set nz [$next get z]
		set vect_x [expr ($cx-($px + $nx)/2)]
                set vect_y [expr ($cy-($py + $ny)/2)]
		set vect_z [expr ($cz-($pz + $nz)/2)]
        	#puts "setting normal componants"
        	#set normal_vector [list 0 0 1]
        	set normal_vector [normal [$phos get index] $radius $i]
        	set norm_x [lindex $normal_vector 0]
       		set norm_y [lindex $normal_vector 1]
        	set norm_z [lindex $normal_vector 2]
        	#puts "normal: ($norm_x $norm_y $norm_z)"
		
		#"(cos(theta))^2 = (v.n)/(|v|*|n|)^2"
		set mvect2 [expr ($vect_x*$vect_x + $vect_y*$vect_y+ $vect_z*$vect_z)]
		set mnorm2 [expr ($norm_x*$norm_x + $norm_y*$norm_y+ $norm_z*$norm_z)]
		set v_dot_n [expr ($vect_x*$norm_x + $vect_y*$norm_y + $vect_z*$norm_z)]
		set cos2 [expr (($v_dot_n*$v_dot_n)/($mvect2*$mnorm2))]

		#"Scd            = -1*(3*(cos(theta))^2 - 1)/2"
		set scd_frame [expr (-1*(3*$cos2 - 1)/2)]

                set scd_sum [expr $scd_sum + $scd_frame]
                set count [expr $count + 1]
                #puts "angle: frame$i order: $scd_frame sum: $scd_sum"
                $phos delete
                $carb delete
                $prev delete
                $next delete
        }

        #puts "computing scd_sum / count: $scd_sum / $count"
        set scd [expr $scd_sum/$count]
	set avePx [expr $sumPx/$count]
        set avePy [expr $sumPy/$count]
        set avePz [expr $sumPz/$count]
	set leaftext "top"
	if {$avePz < $cz} {
		set leaf -1
		set leaftext "bot"
	} 
        #puts "angle order parameter from frame $firstf to frame $lastf: $scd"

	puts "$avePx $avePy $avePz $leaf $scd"
	exec echo "$avePx $avePy $avePz $leaf $scd" >> $output.$leaftext.dat

        unset count
        unset scd_sum
        unset c
        unset firstf
        unset lastf
        unset step
        return $scd
}



#*************************************************************************#
# order_parameter_selector calls order_parameter to return the order of   #
# each carbon in the selection.                                           #
# Outputs the average orderparameter over the trajectory for each         #
# specified carbon location on each chain.                                #
# The bilayer normal is assumed to be in the z direction.                 #
#       lipids     # 
#       chainl     #
#		   	      H
#		   	      |
#		   	head--C--H					  #
#		   	      | 					  #
#			   H--C--fatty-acid(2) 				  #
#			      |						  #
#			   H--C--fatty-acid(1) 				  #
#			      |						  #
#			      H 					  #
#									  #
#       carbonNumL #
#       firstf     first trajectory frame                                 #
#       lastf      last trajectory frame                                  #
#       step       frequency for reading frames                           #
#*************************************************************************#
proc order_parameter_selector {lipids chainL carbonNumL firstf lastf step output} {

	exec echo "SD orderparameter file" > $output.chain1.dat
	exec echo "SD orderparameter file" > $output.chain2.dat
	set chain1 "false"
	set chain2 "false"
	foreach chain $chainL {
		if {$chain == 1} {
			set chain1 "true"
		}
		if {$chain == 2} {
                        set chain2 "true"
                } 
	}
	foreach carbon $carbonNumL {
		exec echo "Px Py Pz leaf scd" > $output.c$carbon.chain1.top.dat
                exec echo "Px Py Pz leaf scd" > $output.c$carbon.chain2.top.dat
		exec echo "Px Py Pz leaf scd" > $output.c$carbon.chain1.bot.dat
                exec echo "Px Py Pz leaf scd" > $output.c$carbon.chain2.bot.dat
		set pcarb [expr ($carbon - 1)]
		set ncarb [expr ($carbon + 1)]
		set values1 0
		set values3 0
		set numlip1 0
		set numlip3 0
		foreach lipid $lipids {
			set residues [lsort -unique [$lipid get resid]]
			set seg [lsort -unique [$lipid get segname]]
      	        	#puts "segname: $seg"
			set chaincount 0
			foreach residue $residues {
				#puts "resid: $residue"
				set carb [atomselect top "segname $seg and resid $residue and name C$carbon"]
				set prev [atomselect top "segname $seg and resid $residue and name C$pcarb"]
				set next [atomselect top "segname $seg and resid $residue and name C$ncarb"]
                                #set hyd1 [atomselect top "segname $seg and resid $residue and name H$carbon\A"]
				#set hyd2 [atomselect top "segname $seg and resid $residue and name H$carbon\B"]
				set resname [$carb get resname]
				#puts "resname: $resname$residue"
				if {$resname=="OLEO"} {
					set c [$carb get index]
					set p [$prev get index]
					set n [$next get index]
                                	#set h1 [$hyd1 get index]
                                	#set h2 [$hyd2 get index]
					#puts "$c $h1 $h2"
					if {$chaincount == 0 && $chain2} {
						#set values1 [expr $values1+[order_parameter_zpos $c $h1 $h2 $firstf $lastf $step $output.chain2.dat]]
						set values1 [expr $values1+[order_parameter_angles $c $p $n $firstf $lastf $step $output.c$carbon.chain2]]
						set numlip1 [expr $numlip1 + 1]		
			}
					if {$chaincount == 2 && $chain1} {
                                                #set values3 [expr $values3+[order_parameter_zpos $c $h1 $h2 $firstf $lastf $step $output.chain1.dat]]
						set values3 [expr $values3+[order_parameter_angles $c $p $n $firstf $lastf $step $output.c$carbon.chain1]]
						set numlip3 [expr $numlip3 + 1]
                                        }
					unset c
					#unset h1
					#unset h2
					unset p
					unset n
				}
				set chaincount [expr ($chaincount + 1)%3]
				$carb delete
				$prev delete
				$next delete
				#$hyd1 delete
				#$hyd2 delete
				unset carb
				#unset hyd1
				#unset hyd2
				unset resname
			}
		}
		#puts "values1-$carbon: $values1"
		#puts "values2-$carbon: $values3"
		if {$numlip3 > 0} {
                        set order1 [expr $values3/$numlip3]

                        puts "chain 1 carbon $carbon order: $order1"

                        exec echo "chain 1 carbon $carbon order: $order1" >> $output.chain1.dat
                	unset order1
		}
		if {$numlip1 > 0} {
			set order2 [expr $values1/$numlip1]

                        puts "chain 2 carbon $carbon order: $order2"

			exec echo "chain 2 carbon $carbon order: $order2" >> $output.chain2.dat
			unset order2
		} 
		if {$numlip1 <= 0 && $numlip3 <= 0} {
			puts "No lipids selected!!"
		}
		unset values1
		unset values3
		unset numlip1
		unset numlip3
	}
	#puts "unsetting in function"
	unset lipids 
	unset chainL 
        unset carbonNumL 
        unset firstf 
        unset lastf 
        unset step 
        unset output
}

#**************************************************************************************
# main 
#**************************************************************************************

package require psfgen


set psf "nbar_dops.10x2.xplor.psf"  
set dcd "nbar_dops.10x2.RS0.dcd"  
#set output "orderparameter"

#mol load psf $psf dcd $dcd

#set S001 [atomselect top "segname S001 and (resname OLEO or resname PCGL)"]
#set C001 [atomselect top "segname C001 and (resname OLEO or resname PCGL)"]
#set BOTH [atomselect top "resname OLEO or resname PCGL"]

#animate write dcd S001.dcd sel $S001
#animate write psf S001.psf sel $S001
#animate write dcd C001.dcd sel $C001
#animate write psf C001.psf sel $C001
#animate write dcd BOTH.dcd sel $BOTH
#animate write psf BOTH.psf sel $BOTH


set files [list BOTH] 
#S001 C001]
foreach file $files {
mol delete all
mol load psf $file.psf dcd $file.dcd



#carbon 9 double bonded to carbon 10
set carbonNumbers [list 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17]
set carbonNumbers [list 12]
set chains [list 1 2]
#set chains [list 2]

puts "carbonNumbers: $carbonNumbers"

set firstf 1
set lastf [molinfo top get numframes]
#set lastf 10
puts "lastf: $lastf"
set step 1

#set val [order_parameter 37187 37188 37189 $firstf $lastf $step $output]
#puts "15: $val"
#set lipids [atomselect top "(same fragment as index 37187)"]
#order_parameter_selector $lipids $chains $carbonNumbers  $firstf $lastf $step $output

#set val [order_parameter 37321 37322 37323 $firstf $lastf $step $output]
#puts "16: $val"
#set lipids [atomselect top "(same fragment as index 37321)"]
#order_parameter_selector $lipids $chains $carbonNumbers  $firstf $lastf $step $output

puts "Extracting data from $file.dcd"
set output "Scd.$file"
#set lipids [atomselect top "segname $file and resname OLEO"]
set lipids [atomselect top "resname OLEO"]
#set lipids [atomselect top "same fragment as index 37271" ]

puts "calculating..."
order_parameter_selector $lipids $chains $carbonNumbers  $firstf $lastf $step $output

puts "unsetting in main"
unset firstf
unset lastf
unset step
unset output
$lipids delete
unset lipids
}

puts "Finished!"

puts "Exiting"
exit