#!/bin/env perl

use strict;
use warnings;
use Cwd 'abs_path';

## Get base dir
my @spath = split(/\//, abs_path($0));
my $basedir = join('/', @spath[0..$#spath-2]);

## Input files
my ($pidf, $indf, $ensf) = (shift, shift, shift);

## Read in decision tree node map
open my $nfh, '<', "$basedir/flat/nodemap.tsv" or die $!;
my %node = ();
while(<$nfh>){
    next if($_ =~ m/^#/);
    chomp;
    my($i, $parent, $parent_call, $dec, $thresh) = split(/\t/, $_);
    $node{$i} = {
	'parent' => $parent,
	'decision' => $dec,
	'parent_call' => $parent_call,
	'thresh' => $thresh
    };
}
close $nfh;

## Define paths
my @path = ();
foreach my $n (sort {$a <=> $b} keys %node){
    if($node{$n}{'decision'} =~ m/^LEAF_NODE/){
	my $p = $node{$n}{'parent'};
	my $traceback = $node{$p}{'decision'}.'%'.$node{$p}{'thresh'}.'-'.$node{$n}{'parent_call'}.'&LEAF_NODE-'.$n;
	while($p != 0){
	    my $t = '';
	    ($p, $t) = getparent($p);
	    $traceback = $t.'&'.$traceback;
	}
	push @path, $traceback;
    }
}
my %m2p = ();
open my $cmfh, '<', "$basedir/flat/fullset0_smiles.faa" or die $!;
while(<$cmfh>){
    chomp;
    if($_ =~ m/^>(.+)/){
	my @a = split(/_/, $1);
	my $q = join('_', @a[0..$#a-1]);
	$q =~ s/[\(\)]/-/g;
	$m2p{$q} = $a[-1];
    }
}
close $cmfh;

## Parse pid and individual results
my %p = ();
open my $pfh, '<', $pidf or die $!;
while(<$pfh>){
    chomp;
    my($query, $pid) = split(/\t/, $_);
    $p{$query}{'pid'} = $pid;
}
close $pfh;
open my $afh, '<', $indf or die $!;
while(<$afh>){
    chomp;
    next if($_ =~ m/^shuffle/);
    my($query, $method, $called_spec, @rest) = split(/\t/, $_);
    $called_spec = 'nocall' unless(defined $called_spec && $called_spec =~ m/\S/);
    $called_spec =~ s/-/\|/g if($method eq 'pHMM');
    if($method eq 'prediCAT_MP'){
	if(exists $m2p{$called_spec}){
	    $called_spec = $m2p{$called_spec};
	}else{
	    $called_spec = 'nocall';
	}
    }
    if($method eq 'prediCAT_SNN' && $called_spec eq 'no_force_needed'){
	$called_spec = $p{$query}{'prediCAT_MP'};
    }
    $called_spec = cleannc($called_spec);
    #print STDERR join("\t", $query, $method, $called_spec)."\n";
    $p{$query}{$method} = $called_spec;
}
close $afh;
## Loop through all results
my $totpassed = 0;
foreach my $q (sort keys %p){
    foreach my $pa (@path){
	my $pass = 1;
	my @dec = split(/&/, $pa);
	foreach my $d (@dec){
	    last if($d =~ m/LEAF_NODE/);
	    my($decision, $threshchoice) = split(/%/, $d);
	    my($thresh, $choice) = split(/-/, $threshchoice);
	    if($decision eq 'pid'){
		if($choice eq 'T'){ ## need greater than thresh to pass
		    if($thresh <= $p{$q}{'pid'}){
			$pass = 0;
			last;
		    }
		}else{ ## need less than or eq thresh to pass
		    if($thresh > $p{$q}{'pid'}){
			$pass = 0;
			last;
		    }
		}
	    }else{ ## not pid
		$decision = cleannc($decision);
		my @a = split(/_/, $decision);
		my $m = join('_', @a[0..$#a-1]);
		my $sp = $a[-1];
		if($choice eq 'T'){ ## less than 0.5, so NOT sp
		    die join("\t", $q, $m)."\n" unless(exists $p{$q}{$m});
		    if($p{$q}{$m} eq $sp){
			$pass = 0;
			last;
		    }
		}else{ ## matches sp
		    if($p{$q}{$m} ne $sp){
			$pass = 0;
			last;
		    }
		}
	    }
	}
	$p{$q}{'pass'} = $pass;
	$p{$q}{'path'} = $pa if($pass==1);
	$totpassed += $pass;
    }
}
#print STDERR "$totpassed passed total\n";

my %ur = ();
open my $ufh, '<', "$basedir/flat/unreliable_paths.txt" or die $!;
while(<$ufh>){
    chomp;
    $ur{$_} = 1;
}
close $ufh;

my %e = ();
my $changes = 0;
open my $efh, '<', $ensf or die $!;
my $pcutoff = 97;
while(<$efh>){
    chomp;
    my($query, $method, $called_spec)=split(/\t/, $_);
    my $path = $p{$query}{'path'};
    if($p{$query}{'pid'} >= $pcutoff){
	## Get consensus
	my %con = ();
	foreach my $meth (sort keys %{$p{$query}}){
	    next if($meth eq 'pid');
	    next if($meth eq 'path');
	    next if($meth eq 'pass');
	    next if($p{$query}{$meth} =~ m/^nocall$/);
	    my @a = split(/\|/, $p{$query}{$meth});
	    foreach my $c (@a){
		$con{$c} += 1;
	    }
	}
        my $highscore = undef;
        my @specs = ();
	foreach my $s (sort {$con{$b} <=> $con{$a}} keys %con){
	    unless (defined $highscore) {
	        $highscore = $con{$s};
            }
            last if ($con{$s} < $highscore);
	    push @specs, $s;
	    $changes += 1;
	}
	$called_spec = join('|', sort @specs);
    }else{
	## Report no call for unreliable path
	if(exists $ur{$path}){
	    $called_spec = 'no_call';
	    $changes += 1;
	}
    }
    unless (defined $called_spec and length $called_spec > 0) {
        $called_spec = 'no_call';
    }
    print join("\t", $query, 'SANDPUMA', $called_spec)."\n";
}
close $efh;
#print STDERR "$changes changes made\n";

sub getparent{
    my $n = shift;
    my $p = $node{$n}{'parent'};
    #my $c = $node{$n}{'parent_call'};
    my $c = $node{$p}{'thresh'}.'-'.$node{$n}{'parent_call'};# if($node{$n}{'decision'} eq 'pid');
    return ($p, $node{$p}{'decision'}.'%'.$c);
}
sub cleannc{
    my $s = shift;
    $s =~ s/_result//;
    $s =~ s/no_confident/nocall/;
    $s =~ s/N\/A/nocall/;
    $s =~ s/no_call/nocall/;
    return($s);
}
