#!/usr/bin/env ruby

# Copyright (C) 2007-2011 Martin A. Hansen.

# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

# http://www.gnu.org/copyleft/gpl.html

# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# This program is part of the Biopieces framework (www.biopieces.org).

# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> DESCRIPTION <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# Remove adaptors or parts thereof from sequences in the stream.

# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<


require 'biopieces'
require 'seq'

VERBOSE = true

def disambiguate(adaptor)
  adaptor_disamb = adaptor.dup
  adaptor_disamb.gsub!('U', 'T')
  adaptor_disamb.gsub!('R', '[AG]')
  adaptor_disamb.gsub!('Y', '[CT]')
  adaptor_disamb.gsub!('S', '[GC]')
  adaptor_disamb.gsub!('W', '[AT]')
  adaptor_disamb.gsub!('M', '[AC]')
  adaptor_disamb.gsub!('K', '[GT]')
  adaptor_disamb.gsub!('V', '[ACG]')
  adaptor_disamb.gsub!('H', '[ACT]')
  adaptor_disamb.gsub!('D', '[AGT]')
  adaptor_disamb.gsub!('B', '[CGT]')
  adaptor_disamb.gsub!('N', '.')
  adaptor_disamb
end

class Seq
  # Method that finds an adaptor or part thereof in the sequence of a Seq object.
  # Returns a Match object if the adaptor was found otherwise nil. The ed_percent
  # indicates the maximum edit distance allowed in all possible overlaps.
  def adaptor_find(adaptor, adaptor_disamb, pos = 0, ed_percent = 0)
    raise SeqError, "Edit distance percent out of range #{ed_percent}" unless (0 .. 100).include? ed_percent

    if pos < 0
      pos = self.length + pos # pos offset from the right end
    end

    if pos < self.length
      if match = adaptor_find_simple(adaptor_disamb, pos)
        return match
      elsif match = adaptor_find_complex(adaptor, pos, ed_percent)
        return match
      elsif match = adaptor_partial_find_complex(adaptor, pos, ed_percent)
        return match
      end
    end
  end

  private

  # Method to find an adaptor in a sequence taking into account ambiguity
  # codes, but not considering mismatches, insertions, and deletions.
  def adaptor_find_simple(adaptor, pos)
    self.seq.upcase.match(adaptor, pos) do |m|
      return Match.new($`.length, m, m.to_s.length, 0, 0, 0, m.to_s.length)
    end
  end

  # Method to find an adaptor in a sequence taking into account ambiguity
  # codes, mismatches, insertions, and deletions.
  def adaptor_find_complex(adaptor, pos, ed_percent)
    ed_max = (adaptor.length * ed_percent * 0.01).round

    match = self.match(adaptor, pos, ed_max)

    match
  end

  # Method to find part of an adaptor at the right end of a sequence taking
  # into account ambiguity codes, mismatches, insertions, and deletions.
  def adaptor_partial_find_complex(adaptor, pos, ed_percent)
    if pos > self.length - adaptor.length
      adaptor = adaptor[0 ... self.length - pos]
    else
      adaptor = adaptor[0 ... adaptor.length - 1]

      pos = self.length - adaptor.length
    end

    #puts self.seq if VERBOSE

    while adaptor.length > 0
      #puts (" " * pos) + adaptor if VERBOSE

      ed_max = (adaptor.length * ed_percent * 0.01).round

      if ed_max == 0
        self.seq.upcase.match(adaptor, pos) do |m|
          return Match.new($`.length, m, m.to_s.length, 0, 0, 0, m.to_s.length)
        end
      else
        self.scan(adaptor, pos, ed_max).each do |match|
          return match
        end
      end

      adaptor = adaptor[0 ... -1]

      pos += 1
    end
  end
end

casts = []
casts << {:long=>'adaptor',       :short=>'r', :type=>'string', :mandatory=>false, :default=>nil, :allowed=>nil, :disallowed=>nil}
casts << {:long=>'edit_distance', :short=>'e', :type=>'uint',   :mandatory=>false, :default=>20,  :allowed=>nil, :disallowed=>nil}
casts << {:long=>'pos',           :short=>'p', :type=>'int',    :mandatory=>false, :default=>1,   :allowed=>nil, :disallowed=>"0"}
casts << {:long=>'cache',         :short=>'c', :type=>'flag',   :mandatory=>false, :default=>nil, :allowed=>nil, :disallowed=>nil}

bp = Biopieces.new

options = bp.parse(ARGV, casts)

adaptor        = options[:adaptor].to_s.upcase
adaptor_disamb = disambiguate(adaptor)

pos  = options[:pos]
pos -= 1 if pos > 0  # pos was 1-based

cache = {}

bp.each_record do |record|
  if record.has_key? :SEQ
    entry = Seq.new(record[:SEQ_NAME], record[:SEQ], "dna", record[:SCORES])

    if cache[entry.seq.upcase] and options[:cache]
      match = cache[entry.seq.upcase]
    else
      match = entry.adaptor_find(adaptor, adaptor_disamb, pos, options[:edit_distance])

      cache[entry.seq.upcase] = match if match and options[:cache]
    end

    if match
      record[:ADAPTOR_POS]   = match.pos
      record[:ADAPTOR_LEN]   = match.length
      record[:ADAPTOR_MATCH] = match.match
    end
  end

  bp.puts record
end


# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<


__END__