# -*- coding: utf-8 -*-
from __future__ import print_function, division

import os
import numpy as np
import numpy.random as npr

import spacy
import util

DETERMINANTS = ['', 'A ', 'The ']
P_DETERMINANTS = [0.4, 0.3, 0.3]


class TitleGenerator(object):
    """Generates the title of a poem"""

    def __init__(self):
        # Init spacy
        self.nlp = spacy.load('en')

    def initialize(self, nouns_file, adjs_file, determinants=DETERMINANTS, tau=1.0):
        # Save temperature
        self.tau = tau
        # Load nouns
        self.nouns = util.load_wordlist(nouns_file)
        # Load adjectives
        self.adjs = util.load_wordlist(adjs_file)
        # Noun vectors
        self.noun_vectors = np.zeros((len(self.nouns), self.nlp.vocab.vectors_length))
        for i, noun in enumerate(self.nouns):
            tok = self.nlp(noun.decode('utf-8'))
            self.noun_vectors[i] = tok.vector / tok.vector_norm if tok.vector_norm > 0 else 0
        # Adjectives vectors
        self.adj_vectors = np.zeros((len(self.adjs), self.nlp.vocab.vectors_length))
        for i, adj in enumerate(self.adjs):
            tok = self.nlp(adj.decode('utf-8'))
            self.adj_vectors[i] = tok.vector / tok.vector_norm if tok.vector_norm > 0 else 0

    def poem_vector(self, poem):
        """Creates a poem vector"""
        doc = self.nlp(poem.decode('utf-8'))
        vectors = [w.vector for w in doc if w.tag_.startswith('NN')]
        vector = sum(vectors) / len(vectors)
        return vector / np.linalg.norm(vector)

    def sample_noun(self, vector):
        p = util.softmax(self.noun_vectors.dot(vector) / self.tau)
        return npr.choice(self.nouns, p=p)

    def sample_adjective(self, vector):
        p = util.softmax(self.adj_vectors.dot(vector) / self.tau)
        return npr.choice(self.adjs, p=p)

    def sample_det(self, vector):
        return npr.choice(DETERMINANTS, p=P_DETERMINANTS)

    def __call__(self, poem):
        """Generates a title for the input"""
        poem_vector = self.poem_vector(poem)
        noun = self.sample_noun(poem_vector).capitalize()
        adj = self.sample_adjective(poem_vector).capitalize()
        det = self.sample_det(poem_vector).capitalize()
        return det + adj + " " + noun

    def save(self, filename):
        np.savez(filename,
                 nouns=self.nouns,
                 noun_vectors=self.noun_vectors,
                 adjs=self.adjs,
                 adj_vectors=self.adj_vectors,
                 tau=self.tau)

    def load(self, filename):
        self.__dict__.update(np.load(filename).items())


def get_title_generator(options):
    tg = TitleGenerator()
    if options.filename is not None and os.path.isfile(options.filename):
        tg.load(options.filename)
    else:
        tg.initialize(options.nouns_file, options.adjs_file, DETERMINANTS, options.tau)
        tg.save(options.filename)
    return tg

@@ -1846,30 +1846,27 @@
 e_det(self, 
-vector
+adj
 ):%0A        r
@@ -1856,38 +1856,37 @@
 , adj):%0A        
-return
+det =
  npr.choice(DETE
@@ -1912,16 +1912,151 @@
 MINANTS)
+%0A        # Edge case of a/an (could be better)%0A        if det == %22A %22 and adj%5B0%5D in 'aeiou':%0A            det = %22An %22%0A        return det
 %0A%0A    de
@@ -2320,27 +2320,19 @@
 ple_det(
-poem_vector
+adj
 ).capita

Address a/an in titles (not perfect but works)