import click
from . import wordplay

SIZES = {
    'small':  'GloVe Wikipedia 100-dim  (~128 MB)',
    'medium': 'GloVe Wikipedia 300-dim  (~376 MB)',
    'large':  'Word2Vec Google News 300-dim  (~1.6 GB)',
}


@click.group()
@click.option(
    '--size',
    type=click.Choice(list(SIZES)),
    default='small',
    show_default=True,
    help='Embedding model to use:\n\n' + '\n'.join(f'  {k}: {v}' for k, v in SIZES.items()),
)
@click.pass_context
def cli(ctx, size):
    """Wordplay - word embedding experiments."""
    ctx.ensure_object(dict)
    ctx.obj['model'] = wordplay.load_embeddings(size)


@cli.command()
@click.argument('word')
@click.option('-n', default=5, show_default=True, help='Number of synonyms to return.')
@click.pass_obj
def synonyms(obj, word, n):
    """Find the n most similar words to WORD."""
    click.echo(wordplay.synonyms(obj['model'], word, n))


@cli.command()
@click.argument('word1')
@click.argument('word2')
@click.pass_obj
def average(obj, word1, word2):
    """Return a word which is the average of the two given words."""
    click.echo(wordplay.average(obj['model'], word1, word2))


@cli.command()
@click.argument('words', nargs=-1, required=True)
@click.pass_obj
def outlier(obj, words):
    """Given a list of words, identifies the word which is the outlier."""
    click.echo(wordplay.outlier(obj['model'], words))


@cli.command()
@click.argument('words', nargs=-1, required=True)
@click.pass_obj
def sort(obj, words):
    """Given a list of words, sorts them in order."""
    click.echo(wordplay.sort(obj['model'], words))


if __name__ == '__main__':
    cli()