from optparse import OptionParser

import os
import sys
import gensim
import random
from collections import OrderedDict

import numpy as np
import theano
from theano import tensor as T
from theano.tensor.shared_randomstreams import RandomStreams

from ..util import file_handling as fh
from ..dataset_scripts.sst.read_data import load_ptb_bitrees

# This is the same as constituency_tree_lstm, but with only one mask for L/R

# Otherwise the deepcopy fails
sys.setrecursionlimit(5000)

THEANO_FLAGS='floatX=float32'
#THEANO_FLAGS='floatX=float32, optimizer=None'

#DYLD_FALLBACK_LIBRARY_PATH='/Applications/anaconda/lib/'


class DyTreeLSTM(object):

    def __init__(self, dv, dh, dx, de, ne, nc, alpha=1.0, init_scale=0.2, initial_embeddings=None, edge_embeddings=None,
                 params_init=None,
                 update='adagrad', seed=None, drop_p=0.5, momentum=0.9, train_embeddings=False):

        self.dv = dv  # vocabulary size
        self.dh = dh  # hidden vector size
        self.dx = dx  # word embedding size
        self.de = de  # edge embedding size
        dxe = dx + de
        self.dxe = dxe  # dimension of word embeddings concatenated with edge embeddings
        self.ne = ne  # number of types of edges (edge vocab size)
        self.nc = nc  # number of classes
        self.alpha = alpha  # regularization strength
        self.drop_p = drop_p  # probability of dropping an input with dropout

        # adagrad parameters
        self.epsilon = 0.00001

        if initial_embeddings is None:
            self.emb = theano.shared(name='embeddings',
                                     value=init_scale * np.random.uniform(-1.0, 1.0,
                                                                          (dv, dx)).astype(theano.config.floatX))
        else:
            self.emb = theano.shared(name='embeddings', value=initial_embeddings.astype(theano.config.floatX))

        if edge_embeddings is None:
            self.edge_emb = theano.shared(name='edge_embeddings',
                                          value=init_scale * np.random.uniform(-1.0, 1.0,
                                                                               (ne, de)).astype(theano.config.floatX))
        else:
            self.edge_emb = theano.shared(name='edge_embeddings', value=edge_embeddings.astype(theano.config.floatX))

        self.W_x_i = theano.shared(name='W_x_i', value=init_scale * np.random.uniform(-1.0, 1.0, (dxe, dh))
                                   .astype(theano.config.floatX))
        self.W_h_i = theano.shared(name='W_h_i', value=init_scale * np.random.uniform(-1.0, 1.0, (dh, dh))
                                   .astype(theano.config.floatX))
        self.b_h_i = theano.shared(name='b_h_i', value=np.array(np.zeros(dh),
                                                                dtype=theano.config.floatX))

        self.W_x_f = theano.shared(name='W_x_f', value=init_scale * np.random.uniform(-1.0, 1.0, (dxe, dh))
                                   .astype(theano.config.floatX))
        self.W_h_f = theano.shared(name='W_h_f', value=init_scale * np.random.uniform(-1.0, 1.0, (dh, dh))
                                   .astype(theano.config.floatX))
        self.b_h_f = theano.shared(name='b_h_f', value=np.array(np.random.uniform(0.0, 1.0, dh),
                                                                dtype=theano.config.floatX))

        self.W_x_o = theano.shared(name='W_x_o', value=init_scale * np.random.uniform(-1.0, 1.0, (dxe, dh))
                                   .astype(theano.config.floatX))
        self.W_h_o = theano.shared(name='W_h_o', value=init_scale * np.random.uniform(-1.0, 1.0, (dh, dh))
                                   .astype(theano.config.floatX))
        self.b_h_o = theano.shared(name='b_h_o', value=np.array(np.zeros(dh),
                                                                dtype=theano.config.floatX))

        self.W_x_u = theano.shared(name='W_x_u', value=init_scale * np.random.uniform(-1.0, 1.0, (dxe, dh))
                                   .astype(theano.config.floatX))
        self.W_h_u = theano.shared(name='W_h_u', value=init_scale * np.random.uniform(-1.0, 1.0, (dh, dh))
                                   .astype(theano.config.floatX))
        self.b_h_u = theano.shared(name='b_h_u', value=np.array(np.zeros(dh),
                                                                dtype=theano.config.floatX))

        self.W_x_z = theano.shared(name='W_x_z', value=init_scale * np.random.uniform(-1.0, 1.0, (dx, nc))
                                   .astype(theano.config.floatX))
        self.W_h_z = theano.shared(name='W_h_z', value=init_scale * np.random.uniform(-1.0, 1.0, (dh, nc))
                                   .astype(theano.config.floatX))
        self.b_z = theano.shared(name='b_z', value=np.array(np.zeros(nc),
                                                            dtype=theano.config.floatX))

        self.W_x_a = theano.shared(name='W_x_a', value=init_scale * np.random.uniform(-1.0, 1.0, (dxe, dh))
                                   .astype(theano.config.floatX))
        self.W_h_a = theano.shared(name='W_h_a', value=init_scale * np.random.uniform(-1.0, 1.0, (dh, dh))
                                   .astype(theano.config.floatX))
        self.v_a = theano.shared(name='v_a', value=init_scale * np.random.uniform(-1.0, 1.0, (dh,))
                                 .astype(theano.config.floatX))


        self.params = [self.W_x_i, self.W_h_i, self.b_h_i]
        self.params += [self.W_x_f, self.W_h_f, self.b_h_f]
        self.params += [self.W_x_o, self.W_h_o, self.b_h_o]
        self.params += [self.W_x_u, self.W_h_u, self.b_h_u]

        self.params += [self.W_x_z, self.W_h_z, self.b_z]
        self.params += [self.W_h_a, self.v_a]

        self.param_shapes = [(dxe, dh), (dh, dh), (dh,),
                             (dxe, dh), (dh, dh), (dh,),
                             (dxe, dh), (dh, dh), (dh,),
                             (dxe, dh), (dh, dh), (dh,),

                             (dx, nc), (dh, nc), (nc,),
                             (dh, dh), (dh,),
                             ]


        if update == 'adagrad':

            self.grad_histories = [
                theano.shared(
                    value=np.zeros(param_shape, dtype=theano.config.floatX),
                    borrow=True,
                    name="grad_hist:" + param.name
                )
                for param_shape, param in zip(self.param_shapes, self.params)
                ]
        elif update == 'sgdm':

            self.velocity = [
                theano.shared(
                    value=np.zeros(param_shape, dtype=theano.config.floatX),
                    borrow=True,
                    name="momentum:" + param.name
                )
                for param_shape, param in zip(self.param_shapes, self.params)
                ]
            self.momentum = momentum

        self.theano_rng = RandomStreams(seed)

        input_idx = T.ivector('input_idx')
        edge_mask = T.imatrix('edge_mask')
        output_mask = T.imatrix('output_mask')
        y = T.ivector('y')

        n_nodes, n_edges = T.shape(output_mask)

        idxs = T.ivector('idxs')
        temp = self.emb[idxs]
        x = temp.reshape([n_nodes, dx])

        edges_idx = T.ivector('edges_idx')
        temp = self.edge_emb[edges_idx]
        edges = temp.reshape([n_edges, de])

        lr = T.scalar('lr', dtype=theano.config.floatX)
        emb_lr = T.scalar('emb_lr', dtype=theano.config.floatX)
        is_train = T.iscalar('is_train')
        drop_x = T.iscalar('drop_x')

        # This is a bit annoying; the 0th dimension of x needs to be sequence, so we can iterate over it
        # but the 0th dimension of the hidden nodes needs to be hidden-node dimension, so that we can broadcast
        # the mask out to it
        def pass_edges(input_idx_t, edge_t, edge_mask_t, counter_t, h_tm1, c_tm1, x):
            h_t = h_tm1
            c_t = c_tm1
            # select the input vector to use for this edge (source)
            x_t_i = x[input_idx_t, :]
            # zero out the input unless this is a leaf node
            x_t_0 = T.switch(T.eq(T.sum(edge_mask_t), 0), x_t_i, x_t_i*0)
            # concatenate with the input edge vector
            x_t_edge = T.concatenate([x_t_0, edge_t])

            # compute attention weights, using a manual softmax
            attention_scores = T.dot(self.v_a, T.tanh(T.dot(self.W_h_a, h_tm1))) # (1, n_edges)
            # find the max of the unmasked values
            max_score = T.max(attention_scores + edge_mask_t * 10000.0) - 10000.0
            # exponentiate the differences, masking first to avoid inf, and then to keep only relevant scores
            exp_scores = T.exp((attention_scores - max_score) * edge_mask_t) * edge_mask_t
            # take the sum, and add one if the mask is all zeros to avoid an inf
            exp_scores_sum = T.sum(exp_scores) + T.switch(T.eq(T.sum(edge_mask_t), 0), 1.0, 0.0)
            # normalize to compute the weights
            weighted_mask = exp_scores / exp_scores_sum

            i_t = T.nnet.sigmoid(T.dot(x_t_edge, self.W_x_i) + T.sum(T.dot(self.W_h_i.T, (weighted_mask * h_tm1)).T, axis=0) + self.b_h_i)
            f_t = T.nnet.sigmoid(T.dot(x_t_edge, self.W_x_f) + T.sum(T.dot(self.W_h_f.T, (weighted_mask * h_tm1)).T, axis=0) + self.b_h_f)
            o_t = T.nnet.sigmoid(T.dot(x_t_edge, self.W_x_o) + T.sum(T.dot(self.W_h_o.T, (weighted_mask * h_tm1)).T, axis=0) + self.b_h_o)
            u_t = T.tanh(T.dot(x_t_edge, self.W_x_u) + T.sum(T.dot(self.W_h_u.T, (weighted_mask * h_tm1)).T, axis=0) + self.b_h_u)