# This script takes takes the various dcds output  from a Replica Exchange simulation and loads the dcds by temperature (not by replica). 
#It is prepared to function with parm7 *top files, not with psf files
# This script requires that a *conf file exists in the directory above it, because some variables are defined there (see file *.conf given in the example provided in NAMD for the replica exchange simulations)
#  This script will not work properly if of the target temperatures is bigger than the number of steps contained in the first dcd frame (to avoid problems, each dcd frame should be composed of 1000 steps at least, because it is unlikely that simulations at T larger than 1000 will be performed.
# This script will not work properly if replica 1 does not experience all possible temperatures (as all replicas should in a replica exchange simulation)

# To run, put this script in the same directory as the *dcd and *dat files output by the RE simulation, in a Linux machine that runs VMD.  Open it and define the name of the relevant *targtemp.dat  and *conf files.  Then type at the prompt
#                     vmd -e REselectByTemp.tcl
#############################################################################
#   User input starts here

set inFile1 "REstr2_2ss.out.job0.targtemp.dat" ;# put the file with the target temperature

source ../RE_str2_2ssVMD.conf ;# put *conf file used for the run; comment out the lines used to run with PBS, otherwise VMD will give an error;  the command "source" reads this file into VMD

# User input stops here
##############################################################################
set tFrame [expr $runs_per_frame * $steps_per_run] ;# the number of timesteps between each frame in each dcd

if {! [info exists i_job]} { set i_job 0 }
set job_output_root "$output_root.job$i_job" ;# set base name for output files
set outputbase $job_output_root
set dcd_filename_format "${outputbase}.%d.dcd" ;# set base name for DCD

set inStream [open $inFile1 r]
set aux1 [split [read $inStream] \n] ;# reads inFile1
set aux2 [lindex $aux1 0] ;# stores first line of inFile1
#puts "aux2 is $aux2"
set aux3 [llength [split $aux2]] ;# tells me the number of columns (= no. of different replicas+1 = no. of different temperatures+1) in inFile1, necessary to create as many lists as different temperatures used in the simulation
#puts "aux3 is $aux3"

set listOfIndex ""
set maxReplicaID [expr $aux3-2] ;# subtracting 1 for the time column and 1 because replica numbering starts in 0 
for {set i 0} {$i <= $maxReplicaID} {incr i} {
    set T$i ""  ;# store info by temperature
    puts "R$i"
    append listOfIndex "$i "  ;# list of replica IDs (same as temperature IDs)
}
#puts "list of index is $listOfIndex"

# get the target temperature values used in this simulation
set R1 ""
foreach line $aux1 {
    set aux4 [split $line]
    append R1 "[lindex $aux4 2] "  ;# append the time-sequence of temperatures experienced by replica 1 in list R1
}
set aux4 [lsort -integer $R1] ;# sort the elements in the list in increasing order
set noFrames [llength $aux4] ;# get number of frames
#puts "sorted list $aux4"
set Temperatures [lindex $aux4 0]
for {set i 1} {$i <= $noFrames} {incr i} {
   set diff [expr [lindex $aux4 $i] - [lindex $aux4 [expr $i-1 ]]]
   if {$diff != 0} {
     append Temperatures " [lindex $aux4 $i]"
     set Temperatures [split [string trim $Temperatures]] ;# contains the list of target temperatures probed by the simulation
   }
}

if { [llength $Temperatures] != [llength $listOfIndex]} {
   error "ERROR! Replica 1 does not experience all temperatures!"
}
puts "temperatures $Temperatures"

# obtain a list of replica IDs at each target temperature at each dcd frame
set frmNo 1
foreach line $aux1 {
    set lineList [split $line]
    if {[lindex $lineList 0] == $frmNo*$tFrame} { ;# if the number of time-steps matches the no. of time-steps for each dcd frame
       foreach T $Temperatures temp_id $listOfIndex {
          lappend T$temp_id [expr [lsearch $lineList $T]-1]
       }
       incr frmNo
    } 
}
#puts "T2 is $T2"

# getting rid of the last term in these lists if it takes the value -1, which is not a replica ID.
foreach temp_id $listOfIndex { 
    set lastElement [lindex [set T$temp_id] end]
#    puts "last element is $lastElement"
    if {$lastElement == -1} { 
        set T$temp_id  [lrange [set T$temp_id] 0 end-1 ] ;# getting rid of the last term in these lists if it takes the value -1, which is not a replica ID.
    }
}
       
# loading data into VMD; each molecule in VMD contains frames from a single temperature
#puts "T2 is $T2"
foreach tempID $listOfIndex temp $Temperatures {
       mol new $parm7_file type parm7 
       mol rename top T$temp
       set frm 0
       foreach replica_id [set T$tempID] {
           set  dcdfile [format $dcd_filename_format $replica_id]
           mol addfile $dcdfile type dcd first $frm last $frm 
           incr frm
       }
}
close $inStream