# output water dipole orientation, averaged along the nanotube.
# water's fallin-region:
# (z>0.0)&&(z<33.0)&&((x-0.12)*(x-0.12)+(y+0.13)*(y+0.13)<36)
#
####################################################################
 
set tcl_precision 6

####################################################################
# calculate cosine value from a dipole vector
proc get_cos id {
    
    # create one selection of atoms from the resid number
    set sel [atomselect top "resid $id"]
    
    # get dipole vector
    set dipole_vec [measure dipole $sel]
    unset sel
    
    # get z component of dipole_vector
    set z [lindex $dipole_vec 2]
    
    # calculate the cosine value from the length and z component
    set scale [veclength $dipole_vec]
    set cos [expr $z / $scale]
    unset dipole_vec
    return $cos
}

####################################################################
# in oder to average the data of water dipole orientations along the
# tube, the nanotube is divided into 50 slices, refer to 50 slots 
# in which cosine values of water dipole orentation should be filled.
# return 50 lists, each list corresponds to a specific region along 
# the tube.
#
# determine which slots should be filled in, return slots' id list

proc allocate z {
    if {($z >= 0.66) && ($z < 1.32)} {
        return {1 2 3} 
    } 
    if {($z >= 1.32) && ($z < 1.98)} {
        return {2 3 4}
    } 
    if {($z >= 1.98) && ($z < 2.64)} {
        return {3 4 5}
    } 
    if {($z >= 2.64) && ($z < 3.3)} {
        return {4 5 6}
    } 
    if {($z >= 3.3) && ($z < 3.96)} {
        return {5 6 7}
    } 
    if {($z >= 3.96) && ($z < 4.62)} {
        return {6 7 8}
    } 
    if {($z >= 4.62) && ($z < 5.28)} {
        return {7 8 9}
    } 
    if {($z >= 5.28) && ($z < 5.94)} {
        return {8 9 10}
    } 
    if {($z >= 5.94) && ($z < 6.6)} {
        return {9 10 11}
    } 
    if {($z >= 6.6) && ($z < 7.26)} {
        return {10 11 12}
    } 
    if {($z >= 7.26) && ($z < 7.92)} {
        return {11 12 13}
    } 
    if {($z >= 7.92) && ($z < 8.58)} {
        return {12 13 14}
    } 
    if {($z >= 8.58) && ($z < 9.24)} {
        return {13 14 15}
    } 
    if {($z >= 9.24) && ($z < 9.9)} {
        return {14 15 16}
    } 
    if {($z >= 9.9) && ($z < 10.56)} {
        return {15 16 17}
    } 
    if {($z >= 10.56) && ($z < 11.22)} {
        return {16 17 18}
    } 
    if {($z >= 11.22) && ($z < 11.88)} {
        return {17 18 19}
    } 
    if {($z >= 11.88) && ($z < 12.54)} {
        return {18 19 20}
    } 
    if {($z >= 12.54) && ($z < 13.2)} {
        return {19 20 21}
    } 
    if {($z >= 13.2) && ($z < 13.86)} {
        return {20 21 22}
    } 
    if {($z >= 13.86) && ($z < 14.52)} {
        return {21 22 23}
    } 
    if {($z >= 14.52) && ($z < 15.18)} {
        return {22 23 24}
    } 
    if {($z >= 15.18) && ($z < 15.84)} {
        return {23 24 25}
    } 
    if {($z >= 15.84) && ($z < 16.5)} {
        return {24 25 26}
    } 
    if {($z >= 16.5) && ($z < 17.16)} {
        return {25 26 27}  
    } 
    if {($z >= 17.16) && ($z < 17.82)} {
        return {26 27 28}
    } 
    if {($z >= 17.82) && ($z < 18.48)} {
        return {27 28 29}
    } 
    if {($z >= 18.48) && ($z < 19.14)} {
        return {28 29 30}
    } 
    if {($z >= 19.14) && ($z < 19.8)} {
        return {29 30 31}
    } 
    if {($z >= 19.8) && ($z < 20.46)} {
        return {30 31 32}
    } 
    if {($z >= 20.46) && ($z < 21.12)} {
        return {31 32 33}
    } 
    if {($z >= 21.12) && ($z < 21.78)} {
        return {32 33 34}
    } 
    if {($z >= 21.78) && ($z < 22.44)} {
        return {33 34 35}
    } 
    if {($z >= 22.44) && ($z < 23.1)} {
        return {34 35 36}
    } 
    if {($z >= 23.1) && ($z < 23.76)} {
        return {35 36 37}
    } 
    if {($z >= 23.76) && ($z < 24.42)} {
        return {36 37 38}
    } 
    if {($z >= 24.42) && ($z < 25.08)} {
        return {37 38 39}
    } 
    if {($z >= 25.08) && ($z < 25.74)} {
        return {38 39 40}
    } 
    if {($z >= 25.74) && ($z < 26.4)} {
        return {39 40 41}
    } 
    if {($z >= 26.4) && ($z < 27.06)} {
        return {40 41 42}
    } 
    if {($z >= 27.06) && ($z < 27.72)} {
        return {41 42 43}
    } 
    if {($z >= 27.72) && ($z < 28.38)} {
        return {42 43 44}
    } 
    if {($z >= 28.38) && ($z < 29.04)} {
        return {43 44 45}
    } 
    if {($z >= 29.04) && ($z < 29.7)} {
        return {44 45 46}
    } 
    if {($z >= 29.7) && ($z < 30.36)} {
        return {45 46 47}
    } 
    if {($z >= 30.36) && ($z < 31.02)} {
        return {46 47 48}
    } 
    if {($z >= 31.02) && ($z < 31.68)} {
        return {47 48 49}
    } 
    if {($z >= 31.68) && ($z < 32.34)} {
        return {48 49 50}
    } else {
        return {49 50 51}
    }
}

####################################################################
# determine the numerator for averaging
proc sum list {
   set sum 0
   foreach i $list {
       set sum [expr {$sum+$i}]
   }
   set sum
}


####################################################################

set numFrame [molinfo top get numframes]
set wat [atomselect top "name OH2"]

# get index list of water
set index [$wat get index]
unset wat

# create lists for data storage.
for {set i 1} {$i <= 50} {incr i} {
    set list$i {}
}

####################################################################

for {set fr 0} {$fr < $numFrame} {incr fr} {

    echo "\$\$\$\$\$ [expr $numFrame - $fr] frames left! \$\$\$\$\$"
    
    molinfo top set frame $fr
    foreach wt_num $index {
        # get atomselection
        set wt_atom [atomselect top "index $wt_num"]   
        # get {x y z}
        set sx [lindex [$wt_atom get x] 0]
        set sy [lindex [$wt_atom get y] 0]
        set sz [lindex [$wt_atom get z] 0]
           
        # search fallin water molecules
        if {($sz > 0.0) && ($sz < 33.0) && (($sx - 0.12)*($sx - 0.12) + ($sy + 0.13)*($sy + 0.13) < 36)} {
            
            set resid [$wt_atom get resid]
            echo "Water:$resid falls in"
            set cos_value [get_cos $resid]
            
            # refresh lists
            
            foreach i [allocate $sz] {
               lappend list$i $cos_value
               echo $cos_value
            }
            unset resid cos_value
        }
        unset wt_atom sx sy sz
    }
}

#######################################
# lists manipulation, output
set f [open water_profile.data a+]

for {set i 1} {$i <= 50} {incr i} {
    echo "entering writing data!"
    set numerator [sum [set list$i]]
    set denominator [llength [set list$i]]
    if {$denominator != 0} {
        set average [expr $numerator / $denominator]
    } else {
        set average 0
    }
    puts $f [format %8f $average]
}
close $f

puts "done!"