#!/usr/bin/env perl
use warnings; #sed replacement for -w perl parameter
# Copyright 2014 David Snyder
# Apache 2.0.
#
# This script produces a vector used by logistic-regression-copy to 
# rescale the logistic regression model which reduces bias due to unbalanced
# classes. This script relies only on the distribution of the test data;
# alternatively, a uniform prior can be used (see run_logistic_regression.sh). 

# The scale parameter controls how sensitive the priors are to the
# distribution of the test data. Typically this ranges from 0.5
# to 1.0. Smaller values are less reliant on the test data distribution.
my ($train_file, $test_file, $lang_file, $scale, $priors_file) = @ARGV;
open(UTT2LANG_TRAIN, "<$train_file") or die "no utt2lang training file";

%train_count = ();
$train_tot = 0;
while(<UTT2LANG_TRAIN>) {
  $line = $_;
  chomp($line);
  @words = split(" ", $line);
  $lang = $words[1];
  if (not exists($train_count{$lang})) {
    $train_count{$lang} = 1;
  } else {
    $train_count{$lang} += 1;
  }
  $train_tot += 1;
}

open(UTT2LANG_TEST, "<$test_file");

%test_count = ();
$test_tot = 0;
while(<UTT2LANG_TEST>) {
  $line = $_;
  chomp($line);
  @words = split(" ", $line);
  $lang = $words[1];
  if (not exists($test_count{$lang})) {
    $test_count{$lang} = 1;
  } else {
    $test_count{$lang} += 1;
  }
  $test_tot += 1;
}

foreach my $key (keys %train_count) {
  if (not exists($test_count{$key})) {
    $test_count{$key} = 0;
  }
}

# load languages file
open(LANGUAGES, "<$lang_file");
@idx_to_lang = ();

$largest_idx = 0;
while(<LANGUAGES>) {
  $line = $_;
  chomp($line);
  @words = split(" ", $line);
  $lang = $words[0];
  $idx = $words[1];
  $idx_to_lang[$idx + 0] = $lang;
  if ($idx > $largest_idx) {
    $largest_idx = $idx;
  }
}

$priors = " [ ";
foreach $lang (@idx_to_lang) {
 $ratio = ((1.0*$test_count{$lang}) / $train_count{$lang})**($scale);
 $priors .= "$ratio ";
}

$priors .= " ]";
open(PRIORS, ">$priors_file");
print PRIORS $priors;