<reponame>raphaelsulzer/dgnn
import sys, os
import torch
import torch.nn as nn
import torch.nn.functional as F

from tqdm import tqdm
from torch_geometric.nn.conv import SAGEConv

class SurfaceNet(nn.Module):

    def __init__(self, n_node_features, clf):

        super(SurfaceNet, self).__init__()  # necessary for all classes extending the module class

        self.clf = clf
        self.n_classes = 2
        self.n_node_feat = n_node_features

        self.sage_convs = torch.nn.ModuleList()

        self.sage_convs.append(SAGEConv(self.n_node_feat, self.clf.training.model_params[0]))
        for i in range(len(self.clf.training.model_params)-1):
            self.sage_convs.append(SAGEConv(self.clf.training.model_params[i], self.clf.training.model_params[i+1], normalize=False))

        self.num_layers = len(self.sage_convs)

        if(self.clf.training.model_name == "sage+"):
            self.vertex_net = nn.Sequential(nn.ReLU(True), nn.Linear(self.clf.training.model_params[-1], 128), nn.ReLU(True), nn.Linear(128, 2))

        if(self.clf.training.model_name == "sage++"):
            self.vertex_net = nn.Sequential(nn.ReLU(True), nn.Linear(self.clf.training.model_params[-1], 128), nn.ReLU(True), nn.Linear(128, 2))
            self.edge_net = nn.Sequential(nn.Linear(self.clf.training.model_params[-1]*2, 128), nn.ReLU(True), nn.Linear(128, 1))


    #######################################################
    ##################### TRAIN FORWARD ###################
    #######################################################
    def forward(self, data_all):
        # produces final embedding batch per batch


        """
        the forward function producing the embeddings for each cell of 'input'
        input = [n_batch, input_feat, n_channels=1] float array: input features
        output = [n_batch, n_class, n_channels=1] float array: cell class scores
        """

        if(self.clf.features.normalization_feature and not self.clf.features.keep_normalization_feature):
            x=data_all.x[data_all.n_id, 1:].to(self.clf.temp.device) # put this batch on gpu
        else:
            x=data_all.x[data_all.n_id, :].to(self.clf.temp.device) # put this batch on gpu

        # for i, (edge_index, _, size) in enumerate(data_all.adjs):
        for i in range(self.num_layers):
            edge_index, e_id, size = data_all.adjs[i]
            x = self.sage_convs[i]((x, x[:size[1]]), edge_index.to(self.clf.temp.device))
            if i != self.num_layers - 1:
                x = F.relu(x)
                # x = F.dropout(x, p=0.5, training=self.training)

            # torch.cuda.empty_cache()

        x_edge = None

        if(self.clf.training.model_name[-1] == "+"):
            x = F.relu(x)
            x = self.vertex_net(x)

        if(self.clf.training.model_name[:-1] == "epsage"):
            x = F.relu(x)
            # concatenate left and right cell embeddings
            # take the edge_indices from self.num_layers + 1, for those I have the final embeddings calculated
            x_edge = self.edge_net(torch.cat((x[data_all.adjs[i+1][0][0]],x[data_all.adjs[i+1][0][1]]),dim=1))
            x = self.vertex_net(x)





        # self.clf.temp.memory.append(lh.get_gpu_memory(self.clf.temp.device))

        return x, x_edge

    #######################################################
    ################## INFERENCE FORWARDS #################
    #######################################################
    def inference_batch_layer(self, x_all, batch_loader):
        # produces embeddings layer by layer, batch per batch
        # subgraph sampling necessary
        # needed when full graph does not fit in memory


        if (self.clf.training.loss == "kl"):
            x_out = torch.zeros([x_all.size()[0], 2], dtype=torch.float32, requires_grad=False).to('cpu')
        elif (self.clf.training.loss == "mse"):
            x_out = torch.zeros([x_all.size()[0], 1], dtype=torch.float32, requires_grad=False).to('cpu')
        else:
            print("{} is not a valid loss. choose either kl or mse".format(self.clf.training.loss))
            sys.exit(1)


        if (self.clf.features.normalization_feature and not self.clf.features.keep_normalization_feature):
            x_all = x_all[:, 1:]  # do not put it on gpu yet, because they will be used batch by batch

        # Compute representations of nodes layer by layer, using *all*
        # available edges. This leads to faster computation in contrast to
        # immediately computing the final representations of each batch.
        for batch_size, n_id, adjs in tqdm(batch_loader, ncols=50):
            x = x_all[n_id, :].to(self.clf.temp.device)
            for i in range(self.num_layers):
                edge_index, _, size = adjs[i]
                x = self.sage_convs[i]((x, x[:size[1]]), edge_index.to(self.clf.temp.device))
                if i != self.num_layers - 1:
                    x = F.relu(x)
                # torch.cuda.empty_cache()
                # self.clf.temp.memory.append(lh.get_gpu_memory(self.clf.temp.device))
            x_out[n_id[:batch_size]] = x.to('cpu')


        if (self.clf.training.model_name == "sage+"):
            x = F.relu(x_out.to(self.clf.temp.device))
            x = self.vertex_net(x)
        else:
            x = x_out.to(self.clf.temp.device)

        self.clf.temp.memory.append(lh.get_gpu_memory(self.clf.temp.device))
        return x



    def inference_layer_batch(self, x_all, batch_loader):
        # produces embeddings layer by layer, batch per batch
        # subgraph sampling necessary
        # needed when full graph does not fit in memory

        if (self.clf.features.normalization_feature and not self.clf.features.keep_normalization_feature):
            x_all = x_all[:, 1:]  # do not put it on gpu yet, because they will be used batch by batch

        # Compute representations of nodes layer by layer, using *all*
        # available edges. This leads to faster computation in contrast to
        # immediately computing the final representations of each batch.
        for i in tqdm(range(self.num_layers), ncols=50):
            xs = []
            for batch_size, n_id, adj in batch_loader:
                edge_index, _, size = adj  # get adjacencies of current layer / hop
                x = x_all[n_id].to(self.clf.temp.device)
                x_target = x[:size[1]]
                x = self.sage_convs[i]((x, x_target), edge_index.to(self.clf.temp.device))
                if i != self.num_layers - 1:
                    x = F.relu(x)
                # torch.cuda.empty_cache()
                # self.clf.temp.memory.append(lh.get_gpu_memory(self.clf.temp.device))
                xs.append(x.cpu())

            x_all = torch.cat(xs, dim=0)

        if (self.clf.training.model_name == "sage+"):
            x = F.relu(x_all.to(self.clf.temp.device))
            x = self.vertex_net(x)
        else:
            x = x_all.to(self.clf.temp.device)

        self.clf.temp.memory.append(lh.get_gpu_memory(self.clf.temp.device))
        return x


    def inference_layer(self, data_all):
        # produces embeddings layer by layer directly for the whole graph
        # no subgraph sampling necessary
        # only works if full graph fits in memory


        if(self.clf.features.normalization_feature and not self.clf.features.keep_normalization_feature):
            x=data_all.x[:, 1:].to(self.clf.temp.device) # put all on gpu, because they will all be used directly
        else:
            x=data_all.x[:, :].to(self.clf.temp.device)

        edge_index = data_all.edge_index.to(torch.long).to(self.clf.temp.device) # put all on gpu, because they will all be used directly
        # Compute representations of nodes layer by layer, using *all*
        # available edges. This leads to faster computation in contrast to
        # immediately computing the final representations of each batch.
        for i in range(self.num_layers):
            x = self.sage_convs[i]((x, x), edge_index)
            if i != self.num_layers - 1:
                x = F.relu(x)
            torch.cuda.empty_cache()

            # todo run on cnes with 36gb memory to achieve faster gpu time

        if(self.clf.training.model_name == "sage+"):
            x = F.relu(x)
            x = self.vertex_net(x)

        self.clf.temp.memory.append(lh.get_gpu_memory(self.clf.temp.device))

        return x
