######################################################################
# REQUIREMENTS:
#   VMD Version 1.8.2
#
# DESCRIPTION:
#   Wrap all atoms of a given selection into the central image of the
#  PBC unit cell (as defined by [molinfo get {a b c}]). 
#  Handling of sheared boxes (using alpha, beta and gamma) is not
#  implemented so far.
#
# PROCEDURES:
#   pbcwrap [molid $molid] [sel $selection] [first $first] [last $last]
#           [keepresidues {yes|no}]
#
#     Wrap all atoms in molecule $molid of the selection $selection into
#   the central image of the PBC unit cell. Unit cell geometry is read
#   from [ molinfo get { a b c } ]. 
#
#   Options:
#     molid 		the molecule to wrap (default: "top")
#     sel 		the selection of atoms to wrap (default: "all")
#     first 		the first frame to wrap (default: "now")
#     last 		the last frame to wrap (default: "now")
#     keepresidues 	denotes if residues should be kept together 
#			(default: "no")
#
# EXAMPLE USAGE:
#   source olpbcwrap.tcl
#   pbcwrap sel "segname OXY" first 17 last "last" keepresidues yes
#
# AUTHOR:
#         Olaf Lenz
#         Condensed Matter Theory
#         Fakultät für Physik
#         D-33501 Bielefeld
#         Germany
#         olenz@Physik.Uni-Bielefeld.DE
#         +49(0)-521-106-6198
#
#
# Feel free to send comments, bugs, etc.
# This script is inspired by Jan Saam's pbcwrap script.
######################################################################
package provide olpbcwrap 1.0

proc pbcwrap { args } {

    # Set the defaults
    set molid		"top"
    set sel		"all"
    set first		"now"
    set last		"now"
    set keepresidues	"no"

    # Parse the arguments
    foreach {argname argval} $args {
	switch $argname {
	    "molid" { set molid $argval }
	    "sel" { set sel $argval }
	    "first" { set first $argval }
	    "last" { set last $argval }
	    "keepresidues" { set keepresidues $argval }
	    default { puts "olpbcwrap: Argument $argname not understood!" }
	}
    }
    if { $molid == "top" } then { set molid [ molinfo top ] }
    if { $keepresidues == "no" || $keepresidues == 0 } \
	then { set keepresidues "" } \
	else { set keepresidues "same residue as" }

    # save the frame number
    set frame_before [ molinfo $molid get frame ]
    
    if { $first == "now" } then { set first $frame_before }
    if { $first == "first" } then { set first 0 }
    if { $last == "now" } then { set last $frame_before }
    if { $last == "last" } then {
        set numframes [molinfo $molid get numframes]
	set last [expr $numframes-1]
    }
    
    display update off

    # Now loop over all frames
    for { set frame $first } { $frame <= $last } { incr frame } {

	# switch to the next frame
	molinfo $molid set frame $frame

	# get the unit cell boundaries
	set A [ molinfo $molid get a ]
	set B [ molinfo $molid get b ]
	set C [ molinfo $molid get c ]

	# wrap the coordinates
	set outseltext "($sel) and $keepresidues x < 0"
	set outsel [atomselect $molid $outseltext]
	while { [$outsel num] > 0 } { 
	    $outsel moveby "$A 0.0 0.0"
	    set outsel [atomselect $molid $outseltext]
	}
	
	set outseltext "($sel) and $keepresidues x > $A"
	set outsel [atomselect $molid $outseltext]
	while { [$outsel num] > 0 } { 
	    $outsel moveby "-$A 0.0 0.0"
	    set outsel [atomselect $molid $outseltext]
	}
	
	set outseltext "($sel) and $keepresidues y < 0"
	set outsel [atomselect $molid $outseltext]
	while { [$outsel num] > 0 } { 
	    $outsel moveby "0.0 $B 0.0"
	    set outsel [atomselect $molid $outseltext]
	}
	
	set outseltext "($sel) and $keepresidues y > $B"
	set outsel [atomselect $molid $outseltext]
	while { [$outsel num] > 0 } { 
	    $outsel moveby "0.0 -$B 0.0"
	    set outsel [atomselect $molid $outseltext]
	}
	
	set outseltext "($sel) and $keepresidues z < 0"
	set outsel [atomselect $molid $outseltext]
	while { [$outsel num] > 0 } { 
	    $outsel moveby "0.0 0.0 $C"
	    set outsel [atomselect $molid $outseltext]
	}
	
	set outseltext "($sel) and $keepresidues z > $C"
	set outsel [atomselect $molid $outseltext]
	while { [$outsel num] > 0 } { 
	    $outsel moveby "0.0 0.0 -$C"
	    set outsel [atomselect $molid $outseltext]
	}
	
    }

    # go back to the frame where we started from
    molinfo $molid set frame $frame_before

    display update on
}