#!/usr/bin/python3

import argparse
import collections
import sys

import nltk
import nltk.corpus
import nltk.stem
import nltk.tokenize

# call this once before executing this script
# pip3 install --user nltk
# nltk.download('averaged_perceptron_tagger')
# nltk.download('punkt')
# nltk.download('wordnet')

ap = argparse.ArgumentParser()
ap.add_argument('inputs', metavar='INPUT', nargs='*', help='input files')
ap.add_argument('--output', '-o', metavar='OUTPUT', help='output file')
args = vars(ap.parse_args())

# if no inputs are given, read from standard input
if not args['inputs']:
    args['inputs'] = '-'

# prepare stuff
lemmatizer = nltk.stem.WordNetLemmatizer()
tokenizer = nltk.tokenize.RegexpTokenizer("[\w']+")

# return a list with normalized words
def lemmatize(words):
    # NLTK’s tagger and lemmatizer use different part-of-speech tags, this converts
    tags = {
        'J' : nltk.corpus.wordnet.ADJ,
        'R': nltk.corpus.wordnet.ADV,
        'V': nltk.corpus.wordnet.VERB,
    }
    def pos2tag(pos):
        return tags.get(pos[0], nltk.corpus.wordnet.NOUN)

    # add part-of-speech (noun, verb, adjective…) tags to words
    tagged_words = nltk.pos_tag(list(words))
    return (lemmatizer.lemmatize(word, pos2tag(pos)) for word, pos in tagged_words)

# count words in all the files
counts = collections.Counter()
for filename in args['inputs']:
    # special filename - reads from standard input
    with (sys.stdin if filename == '-' else open(filename)) as stream: 
        # read text from file and clean it up: convert to lowercase and remove HTML tags
        text = stream.read().lower().replace('<br />', ' ')

    # split text into words and normalize them
    tokens = tokenizer.tokenize(text)
    words = lemmatize(tokens)

    # add words to counter
    counts.update(words)

# write word counts to output file (or standard output if not given)
output = open(args['output'], 'w') if args['output'] else sys.stdout
for word, count in counts.most_common():
    print(f'{word} {count}', file=output)