"""
Examine models generated by kenlm
"""
import kenlm
from statistics import mean
import argparse
import csv

def parse() -> argparse.Namespace:
    parser = argparse.ArgumentParser(\
        description="Gets the perplexity of the model named 'bigram'")

    parser.add_argument('--csv', default=None, \
        help='Whether to store the output to the given csv instead')

    return parser.parse_args()

def main():
    args = parse()

    model = kenlm.Model('bigram')

    perplexities = []
    sample_size = 1000

    with open('/u/demorali/corpora/1g-word-lm-benchmark-r13output/heldout-monolingual.tokenized.shuffled/news.en-00000-of-00100', 'r') as f:
        for line in f:
            perplexities.append(model.perplexity(line))
            if len(perplexities) == sample_size:
                break

    if args.csv is None:
        print(
            """Test sentence count: {size}
            mean is {mean}
            max is {max}
            min is {min}
            """.format(mean=mean(perplexities), max=max(perplexities), min=min(perplexities),
                    size=sample_size))
    else:
        with open(args.csv, 'a') as f:
            writer = csv.writer(f)
            writer.writerow([mean(perplexities), max(perplexities), min(perplexities)])


if __name__ == "__main__":
    main()
