import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

import torch_geometric.nn as pyg_nn
import torch_geometric.utils as pyg_utils

import rlkit.torch.pytorch_util as ptu
from rlkit.torch.networks.network_utils import get_activation, build_gnn_conv_model

class R2GQNetBase(torch.nn.Module):
    def __init__(self, 
                pre_graph_builder, 
                obs_dim,
                action_dim, 
                node_dim,
                post_mlp,
                conv_type='GSage',
                skip_connect=False,
                concat_emb=False,
                num_conv_layers=3,
                hidden_activation=None,
                output_activation=None,
                ):
        super(R2GQNetBase, self).__init__()
        print(conv_type)
        # graph builder
        self.pre_graph_builder = pre_graph_builder

        # convs
        self.obs_dim = obs_dim
        self.action_dim = action_dim
        self.node_input_dim = pre_graph_builder.output_node_dim+2*self.action_dim
        self.node_dim = node_dim
        self.conv_type = conv_type
        self.num_conv_layers = num_conv_layers
        self.convs = self.build_convs(self.node_input_dim, self.node_dim, self.num_conv_layers)
        self.skip_connect = skip_connect
        self.concat_emb = concat_emb

        self.post_mlp = post_mlp

        if self.skip_connect:
            print('skip_connect')
            self.mlps = nn.ModuleList()
            self.mlps.append(nn.Linear(self.node_input_dim+node_dim, node_dim))
            for l in range(1,num_conv_layers):
                self.mlps.append(nn.Linear(node_dim+node_dim, node_dim))
        if self.concat_emb:
            print('concat_emb')
        self.hidden_activations = nn.ModuleList([get_activation(hidden_activation) for l in range(num_conv_layers-1)])
        self.output_activation = get_activation(output_activation)
        

    def build_convs(self, node_input_dim, node_dim, num_conv_layers):
        convs = nn.ModuleList()
        conv = build_gnn_conv_model(self.conv_type, node_input_dim, node_dim)
        convs.append(conv)
        for l in range(1,num_conv_layers):
            conv = build_gnn_conv_model(self.conv_type, node_dim, node_dim)
            convs.append(conv)
        return convs

    def forward(self, obs, context_actions, eval_actions):
        batch_size = obs.shape[0]
        batch_size, num_agent, a_dim = context_actions.shape

        x_obs, edge_index = self.pre_graph_builder(obs)
        x_obs = x_obs.reshape(batch_size,-1,x_obs.shape[-1])

        x = torch.zeros(batch_size, x_obs.shape[-2], x_obs.shape[-1]+2*a_dim).to(ptu.device)
        x[:,:,:x_obs.shape[-1]] = x_obs[:,:,:]
        x[:,:num_agent,x_obs.shape[-1]:x_obs.shape[-1]+a_dim] = context_actions[:,:,:]
        x[:,:num_agent,x_obs.shape[-1]+a_dim:] = eval_actions[:,:,:]
        x = x.reshape(-1,x.shape[-1])

        if self.concat_emb:
            x_list = [x.clone()]
            concat_length = x.shape[-1]

        for l, conv in enumerate(self.convs):
            # self.check_input(x, edge_index)
            x_new = conv(x, edge_index)
            if self.skip_connect:
                x = self.mlps[l](torch.cat([x,x_new],dim=-1))
            else:
                x = x_new
            if l < len(self.convs)-1:
                x = self.hidden_activations[l](x)

            if self.concat_emb:
                x_list.append(x.clone())
                concat_length += x.shape[-1]

        if self.concat_emb:
            x = torch.cat(x_list, dim=-1)
            x = x.reshape(batch_size, -1, concat_length)
        else:
            x = x.reshape(batch_size, -1, self.node_dim)                
        x = self.output_activation(x)
        x = self.post_mlp(x)
        return x

    def check_input(self, xs, edge_index):
        Os = {}
        for indx in range(edge_index.shape[1]):
            i=edge_index[1,indx].detach().numpy()
            j=edge_index[0,indx].detach().numpy()
            xi=xs[i].detach().numpy()
            xj=list(xs[j].detach().numpy())
            if str(i) not in Os.keys():
                Os[str(i)] = {'x_j':[]}
            Os[str(i)]['x_i'] = xi
            Os[str(i)]['x_j'] += xj

        import matplotlib.pyplot as plt
        plt.figure()
        plt.subplot(1,2,1)
        for i in Os.keys():
            plt.plot(Os[str(i)]['x_i'],label=str(i))
            plt.title('x_i')
        plt.legend()
        plt.subplot(1,2,2)
        for i in Os.keys():
            plt.plot(Os[str(i)]['x_j'],label=str(i))
            plt.title('x_j')
        plt.legend()
        plt.show()




