# coding=utf-8
r"""Estimates entropy from an SGNMT timecsv directory.
"""

import logging
import argparse
import sys
import os
from collections import defaultdict
import operator
import math

parser = argparse.ArgumentParser(description='Estimates entropy from a timecsv directory generated by SGNMT.')
parser.add_argument('-d','--dir', help='Path to timecsv directory', required=True)
args = parser.parse_args()

n_sentences = defaultdict(lambda: 0)
n_words = defaultdict(lambda: 0)
rank_freqs = defaultdict(lambda: defaultdict(lambda: 0))
acc_probs = defaultdict(lambda: 0.0)

for file_name in os.listdir(args.dir):
  with open("%s/%s" % (args.dir, file_name)) as f:
    lines = [l.split()[1:] for l in f]
    prob_fields = {}
    rank_fields = {}
    for col, header in enumerate(lines[0]):
      if header[-7:] != "_weight":
        n_sentences[header] += 1
        if header[:5] == "rank_":
          rank_fields[header] = col
        else:
          prob_fields[header] = col
    prev_acc_rank = {k: 0 for k in rank_fields.iterkeys()}
    sentence_probs = {k: 0.0 for k in prob_fields.iterkeys()}
    for line in lines[1:]:
      for field_name, col in prob_fields.iteritems():
        if line[col] != "-":
          sentence_probs[field_name] = float(line[col])
          n_words[field_name] += 1
      for field_name, col in rank_fields.iteritems():
        if line[col] != "-":
          acc_rank = -int(float(line[col]))
          rank = max(1, acc_rank - prev_acc_rank[field_name])
          prev_acc_rank[field_name] = acc_rank
          rank_freqs[field_name][rank] += 1
          n_words[field_name] += 1
    for field_name, prob in sentence_probs.iteritems():
      acc_probs[field_name] += prob

print("RANK FREQUENCIES")
for field_name, freqs in rank_freqs.iteritems():
  print("%s:" % field_name)
  n = float(n_words[field_name])
  for rank, abs_freq in sorted(freqs.iteritems()):
    print("%d: %d (%f)" % (rank, abs_freq, abs_freq / n))
  print("")

print("")
print("RANK STATISTICS")
for field_name, freqs in rank_freqs.iteritems():
  print("%s:\nn_sentences: %d\nn_words: %d" % (field_name, n_sentences[field_name], n_words[field_name]))
  n = float(n_words[field_name])
  acc = 0.0
  log_acc = 0.0
  for rank, abs_freq in sorted(freqs.iteritems()):
    acc += rank * abs_freq
    log_acc += math.log(rank) * abs_freq
  print("average_rank: %f\naverage_log_rank: %f" % (acc / n, log_acc / n))
  print("")

print("")
print("PROB STATISTICS")
for field_name, acc_prob in acc_probs.iteritems():
  acc_prob_base_2 = acc_prob / math.log(2.0)
  print("%s:\nn_sentences: %d\nn_words: %d\nentropy_per_sentence_base_e: %f\nentropy_per_word_base_e: %f\nentropy_per_word_base_2: %f\nperplexity_per_word_base_2: %f" % 
    (field_name, 
     n_sentences[field_name],
     n_words[field_name], 
     -acc_prob / n_sentences[field_name],
     -acc_prob / n_words[field_name],
     -acc_prob_base_2 / n_words[field_name],
     math.pow(2.0, -acc_prob_base_2 / n_words[field_name])
    ))
  print("")

