#!/usr/bin/env python3
from argparse import ArgumentParser, RawTextHelpFormatter as RawText

from gensim.models import KeyedVectors

def run(model, tests, count=10):
    tests = read_tests(tests)
    model = KeyedVectors.load_word2vec_format(model, binary=False)
    print_model(model, tests, count)

def read_tests(tests):
    read_tests = []
    for test in tests:
        read_tests.extend(read_test(test))
    return read_tests

def read_test(test):
    with open(test) as file_:
        return [tuple(line.strip().split(" "))
                for line in file_]

def print_model(model, tests, count=10):
    for analogy in tests:
        found = get_analogy(model, analogy, count)
        if found:
            print("{}\t{}".format("\t".join(analogy), "\t".join(found)))

def get_analogy(model, analogy, count=10):
    (a, b, c, d) = analogy
    try:
        result = model.most_similar([b, c], [a], count)
        result = [word for (word, _) in result]
        return result
    except KeyError:
        return None

def test_analogy(model, analogy, count=10):
    result = get_analogy(model, analogy, count)
    if not result:
        return None
    target = analogy[3]
    ranking = result.index(target) if target in result else -1
    return (ranking, result[0])

def _main():
    args = _get_args()
    run(**vars(args))

def _get_args():
    parser = ArgumentParser(description=__doc__, formatter_class=RawText)
    parser.add_argument("model", help="Word2vec text models")
    parser.add_argument("tests", nargs="+", help="Analogy test files")
    parser.add_argument("-c", "--count", type=int, default=10,
                        help="Number of analogies to find")
    return parser.parse_args()

if __name__ == "__main__":
    _main()

