import re
import numpy as np
import tensorflow as tf

from model import *

class headGAN(object):
    def __init__(self, d_net, g_net, wordvec, article, title, wd_list):
        print 'GAN headline'
        self.wordvec = wordvec
        self.article = article
        self.title = title
        self.d_net = d_net
        self.g_net = g_net
        self.wd_list = wd_list

        self.sess = tf.InteractiveSession()

        self.build_model()
        self.train_model()

    def build_model(self):
        art_len = self.article.shape[1]
        ttl_len = self.title.shape[1]
        wd_dim = self.wordvec.shape[1]

        self.in_art = tf.placeholder(tf.int32,[None, art_len])
        self.in_ttl = tf.placeholder(tf.int32,[None, ttl_len])

        r_art = tf.nn.embedding_lookup(self.wordvec, self.in_art)
        r_art = tf.expand_dims(r_art, -1)
        r_art = tf.transpose(r_art, [0,2,1,3], name='r_art')
        self.r_art = r_art

        self.r_ttl = tf.nn.embedding_lookup(self.wordvec, self.in_ttl)
        self.r_ttl = tf.expand_dims(self.r_ttl, -1)
        self.r_ttl = tf.transpose(self.r_ttl, [0,2,1,3], name='r_ttl')

        self.f_ttl = self.g_net(input_data=r_art)
        r_logits = self.d_net(input_data=self.r_ttl, reuse=False)
        f_logits = self.d_net(input_data=self.f_ttl, reuse=True)

        r_ent = tf.nn.sigmoid(r_logits)
        f_ent = tf.nn.sigmoid(f_logits)
        self.d_loss = tf.reduce_mean(r_ent) - tf.reduce_mean(f_ent)

        self.g_loss = tf.reduce_mean(f_ent, name='g_loss')
        
        self.g_var = self.g_net.vars
        self.d_var = self.d_net.vars

        self.opt_method = 'rmsprop'
        if self.opt_method == 'rmsprop':
            self.d_opt = tf.train.RMSPropOptimizer(0.01,decay=0.9).minimize(self.d_loss,var_list=self.d_var)
            self.g_opt = tf.train.RMSPropOptimizer(0.01,decay=0.9).minimize(self.g_loss,var_list=self.g_var)
        else:
            self.d_opt = tf.train.AdamOptimizer().minimize(self.d_loss,var_list=self.d_var)
            self.g_opt = tf.train.AdamOptimizer().minimize(self.g_loss,var_list=self.g_var)

    def train_model(self):
        self.sess.run(tf.global_variables_initializer())

        k = 0
        while k < 20000:
            feed_ = {self.in_art: [self.article[k]], self.in_ttl: [self.title[k]]}
            self.sess.run(self.g_opt, feed_dict=feed_)
            self.sess.run(self.d_opt, feed_dict=feed_)
            if k%1000 == 0:
                print k, self.sess.run([self.d_loss, self.g_loss], feed_dict=feed_)
                tt = self.sess.run(self.f_ttl, feed_dict=feed_)
                xx = self.sess.run(self.r_art[0,:,:,0], feed_dict=feed_)
                dd = self.sess.run(tf.matmul(tt[0,0], xx)))
                idx = np.argmax(dd, axis=1).tolist()
		for kk in idx:
		   print self.wd_list[self.article[k,kk]],
		print
		for kk in self.title[k]:
		   print self.wd_list[kk],
		print
                for kk in self.article[k]:
                   print self.wd_list[kk],
                print
                print
            k += 1
            if k == 19999:
                k = 0
        
