# Contacts amongs molecules in trajectories
# Needs basics.tcl

proc conseg2 { mol1 mol2 {seltext1 "backbone"} {seltext2 "backbone"} {cutoff 5.0} {file ""} {atomlist 0} } {
# Two molecules
  if {$file != ""} {
    set file_id [open $file w]
  }

  set mol1sel [atomselect $mol1 $seltext1 frame 0]
  set mol2sel [atomselect $mol2 $seltext2 frame 0]

  set n1 [molinfo $mol1 get numframes]
  set n2 [molinfo $mol2 get numframes]

  for {set i 0} {$i < $n1} {incr i} {
    #puts -nonewline "mol1 $i"
    $mol1sel frame $i
    set output ""
    for {set j 0} {$j < $n2} {incr j} {
      if {$atomlist > 1} {
	puts "Frames $i $j"
      }
      if {$mol1 == $mol2 && $j < $i} {
	set ncont 0
      } else {
	$mol2sel frame $j
	set ncont [ncontacts $mol1sel $mol2sel $cutoff $atomlist]
      }
      append output [format "%5d" $ncont]
    }
    if {$file != ""} {
      puts $file_id $output
    } else {
      puts $output
    }
  }

  if {$file != ""} {
    close $file_id
  }
}

proc conseg1 { {seltext1 "backbone"} {seltext2 "backbone"} {cutoff 5.0} {file ""} {atomlist 0} } {
# One molecule
  if {$file != ""} {
    set file_id [open $file w]
  }

  set sel1 [atomselect top $seltext1 frame 0]
  set sel2 [atomselect top $seltext2 frame 0]

  set n [molinfo top get numframes]

  set header ""
  if {$atomlist == 1} {
    set numres 0
    for {set a 0} {$a < [molinfo top get numatoms]} {incr a} {
      set new [[atomselect top "index $a"] get resid]
      if {$numres < $new} {
	set numres $new
      }
    }
    set b 0
    for {set a 1} {$a <= $numres} {incr a} {
      set l [expr $a%10]
      if {$l == 0} {
	incr b
	if {$b > 9} {set b 0}
	set label $b
      } else {
	set label " "
      }
      append header $label

    }
    append header "\n"
    for {set a 1} {$a <= $numres} {incr a} {
      set label [expr $a%10]
      append header $label
    }
    if {$file != ""} {
      puts $file_id $header
    } else {
      puts $header
    }
  }

  for {set i 0} {$i < $n} {incr i} {
    #puts -nonewline "mol1 $i"
    #puts "DEBUG: frame $i"
    $sel1 frame $i
    $sel2 frame $i
    set data ""
    if {$atomlist == 0} {
      set ncont [ncontacts $sel1 $sel2 $cutoff $atomlist]
      append data "$ncont"
    } elseif {$atomlist == 1} {
      set reslist [ncontacts $sel1 $sel2 $cutoff $atomlist]
      for {set j 1} {$j <= $numres} {incr j} {
	if {[lsearch -exact $reslist "$j"] >= 0} {
	  set label "x"
	} else {
	  set label " "
	}
	append data $label
      }
    } elseif {$atomlist > 1} {
      puts "Frame $i"
      append data [ncontacts $sel1 $sel2 $cutoff $atomlist]
    }

    if {$file != ""} {
      puts $file_id $data
    } else {
      puts $data
    }
    
  }
  
  if {$file != ""} {
    close $file_id
  }
}

proc ncontacts { sel1 sel2 {cutoff 5.0} {atomlist 0} } {
  # Number of contacts between two selections (same or diff molecule)
  # and contact info

  set caller [lindex [info level 1] 0]

  set count 0

  set maxi [$sel1 num]
  set maxj [$sel2 num]
  set coords1 [$sel1 get {x y z}]
  set coords2 [$sel2 get {x y z}]
  set mol1 [$sel1 molid]
  set mol2 [$sel2 molid]

  set output ""
  set reslist {}
  for {set i 0} {$i < $maxi} {incr i} {
    set c1 [lindex $coords1 $i]
    if {$caller eq "conseg1"} {
      set minj [expr $i+1]
    } else {
      set minj 0
    }
    for {set j $minj} {$j < $maxj} {incr j} {
      set at1 [lindex [$sel1 get {resname resid segname name index}] $i]
      set at2 [lindex [$sel2 get {resname resid segname name index}] $j]
      if {$caller eq "conseg1" && $at1 eq $at2} {
	#puts "DEBUG: igual $at1 $at2"
	continue
      }
      if {$caller eq "conseg2" && $at1 eq $at2 && [$sel1 molid] == [$sel2 molid] && [$sel1 frame] == [$sel2 frame]} {
	#puts "DEBUG: igual $at1 $at2"
	continue
      }
      set c2 [lindex $coords2 $j]
      set d [dist $c1 $c2]
      if { $d <= $cutoff } {
	incr count
	if {$atomlist == 1} {
	  lappend reslist [lindex [$sel2 get resid] $j]
	} elseif {$atomlist > 1} {
	  set at1_rn [lindex $at1 0]
	  set at1_ri [lindex $at1 1]
	  set at1_s  [lindex $at1 2]
	  set at1_n  [lindex $at1 3]
	  set at1_i  [lindex $at1 4]
	  set at2_rn [lindex $at2 0]
	  set at2_ri [lindex $at2 1]
	  set at2_s  [lindex $at2 2]
	  set at2_n  [lindex $at2 3]
	  set at2_i  [lindex $at2 4]
#	  puts [format "%3s %5d %4s %4s %5d -- %3s %5d %4s %4s %5d = %5.2f" $at1_rn $at1_ri $at1_s $at1_n $at1_i $at2_rn $at2_ri $at2_s $at2_n $at2_i $d]
	  append output [format "%3s %5d %4s %4s %5d -- %3s %5d %4s %4s %5d = %5.2f" $at1_rn $at1_ri $at1_s $at1_n $at1_i $at2_rn $at2_ri $at2_s $at2_n $at2_i $d]
	  append output "\n"
	}
      }
      #puts "$d $count"
    }
  }


  if {$caller eq "conseg1"} {
    if {$atomlist == 0} {
      return $count
    } elseif {$atomlist == 1} {
      return $reslist
    } elseif {$atomlist > 1} {
      return $output
    }
  } elseif {$caller eq "conseg2"} {
    if {$atomlist > 1} {
      puts $output
    }
    return $count
  }
}