#!/usr/bin/env perl

# Copyright (C) 2007-2009 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


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

# Vmatch sequences in the stream against a specified database.

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


use warnings;
use strict;
use Data::Dumper;
use Maasha::Biopieces;
use Maasha::Common;
use Maasha::Filesys;
use Maasha::Seq;
use Maasha::Fasta;


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


my ( $options, $in, $out, @index_files, $tmp_dir, $query_file, $fh_out, $fh_in, $record, $entry,
     $vmatch_args, $result_file, $count_hash );

$options = Maasha::Biopieces::parse_options(
    [
        { long => 'genome',       short => 'g', type => 'genome', mandatory => 'no', default => undef, allowed => undef, disallowed => undef },
        { long => 'index_name',   short => 'i', type => 'file',   mandatory => 'no', default => undef, allowed => undef, disallowed => undef },
        { long => 'count',        short => 'c', type => 'flag',   mandatory => 'no', default => undef, allowed => undef, disallowed => undef },
        { long => 'max_hits',     short => 'm', type => 'uint',   mandatory => 'no', default => undef, allowed => undef, disallowed => 0 },
        { long => 'hamming_dist', short => 'h', type => 'uint',   mandatory => 'no', default => undef, allowed => undef, disallowed => 0 },
        { long => 'edit_dist',    short => 'e', type => 'uint',   mandatory => 'no', default => undef, allowed => undef, disallowed => 0 },
    ]   
);

$options->{ 'count' } = 1 if $options->{ 'max_hits' };

Maasha::Common::error( qq(both --index_hame and --genome specified) ) if     $options->{ "genome" } and     $options->{ "index_name" };
Maasha::Common::error( qq(no --index_name or --genome specified) )    if not $options->{ "genome" } and not $options->{ "index_name" };

$in  = Maasha::Biopieces::read_stream( $options->{ "stream_in" } );
$out = Maasha::Biopieces::write_stream( $options->{ "stream_out" } );

@index_files = get_index_files( $options );

$tmp_dir = Maasha::Biopieces::get_tmpdir();

$query_file  = "$tmp_dir/query.seq";
$result_file = "$tmp_dir/vmatch.out";

$fh_out = Maasha::Filesys::file_write_open( $query_file );

while ( $record = Maasha::Biopieces::get_record( $in ) ) 
{
    if ( $entry = Maasha::Fasta::biopiece2fasta( $record ) )
    {
        next if length $record->{ 'SEQ' } < 12; # assuming that the index is created for 12 as minimum length

        Maasha::Fasta::put_entry( $entry, $fh_out, 100 );
    }

    Maasha::Biopieces::put_record( $record, $out ); # must this be here?
}

close $fh_out;

$vmatch_args = vmatch_args( $options, $query_file );

map { Maasha::Common::run( "vmatch", "$vmatch_args $_ >> $result_file" ) } @index_files;

unlink $query_file;

$fh_in = Maasha::Filesys::file_read_open( $result_file );

if ( $options->{ "count" } )
{
    $count_hash = vmatch_count_hits( $result_file ) if ( $options->{ "count" } );

    if ( $options->{ "max_hits" } )
    {
        while ( $record = vmatch_get_entry( $fh_in ) )
        {
            $record->{ 'SCORE' } = $count_hash->{ $record->{ 'Q_ID' } };

            next if $record->{ 'SCORE' } > $options->{ 'max_hits' };

            Maasha::Biopieces::put_record( $record, $out );
        }
    
    }
    else
    {
        while ( $record = vmatch_get_entry( $fh_in ) )
        {
            $record->{ 'SCORE' } = $count_hash->{ $record->{ 'Q_ID' } };

            Maasha::Biopieces::put_record( $record, $out );
        }
    }
}
else
{
    while ( $record = vmatch_get_entry( $fh_in ) ) {
        Maasha::Biopieces::put_record( $record, $out );
    }
}

close $fh_in;

unlink $result_file;


# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> SUBROUTINES <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<


sub get_index_files
{
    # Martin A. Hansen, May 2009.

    # Get the index_files from the index_name or genome options.

    my ( $options,   # hashref
       ) = @_;

    # Returns a list.

    my ( @index_files, %hash );

    if ( $options->{ "index_name" } )
    {
        @index_files = $options->{ "index_name" };
    }
    else
    {
        @index_files = Maasha::Filesys::ls_files( "$ENV{ 'BP_DATA' }/genomes/$options->{ 'genome' }/vmatch" );

        map { $_ =~ /^(.+)\.[a-z1]{3,4}$/; $hash{ $1 } = 1 } @index_files;

        @index_files = sort keys %hash;
    }

    return wantarray ? @index_files : \@index_files;
}


sub vmatch_args
{
    # Martin A. Hansen, May 2009.

    # Given an option hashref and query file, compiles the arguments for running
    # Vmatch on the commandline.

    my ( $options,      # hashref
         $query_file,   # path to query file
       ) = @_;

    # Returns a string.

    my ( $vmatch_args );

    $vmatch_args  = "-complete -d -p -showdesc 100 -q $query_file";
    
    $vmatch_args .= " -h " . $options->{ "hamming_dist" } if $options->{ "hamming_dist" };
    $vmatch_args .= " -e " . $options->{ "edit_dist" }    if $options->{ "edit_dist" };

    return $vmatch_args;
}


sub vmatch_count_hits
{
    # Martin A. Hansen, April 2008.

    # Given a Vmatch result files, count duplications based
    # on q_id. The counts are returned in a hash.

    my ( $file,   # vmatch result file
       ) = @_;

    # Returns a list.

    my ( $fh_in, $line, @fields, %count_hash );

    $fh_in = Maasha::Filesys::file_read_open( $file );

    while ( $line = <$fh_in> )
    {
        chomp $line;

        next if $line =~ /^#/;

        @fields = split " ", $line;

        $count_hash{ $fields[ 5 ] }++;
    }

    close $fh_in;

    return wantarray ? %count_hash : \%count_hash;
}


sub vmatch_get_entry
{
    # Martin A. Hansen, January 2008.

    # Parses vmatch output records.

    my ( $fh,   # file handle to vmatch result file.
       ) = @_;

    # Returns a hash.

    my ( $line, @fields, %record );

    while ( $line = <$fh> )
    {
        chomp $line;

        next if $line =~ /^#/;
        $line =~ s/^\s+//;

        @fields = split " ", $line;

        $record{ "REC_TYPE" } = "VMATCH";

        $record{ "S_LEN" }      = $fields[ 0 ];
        $record{ "S_ID" }       = $fields[ 1 ];
        $record{ "S_BEG" }      = $fields[ 2 ];

        if ( $fields[ 3 ] eq "D" ) {
            $record{ "STRAND" } = "+";
        } else {
            $record{ "STRAND" } = "-";
        }

        $record{ "Q_LEN" }      = $fields[ 4 ];
        $record{ "Q_ID" }       = $fields[ 5 ];
        $record{ "Q_BEG" }      = $fields[ 6 ];
        $record{ "MATCH_DIST" } = $fields[ 7 ];
        $record{ "E_VAL" }      = $fields[ 8 ];
        $record{ "SCORE" }      = $fields[ 9 ];
        $record{ "IDENT" }      = $fields[ 10 ];

        $record{ "Q_END" }      = $record{ "Q_BEG" } + $record{ "Q_LEN" } - 1;
        $record{ "S_END" }      = $record{ "S_BEG" } + $record{ "S_LEN" } - 1;

        return wantarray ? %record : \%record;
    }
}

Maasha::Biopieces::close_stream( $in );
Maasha::Biopieces::close_stream( $out );


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


BEGIN
{
    Maasha::Biopieces::status_set();
}


END
{
    Maasha::Biopieces::status_log();
}


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


__END__