#!/usr/bin/perl


#####################################################
#
# Program: MetaPyler Version 1.07
#          Taxonomic profiling for metagenomic sequences
#          
#
# Author: Bo Liu, boliu@umiacs.umd.edu
#
# Last Update: Fri Jul 16 15:59:20 EDT 2010
#
#####################################################


use strict;
use warnings;


#---------------------------------------------------#
# command line parameters
#---------------------------------------------------#
my $blastfile = "";       # input blastx file
my $prefix = "";          # output files prefix
if (scalar @ARGV == 2) {
    $blastfile = $ARGV[0];
    $prefix = $ARGV[1];
} else {
    Usage();
}
#---------------------------------------------------#


#---------------------------------------------------#
# taxonomy information for reference marker genes 
#---------------------------------------------------#

# $gtax{gene id}{tax level} = tax name
my %gtax = ();   

# taxonomic info. for reference marker genes
my $gtaxfile = "./data/markerstax.tab";
open(FH, $gtaxfile) or die("Could not open file $gtaxfile\n");
foreach my $line (<FH>) {
    $line =~ /^(\S+)\t(\S+)\t(.+?)\t(\w+)$/;
    my ($acc, $tid, $tname, $lev) = ($1, $2, $3, $4);
    $gtax{$acc}{$lev} = $tname;
}
#---------------------------------------------------#


#---------------------------------------------------#
# load classifiers
#---------------------------------------------------#

# $g2m{$gene} = marker name
my %g2m = ();

# $abund{tax level}{tax name} = number of reads mapped
my %abund = ();

# total number of reads mapped
my $numMapped = 0;

# classifier for each gene
my %gcut = ();

# classifier file
open(FH, "./data/lr.tab");
foreach my $line (<FH>) {
    chomp $line;
    
    # @values: $a, $b, $minlen, $level, $tid, $tname, $m
    my ($acc, @values) = split("\t", $line);
    my $a = $values[0];
    my $b = $values[1];
    my $lev = $values[3];
    $gcut{$acc}{$lev} = [$a, $b];
    $g2m{$acc} = $values[6];
}
close FH;
#----------------------------------------------------#


#---------------------------------------------------#
# taxonomic classification for query reads
#---------------------------------------------------#

# classify at these tax levels
my @tlevs = ("genus", "family", "order", "class", "phylum");

# $hits{query id} = [hit ids]
my %hits = ();

# $hitsbit{query id} = [hit bit scores]
my %hitsbit = ();

# $hitspct{query id} = [hit percentages]
my %hitspct = ();

# $hitshspl{query id} = [hsp lengths]
my %hitshspl = ();

# open input blastx file in -m 8 format
open(FH, $blastfile);
foreach my $line (<FH>) {

    # all lines should match
    unless ($line =~ /^(\S+)\t(\S+)\t(\S+)\t(\S+).*\s(\S+)$/) {
	print "This line does not match:\n";
	print $line;
	exit;
    }
    my ($qacc, $hacc, $pct, $hspl, $bits) = ($1, $2, $3, $4, $5);

    # ignore if HSP length < 60bp
    if ($hspl < 20) { next;}

    # ignore if bit score <= 35
    if ($bits <= 35) { next;}

    # ignore if previous hits have been used
    # and % identity of current hit < 98%
    if (exists $hits{$qacc} && $pct < 98) { next;}

    # store blast information
    push(@{$hits{$qacc}}, $hacc);
    push(@{$hitspct{$qacc}}, $pct);
    push(@{$hitsbit{$qacc}}, $bits);
    push(@{$hitshspl{$qacc}}, $hspl);
}
close FH;
#---------------------------------------------------#


#---------------------------------------------------#
# classify query reads
#---------------------------------------------------#

# output classification file
open(OUTPUT, ">./$prefix.classify");
select OUTPUT;

foreach my $rid (keys %hits) {

    # accession number of the best hit
    my $hacc = $hits{$rid} -> [0];

    # bit score of the best hit
    my $bestbit = $hitsbit{$rid} -> [0];

    # HSP length of the best hit
    my $hspl = $hitshspl{$rid} -> [0];

    # classify at each tax level
    foreach my $lev (@tlevs) {

	my $assign = 0;
	
	##################################
	# if top hits are equally good
	# then see if their taxonomic labels
	# agree with each other
	my $num = scalar @{$hits{$rid}};
	my @tids = ();
	for (my $i = 0; $i < $num; $i++) {
	    if ($hitspct{$rid} -> [$i] >= 98) {
		my $refid = $hits{$rid} -> [$i];
		if (exists $gtax{$refid}{"$lev"}) {
		    push(@tids, $gtax{$refid}{$lev});
		}
	    }
	    else {
		last;
	    }
	}

	my $num2 = scalar @tids;
	for (my $i = 1; $i < $num2; $i++) {
	    my $preid = $tids[$i - 1];
	    my $id = $tids[$i];

	    # agree
	    if ($preid eq $id) {
		$assign = 1;
	    }

	    # do not agree
	    else {
		$assign = 2;
		last;
	    }
	}
	##################################
	
	# only one hit >= 98%
	# then classify to it
	if ($num2 == 1) {
	    $assign = 1;
	}

	# top hits do not agree
	# do not classify
	if ($assign == 2) { next;}

	# use bit score to classify
	if ($assign != 1) {
	    if (exists $gcut{$hacc}{"$lev"}) {
		my $a = $gcut{$hacc}{"$lev"} -> [0];
		my $b = $gcut{$hacc}{"$lev"} -> [1];
		my $cutvalue = $a + $b * $hspl;
		if ($bestbit >= $cutvalue) {
		    $assign = 1;
		}
	    }
	}

	# read is classified at $lev level
	if ($assign == 1) {

	    # total number of reads mapped
	    $numMapped++;
	    
	    print "$rid\t$hacc\t$g2m{$hacc}\t";
	    my $levtag = 0;
	    foreach my $tlev (@tlevs) {

		# classified at this level
		if ($tlev eq $lev) {
		    $levtag = 1;
		}

		# print out taxonomic information
		if ($levtag == 1) {

		    # has name at this tax level
		    if (exists $gtax{$hacc}{$tlev}) {
			my $tname = $gtax{$hacc}{$tlev};
			print "$tname\t";
			$abund{$tlev}{$tname}++;
		    }

		    # does not have name at this tax level
		    else {
			print "NA\t";
			$abund{$tlev}{"NA"}++;
		    }
		}

		# not classified at this level
		else {
		    print "NA\t";
		    $abund{$tlev}{"NA"}++;
		}
	    }
	    
	    print "\n";
	    last;
	}
    }
}
close OUTPUT;
select STDOUT;


#---------------------------------------------------#
# taxonomic profile at each level
#---------------------------------------------------#
# output taxonomic profile file
open(OUTPUT, ">./$prefix.tax");
select OUTPUT;

foreach my $lev (@tlevs) {
    
    print ">$lev\n";

    # sort by abundances from high to low
    foreach my $tname (sort {$abund{$lev}{$b} <=> $abund{$lev}{$a}} keys %{$abund{$lev}}) {
	my $abund = $abund{$lev}{$tname};

	# % abundance
	my $pct = $abund / $numMapped;

	# if % abundance <= 0.0005, ignore
	# if <= 2 reads classified at this tax level, ignore
	if ($tname eq "NA") {
	    if ($lev eq "phylum") { next;}
	}
	else {
	    if ($pct <= 0.0005) { next;}
	    if ($abund <= 2) { last;}
	}
	
	print "$tname\t";
	printf("%.2f\n", $pct * 100);
    }
    print "\n";
}
close OUTPUT;
select STDOUT;
#---------------------------------------------------#


exit;


sub Usage {
    die("
Usage:
       perl metaphyler.pl <BLAST file> <output prefix>

Output:
       \"prefix\".classify
           column 1: query read id
           column 2: reference gene id
           column 3: reference gene name
           column 4: taxonomic label at genus level
           column 5: taxonomic label at family level
           column 6: taxonomic label at order level
           column 7: taxonomic label at class level
           column 8: taxonomic label at phylum level

       \"prefix\".taxprof
           >taxonomic_level
           taxonomic_name percent_abundance

Notes:
       1, The input file is the BLASTP or BLASTX file (-m 8 format)
          of the metagenomic sequences against the reference marker genes.

Example:
       perl metaphyler.pl ./test/test.blastx test

");
}
