#!/bin/bash

# Copyright  2017  Johns Hopkins University (author: Daniel Povey)
# Apache 2.0.


# This script works out how many pieces we want to split the data into for a
# particular training run, based on how many words are in the data directory
# (excluding dev), and the target words-per-split.


if [ $# != 3 ]; then
  (
    echo "Usage: rnnlm/get_num_splits.sh <target-words-per-split> <data-dir> <weights-file>"
    echo "e.g.: rnnlm/get_num_splits.sh 200000 data/text exp/rnnlm/data_weights.txt"
    echo "This works out how many pieces to split a data directory into, and"
    echo "(if just one piece) how many times that piece should be repeated to"
    echo "get the target words-per-split.  A number is printed to the standard"
    echo "output.  If no repeats are necessary it will be the number of splits,"
    echo "a positive number.  If repeats are necessary, then a negative number,"
    echo "interpretable as the negative of the number of times we should repeat"
    echo "the data, is echoed, and the number of splits should be taken to be 1."
    echo "To compute the number of words of training data"
    echo "this script uses <data-dir>/*.counts; they are scaled by the data-multiplicities"
    echo "given as the second field of <weights-file> for each data source."
  ) 1>&2
  exit 1
fi


words_per_split=$1
text=$2
weights_file=$3

! [ $words_per_split -eq $words_per_split ] && \
  echo "$0: first arg must be an integer" 1>&2 && exit 1;

[ ! -d $text ] && \
  echo "$0: no such directory $text" 1>&2 && exit 1;

[ ! -f $weight ] && \
  echo "$0: expected weights file in $weight" 1>&2 && exit 1;

rnnlm/ensure_counts_present.sh $text 1>&2


set -e -o pipefail -u

export LC_ALL=C





multiplicities=$(mktemp tmp.XXXX)
trap "rm $multiplicities" EXIT

if ! awk '{if(NF!=3){ exit(1); } print $1, $2; } END{if(NR==0) exit(1);}' <$weights_file > $multiplicities; then
  echo "$0: weights file $weights_file has the wrong format."
fi

tot_orig=0
tot_with_multiplicities=0


for f in $text/*.counts; do
  if [ "$f" != "$text/dev.counts" ]; then
    this_tot=$(cat $f | awk '{tot += $2} END{printf("%d", tot)}')
    if ! [ $this_tot -gt 0 ]; then
      echo "$0: there were no counts in counts file $f" 1>&2
      exit 1
    fi
    # weight by the data multiplicity which is the second field of the weights file.
    multiplicity=$(basename $f | sed 's:.counts$::' | utils/apply_map.pl $multiplicities)
    if ! [ $multiplicity -eq $multiplicity ]; then
      echo "$0: error getting multiplicity for data-source $f, check weights file $weights_file"
      exit 1
    fi
    tot_orig=$[tot_orig+this_tot]
    tot_with_multiplicities=$[tot_with_multiplicities+(this_tot*multiplicity)]
  fi

done

if ! [ $tot_orig -gt 0 ]; then
  echo "$0: there was a problem getting counts from directory $text (no counts present?)" 1>&2
  exit 1
fi
if ! [ $tot_with_multiplicities -gt 0 ]; then
  echo "$0: there was a problem getting counts from directory $text (check data-weights file $weights_file)" 1>&2
  exit 1
fi


# adding words_per_split-1 below causes us to round up the number of splits.
num_splits=$[(tot_with_multiplicities+words_per_split-1)/words_per_split]

actual_words_per_split=$[tot_with_multiplicities/num_splits]

if ! [ $num_splits -gt 0 ]; then
  echo "$0: there was a problem getting the number of splits" 1>&2
  exit 1
fi


num_repeats=$[words_per_split/actual_words_per_split]
if ! [ $num_repeats -ge 1 ]; then
  echo "$0: error computing the number of repeats, got $num_repeats." 1>&2
  exit 1
fi

if [ $num_repeats -gt 1 -a $num_splits -gt 1 ]; then
  echo "$0: script error: both num-repeats and num-splits are over 1." 1>&2
  exit 1
fi

echo -n "get_num_splits.sh: based on tot-words=$tot_orig (with multiplicities: $tot_with_multiplicities)" 1>&2
echo " and target-words-per-split=$words_per_split, got $num_splits splits, actual words-per-split is $actual_words_per_split" 1>&2
if [ $num_repeats -gt 1 ]; then
  echo " ... and num-repeats is $num_repeats" 1>&2
fi


if [ $num_repeats -eq 1 ]; then
  echo $num_splits
else
  echo -$num_repeats
fi