class PlayfairCipher
  AZ = ('A'..'Z').to_a - ['J']
  def initialize keyword
    @square = keyword.upcase.gsub(/[^A-Z]/, '').gsub('J', 'I').split('').uniq
    AZ.each{|x| @square.push x unless @square.include? x}
    @square = 0.step(20, 5).collect{|x| @square.slice x,5}
    @h = AZ.collect{|x| AZ.collect{|y| "'#{x+y}' => '#{convert x, y}'"}}.flatten.join ', '
    @h = eval "{#{@h}}"
    @rh = @h.invert
  end

  attr_reader :square

  def convert a, b
    ax, ay, bx, by, i = nil
    @square.each.with_index do |line, j|
      (ax, ay = i, j) if i = line.index(a)
      (bx, by = i, j) if i = line.index(b)
    end
    case
      when ay == by
        @square[ay][(ax + 1) % 5] + @square[by][(bx + 1) % 5]
      when ax == bx
        @square[(ay + 1) % 5][ax] + @square[(by + 1) % 5][bx]
      else
        @square[ay][bx] + @square[by][ax]
    end
  end

  def insert_xz s
    return [s, nil] if s.length < 2
    xz = @flag ? 'Z' : 'X'
    s[0] == s[1] ? (@flag = !@flag; [s[0, 1] + xz, s[1..-1]]) : (@flag = false; [s[0, 2], s[2..-1]])
  end

  def encrypt str
    s = str.upcase.gsub(/[^A-Z]/, '').gsub('J', 'I')
    str = ''
    (pair, s = insert_xz(s); str << pair) while s
    str += 'X' unless (str.length % 2).zero?
    str.gsub(/(..)/){$1+' '}.split.collect{|s| @h[s]}.join
  end
end