#!/usr/bin/perl -w

$|=1;

use strict;
use Getopt::Std;
use Statistics::Basic qw(mean stddev median);
use List::Util qw(max);

my $EXPRTABEXT = '.expr.tab';
my $PLOTFORMAT = 'pdf'; # png
my $RVERBOSE = 0;

my $Rparams = '';
if(!$RVERBOSE){ $Rparams = '-q 2>&1 > /dev/null' }

my ($INP_matrix,$INP_taxa,$INP_length,$INP_nucleotides,$INP_outmatrix,$INP_outplot,$INP_expr_dir,%opts) = ('',1,0,0);
my ($INP_identity,$INP_cover) = (90,75);

getopts('ho:m:e:p:', \%opts);

if(($opts{'h'})||(scalar(keys(%opts))==0))
{
  print   "\n[options]: \n";
  print   "-h \t this message\n";
  print   "-m \t input pangenome matrix .tab            (required, expects source: in 1st cell)\n"; 
  print   "-e \t folder with $EXPRTABEXT files            (required, expects files with 2 cols: id expression)\n"; 
  print   "-o \t output matrix filename                 (optional, adds column with mean expression)\n"; 
  print   "-p \t output plot filename                   (required, $PLOTFORMAT format)\n"; 
  exit(-1);
}

if(defined($opts{'m'})){ $INP_matrix = $opts{'m'} }
else{ die "\n# $0 : need -m parameter, exit\n"; }

if(defined($opts{'e'})){ $INP_expr_dir = $opts{'e'} }
else{ die "\n# $0 : need -e parameter, exit\n"; }

if(defined($opts{'o'})){ $INP_outmatrix = $opts{'o'} }

if(defined($opts{'p'})){ $INP_outplot = $opts{'p'} }
else{ die "\n# $0 : need -p parameter, exit\n"; }


#################################### MAIN PROGRAM  ################################################

my (@expr_files,@cluster_names,@occup,%taxa,%panexpr,%expr_values);
my ($col,$cluster_name,$seq,$expr,$cluster_expr_line,$cluster_expr,$e,$n_of_taxa);
my ($n_of_clusters,$max_occup,$maxsize,$ecotype,$cluster_path,$name,$taxon) = (0,0,0);

## 1) read expression values
opendir(EXPRDIR,$INP_expr_dir);
@expr_files = grep{/$EXPRTABEXT/} readdir(EXPRDIR);
closedir(EXPRDIR);

foreach my $expr_file (@expr_files)
{
  $taxon = $expr_file;
  $taxon =~ s/$EXPRTABEXT//; 

  open(EXPRFILE,"$INP_expr_dir/$expr_file") || die "# cannot read $INP_expr_dir/$expr_file\n";
  while(<EXPRFILE>)
  {
    chomp;
    ($name,$expr) = split(/\s+/,$_,2); #print "|$taxon|$name| $expr\n"; #|SBCC073_fLF.Trinity|TR1|c0_g1_i1| 4.3369
    $expr_values{$taxon}{$name} = $expr; #Scarlett_fLF.Trinity            SBCC073_fLF.Trinity
  }
  close(EXPRFILE);
  
  print "# read $INP_expr_dir/$expr_file -> $taxon\n" if($expr_values{$taxon});
} print "\n"; 

## 2) parse pangenome matrix
if($INP_outmatrix)
{
    open(OUTMAT,'>',$INP_outmatrix) || die "#EXIT: cannot create $INP_outmatrix\n";
}

open(MAT,$INP_matrix) || die "# EXIT : cannot read $INP_matrix\n";
while(<MAT>)
{
  next if(/^#/ || /^$/);
  chomp;
  my @data = split(/\s+/,$_);
  #source:/data2/Brachypodium/genome/genomes/clusters_pep_cov75	24_Brdisv1ABR21043430m.faa	297_Brdisv1ABR21026756m.faa	...
  if($data[0] =~ /^source:(\S+)/) 
  {
    $cluster_path = $1;
    foreach $col (1 .. $#data)
    { 
      $data[$col] =~ s/\+\d$//;
      $cluster_names[$col-1] = $data[$col];
    }
    $n_of_clusters = $#data-1;

    print OUTMAT "$_\n" if($INP_outmatrix);
  }
  else 
  {
    #_Escherichia_coli_ETEC_H10407_uid42749.gbk     1       1	...
    #Bd3-1_r.1.cds.fna.nucl	1	2 ...
    foreach $col (1 .. $#data)
    {
      $occup[$col-1]++ if($data[$col] > 0);
    }

    print OUTMAT "$_\n" if($INP_outmatrix);
  }
}
close(MAT);
printf( STDERR "# matrix contains %d clusters\n",$n_of_clusters+1);

## 3) parse clusters and get expression values of matched sequences
print STDERR "\n# parsing clusters...\n";
foreach $col (0 .. $n_of_clusters)
{
  #print "$cluster_names[$col] $occup[$col]\n";
  my (@expression);
  $cluster_expr = 0;
  open(CLUSTER,"$cluster_path/$cluster_names[$col]");
  while(<CLUSTER>)
  {
      #>Brdisv1S8iiC1042528m ... [S8iiC.1.cds.fna]
      #>TR20167|c1_g1_i1 len=1493 ... [SBCC073_fLF.Trinity.fna.gz]
      if(/^>(\S+).*\[(\S*?).fn*a.*?\]$/) # works for Athaliana and barley if(/^>(\S+).*?\] \[(\S*?).fna/)
      {
        ($name,$taxon) = ($1,$2); #print "|$taxon|$name|\n";
        #$taxon = lc($taxon); # comment for barley
        #next if($strains2skip{$taxon});
        
        if($expr_values{$taxon}{$name})
        {
          $expr = abs(log($expr_values{$taxon}{$name})/log(2));
          push(@expression,$expr);
          #print "$col $occup[$col] $cluster_names[$col] $name,$taxon,$expr_values{$taxon}{$name}\n";
        }
        else
        {
          push(@expression,0); # unpressed sequence
        }
      }
      
  } 
  close(CLUSTER);
  
  #$cluster_expr = mean(\@expression); # Athaliana
  $cluster_expr = max(@expression); # barley
  $cluster_expr_line .= "\t$cluster_expr";
  
  print "# cluster $col $occup[$col] $cluster_expr\r";
  
  push(@{$panexpr{$occup[$col]}{'data'}},$cluster_expr);
  $panexpr{$occup[$col]}{'n'}++;
  
  if($panexpr{$occup[$col]}{'n'} > $maxsize){ $maxsize = $panexpr{$occup[$col]}{'n'} }
  
  #last if($col>500); # debug
} print "\n";

$max_occup = (sort {$b<=>$a}(keys(%panexpr)))[0]; 

print OUTMAT "mean_expression".$cluster_expr_line."\n" if($INP_outmatrix);
close(OUTMAT);

my $CORRFILE = '_corrdata.tab';
#open(CORRDATA,'>',$CORRFILE);
print "#occupancy\tn\tmean\tstddev\tmedian\n";
foreach $n_of_taxa ( 1 .. $max_occup )
{
  next if(!$panexpr{$n_of_taxa}{'n'});
  
  printf("%d\t%d\t%1.2f\t%1.2f\t%1.2f\n",
    $n_of_taxa,
    $panexpr{$n_of_taxa}{'n'},
    mean($panexpr{$n_of_taxa}{'data'}), 
    stddev($panexpr{$n_of_taxa}{'data'}),
    median($panexpr{$n_of_taxa}{'data'}), 
    );
    
  #foreach $e ( 0 .. $panexpr{$n_of_taxa}{'n'}-1){print CORRDATA "$n_of_taxa\t$panexpr{$n_of_taxa}{'data'}->[$e]\n";}  
}
#close(CORRDATA);

## 4) make boxplot and calculate correlations in R
my $TMPBOXPLOT = '_boxdata.tab';
open(TAB,">$TMPBOXPLOT");
foreach $n_of_taxa ( 1 .. $max_occup )
{  
  print TAB "$n_of_taxa\t";
}
print TAB "\n";

for($e=0;$e<$maxsize;$e++)
{
   foreach $n_of_taxa ( 1 .. $max_occup )
   {
      $expr = $panexpr{$n_of_taxa}{'data'}->[$e] || 'NA';
      print TAB "$expr\t";
   }
   print TAB "\n";
}
close(TAB);


open(RSHELL,"|R --no-save $Rparams ") || die "# $0 : cannot call R: $!\n";
print RSHELL<<EOR;

#png("$INP_outplot",width=800);
pdf("$INP_outplot",width=10);  

#pancolors = c(rep('red',2),rep('orange',15),rep('yellow',1),'white') # Athaliana
pancolors = c(rep('red',2),rep('orange',12),rep('yellow',1),'white') # barley

#raw  = read.table(file="$CORRFILE",header=F)
expr = read.table(file="$TMPBOXPLOT",header=T,check.names=F); 
boxplot( expr, ylab='|log2(transcripts per million)|',xlab='genome occupancy', 
         varwidth=F, show.names=T, outline=F, notch=T, #outpch=20,
         col=pancolors, cex.axis=1.5, cex.names=1.5, cex.lab=1.5 );

grid(nx=NA,ny=20,lwd=2) #grid over boxplot
par(new=TRUE)

boxreport = boxplot( expr, ylab='|log2(transcripts per million)|',xlab='genome occupancy', 
         varwidth=F, show.names=T, outline=F, notch=T, #outpch=20,
         col=pancolors, cex.axis=1.5, cex.names=1.5, cex.lab=1.5 );
         
#lines(lowess(raw\$V1,raw\$V2),col="blue")        


legend("topleft", inset=.05, c("cloud","shell","soft-core","core"), 
   fill=c('red','orange','yellow','white'), horiz=TRUE, cex=1.5)

dev.off()

q()
EOR
close RSHELL;

unlink($TMPBOXPLOT);


sub read_FASTA_file
{

# in FASTA format
# returns a reference to a 2D array for 4 secondary keys: NAME,SEQ,IDENTICALS,NUMBER_IDENTICALS
# first valid index (first sequence) is '0'
  my ( $infile  ) = @_;
  my (%FASTA,$name,$seq,$n_of_sequences,$magic);
  $n_of_sequences = -1;

  # check input file format and open it accordingly
  open(INFILE,$infile) || die "# read_FASTA_sequence_array: cannot read $infile, exit\n";
  sysread(INFILE,$magic,2);
  close(INFILE);

  if($infile =~ /\.gz$/ || $magic eq "\x1f\x8b") # GZIP compressed input
  {
    if(!open(FASTA,"gzip -dc $infile |"))
    {
      die "# read_FASTA_sequence_array: cannot read GZIP compressed $infile $!\n"
        ."# please check gzip in installed\n";
    }
  }
  elsif($infile =~ /\.bz2$/ || $magic eq "BZ") # BZIP2 compressed input
  {
    if(!open(FASTA,"bzip2 -dc $infile |"))
    {
      die "# read_FASTA_sequence_array: cannot read BZIP2 compressed $infile $!\n"
        ."# please check bzip2 in installed\n";
    }
  }
  else{ open(FASTA,"<$infile") || die "# read_FASTA_sequence_array: cannot read $infile $!\n"; }

  while(<FASTA>)
  {
    next if(/^$/);
    next if(/^#/);
    if(/^\>(\S+)/)
    {
      $n_of_sequences++;
      $name = $1; #print "# $name\n";
    }
    elsif($n_of_sequences>-1)
    {
      $_ =~ s/[\s|\n|\-|\.]//g;
      $FASTA{$name} .= $_; # conserve case
    }
  }
  close(FASTA);

  return \%FASTA;  
}
