# perl /home/xiangyang/土壤银120个样本数据/scripts_meta/filter_taxonomy_bracken.pl kreports bracken_outputs $include_tax
# usage
# perl /home/xiangyang/土壤银120个样本数据/scripts_meta/filter_taxonomy_bracken.pl /home/xiangyang/土壤银120个样本数据/RESULT_120/kreports /home/xiangyang/土壤银120个样本数据/RESULT_120/bracken_outputs 'd__Bacteria|d__Archaea'


use strict;
use warnings;
use File::Basename;
use FindBin;

#Author: Xiangyang Li
#Function: filter the *.bracken by d__Bacteria, d__Archaea, d__Eukaryota or k__Fungi, d__Viruses


my $k2report_dir = $ARGV[0];
my $bracken_dir = $ARGV[1];
my $tax_include = $ARGV[2]; #d__Bacteria, d__Archaea, d__Eukaryota (d__Eukaryota|k__Fungi), d__Viruses
my @tax_include = split /\|/, $tax_include;

my $metadata = "/home/xiangyang/土壤银120个样本数据/meta_data.txt";
my %order;
open(META, "$metadata");
<META>; #delete header
while(<META>){
    chomp;
    my @meta = split /\t/, $_;
    $order{$meta[3]} = $meta[0];
}
opendir (K, $k2report_dir);
my @k = readdir K;
@k = grep($_ !~ /^\./, @k);
closedir K;

my $result_dir = dirname($bracken_dir);
my $bracken_filter_dir = dirname($bracken_dir)."/bracken_filter_outputs";
mkdir $bracken_filter_dir;
my $mpa_dir = dirname($bracken_dir)."/mpas";
mkdir $mpa_dir;
my $home_dir = $FindBin::Bin;
my @fname;
my @bracken;
my @indexs=("BP", "Sh", "F", "Si", "ISi");

open(A, ">>$result_dir/αdiversity.txt");

foreach(sort{$a<=>$b} keys %order){
    my $fname = $order{$_};
    system("python $home_dir/KrakenTools/kreport2mpa_Li.py -r $k2report_dir/$order{$_}.k2report --display-header -o $mpa_dir/$fname.mpa.tmp");
    open(T, "$mpa_dir/$fname.mpa.tmp");
    open(M, ">$mpa_dir/$fname.mpa");
    print M "#Classification\t$order{$_}.k2report\n"; #print header
    my %taxid;
    while(<T>){
        chomp;
        next unless $_ !~ /^#/; #skip header
        my @m = split /\t/, $_;
        print M "$m[2]\t$m[3]\n" if grep($m[2] =~ /\Q$_\E/, @tax_include);
        $taxid{$m[0]} = $m[2];   # taxid vs. full taxonomy information
    }
    close T;
    close M;
    system("rm $mpa_dir/$fname.mpa.tmp");
    push @fname, "$mpa_dir/$fname.mpa";
    
    open(B, "$bracken_dir/$fname.bracken");
    open(F, ">$bracken_filter_dir/$fname.bracken");
    while(<B>){
        chomp;
        print F "$_\n" if $_ =~ /^name\ttaxonomy_id/;
        my @b = split /\t/, $_;
        print F "$_\n" if defined $taxid{$b[1]} && grep($taxid{$b[1]} =~ /\Q$_\E/, @tax_include);
    }
    close F;
    close B;
    
    ###计算αdiversity 
    
    print A "$fname\n";
    foreach(@indexs){
        system("python $home_dir/KrakenTools/DiversityTools/alpha_diversity.py -f $bracken_filter_dir/$fname.bracken -a $_ >> $result_dir/αdiversity.txt");
    }
    push @bracken, "$bracken_filter_dir/$fname.bracken";
}

close A; #"$result_dir/αdiversity.txt"

### 处理αdiversity.txt文件格式
open (AA, "$result_dir/αdiversity.txt");
open (BB, ">$result_dir/αdiversity_tab.txt");
print BB "sample\t", join("\t", @indexs)."\n";
my $count=0;
my @content;
while(<AA>){
    chomp;
    $count++; 
    $_ =~ s/.*\: //;
    push @content, "$_" if $_ ne "Fisher's alpha...loading";
    print BB join("\t", @content), "\n" if $count/7 == int($count/7);
    @content = () if $count/7 == int($count/7);
}
close AA;
close BB;

#Calculate β-diversity
my $bracken_all_str = join(" ", @bracken);
system("python $home_dir/KrakenTools/DiversityTools/beta_diversity.py -i $bracken_all_str --type bracken > $result_dir/βdiversity.txt");

###合并所有mpa，输出all.mpa
my $fname_all_str = join(" ", @fname);
system("python $home_dir/KrakenTools/combine_mpa.py -i $fname_all_str -o $result_dir/all.mpa");

###将mpa转为spf格式文件
system("perl $home_dir/script/metaphlan_to_stamp.pl $result_dir/all.mpa > $result_dir/all.mpa.spf");


