#!/usr/bin/tclsh
# scripts to create and restore backups
# from simulation data within VMD

# store a give trajectory/molecule into a new molecule
proc mol_backup { args } {
    set savetop [molinfo top]

    # set defaults
    set name "backup"
    set beg 0
    set end -1
    set step 1

    if { ([llength $args] < 1) || ([llength $args] == 2) } {
        puts "usage: mol_backup <molid> \[name $name\] \[beg $beg\] \[end $end\] \[step $step\]"
        return -1;
    }

    # handle arguments
    set mol [lindex $args 0]
    set args [lrange $args 1 end]
    
    while {[llength $args] > 1} {
        set tag [lindex $args 0]
        set value [lindex $args 1]
        set args [lrange $args 2 end]
        
        switch $tag {
            name    { set name $value }
            beg     { set beg  $value }
            end     { set end  $value }
            step    { set step $value }
            default { 
                puts "mol_backup: unknown flag '$tag'"
                return -2 
            }
        }
    }

    if {[lsearch -exact [molinfo list] $mol]} {
        puts "mol_backup: error molecule $mol does not exist."
        return -3
    }
    if {$end < 0} { set end [molinfo $mol get numframes] }
    if {($beg < 0) || ($beg > $end)} {
        puts "mol_backup: illegal frame range: $beg - $end"
        return -4
    }

    set tmpdir "/tmp"
    #  if {[info exists env(TMPDIR)]} {set tmpdir $env(TMPDIR)}

    set old [atomselect $mol all]
    animate write psf "${tmpdir}/${name}.psf" beg 0 end 0 waitfor all sel $old $mol

    set newmol [mol new "${tmpdir}/${name}.psf" type psf waitfor all]
    exec rm ${tmpdir}/${name}.psf
    mol rename $newmol $name
    set new [atomselect $newmol all]

    $new set {name type resname resid } \
        [ $old get {name type resname resid } ]
    for {set i $beg} {$i < $end} {incr i $step} {
        animate dup $newmol
        $old frame $i
        $new frame [expr $i/$step]
        $new set {x y z user} [ $old get {x y z user} ]
        molinfo $mol set frame $i
        molinfo $newmol set frame [expr $i/$step]
        molinfo $newmol set {a b c} [molinfo $mol get {a b c}]
    }

    # clean up
    mol off $newmol
    mol top $savetop
    $new delete
    $old delete
    return 0
}

proc mol_restore { args } {
    set savetop [molinfo top]

    # set defaults
    set name "backup"
    set beg 0
    set end -1
    set step 1
    set append 0

    if { ([llength $args] < 1) || ([llength $args] == 2) } {
        puts "usage: mol_restore <molid> \[name $name\] \[beg $beg\] \[end $end\] \[step $step\] \[append $append\]"
        return -1;
    }

    # handle arguments
    set mol [lindex $args 0]
    set args [lrange $args 1 end]
    
    while {[llength $args] > 1} {
        set tag [lindex $args 0]
        set value [lindex $args 1]
        set args [lrange $args 2 end]
        
        switch $tag {
            name    { set name $value }
            beg     { set beg  $value }
            end     { set end  $value }
            step    { set step $value }
            append  {set append $value}
            default { 
                puts "mol_restore: unknown flag '$tag'"
                return -2 
            }
        }
    }

    if {[lsearch -exact [molinfo list] $mol]} {
        puts "mol_restore: error target molecule $mol does not exist."
        return -3
    }
    set newmol -1
    foreach m [molinfo list] {
        if {[string equal $name [molinfo $m get name]]} {
            set newmol $m
        }
    }
    if {$newmol < 0} {
        puts "mol_restore: error source molecule name $name does not exist."
        return -5
    }


    if {$end < 0} { set end [molinfo $newmol get numframes] }
    if {($beg < 0) || ($beg > $end)} {
        puts "mol_restore: illegal frame range: $beg - $end"
        return -4
    }

    set old [atomselect $mol all]
    set new [atomselect $newmol all]
    if {$append == 0} {
        animate delete beg 0 end -1 $mol
    }
    set append [molinfo $mol get numframes]

    $old set {name type resname resid } \
        [ $new get {name type resname resid } ]
    
    for {set i $beg} {$i < $end} {incr i $step} {
        animate dup $mol
        $old frame [expr (($i - $beg)/$step) + $append]
        $new frame $i
        $old set {x y z user} [ $new get {x y z user} ]
        molinfo $mol set frame [expr (($i - $beg)/$step)  + $append]
        molinfo $newmol set frame $i
        molinfo $mol set {a b c} [molinfo $newmol get {a b c}]
    }
    mol top $savetop
    $new delete
    $old delete
    return 0
}