import logging
from typing import Dict, List, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from networkx import DiGraph
from torch import Tensor, nn as nn
from torch.autograd.variable import Variable

from binlin.data.ud import index_data
from binlin.model.nn_utils import get_embed_matrix, pad_seq
from binlin.model.syn.sympairs import SymModel
from binlin.model.syn.utils.bintree import BinTreeBase
from binlin.utils.combinatorics import flatten_nested_lists
from binlin.utils.constants import VocabSymbols

logger = logging.getLogger('main')


class SymGraphEncModel(SymModel):

    @property
    def _data_fields(self):
        return ['LEMMA', 'UPOS', 'XPOS', 'DEPREL']

    @property
    def _num_embedding_feats(self):
        # todo: need to avoid hard-coding
        # 6 = 2 * 3
        # 2 because we consider two nodes at a time,
        # 3 because for a node we consider itself + its context
        # consisting of its parent and a random child
        return 2 * 3 * self._data_fields_num

    def _init_weights(self):
        self._dim_emb = self.config["embedding_dim"]
        self._mat_emb = get_embed_matrix(len(self._id2tok), self._dim_emb, padding_idx=self.PAD_ID)

        self._dim_emb_proj_in = self._num_embedding_feats * self._dim_emb
        self._dim_emb_proj_out = self._num_embedding_feats * self.config['embedding_proj_dim']
        self._mat_emb_proj = nn.Linear(self._dim_emb_proj_in, self._dim_emb_proj_out)

        self._dim_dense_out = self.config["dense_dim"]
        self._mat_dense = nn.Linear(self._dim_emb_proj_out, self._dim_dense_out)

        self._mat_attn = nn.Linear(self._dim_emb, 1)

        self._dim_concat_in = self._dim_dense_out + self._dim_emb
        self._dim_concat_out = self._dim_dense_out
        self._mat_concat = nn.Linear(self._dim_concat_in, self._dim_concat_out)

        self._dim_out = 1
        self._mat_out = nn.Linear(self._dim_dense_out, self._dim_out)

    def forward(self, batch_data: Tuple[Tensor, Tensor, Union[Tensor, None]]) -> Dict:
        other_nodes_var, head_child_var, _ = batch_data

        # head-child pair is encoded using an MLP
        x_head_child = self._mat_emb(head_child_var).view(-1,
                                                          self._dim_emb_proj_in)  # size(num_node_pairs, self.emb_proj_in)
        x_head_child = F.leaky_relu(self._mat_emb_proj(x_head_child))  # size (num_node_pairs, self.emb_proj_out)
        x_head_child = F.leaky_relu(self._mat_dense(x_head_child))  # size (num_node_pairs, self.dense1_dim)
        # x_head_child = F.leaky_relu(x_head_child) # size (num_node_pairs, self.dense1_dim)

        # graph nodes are encoded using embedding lookup --> summing
        node_number = other_nodes_var.shape[-1]
        x_graph = self._mat_emb(other_nodes_var)  # size (max_num_nodes, batch_size, self.emb_dim)

        # variant1: sum over all vecs
        # x_graph = torch.sum(x_graph, dim=[0], keepdim=True).view(-1, self._dim_emb) # size (batch_size, emb_dim)

        # variant2: use attn scores
        attn_unnorm_scores = self._mat_attn(x_graph.view(-1, self._dim_emb))  # num_edges x 1
        attn_weights = F.leaky_relu(attn_unnorm_scores).view(-1, 1, node_number)

        # TODO: find a simpler way to do it w/o squeezing and unsqueezing?
        # apply attention weights to the graph vectors to get weighted average
        # size (1, dense)
        x_graph = torch.bmm(attn_weights, x_graph).squeeze(1)  # size: (bsize, emb_size)

        # Concat head, child and graph representations
        x_combined = torch.cat((x_head_child, x_graph), 1)  # size (bs, emb_dim + self.dense1_dim)

        # size (batch_size, self.dense1_dim)
        x_combined = self._mat_concat(x_combined)
        x_combined = F.leaky_relu(x_combined)
        x_combined = self._mat_out(x_combined)
        logits = torch.sigmoid(x_combined)

        return {'logits': logits}

    def extract_features(self, bt, new_node_nxid, dg, feats_d):

        # pair-level features
        head_nxid = bt.nxid
        head_deptree_feats = feats_d[head_nxid]
        child_deptree_feats = feats_d[new_node_nxid]
        x_pair_ids_l = head_deptree_feats + child_deptree_feats

        # extracting graph-level features
        graph_ids_l = self.extract_graph_level_feats(dg, new_node_nxid, bt, feats_d)

        return (graph_ids_l, x_pair_ids_l)

    def extract_graph_level_feats(self, dg, new_node_nxid, bt, feats_d):
        head_sbl = dg.node[bt.nxid]['sbl']
        if head_sbl is None:
            head_sbl_feats_l = self._dummy_node_feats_vec
        else:
            head_sbl_feats_l = flatten_nested_lists([feats_d[ch] for ch in head_sbl])

        child_sbl = dg.node[new_node_nxid]['sbl']
        if child_sbl is None:
            ch_sbl_feats_l = self._dummy_node_feats_vec
        else:
            ch_sbl_feats_l = flatten_nested_lists([feats_d[ch] for ch in child_sbl])

        child_children = dg[new_node_nxid]
        if len(child_children) == 0:
            ch_ch_feats_l = self._dummy_node_feats_vec
        else:
            ch_ch_feats_l = flatten_nested_lists([feats_d[ch] for ch in child_children])

        graph_ids_l = head_sbl_feats_l + ch_sbl_feats_l + ch_ch_feats_l
        return graph_ids_l

    def init_data_containers(self):
        return {'X': [],
                'Y': [],
                'Xg': []}

    def add_xy_pairs(self, data_containers: Dict, y: int, model_inputs: Tuple[List[int], List[int]]):
        x_graph, x_head_child = model_inputs
        data_containers['X'].append(x_head_child)
        data_containers['Xg'].append(x_graph)
        data_containers['Y'].append(y)

    def _batchify(self, data_containers: Dict, batch_size: int):

        # sort according to the lemma length
        sorted_data = sorted(zip(*(data_containers['Xg'],
                                   data_containers['X'],
                                   data_containers['Y'])), key=lambda p: len(p[0]), reverse=True)
        data_size = len(sorted_data)
        num_batches = data_size // batch_size
        data_indices = index_data(data_size, mode='no_shuffling')
        batch_pairs = []

        for bi in range(num_batches + 1):  # including the last (smaller) batch

            batch_x_pair_feats = []
            batch_x_graph_feats = []
            batch_x_lens = []
            batch_y = []

            curr_batch_indices = data_indices[bi * batch_size: (bi + 1) * batch_size]

            if len(curr_batch_indices) == 0:
                break

            for idx in curr_batch_indices:
                graph_f_ids, node_pairs_f_ids, y_ids = sorted_data[idx]
                batch_x_graph_feats.append(graph_f_ids)
                batch_x_lens.append(len(graph_f_ids))
                batch_x_pair_feats.append(node_pairs_f_ids)
                batch_y.append(y_ids)

            max_graph_f_len = max(batch_x_lens)
            batch_x_graph_feats_padded = [pad_seq(x, max_graph_f_len, pad_id=self.PAD_ID) for x in batch_x_graph_feats]

            # size: (num_nodes, batch_size)
            batch_x_graph_feats_var = Variable(torch.LongTensor(batch_x_graph_feats_padded)).to(self.device)
            # size: (batch_size, 2 * num_node_feats)
            batch_x_pair_feats_var = Variable(torch.LongTensor(batch_x_pair_feats)).to(self.device)
            # size: (batch_size, 1)
            batch_y_var = Variable(torch.FloatTensor(batch_y)).unsqueeze(1).to(self.device)
            batch_pairs.append((batch_x_graph_feats_var, batch_x_pair_feats_var, batch_y_var))

        return batch_pairs

    def make_decision(self,
                      bt: BinTreeBase,
                      new_node_nxid: str,
                      dg: DiGraph,
                      feats_d: Dict, *other_inputs) -> int:

        x_graph_ids_l, x_ids_l = self.extract_features(bt, new_node_nxid, dg, feats_d)
        x_ids_np = np.asarray(x_ids_l)
        x_graph_ids_l = np.asarray([x_graph_ids_l])
        outputs = self.__call__(
            (
                torch.from_numpy(x_graph_ids_l).to(self.device),
                torch.from_numpy(x_ids_np).to(self.device),
                None)
        )
        logit_val = outputs['logits'].cpu().data[0].numpy()

        if logit_val >= 0.5:
            decision = VocabSymbols.RIGHT
        else:
            decision = VocabSymbols.LEFT

        return decision


component = SymGraphEncModel

import logging
from typing import Dict, List, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from networkx import DiGraph
from torch import Tensor, nn as nn
from torch.autograd.variable import Variable

from binlin.data.ud import index_data
from binlin.model.nn_utils import get_embed_matrix, pad_seq
from binlin.model.syn.sympairs import SymModel
from binlin.model.syn.utils.bintree import BinTreeBase
from binlin.utils.combinatorics import flatten_nested_lists
from binlin.utils.constants import VocabSymbols

logger = logging.getLogger('main')


class SymGraphEncModel(SymModel):

    @property
    def _data_fields(self):
        return ['LEMMA', 'UPOS', 'XPOS', 'DEPREL']

    @property
    def _num_embedding_feats(self):
        # todo: need to avoid hard-coding
        # 6 = 2 * 3
        # 2 because we consider two nodes at a time,
        # 3 because for a node we consider itself + its context
        # consisting of its parent and a random child
        return 2 * 3 * self._data_fields_num

    def _init_weights(self):
        self._dim_emb = self.config["embedding_dim"]
        self._mat_emb = get_embed_matrix(len(self._id2tok), self._dim_emb, padding_idx=self.PAD_ID)

        self._dim_emb_proj_in = self._num_embedding_feats * self._dim_emb
        self._dim_emb_proj_out = self._num_embedding_feats * self.config['embedding_proj_dim']
        self._mat_emb_proj = nn.Linear(self._dim_emb_proj_in, self._dim_emb_proj_out)

        self._dim_dense_out = self.config["dense_dim"]
        self._mat_dense = nn.Linear(self._dim_emb_proj_out, self._dim_dense_out)

        self._mat_attn = nn.Linear(self._dim_emb, 1)

        self._dim_concat_in = self._dim_dense_out + self._dim_emb
        self._dim_concat_out = self._dim_dense_out
        self._mat_concat = nn.Linear(self._dim_concat_in, self._dim_concat_out)

        self._dim_out = 1
        self._mat_out = nn.Linear(self._dim_dense_out, self._dim_out)

    def forward(self, batch_data: Tuple[Tensor, Tensor, Union[Tensor, None]]) -> Dict:
        other_nodes_var, head_child_var, _ = batch_data

        # head-child pair is encoded using an MLP
        x_head_child = self._mat_emb(head_child_var).view(-1,
                                                          self._dim_emb_proj_in)  # size(num_node_pairs, self.emb_proj_in)
        x_head_child = F.leaky_relu(self._mat_emb_proj(x_head_child))  # size (num_node_pairs, self.emb_proj_out)
        x_head_child = F.leaky_relu(self._mat_dense(x_head_child))  # size (num_node_pairs, self.dense1_dim)
        # x_head_child = F.leaky_relu(x_head_child) # size (num_node_pairs, self.dense1_dim)

        # graph nodes are encoded using embedding lookup --> summing
        node_number = other_nodes_var.shape[-1]
        x_graph = self._mat_emb(other_nodes_var)  # size (max_num_nodes, batch_size, self.emb_dim)

        # variant1: sum over all vecs
        # x_graph = torch.sum(x_graph, dim=[0], keepdim=True).view(-1, self._dim_emb) # size (batch_size, emb_dim)

        # variant2: use attn scores
        attn_unnorm_scores = self._mat_attn(x_graph.view(-1, self._dim_emb))  # num_edges x 1
        attn_weights = F.leaky_relu(attn_unnorm_scores).view(-1, 1, node_number)

        # TODO: find a simpler way to do it w/o squeezing and unsqueezing?
        # apply attention weights to the graph vectors to get weighted average
        # size (1, dense)
        x_graph = torch.bmm(attn_weights, x_graph).squeeze(1)  # size: (bsize, emb_size)

        # Concat head, child and graph representations
        x_combined = torch.cat((x_head_child, x_graph), 1)  # size (bs, emb_dim + self.dense1_dim)

        # size (batch_size, self.dense1_dim)
        x_combined = self._mat_concat(x_combined)
        x_combined = F.leaky_relu(x_combined)
        x_combined = self._mat_out(x_combined)
        logits = torch.sigmoid(x_combined)

        return {'logits': logits}

    def extract_features(self, bt, new_node_nxid, dg, feats_d):

        # pair-level features
        head_nxid = bt.nxid
        head_deptree_feats = feats_d[head_nxid]
        child_deptree_feats = feats_d[new_node_nxid]
        x_pair_ids_l = head_deptree_feats + child_deptree_feats

        # extracting graph-level features
        graph_ids_l = self.extract_graph_level_feats(dg, new_node_nxid, bt, feats_d)

        return (graph_ids_l, x_pair_ids_l)

    def extract_graph_level_feats(self, dg, new_node_nxid, bt, feats_d):
        head_sbl = dg.node[bt.nxid]['sbl']
        if head_sbl is None:
            head_sbl_feats_l = self._dummy_node_feats_vec
        else:
            head_sbl_feats_l = flatten_nested_lists([feats_d[ch] for ch in head_sbl])

        child_sbl = dg.node[new_node_nxid]['sbl']
        if child_sbl is None:
            ch_sbl_feats_l = self._dummy_node_feats_vec
        else:
            ch_sbl_feats_l = flatten_nested_lists([feats_d[ch] for ch in child_sbl])

        child_children = dg[new_node_nxid]
        if len(child_children) == 0:
            ch_ch_feats_l = self._dummy_node_feats_vec
        else:
            ch_ch_feats_l = flatten_nested_lists([feats_d[ch] for ch in child_children])

        graph_ids_l = head_sbl_feats_l + ch_sbl_feats_l + ch_ch_feats_l
        return graph_ids_l

    def init_data_containers(self):
        return {'X': [],
                'Y': [],
                'Xg': []}

    def add_xy_pairs(self, data_containers: Dict, y: int, model_inputs: Tuple[List[int], List[int]]):
        x_graph, x_head_child = model_inputs
        data_containers['X'].append(x_head_child)
        data_containers['Xg'].append(x_graph)
        data_containers['Y'].append(y)

    def _batchify(self, data_containers: Dict, batch_size: int):

        # sort according to the lemma length
        sorted_data = sorted(zip(*(data_containers['Xg'],
                                   data_containers['X'],
                                   data_containers['Y'])), key=lambda p: len(p[0]), reverse=True)
        data_size = len(sorted_data)
        num_batches = data_size // batch_size
        data_indices = index_data(data_size, mode='no_shuffling')
        batch_pairs = []

        for bi in range(num_batches + 1):  # including the last (smaller) batch

            batch_x_pair_feats = []
            batch_x_graph_feats = []
            batch_x_lens = []
            batch_y = []

            curr_batch_indices = data_indices[bi * batch_size: (bi + 1) * batch_size]

            if len(curr_batch_indices) == 0:
                break

            for idx in curr_batch_indices:
                graph_f_ids, node_pairs_f_ids, y_ids = sorted_data[idx]
                batch_x_graph_feats.append(graph_f_ids)
                batch_x_lens.append(len(graph_f_ids))
                batch_x_pair_feats.append(node_pairs_f_ids)
                batch_y.append(y_ids)

            max_graph_f_len = max(batch_x_lens)
            batch_x_graph_feats_padded = [pad_seq(x, max_graph_f_len, pad_id=self.PAD_ID) for x in batch_x_graph_feats]

            # size: (num_nodes, batch_size)
            batch_x_graph_feats_var = Variable(torch.LongTensor(batch_x_graph_feats_padded)).to(self.device)
            # size: (batch_size, 2 * num_node_feats)
            batch_x_pair_feats_var = Variable(torch.LongTensor(batch_x_pair_feats)).to(self.device)
            # size: (batch_size, 1)
            batch_y_var = Variable(torch.FloatTensor(batch_y)).unsqueeze(1).to(self.device)
            batch_pairs.append((batch_x_graph_feats_var, batch_x_pair_feats_var, batch_y_var))

        return batch_pairs

    def make_decision(self,
                      bt: BinTreeBase,
                      new_node_nxid: str,
                      dg: DiGraph,
                      feats_d: Dict, *other_inputs) -> int:

        x_graph_ids_l, x_ids_l = self.extract_features(bt, new_node_nxid, dg, feats_d)
        x_ids_np = np.asarray(x_ids_l)
        x_graph_ids_l = np.asarray([x_graph_ids_l])
        outputs = self.__call__(
            (
                torch.from_numpy(x_graph_ids_l).to(self.device),
                torch.from_numpy(x_ids_np).to(self.device),
                None)
        )
        logit_val = outputs['logits'].cpu().data[0].numpy()

        if logit_val >= 0.5:
            decision = VocabSymbols.RIGHT
        else:
            decision = VocabSymbols.LEFT

        return decision


component = SymGraphEncModel

# todo: need to avoid hard-coding # 6 = 2 * 3 # 2 because we consider two nodes at a time, # 3 because for a node we consider itself + its context # consisting of its parent and a random child # head-child pair is encoded using an MLP # size(num_node_pairs, self.emb_proj_in) # size (num_node_pairs, self.emb_proj_out) # size (num_node_pairs, self.dense1_dim) # x_head_child = F.leaky_relu(x_head_child) # size (num_node_pairs, self.dense1_dim) # graph nodes are encoded using embedding lookup --> summing # size (max_num_nodes, batch_size, self.emb_dim) # variant1: sum over all vecs # x_graph = torch.sum(x_graph, dim=[0], keepdim=True).view(-1, self._dim_emb) # size (batch_size, emb_dim) # variant2: use attn scores # num_edges x 1 # TODO: find a simpler way to do it w/o squeezing and unsqueezing? # apply attention weights to the graph vectors to get weighted average # size (1, dense) # size: (bsize, emb_size) # Concat head, child and graph representations # size (bs, emb_dim + self.dense1_dim) # size (batch_size, self.dense1_dim) # pair-level features # extracting graph-level features # sort according to the lemma length # including the last (smaller) batch # size: (num_nodes, batch_size) # size: (batch_size, 2 * num_node_feats) # size: (batch_size, 1)