# Yo Emacs this is  -*-Tcl-*-
# Last changed Time-stamp: <96/03/25 00:44:21 ivo>
# A script to save current viewpoint, representations etc to a file.
# By Ivo Hofacker <ivo@ks.uiuc.edu>
# Usage:
# In the vmd console type 
#       save_state foo
# to save your work in file foo. Restarting vmd as
#       vmd -e foo
# should restore most of your previous session.
#
# Essentially everything that can be changed from the graphics and mol and
# color forms will be saved, but nothing else.
#  
# In particular the script will restore
#     molecules read in from files
#     graphics objects (spheres, cylenders...)
#     all representations
#     the transformation matrices for each molecule
#     which color to use for which representation, definitions of colors 0..33
# It will not restore:
#     trajectories loaded via the edit menu
#     interactive MD sessions
#     anything in the display form, such as lighting, axes, stage, ...


proc save_viewpoint {} {
   global viewpoints
   if [info exists viewpoints] {unset viewpoints}
   # get the current matricies
   foreach mol [molinfo list] {
      set viewpoints($mol) [molinfo $mol get {
        center_matrix rotate_matrix scale_matrix global_matrix}]
   }
}
proc restore_viewpoint {} {
   global viewpoints
   foreach mol [molinfo list] {
      if [info exists viewpoints($mol)] {
         molinfo $mol set {center_matrix rotate_matrix scale_matrix
           global_matrix} $viewpoints($mol)
      }
   }
}

proc save_reps {} {
    global representations
    foreach mol [molinfo list] {
	set representations($mol) ""
	for {set i 0} {$i < [molinfo $mol get numreps]} {incr i} {
	    lappend representations($mol) \
		[molinfo $mol get "{rep $i} {selection $i} {color $i} {material $i}"]
	}
    }
}

proc restore_reps {} {
    global representations
    foreach mol [molinfo list] {
	if [info exists representations($mol)] {
	    #delete current representations
	    for {set i 0} {$i < [molinfo $mol get numreps]} {incr i} {
		mol delrep $i $mol
	    }
	    #restore saved representations
	    foreach rep $representations($mol) {
		puts $rep
		lassign $rep r s c m
		mol representation $r
		mol color $c
		mol selection $s
                mol material $m
		mol addrep $mol
	    }
	}
    }
}
	    
proc save_state {{file EMPTYFILE}} {
  global representations
  global viewpoints
  save_viewpoint
  save_reps

  # If no file was given, get a filename.  Use the Tk file dialog if 
  # available, otherwise get it from stdin. 
  if {![string compare $file EMPTYFILE]} {
    set title "Enter filename to save current VMD state:"
    set filetypes [list {{VMD files} {.vmd}} {{All files} {*}}]
    if { [info commands tk_getSaveFile] != "" } {
      set file [tk_getSaveFile -defaultextension ".vmd" \
	        -title $title -filetypes $filetypes]
    } else {
      puts "Enter filename to save current VMD state:"
      set file [gets stdin]
    }
  }
  if { ![string compare $file ""] } {
    return
  }

  set fildes [open $file w]
  puts $fildes "\#!/usr/local/bin/vmd"
  puts $fildes "\# vmd script saved by save_state"
  puts $fildes "set viewplist {}"
  puts $fildes "set fixedlist {}"
  save_materials $fildes
  save_atomselmacros $fildes
  foreach mol [molinfo list] {
    set files [lindex [molinfo $mol get filename] 0]
    set types [lindex [molinfo $mol get filetype] 0]
    set nfiles [llength $files]
    if { $nfiles >= 1 } {
      puts $fildes "mol new {[lindex $files 0]} type [lindex $types 0]"
    } else {
      puts $fildes "mol new"
    }
    for { set i 1 } { $i < $nfiles } { incr i } {
      puts $fildes "mol addfile {[lindex $files $i]} type [lindex $types $i]"
    }
    foreach g [graphics $mol list] {
      puts $fildes "graphics top [graphics $mol info $g]"
    }
    puts $fildes "mol delrep 0 top"
    if [info exists representations($mol)] {
      foreach rep $representations($mol) {
        lassign $rep r s c m
        puts $fildes "mol representation $r"
        puts $fildes "mol color $c"
        puts $fildes "mol selection {$s}"
        puts $fildes "mol material $m"
        puts $fildes "mol addrep top"
      } 
    }
    puts $fildes "mol rename top \{[lindex [molinfo $mol get name] 0]\}"
	  if {[molinfo $mol get drawn] == 0} {
	      puts $fildes "molinfo top set drawn 0"
	  }
	  if {[molinfo $mol get active] == 0} {
	      puts $fildes "molinfo top set active 0"
	  }
	  if {[molinfo $mol get fixed] == 1} {
	      puts $fildes "lappend fixedlist \[molinfo top\]"
	  }

	  puts $fildes "set viewpoints(\[molinfo top\]) [list $viewpoints($mol)]"
	  puts $fildes "lappend viewplist \[molinfo top\]"
	  if {$mol == [molinfo top]} {
	    puts $fildes "set topmol \[molinfo top\]"
	  }
	  puts $fildes "\# done with molecule $mol"
  } 
  puts $fildes "foreach v \$viewplist \{"
  puts $fildes "  molinfo \$v set {center_matrix rotate_matrix scale_matrix global_matrix} \$viewpoints(\$v)"
  puts $fildes "\}"
  puts $fildes "foreach v \$fixedlist \{"
  puts $fildes "  molinfo \$v set fixed 1"
  puts $fildes "\}"
  puts $fildes "unset viewplist"
  puts $fildes "unset fixedlist"
  if {[llength [molinfo list]] > 0} {
    puts $fildes "mol top \$topmol"
    puts $fildes "unset topmol"
  }
  save_colors $fildes
  save_labels $fildes
  close $fildes
}

# an up-to-date list of color categories can be produced using
#    foreach c [colorinfo categories] {
#	foreach cc [colorinfo category $c] {
#	    lappend colcatlist "$c,$cc" [colorinfo category $c $cc]
#	}
#    }
# the list of default rgb values is generated by
# for {set c 0} {$c < 2*[colorinfo num]} {incr c} {
#    lappend def_rgb [colorinfo rgb $c]
# }

proc save_colors {fildes} {
    set colcatlist { Display,Background black
	Axes,X red Axes,Y green Axes,Z blue Axes,Origin cyan Axes,Labels white
	Stage,Even blue Stage,Odd green
	Labels,Dihedrals cyan Labels,Angles yellow Labels,Bonds white
	Labels,Atoms green
	Name,H white Name,O red Name,N blue Name,C cyan Name,S yellow Name,P
	tan Name,Z silver
	Type,H white Type,O red Type,N blue Type,C cyan Type,S yellow Type,P
	tan Type,Z silver
	Resname,ALA blue Resname,ARG white Resname,ASN tan Resname,ASP red 
	Resname,CYS yellow Resname,GLY white Resname,GLU pink Resname,GLN
	orange Resname,HIS cyan Resname,ILE green Resname,LEU pink
	Resname,LYS cyan Resname,MET yellow Resname,PHE purple Resname,PRO
	ochre Resname,SER yellow Resname,THR mauve Resname,TRP silver
	Resname,TYR green Resname,VAL tan Resname,ADE blue Resname,CYT orange
	Resname,GUA yellow Resname,THY purple Resname,URA green Resname,TIP
	cyan Resname,TIP3 cyan Resname,WAT cyan Resname,SOL cyan Resname,H2O
	cyan Resname,LYR purple Resname,ZN silver Resname,NA yellow
	Resname,CL green 
	Restype,Unassigned cyan Restype,Solvent yellow Restype,Nucleic_Acid
	purple Restype,Basic blue Restype,Acidic red Restype,Polar green
	Restype,Nonpolar white Restype,Ion tan
	Highlight,Proback green Highlight,Nucback yellow Highlight,Nonback 
	blue
	{Structure,Alpha Helix} purple Structure,3_10_Helix mauve
	Structure,Pi_Helix red Structure,Extended_Beta yellow
	Structure,Bridge_Beta tan Structure,Turn cyan Structure,Coil white
    }

    set def_rgb {
	{0.25 0.25 1.0} {1.0 0.0 0.0} {0.35 0.35 0.35} {0.8 0.5 0.2}
	{0.8 0.8 0.0} {0.5 0.5 0.2} {0.6 0.6 0.6} {0.2 0.7 0.2} {1.0 1.0 1.0}
	{1.0 0.6 0.6} {0.25 0.75 0.75} {0.65 0.3 0.65} {0.5 0.9 0.4}
	{0.9 0.4 0.7} {0.5 0.3 0.0} {0.5 0.75 0.75} {0.0 0.0 0.0}
    } 
    
    array set def_colcat $colcatlist

    puts $fildes "proc vmdrestoremycolors \{\} \{"

    foreach c [colorinfo categories] {
	foreach cc [colorinfo category $c] {
	    set col [colorinfo category $c $cc]
	    if {![info exists def_colcat($c,$cc)] ||
		[string compare $col $def_colcat($c,$cc)]} {
		puts $fildes "  color $c \{$cc\} $col"
	    }
	}
    }
    set cnum [colorinfo num]
    for {set c 0} {$c < $cnum} {incr c} {
	set rgb [colorinfo rgb $c]
	if  {[string compare $rgb [lindex $def_rgb $c]]}  {
	    puts $fildes "  color change rgb $c $rgb"
	}
    }

    puts $fildes "\}"
    puts $fildes "vmdrestoremycolors"
}

proc save_materials { filedes } {
  set mlist [material list]
  # The first two materials are non-modifiable, so start with index 2
  # Note: this only works if no materials have been created yet!
  for { set i 2 } { $i < [llength $mlist] } { incr i } {
    set mat [lindex $mlist $i]
    lassign [material settings $mat] amb spec dif shin opac
    puts $filedes "material add $mat"
    puts $filedes "material change ambient $mat $amb"
    puts $filedes "material change specular $mat $spec" 
    puts $filedes "material change diffuse $mat $dif"  
    puts $filedes "material change shininess $mat $shin"  
    puts $filedes "material change opacity $mat $opac"  
  }
}

proc save_labels { file } {
  set atomlist [label list Atoms]
  set bondlist [label list Bonds]
  set anglelist [label list Angles]
  set dihedrallist [label list Dihedrals]

  foreach atom $atomlist {
    lassign $atom atom1 value show
    lassign $atom1 mol index
    puts $file "label add Atoms ${mol}/${index}"
    if { [string compare $show show] } { 
      # don't show this label.  Do this by simply repeating the add command
      puts $file "label add Atoms ${mol}/${index}"
    } 
  } 

  foreach bond $bondlist {
    lassign $bond atom1 atom2
    lassign $atom1 mol1 index1
    lassign $atom2 mol2 index2
    puts $file "label add Bonds ${mol1}/${index1} ${mol2}/${index2}"
    if { [string compare $show show] } { 
      # don't show this label.  Do this by simply repeating the add command
      puts $file "label add Bonds ${mol1}/${index1} ${mol2}/${index2}"
    } 
  }

  foreach angle $anglelist {
    lassign $angle atom1 atom2 atom3
    lassign $atom1 mol1 index1
    lassign $atom2 mol2 index2
    lassign $atom3 mol3 index3
    puts $file "label add Angles ${mol1}/${index1} ${mol2}/${index2} ${mol3}/${index3}"
    if { [string compare $show show] } { 
      # don't show this label.  Do this by simply repeating the add command
      puts $file "label add Angles ${mol1}/${index1} ${mol2}/${index2} ${mol3}/${index3}"
    } 
  }

  foreach dihedral $dihedrallist {
    lassign $dihedral atom1 atom2 atom3 atom4
    lassign $atom1 mol1 index1
    lassign $atom2 mol2 index2
    lassign $atom3 mol3 index3
    lassign $atom4 mol4 index4
    puts $file "label add Dihedrals ${mol1}/${index1} ${mol2}/${index2} ${mol3}/${index3} ${mol4}/${index4}"
    if { [string compare $show show] } { 
      # don't show this label.  Do this by simply repeating the add command
      puts $file "label add Dihedrals ${mol1}/${index1} ${mol2}/${index2} ${mol3}/${index3} ${mol4}/${index4}"
    } 
  }
}

proc save_atomselmacros { file } {
  set mlist [atomselect macro]
  if { [llength $mlist] } {
    puts $file "# Atom selection macros"
    foreach macro [atomselect macro] {
      puts $file [list atomselect macro $macro [atomselect macro $macro]]
    }
  }
}