from torch_geometric.nn import Linear, SAGEConv, GCNConv, SGConv, Sequential, to_hetero, HeteroConv

import torch

import torch.nn.functional as F

import torch.optim as optim

import torch.nn as nn

from .conv import GATConv

class SimpleMLP(nn.Module):

    def __init__(self, input_dim, hidden_dim, output_dim):

        super(SimpleMLP, self).__init__()

        self.FC_hidden = nn.Linear(input_dim, hidden_dim)

        self.FC_hidden2 = nn.Linear(hidden_dim, hidden_dim)

        self.FC_output = nn.Linear(hidden_dim, output_dim)

        self.ReLU = nn.ReLU()       

    def forward(self, x):

        h     = self.ReLU(self.FC_hidden(x))

        h     = self.ReLU(self.FC_hidden2(h))

        x_hat = self.FC_output(h)

        return x_hat

class HeteroGNN(torch.nn.Module):

    def __init__(self, pyg_data, hidden_channels, out_channels, num_layers, gnn_backbone, gnn_aggr, snp_init_dim_size, gene_init_dim_size, go_init_dim_size, gat_num_head, no_relu = False):

        super().__init__()

        edge_types = pyg_data.edge_types

        self.convs = torch.nn.ModuleList()

        self.snp_feat_mlp = SimpleMLP(snp_init_dim_size, hidden_channels, hidden_channels)

        self.go_feat_mlp = SimpleMLP(go_init_dim_size, hidden_channels, hidden_channels)

        self.gene_feat_mlp = SimpleMLP(gene_init_dim_size, hidden_channels, hidden_channels)

        self.ReLU = nn.ReLU()   

        for _ in range(num_layers):

            conv_layer = {}

            for i in edge_types:

                if gnn_backbone == 'SAGE':

                    conv_layer[i] = SAGEConv((-1, -1), hidden_channels)

                elif gnn_backbone == 'GAT':

                    conv_layer[i] = GATConv((-1, -1), hidden_channels, 

                                            heads = gat_num_head, 

                                            add_self_loops = False)

                elif gnn_backbone == 'GCN':

                    conv_layer[i] = GCNConv(-1, hidden_channels, add_self_loops = False)

                elif gnn_backbone == 'SGC':

                    conv_layer[i] = SGConv(-1, hidden_channels, add_self_loops = False)

            conv = HeteroConv(conv_layer, aggr=gnn_aggr)

            self.convs.append(conv)

        self.lin = Linear(hidden_channels, out_channels)

        self.no_relu = no_relu

    def forward(self, x_dict, edge_index_dict, batch_size, genotype = None, return_h = False, 

                return_attention_weights = False):

        x_dict['SNP'] = self.snp_feat_mlp(x_dict['SNP'])

        x_dict['Gene'] = self.gene_feat_mlp(x_dict['Gene'])

        x_dict['CellularComponent'] = self.go_feat_mlp(x_dict['CellularComponent'])

        x_dict['BiologicalProcess'] = self.go_feat_mlp(x_dict['BiologicalProcess'])

        x_dict['MolecularFunction'] = self.go_feat_mlp(x_dict['MolecularFunction'])

        attention_all_layers = []

        for conv in self.convs:

            if return_attention_weights:

                out = conv(x_dict, edge_index_dict, 

                              return_attention_weights_dict = dict(zip(list(edge_index_dict.keys()), 

                                                            [True] * len(list(edge_index_dict.keys())))))

                #attention_layer = {i: [x[1] for x in j[1]] for i,j in out.items()}

                mean_attention = torch.mean(torch.vstack([torch.vstack([x[1] for x in j[1]]) for i,j in out.items()]))

                x_dict = {i: j[0] for i,j in out.items()}

                attention_all_layers.append(mean_attention)

            else:

                x_dict = conv(x_dict, edge_index_dict)

            x_dict = {key: x.relu() for key, x in x_dict.items()}

        if return_h:

            return self.ReLU(self.lin(x_dict['SNP']))[:batch_size], x_dict['SNP'][:batch_size] 

        if return_attention_weights:

            return self.ReLU(self.lin(x_dict['SNP']))[:batch_size], attention_all_layers

        else:

            if self.no_relu:

                return self.lin(x_dict['SNP'])[:batch_size]

            else:

                return self.ReLU(self.lin(x_dict['SNP']))[:batch_size]