#include "dynet/treelstm.h"

#include "dynet/param-init.h"

#include <string>
#include <vector>
#include <iostream>

using namespace std;
using namespace dynet;

enum { X2I, BI, X2F, BF, X2O, BO, X2C, BC };
enum { H2I, H2F, H2O, H2C, C2I, C2F, C2O };

Expression TreeLSTMBuilder::add_input_impl(int prev, const Expression& x) { throw std::runtime_error("add_input_impl() not a valid function for TreeLSTMBuilder"); }
Expression TreeLSTMBuilder::back() const { throw std::runtime_error("back() not a valid function for TreeLSTMBuilder"); }
std::vector<Expression> TreeLSTMBuilder::final_h() const { throw std::runtime_error("final_h() not a valid function for TreeLSTMBuilder"); }
std::vector<Expression> TreeLSTMBuilder::final_s() const { throw std::runtime_error("final_s() not a valid function for TreeLSTMBuilder"); }
unsigned TreeLSTMBuilder::num_h0_components() const { throw std::runtime_error("num_h0_components() not a valid function for TreeLSTMBuilder"); }
void TreeLSTMBuilder::copy(const RNNBuilder&) { throw std::runtime_error("copy() not a valid function for TreeLSTMBuilder"); }

// See "Improved Semantic Representations From Tree-Structured Long Short-Term Memory Networks"
// by Tai, Nary, and Manning (2015), section 3.2, for details on this model.
// http://arxiv.org/pdf/1503.00075v3.pdf
NaryTreeLSTMBuilder::NaryTreeLSTMBuilder(unsigned N,
                         unsigned layers,
                         unsigned input_dim,
                         unsigned hidden_dim,
                         ParameterCollection& model) : layers(layers), N(N), cg(nullptr) {
  unsigned layer_input_dim = input_dim;
  local_model = model.add_subcollection("nary-tree-lstm-builder");
  for (unsigned i = 0; i < layers; ++i) {
    // i
    Parameter p_x2i = local_model.add_parameters({hidden_dim, layer_input_dim});
    LookupParameter p_h2i = local_model.add_lookup_parameters(N, {hidden_dim, hidden_dim});
    LookupParameter p_c2i = local_model.add_lookup_parameters(N, {hidden_dim, hidden_dim});
    Parameter p_bi = local_model.add_parameters({hidden_dim}, ParameterInitConst(0.f));

    // f
    Parameter p_x2f = local_model.add_parameters({hidden_dim, layer_input_dim});
    LookupParameter p_h2f = local_model.add_lookup_parameters(N*N, {hidden_dim, hidden_dim});
    LookupParameter p_c2f = local_model.add_lookup_parameters(N*N, {hidden_dim, hidden_dim});
    Parameter p_bf = local_model.add_parameters({hidden_dim}, ParameterInitConst(0.f));

    // o
    Parameter p_x2o = local_model.add_parameters({hidden_dim, layer_input_dim});
    LookupParameter p_h2o = local_model.add_lookup_parameters(N, {hidden_dim, hidden_dim});
    LookupParameter p_c2o = local_model.add_lookup_parameters(N, {hidden_dim, hidden_dim});
    Parameter p_bo = local_model.add_parameters({hidden_dim}, ParameterInitConst(0.f));

    // c (a.k.a. u)
    Parameter p_x2c = local_model.add_parameters({hidden_dim, layer_input_dim});
    LookupParameter p_h2c = local_model.add_lookup_parameters(N, {hidden_dim, hidden_dim});
    Parameter p_bc = local_model.add_parameters({hidden_dim}, ParameterInitConst(0.f));

    layer_input_dim = hidden_dim;  // output (hidden) from 1st layer is input to next

    vector<Parameter> ps = {p_x2i, p_bi, p_x2f, p_bf, p_x2o, p_bo, p_x2c, p_bc};
    vector<LookupParameter> lps = {p_h2i, p_h2f, p_h2o, p_h2c, p_c2i, p_c2f, p_c2o};
    params.push_back(ps);
    lparams.push_back(lps);
  }  // layers
}

void NaryTreeLSTMBuilder::new_graph_impl(ComputationGraph& cg, bool update) {
  this->cg = &cg;
  param_vars.clear();
  lparam_vars.clear();
  param_vars.reserve(layers);
  lparam_vars.reserve(layers);

  for (unsigned i = 0; i < layers; ++i){
    auto& p = params[i];
    auto& lp = lparams[i];

    //i
    Expression i_x2i = update ? parameter(cg, p[X2I]) : const_parameter(cg, p[X2I]);
    Expression i_bi = update ? parameter(cg, p[BI]) : const_parameter(cg, p[BI]);
    //f
    Expression i_x2f = update ? parameter(cg, p[X2F]) : const_parameter(cg, p[X2F]);
    Expression i_bf = update ? parameter(cg, p[BF]) : const_parameter(cg, p[BF]);
    //o
    Expression i_x2o = update ? parameter(cg, p[X2O]) : const_parameter(cg, p[X2O]);
    Expression i_bo = update ? parameter(cg, p[BO]) : const_parameter(cg, p[BO]);
    //c
    Expression i_x2c = update ? parameter(cg, p[X2C]) : const_parameter(cg, p[X2C]);
    Expression i_bc = update ? parameter(cg, p[BC]) : const_parameter(cg, p[BC]);

    vector<Expression> vars = {i_x2i, i_bi, i_x2f, i_bf, i_x2o, i_bo, i_x2c, i_bc};
    param_vars.push_back(vars);

    DYNET_ASSERT(lp.size() == C2O + 1, "Dimension mismatch in TreeLSTM");
    vector<vector<Expression>> lvars(lp.size());
    for (unsigned p_type = H2I; p_type <= C2O; p_type++) {
    LookupParameter p = lp[p_type];
      vector<Expression> vals(p.get_storage().values.size());
      for (unsigned k = 0; k < p.get_storage().values.size(); ++k) {
        //vals[k] = lookup(cg, p, k);
        vals[k].i = 0;
      }
      lvars[p_type] = vals;
    }
    lparam_vars.push_back(lvars);
  }
}

Expression NaryTreeLSTMBuilder::Lookup(unsigned layer, unsigned p_type, unsigned value) {
  if (lparam_vars[layer][p_type][value].i == 0) {
    LookupParameter p = lparams[layer][p_type];
    lparam_vars[layer][p_type][value] = lookup(*cg, p, value);
  }
  return lparam_vars[layer][p_type][value];
}

// layout: 0..layers = c
//         layers+1..2*layers = h
void NaryTreeLSTMBuilder::start_new_sequence_impl(const vector<Expression>& hinit) {
  h.clear();
  c.clear();
  if (hinit.size() > 0) {
    DYNET_ARG_CHECK(layers*2 == hinit.size(),
                            "Incorrectly sized initialization in TreeLSTM (" << hinit.size() << "). "
                            "Must be twice the number of layers (which is " << layers<< ")");
    h0.resize(layers);
    c0.resize(layers);
    for (unsigned i = 0; i < layers; ++i) {
      c0[i] = hinit[i];
      h0[i] = hinit[i + layers];
    }
    has_initial_state = true;
  } else {
    has_initial_state = false;
  }
}

void NaryTreeLSTMBuilder::set_num_elements(int length) {
  h.resize(length);
  c.resize(length);
}

Expression NaryTreeLSTMBuilder::add_input(int id, vector<int> children, const Expression& x) {
  DYNET_ASSERT(id >= 0 && h.size() == (unsigned)id, "Failed dimension check in TreeLSTMBuilder");
  DYNET_ASSERT(id >= 0 && c.size() == (unsigned)id, "Failed dimension check in TreeLSTMBuilder");
  h.push_back(vector<Expression>(layers));
  c.push_back(vector<Expression>(layers));
  vector<Expression>& ht = h.back();
  vector<Expression>& ct = c.back();

  Expression in = x;
  for (unsigned i = 0; i < layers; ++i) {
    const vector<Expression>& vars = param_vars[i];
    vector<Expression> i_h_children, i_c_children;
    i_h_children.reserve(children.size() > 1 ? children.size() : 1);
    i_c_children.reserve(children.size() > 1 ? children.size() : 1);

    bool has_prev_state = (children.size() > 0 || has_initial_state);
    if (children.size() == 0) {
      i_h_children.push_back(Expression());
      i_c_children.push_back(Expression());
      if (has_initial_state) {
        // intial value for h and c at timestep 0 in layer i
        // defaults to zero matrix input if not set in add_parameter_edges
        i_h_children[0] = h0[i];
        i_c_children[0] = c0[i];
      }
    }
    else {  // t > 0
      for (int child : children) {
        i_h_children.push_back(h[child][i]);
        i_c_children.push_back(c[child][i]);
      }
    }

    // input
    Expression i_ait;
    if (has_prev_state) {
      vector<Expression> xs = {vars[BI], vars[X2I], in};
      xs.reserve(4 * children.size() + 3);
      for (unsigned j = 0; j < children.size(); ++j) {
        unsigned ej = (j < N) ? j : N - 1;
        xs.push_back(Lookup(i, H2I, ej));
        xs.push_back(i_h_children[j]);
        xs.push_back(Lookup(i, C2I, ej));
        xs.push_back(i_c_children[j]);
      }
      DYNET_ASSERT(xs.size() == 4 * children.size() + 3, "Failed dimension check in TreeLSTMBuilder");
      i_ait = affine_transform(xs);
    }
    else
      i_ait = affine_transform({vars[BI], vars[X2I], in});
    Expression i_it = logistic(i_ait);

    // forget
    vector<Expression> i_ft;
    for (unsigned k = 0; k < children.size(); ++k) {
      unsigned ek = (k < N) ? k : N - 1;
      Expression i_aft;
      if (has_prev_state) {
        vector<Expression> xs = {vars[BF], vars[X2F], in};
        xs.reserve(4 * children.size() + 3);
        for (unsigned j = 0; j < children.size(); ++j) {
          unsigned ej = (j < N) ? j : N - 1;
          xs.push_back(Lookup(i, H2F, ej * N + ek));
          xs.push_back(i_h_children[j]);
          xs.push_back(Lookup(i, C2F, ej * N + ek));
          xs.push_back(i_c_children[j]);
        }
        DYNET_ASSERT(xs.size() == 4 * children.size() + 3, "Failed dimension check in TreeLSTMBuilder");
        i_aft = affine_transform(xs);
      }
      else
        i_aft = affine_transform({vars[BF], vars[X2F], in});
      i_ft.push_back(logistic(i_aft + 1.f));
    }

    // write memory cell
    Expression i_awt;
    if (has_prev_state) {
      vector<Expression> xs = {vars[BC], vars[X2C], in};
      // This is the one and only place that should *not* condition on i_c_children
      // This should condition only on x (a.k.a. in), the bias (vars[BC]) and i_h_children
      xs.reserve(2 * children.size() + 3);
      for (unsigned j = 0; j < children.size(); ++j) {
        unsigned ej = (j < N) ? j : N - 1;
        xs.push_back(Lookup(i, H2C, ej));
        xs.push_back(i_h_children[j]);
      }
      DYNET_ASSERT(xs.size() == 2 * children.size() + 3, "Failed dimension check in TreeLSTMBuilder");
      i_awt = affine_transform(xs);
    }
    else
      i_awt = affine_transform({vars[BC], vars[X2C], in});
    Expression i_wt = tanh(i_awt);

    // compute new cell value
    if (has_prev_state) {
      Expression i_nwt = cmult(i_it, i_wt);
      vector<Expression> i_crts(children.size());
      for (unsigned j = 0; j < children.size(); ++j) {
        i_crts[j] = cmult(i_ft[j], i_c_children[j]);
      }
      Expression i_crt = sum(i_crts);
      ct[i] = i_crt + i_nwt;
    }
    else {
      ct[i] = cmult(i_it, i_wt);
    }

    // output
    Expression i_aot;
    if (has_prev_state) {
      vector<Expression> xs = {vars[BO], vars[X2O], in};
      xs.reserve(4 * children.size() + 3);
      for (unsigned j = 0; j < children.size(); ++j) {
        unsigned ej = (j < N) ? j : N - 1;
        xs.push_back(Lookup(i, H2O, ej));
        xs.push_back(i_h_children[j]);
        xs.push_back(Lookup(i, C2O, ej));
        xs.push_back(i_c_children[j]);
      }
      DYNET_ASSERT(xs.size() == 4 * children.size() + 3, "Failed dimension check in TreeLSTMBuilder");
      i_aot = affine_transform(xs);
    }
    else
      i_aot = affine_transform({vars[BO], vars[X2O], in});
    Expression i_ot = logistic(i_aot);

    // Compute new h value
    Expression ph_t = tanh(ct[i]);
    in = ht[i] = cmult(i_ot, ph_t);
  }
  return ht.back();
}

void NaryTreeLSTMBuilder::copy(const RNNBuilder & rnn) {
  const NaryTreeLSTMBuilder & rnn_treelstm = (const NaryTreeLSTMBuilder&)rnn;
  DYNET_ASSERT(params.size() == rnn_treelstm.params.size(), "Failed dimension check in TreeLSTMBuilder");
  for(size_t i = 0; i < params.size(); ++i) {
    for(size_t j = 0; j < params[i].size(); ++j) {
      params[i][j] = rnn_treelstm.params[i][j];
    }
  }
}

ParameterCollection & NaryTreeLSTMBuilder::get_parameter_collection() {
  return local_model;
}

UnidirectionalTreeLSTMBuilder::UnidirectionalTreeLSTMBuilder(unsigned layers,
                         unsigned input_dim,
                         unsigned hidden_dim,
                         ParameterCollection& model) {
  local_model = model.add_subcollection("unidirectional-tree-lstm-builder");
  node_builder = LSTMBuilder(layers, input_dim, hidden_dim, local_model);
}

void UnidirectionalTreeLSTMBuilder::new_graph_impl(ComputationGraph& cg, bool update) {
   node_builder.new_graph(cg, update);
}

// layout: 0..layers = c
//         layers+1..2*layers = h
void UnidirectionalTreeLSTMBuilder::start_new_sequence_impl(const vector<Expression>& hinit) {
  h.clear();
  node_builder.start_new_sequence(hinit);
}

void UnidirectionalTreeLSTMBuilder::set_num_elements(int length) {
  h.resize(length);
}

Expression UnidirectionalTreeLSTMBuilder::add_input(int id, vector<int> children, const Expression& x) {
  DYNET_ASSERT(id >= 0 && h.size() == (unsigned)id, "Failed dimension check in TreeLSTMBuilder");

  RNNPointer prev = (RNNPointer)(-1);
  Expression embedding = node_builder.add_input(prev, x);
  prev = node_builder.state();

  for (unsigned child : children) {
    embedding = node_builder.add_input(prev, h[child]);
    prev = node_builder.state();
  }
  h.push_back(embedding);
  return embedding;
}

BidirectionalTreeLSTMBuilder::BidirectionalTreeLSTMBuilder(unsigned layers,
                         unsigned input_dim,
                         unsigned hidden_dim,
                         ParameterCollection& model) {
  DYNET_ASSERT(hidden_dim % 2 == 0, "Failed dimension check in TreeLSTMBuilder");
  local_model = model.add_subcollection("bidirectional-tree-lstm-builder");
  fwd_node_builder = LSTMBuilder(layers, input_dim, hidden_dim / 2, local_model);
  rev_node_builder = LSTMBuilder(layers, input_dim, hidden_dim / 2, local_model);
}

void BidirectionalTreeLSTMBuilder::new_graph_impl(ComputationGraph& cg, bool update) {
   fwd_node_builder.new_graph(cg, update);
   rev_node_builder.new_graph(cg, update);
}

// layout: 0..layers = c
//         layers+1..2*layers = h
void BidirectionalTreeLSTMBuilder::start_new_sequence_impl(const vector<Expression>& hinit) {
  h.clear();
  fwd_node_builder.start_new_sequence(hinit);
  rev_node_builder.start_new_sequence(hinit);
}

void BidirectionalTreeLSTMBuilder::set_num_elements(int length) {
  h.resize(length);
}

Expression BidirectionalTreeLSTMBuilder::add_input(int id, vector<int> children, const Expression& x) {
  DYNET_ASSERT(id >= 0 && h.size() >= (unsigned)id, "Failed dimension check in TreeLSTMBuilder");

  RNNPointer prev = (RNNPointer)(-1);
  Expression fwd_embedding = fwd_node_builder.add_input(prev, x);
  prev = fwd_node_builder.state();
  for (unsigned child : children) {
    fwd_embedding = fwd_node_builder.add_input(prev, h[child]);
    prev = fwd_node_builder.state();
  }

  prev = (RNNPointer)(-1);
  Expression rev_embedding = rev_node_builder.add_input(prev, x);
  prev = rev_node_builder.state();
  for (unsigned i = children.size(); i-- > 0;) {
    unsigned  child = children[i];
    rev_embedding = rev_node_builder.add_input(prev, h[child]);
    prev = rev_node_builder.state();
  }

  Expression embedding = concatenate({fwd_embedding, rev_embedding});
  if (h.size() == id)
	h.push_back(embedding);
  else
	h[id] = embedding;

  return embedding;
}