#!/usr/bin/perl -w

use strict;

# TREC 2012 Text Relevance Assessing Task submission Score script. 
# Scores designated run and computes: (given threshold step size)
#   * LAM
#   * AUC (using the Mann-Whitney U statistic

# Each line in the Submission (Results) file should be in the form
#   topic docno judgement probability_of_relevance runtag
# e.g.
#   411 FBIS3-10248 0 0.3 UWatMDSvg
# or (if not reporting probability of relevance)
#   411 FBIS3-10248 0 NA UWatMDSvg

# Author: Gaurav Baruah
# Date: October 2012
# version: 1.0 (October 11, 2012)

my $usage = "Usage: $0 <topic-docno-pairs-file> <qrel-file> <run-file>\nComputes LAM and/or AUC for <run-file>. Prints score to console.\n";
if (($#ARGV + 1) != 3) {
    die $usage;
}
my $tdfile = $ARGV[0];
my $qrelfile = $ARGV[1];
my $runfile = $ARGV[2];
#my $scorefile = $runfile.".SCORE";

# init/declare variables
my %topicLAM;  #topic --> {tp => , tn => , fp => , fn => , fpr => , fnr => , lam => } to be used for computing LAM 
my %topicAUC;  #topic --> ( (score, pos/neg), ... )  array of (score, class) tuples to be used to compute rankings and thereafter the U statistic and then the AUC
my $tp = 0;	# true positive 
my $fp = 0; # false positive
my $tn = 0;
my $fn = 0; 

open TDFILE, "<$tdfile" || die "Cannot open topic-docno-pair-file $tdfile.\n";

# create and load maps
my @tdpairs = <TDFILE>;
chomp($_) for (@tdpairs);

my %groundTruth = map { $_ => 0 } @tdpairs;  #ground truth topic-docno --> relevance labels

close TDFILE || die "Close failed for $tdfile";
   
#load the qrel file into hash
open QRELFILE, "<$qrelfile" || die "Cannot open $qrelfile file. \n";
my $line = "";
while ($line = <QRELFILE>) {
    chomp($line);
    next if ($line =~ /^\s*$/);
    my @fields = split (/\s+/, $line);
    if (@fields != 4) { 
        print "ERROR: $qrelfile has unexpected number of columns at \n$line\nExpecting columns \"<topic> 0 <docno> <relevance-label>\".\nPlease also check for extraneous whitespace.\n";
        exit 255;
    }
    if ($fields[1] ne "0" || ($fields[3] ne "0" && $fields[3] ne "1")) {
        print "ERROR: $qrelfile has invalid values at \n$line\nExpecting \"<topic> 0 <docno> <relevance-label = 0|1>\".\n";
        exit 255;

    }
    my $tdpair = $fields[0]." ".$fields[2];
    my $gtlabel = $fields[3];
    next if (! exists($groundTruth{$tdpair})) ;
    
    $groundTruth{$tdpair} = $gtlabel;    
    if ( !exists($topicLAM{$fields[0]})) {
        $topicLAM{$fields[0]} = {} ; #("tp", 0, "fp", 0, "tn", 0, "fn", 0);
        $topicLAM{$fields[0]}->{"tp"} = 0;
        $topicLAM{$fields[0]}->{"fp"} = 0;
        $topicLAM{$fields[0]}->{"tn"} = 0;
        $topicLAM{$fields[0]}->{"fn"} = 0;
        $topicAUC{$fields[0]} = [ ];        
    }        
}
close QRELFILE || die "Close failed for $qrelfile";

#print "size of topicLAM = ". keys(%topicLAM);

#process the run file 
open RUNFILE, "<$runfile" ||
	die "Unable to open run-file $runfile: $!\n";

my @lines = <RUNFILE>;

if (@lines != keys (%groundTruth))  {
    #die "WARNING: Invalid number of judgements ".@lines.". Expecting ". keys(%groundTruth).".\n";
    print "WARNING: Invalid number of judgements ".@lines.". Expecting ". keys(%groundTruth).".\n";
}

my ($topic,$docno,$judgement,$PoR,$tag);
$line = "";    
foreach my $line (@lines) { 
    chomp $line;
    next if ($line =~ /^\s*$/);
    
    my @fields = split " ", $line;
           
    if (scalar(@fields) == 5) {
        ($topic,$docno,$judgement,$PoR,$tag) = @fields;
    } else {
        printf ("ERROR: Wrong number of columns (expecting 5 columns)");
        exit 255;
    }
     
     
    #after the checks, getting on with the main task
    #get TP, FP, TN, FN
    
    my $tdpair = $topic." ".$docno;
    my $gtlabel = $groundTruth{$tdpair};
    
    if ($PoR ne "NA") { # add to topicAUC
        my @sc = ($PoR, $gtlabel);
        push(@{$topicAUC{$topic}}, \@sc);
    }
    
            
    if ($gtlabel == 1 and $judgement == 1) {
        $topicLAM{$topic}->{"tp"} += 1;
        #$tp += 1;
    }
    if ($gtlabel == 0 and $judgement == 0) {
        $topicLAM{$topic}->{"tn"} += 1;
        #$tn += 1;
    }
    if ($gtlabel == 0 and $judgement == 1) {
        $topicLAM{$topic}->{"fp"} += 1;
        #$fp += 1;
    }
    if ($gtlabel == 1 and $judgement == 0) {
        $topicLAM{$topic}->{"fn"} += 1;
        #$fn += 1;
    }        
}
close RUNFILE || die "Close failed for $runfile";

#TODO: for the topic "all" are we sure we will be computing the average of all topic scores? We are in a position to compute the overall score.

#Compute the scores
#open SCOREFILE, ">$scorefile" || die "Unable to open output file $runfile.LAM: $!\n";
    
#print SCOREFILE "Metric Topic Value\n";
print "Metric Topic Value\n";

my $sumLAM = 0;
my $sumAUC = 0;

for my $topic (sort (keys (%topicLAM))) {
    
    #LAM
    my $topicHash = $topicLAM{$topic};
    $tp = $topicHash->{"tp"};
    $fp = $topicHash->{"fp"};
    $tn = $topicHash->{"tn"};
    $fn = $topicHash->{"fn"};

    my $fpr = ($fp + 0.5)/($fp + $tn + 1);
    my $fnr = ($fn + 0.5)/($fn + $tp + 1);
    my $LAM = &lam($fpr, $fnr);
    #printf SCOREFILE "LAM $topic %0.3f\n", $LAM;    
    printf "LAM $topic %0.3f\n", $LAM;    
    $sumLAM += $LAM;
    
    if($PoR ne "NA" ) {
        #AUC
        my @ranking = &rank(@{$topicAUC{$topic}});
        my @AUC = &UStatArea(@ranking);     
        #printf SCOREFILE "AUC $topic %0.3f\n", $AUC[0];
	if($AUC[0] == -1) {
            printf ("ERROR: Cannot compute AUC.  No negative (0|non-rel) ground-truth values in qrels for topic $topic\n");
            next; #exit 255;
	} elsif ( $AUC[1] == -1) {
            printf ("ERROR: Cannot compute AUC.  No positive (1|rel) ground-truth values in qrels for topic $topic\n");
            next; #exit 255;	
	} else {
            printf "AUC $topic %0.3f\n", $AUC[0];
            $sumAUC += $AUC[0];
        }
    }
    
}   

my @alltopics = keys (%topicLAM);

$sumLAM /= @alltopics;
$sumAUC /= @alltopics;

#printf SCOREFILE "LAM all %0.3f\n", $sumLAM;
printf "LAM all %0.3f\n", $sumLAM;
if($PoR ne "NA" ) {
    #printf SCOREFILE "AUC all %0.3f\n", $sumAUC;
    printf "AUC all %0.3f\n", $sumAUC;
}

#close SCOREFILE || die "Unable to close output file $runfile.LAM: $!\n";


exit 0;

sub logit {
    my $p = $_[0];
    return log($p/(1-$p));
}

sub lam {
    my $p = $_[0];
    my $q = $_[1];
    my $x = (&logit($p) + &logit($q)) / 2.0;
    my $e2x = exp($x);
    return $e2x/(1 + $e2x);
}

sub rank {
    my @arr = @_;    
    my @sorted = sort { @{$a}[0] <=> @{$b}[0] } @arr;
    #for my $a (@sorted) {
    #    print @{$a}[0]." ".@{$a}[1]."\n";
    #}
    #print "size of arr = ".@arr."\n";
    return @sorted;
}

sub UStatArea {
    my @ranking = @_;
    #print "size = ".@ranking."\n";
    #print $ranking[0][0]. " " . $ranking[0][1];
    
    my ($n1, $n2, $R1, $R2) = (0, 0, 0, 0);
    my $sameScoreR1 = 0;    
    my $sameScoreR2 = 0;    
            
    #print "scor, c, r, R1 R2 n1 n2 \n";  
    for(my $i = 0; $i < @ranking; $i++) {
        my $rank = $i + 1;
        my $score = $ranking[$i][0];
        my $class = $ranking[$i][1];
        my ($nextScore, $nextClass) = (-1, -1);
        if ($i < @ranking -1) {
            $nextScore = $ranking[$i + 1][0];
            $nextClass = $ranking[$i + 1][1];
        }
        if ($class == 0) {
            $n1 += 1;
        } else {
            $n2 += 1;
        }
        if ($nextScore >= 0 && $nextScore == $score) {                          
            #print "equal\n";
            if ($class == 0) {
                $sameScoreR1 += 1;
            } else {
                $sameScoreR2 += 1;
            } 
            
        } else { # ($nextScore != $score) { #set the ranks appropriately
            #print "not equal\n";
            if ($sameScoreR1 > 0 || $sameScoreR2 > 0) {
                #print "is same score\n";
                if ($class == 0) {
                    $sameScoreR1 += 1;
                } else {
                    $sameScoreR2 += 1;
                }
                my $numSame = $sameScoreR1 + $sameScoreR2;
                
                my $avgRank = ($rank + $rank-$numSame + 1)/2.0;
                $R1 += ($avgRank * $sameScoreR1);
                $R2 += ($avgRank * $sameScoreR2);
                
                $sameScoreR1 = $sameScoreR2 = 0;
            } else {
                #print "not same score\n";
                if ($class == 0) {
                    $R1 += $rank;
                } else {
                    $R2 += $rank;
                }
            }
        }      
        
        #print "$score, $class, $rank, $R1, $R2, $n1, $n2, \n";          
    }
    my $U2 = ($n1*$n2) + ($n2*($n2+1)/2.0) - $R2;
    my $U1 = ($n1*$n2) + ($n1*($n1+1)/2.0) - $R1;

    my @ustatareas = (-1, -1);

    if($n1 == 0) {
	@ustatareas = ( -1, 0);
    } elsif ($n2 == 0 ) {
	@ustatareas = ( 0, -1);
    } else {
	@ustatareas = ( $U1/($n1*$n2), $U2/($n1*$n2) );
    }
    return @ustatareas;
}
