package Maasha::Fasta;

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


# Routines for manipulation of FASTA files and FASTA entries.


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


use strict;
use Data::Dumper;
use Maasha::Common;
use Maasha::Seq;
use vars qw ( @ISA @EXPORT );

@ISA = qw( Exporter );

use constant {
    SEQ_NAME => 0,
    SEQ  => 1,
};


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


sub fasta_format_ok
{
    # Martin A. Hansen, March 2007.

    # Checks if a given FASTA file is formatted with
    # one header per line and one sequence per line.
    # returns 1 if so, else 0.

    my ( $path,   # full path to FASTA file
       ) = @_;

    # Returns boolean

    my ( $fh, $line, $count );

    $fh = Maasha::Common::read_open( $path );

    $count = 0;

    while ( $line = <$fh> )
    {
        if ( not $count % 2 and substr( $line, 0, 1 ) ne ">" ) {
            return 0;
        }

        $count++;
    }

    close $fh;

    return 1;
}


sub get_entries
{
    # Martin A. Hansen, December 2006.

    # Parses a fasta file and returns a list of headers and sequence tuples.

    my ( $path,   # full path to FASTA file
         $count,  # number of sequences to read - OPTIONAL
       ) = @_;

    # returns list of tuples

    my ( $fh, $entry, @entries );

    $fh = Maasha::Common::read_open( $path );

    while ( $entry = get_entry( $fh ) )
    {
        push @entries, $entry;

        if ( $count and $count == @entries ) {
            last;
        }
    }

    close $fh;

    return wantarray ? @entries : \@entries;
}


sub put_entries
{
    # Martin A. Hansen, March 2004.

    # writes fasta sequences to STDOUT or file

    my ( $entries,   # list of fasta entries
         $path,      # full path to file - OPTIONAL
         $wrap,      # line width        - OPTIONAL
       ) = @_;

    my ( $fh );

    $fh = Maasha::Common::write_open( $path ) if $path;

    map { put_entry( $_, $fh, $wrap ) } @{ $entries };

    close $fh if defined;
}


sub wrap
{
    # Martin A. Hansen, June 2007

    # Wraps the sequence of a given FASTA entry
    # to a given length.

    my ( $entry,   # FASTA entry
         $wrap,    # wrap length
       ) = @_;

    # Returns nothing.

    Maasha::Seq::wrap( \$entry->[ SEQ ], $wrap );
}


sub get_entry
{
    # Martin A. Hansen, January 2007.

    # Given a filehandle to an FASTA file,
    # fetches the next FASTA entry, and returns
    # this as a tuple of [ header, sequence ].

    my ( $fh,    # filehandle to FASTA file
       ) = @_;

    # Returns string.
    
    my ( $block, @lines, $seq_name, $seq, $entry );

    local $/ = "\n>";

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

        last if $block !~ /^\s+$/;
    }

    return if not defined $block;

    # $block =~ />?([^\n]+)\n/m;
    $block =~ /^>?(.+)\n/m;
    $seq_name = $1;
    $seq      = $';

    local $/ = "\n";

    chomp $seq;

    $seq_name =~ tr/\r//d;
    $seq      =~ tr/ \t\n\r//d;

    $entry = [ $seq_name, $seq ];

    return wantarray ? @{ $entry } : $entry;
}


sub put_entry
{
    # Martin A. Hansen, January 2007.

    # Writes FASTA entries to STDOUT or file.

    my ( $entry,     # a FASTA entries
         $fh,        # file handle to output file - OPTIONAL
         $wrap,      # line width                 - OPTIONAL
       ) = @_;

    # Returns nothing.

    Maasha::Common::error( qq(FASTA entry has no header) )   if not defined $entry->[ SEQ_NAME ];
    Maasha::Common::error( qq(FASTA entry has no sequence) ) if not defined $entry->[ SEQ ];

    if ( $wrap ) {
        Maasha::Fasta::wrap( $entry, $wrap );
    }

    if ( defined $fh ) {
        print $fh ">$entry->[ SEQ_NAME ]\n$entry->[ SEQ ]\n";
    } else {
        print ">$entry->[ SEQ_NAME ]\n$entry->[ SEQ ]\n";
    }
}


sub find_shortest
{
    # Martin A. Hansen, June 2007.

    # Given a stack of FASTA entries, find and return
    # the shortest entry.

    my ( $entries,   # list of FASTA entries
       ) = @_;

    # returns tuple

    my ( $min, $entry, $min_entry );

    $min = 99999999999;

    foreach $entry ( @{ $entries } )
    {
        if ( length( $entry->[ SEQ ] ) < $min )
        {
            $min_entry = $entry;
            $min = length $entry->[ SEQ ];
        }
    }

    return wantarray ? @{ $min_entry } : $min_entry;
}


sub find_longest
{
    # Martin A. Hansen, June 2007.

    # Given a stack of FASTA entries, find and return
    # the longest entry.

    my ( $entries,   # list of FASTA entries
       ) = @_;

    # returns tuple

    my ( $max, $entry, $max_entry );

    $max = 0;

    foreach $entry ( @{ $entries } )
    {
        if ( length( $entry->[ SEQ ] ) > $max )
        {
            $max_entry = $entry;
            $max = length $entry->[ SEQ ];
        }
    }

    return wantarray ? @{ $max_entry } : $max_entry;
}


sub fasta_get_headers
{
    # Martin A. Hansen, May 2007.

    # Gets the header names of a FASTA file,
    # and returns these in a list.

    my ( $path,   # full path to FASTA file
       ) = @_;

    # returns list

    my ( $fh, $entry, @list );

    $fh = Maasha::Common::read_open( $path );

    while ( $entry = get_entry( $fh ) ) {
        push @list, $entry->[ SEQ_NAME ];
    }

    close $fh;

    return wantarray ? @list : \@list;
}


sub fasta_reformat
{
    # Martin A. Hansen, December 2004.

    # Given a file of one or more FASTA entries, reformats these so
    # each entry consits of one line with header and one line with sequence.

    my ( $path,   # full path to file with multiple FASTA entries
       ) = @_;

    my ( $fh_in, $fh_out, $entry );

    $fh_in  = Maasha::Common::read_open( $path );
    $fh_out = Maasha::Common::write_open( "$path.temp" );

    while ( $entry = get_entry( $fh_in ) ) {
        put_entry( $entry, $fh_out );
    }

    close $fh_in;
    close $fh_out;

    rename( "$path.temp", $path );
}


sub fasta_index
{
    # Martin A. Hansen, July 2008.

    # Given a path to a FASTA file create a FASTA index
    # and save to the specified path.

    my ( $fasta_path,   # path to FASTA file
         $index_path,   # path to index file
       ) = @_;

    # Returns nothing.

    my ( $index );

    $index = index_create( $fasta_path );

    Maasha::Fasta::index_store( $index_path, $index );
}

sub index_create
{
    # Matin A. Hansen, December 2004.

    # Given a FASTA file formatted with one line of header and one line of sequence,
    # returns a list of header, seq beg and seq length (first nucleotide is 0). Also,
    # the file size of the indexed file is written to the index for checking purposes.

    my ( $path,   #  full path to file with multiple FASTA entries
       ) = @_;

    # returns a hashref

    my ( $file_size, $fh, $entry, $beg, $len, %hash, @index );

    $file_size = Maasha::Common::file_size( $path );

    push @index, "FILE_SIZE=$file_size";

    $fh = Maasha::Common::read_open( $path );

    $beg = 0;
    $len = 0;

    while ( $entry = get_entry( $fh ) )
    {
        warn qq(WARNING: header->$entry->[ SEQ_NAME ] alread exists in index) if exists $hash{ $entry->[ SEQ_NAME ] };

        $beg += $len + 2 + length $entry->[ SEQ_NAME ];
        $len = length $entry->[ SEQ ];

        push @index, [ $entry->[ SEQ_NAME ], $beg, $len ];

        $hash{ $entry->[ SEQ_NAME ] } = 1;

        $beg++;
    }
    
    close $fh;

    return wantarray ? @index : \@index;
}


sub index_search
{
    # Martin A. Hansen, December 2004.

    # Searches the index for matching entries.
    
    my ( $index,   # index list
         $regex,   # regex to match FASTA headers [OPTIONAL]
         $invert,  # invert matching
       ) = @_;

    # returns list

    my ( @results );

    if ( not $regex )
    {
        @results = @{ $index };
    }
    else
    {
        if ( $invert ) {
            @results = grep { $_->[ 0 ] !~ /$regex/ } @{ $index };
        } else {
            @results = grep { $_->[ 0 ] =~ /$regex/ } @{ $index };
        }
    }

    return wantarray ? @results : \@results;
}


sub index_lookup
{
    # Martin A. Hansen, July 2007.

    # Lookup a list of exact matches in the index and returns these

    my ( $index,     # index list
         $headers,   # headers to lookup
       ) = @_;

    # returns a list

    my ( %hash, $head, @results );

    map { $hash{ $_->[ 0 ] } = [ $_->[ 1 ], $_->[ 2 ] ] } @{ $index };

    foreach $head ( @{ $headers } )
    {
        if ( exists $hash{ $head } ) {
            push @results, [ $head, $hash{ $head }->[ 0 ], $hash{ $head }->[ 1 ] ];
        }
    }

    return wantarray ? @results : \@results;
}


sub index_store
{
    # Martin A. Hansen, May 2007.

    # Stores a FASTA index to binary file.

    my ( $path,   # full path to file
         $index,  # list with index
       ) = @_;

    # returns nothing

    Maasha::Common::file_store( $path, $index );
}


sub index_retrieve
{
    # Martin A. Hansen, May 2007.

    # Retrieves a FASTA index from binary file.

    my ( $path,   # full path to file
       ) = @_;

    # returns list

    my $index;

    $index = Maasha::Common::file_retrieve( $path );

    return wantarray ? @{ $index } : $index;
}


sub biopiece2fasta
{
    # Martin A. Hansen, July 2008.

    # Given a biopiece record converts it to a FASTA record.
    # If no generic SEQ or SEQ_NAME is found, the Q_* and S_* are
    # tried in that order.

    my ( $record,    # record
       ) = @_;

    # Returns a tuple.

    my ( $seq_name, $seq );

    $seq_name = $record->{ "SEQ_NAME" } || $record->{ "Q_ID" }  || $record->{ "S_ID" };
    $seq      = $record->{ "SEQ" }      || $record->{ "Q_SEQ" } || $record->{ "S_SEQ" };

    if ( defined $seq_name and defined $seq ) {
        return wantarray ? ( $seq_name, $seq ) : [ $seq_name, $seq ];
    } else {
        return;
    }
}


sub fasta2biopiece
{
    # Martin A. Hansen, May 2009

    # Convert a FASTA record to a Biopiece record.

    my ( $entry,   # FASTA entry
       ) = @_;

    # Returns a hashref.

    my ( $record );

    if ( defined $entry->[ SEQ_NAME ] and $entry->[ SEQ ] )
    {
        $record = {
            SEQ_NAME => $entry->[ SEQ_NAME ],
            SEQ      => $entry->[ SEQ ],
            SEQ_LEN  => length $entry->[ SEQ ],
        };
    }

    return wantarray ? %{ $record } : $record;
}



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