#!/usr/bin/env ruby

#
# This script enables you to plot the hydropathicity of a batch of sequences.
#
# 2008-02-27                   Michael Matschiner <michaelmatschiner@mac.com>
#
#
# This script calculates the hydropathicity of a batch of amino acid sequences 
# according to Kyte & Doolittle (1982). It prepares a tab separated file that 
# can be read by Excel. You are then able to plot all trajectories into one graph,
# or to calculate statistics on your hydropathicity data (e.g. whether one or 
# more of the sequences differs significantly).
#
# References:
# Kyte and Doolittle (1982). A simple method for displaying the hydropathic
# character of a protein. Journal of Molecular Biology 157:105-132


class Array2
  def initialize
    @store = [[]]
  end
  def [](a,b)
    if @store[a]==nil ||
      @store[a][b]==nil
      return nil
    else
      return @store[a][b]
    end
  end
  def []=(a,b,x)
    @store[a] = [] if @store[a]==nil
    @store[a][b] = x
  end
end

class String
  def hydropathy
    hp = 0.0
    self.length.times do |l|
      hp += 4.5 if self[l..l] == "I"
      hp += 4.2 if self[l..l] == "V"
      hp += 3.8 if self[l..l] == "L"
      hp += 2.8 if self[l..l] == "F"
      hp += 2.5 if self[l..l] == "C"
      hp += 1.9 if self[l..l] == "M"
      hp += 1.8 if self[l..l] == "A"
      hp -= 0.4 if self[l..l] == "G"
      hp -= 0.7 if self[l..l] == "T"
      hp -= 0.9 if self[l..l] == "W"
      hp -= 0.8 if self[l..l] == "S"
      hp -= 1.3 if self[l..l] == "Y"
      hp -= 1.6 if self[l..l] == "P"
      hp -= 3.2 if self[l..l] == "H"
      hp -= 3.5 if self[l..l] == "E"
      hp -= 3.5 if self[l..l] == "Q"
      hp -= 3.5 if self[l..l] == "D"
      hp -= 3.5 if self[l..l] == "N"
      hp -= 3.9 if self[l..l] == "L"
      hp -= 4.5 if self[l..l] == "R"
    end
    hp
  end
end

class Fixnum
  def even?
    if (self/2)*2 == self
      true
    else
      false
    end
  end
end



# The input file is the first command line argument. This is the only valid
# command line argument. If no argument is given, the script asks for the input
# file.
#
fileIn = ARGV[0]
windowSize = 9
windowSize = ARGV[1].to_i unless ARGV[1] == nil

if fileIn == nil
  print "Please choose the input file: "
  fileIn = gets.chomp
end
if windowSize == nil
  print "Please choose the window size: "
  windowSize = gets.chomp.to_i
  raise( 'The window size must be an odd number' ) if windowSize.even?
end

seqs = IO.readlines( fileIn )
mat = Array2.new
maxSeqLength = 0

seqs.size.times do |s|
  maxSeqLength = seqs[s].chomp.length if seqs[s].length > maxSeqLength
  ((windowSize-1)/2).upto(seqs[s].chomp.length-((windowSize+1)/2)) do |l|
    mat[s,l] = eval(sprintf("%5.2f",seqs[s][l-((windowSize-1)/2)..l+((windowSize-1)/2)].hydropathy))
  end
end

output = String.new
output << "Position\t"
(seqs.size).times do |s|
  output << "Seqence "
  output << (s+1).to_s
  output << "\t"
end
output << "\n"

((windowSize-1)/2).upto(maxSeqLength-((windowSize+1)/2)) do |m|
  output << (m+1).to_s
  output << "\t"
  seqs.size.times do |s|
    output << mat[s,m].to_s
    output << "\t"
  end
  output << "\n"
end


fileOutName = "hp_plotter.out"
fileOut = File.new(fileOutName, "w")
fileOut.print output.to_s

puts "\nhp_plotter.rb: Program successfully wrote \'hp_plotter.out\'. Please open this file with MS Excel.\n\n"