# Pytorch 

import torch

import torch.nn as nn

import torch.nn.functional as F

from torch_geometric.nn import BatchNorm, LayerNorm, GATv2Conv

# Pytorch Lightning

import pytorch_lightning as pl

from pytorch_lightning.loggers import WandbLogger

# General

import numpy as np

import math

import tqdm

import time

import wandb

# Own

from utils.pretrain_utils import sample_node_for_et, get_batched_data, get_edges, calc_metrics, plot_roc_curve, metrics_per_rel

from decoders import bilinear, trans, dot

# Global variables

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class NodeEmbeder(pl.LightningModule):

    def __init__(self, all_data, edge_attr_dict, hp_dict=None, num_nodes=None, combined_training=False, spl_mat=[]):

        super().__init__()

        # save hyperparameters

        self.save_hyperparameters("hp_dict", ignore=["spl_mat"])

        # Data

        self.all_data = all_data

        self.edge_attr_dict = edge_attr_dict

        # Model parameters

        self.lr = self.hparams.hp_dict['lr']

        self.lr_factor = self.hparams.hp_dict['lr_factor']

        self.lr_patience = self.hparams.hp_dict['lr_patience']

        self.lr_threshold = self.hparams.hp_dict['lr_threshold']

        self.lr_threshold_mode = self.hparams.hp_dict['lr_threshold_mode']

        self.lr_cooldown = self.hparams.hp_dict['lr_cooldown']

        self.min_lr = self.hparams.hp_dict['min_lr']

        self.eps = self.hparams.hp_dict['eps']

        self.wd = self.hparams.hp_dict['wd']

        self.decoder_type = self.hparams.hp_dict['decoder_type']

        self.pred_threshold = self.hparams.hp_dict['pred_threshold']

        self.use_spl = None

        self.spl_mat = []

        self.spl_dim = 0

        self.nfeat = self.hparams.hp_dict['nfeat']

        self.nhid1 = self.hparams.hp_dict['hidden'] * 2

        self.nhid2 = self.hparams.hp_dict['hidden']

        self.output = self.hparams.hp_dict['output']

        self.node_emb = nn.Embedding(num_nodes, self.nfeat)

        self.num_nodes = num_nodes 

        self.num_relations = len(edge_attr_dict)

        self.n_heads = self.hparams.hp_dict['n_heads']

        self.dropout = self.hparams.hp_dict['dropout']

        self.norm_method = self.hparams.hp_dict['norm_method']

        # Select decoder

        if self.decoder_type == "bilinear": self.decoder = bilinear

        elif self.decoder_type == "trans": self.decoder = trans

        elif self.decoder_type == "dot": self.decoder = dot

        self.n_layers = 3

        self.loss_type = self.hparams.hp_dict['loss']

        self.combined_training = combined_training

        # Conv layers

        self.convs = torch.nn.ModuleList()

        self.convs.append(GATv2Conv(self.nfeat, self.nhid1, self.n_heads)) # input = nfeat, output = nhid1*n_heads

        if self.n_layers == 3:

            self.convs.append(GATv2Conv(self.nhid1*self.n_heads, self.nhid2, self.n_heads)) # input = nhid1*n_heads, output = nhid2*n_heads

            self.convs.append(GATv2Conv(self.nhid2*self.n_heads, self.output, self.n_heads)) # input = nhid2*n_heads, output = output*n_heads

        else:

            self.convs.append(GATv2Conv(self.nhid1*self.n_heads, self.output, self.n_heads)) # input = nhid2*n_heads, output = output*n_heads

        # Relation learnable weights

        self.relation_weights = nn.Parameter(torch.Tensor(self.num_relations, self.output * self.n_heads))

        # Normalization (applied after a single conv layer)

        if self.norm_method == "batch":

            self.norms = torch.nn.ModuleList()

            self.norms.append(BatchNorm(self.nhid1*self.n_heads))

            self.norms.append(BatchNorm(self.nhid2*self.n_heads))

        elif self.norm_method == "layer":

            self.norms = torch.nn.ModuleList()

            self.norms.append(LayerNorm(self.nhid1*self.n_heads))

            self.norms.append(LayerNorm(self.nhid2*self.n_heads))

        elif self.norm_method == "batch_layer":

            self.batch_norms = torch.nn.ModuleList()

            self.batch_norms.append(BatchNorm(self.nhid1*self.n_heads))

            if self.n_layers == 3: self.batch_norms.append(BatchNorm(self.nhid2*self.n_heads))

            self.layer_norms = torch.nn.ModuleList()

            self.layer_norms.append(LayerNorm(self.nhid1*self.n_heads))

            if self.n_layers == 3: self.layer_norms.append(LayerNorm(self.nhid2*self.n_heads))

        self.reset_parameters()

    def reset_parameters(self):

        for conv in self.convs:

            conv.reset_parameters()

        nn.init.xavier_uniform_(self.relation_weights, gain = nn.init.calculate_gain('leaky_relu'))

    def forward(self, n_ids, adjs): 

        x = self.node_emb(n_ids)

        gat_attn = []

        assert len(adjs) == self.n_layers

        for i, (edge_index, _, edge_type, size) in enumerate(adjs):

            # Update node embeddings

            x_target = x[:size[1]]  # Target nodes are always placed first. 

            x, (edge_i, alpha) = self.convs[i]((x, x_target), edge_index, return_attention_weights=True)

            edge_i = edge_i.detach().cpu()

            alpha = alpha.detach().cpu()

            edge_i[0,:] = n_ids[edge_i[0,:]]

            edge_i[1,:] = n_ids[edge_i[1,:]]

            gat_attn.append((edge_i, alpha))

            # Normalize

            if i != self.n_layers - 1:

                if self.norm_method in ["batch", "layer"]:

                    x = self.norms[i](x)

                elif self.norm_method == "batch_layer":

                    x = self.layer_norms[i](x)

                x = F.leaky_relu(x)

                if self.norm_method == "batch_layer":

                    x = self.batch_norms[i](x)

                x = F.dropout(x, p=self.dropout, training=self.training)

        return x, gat_attn

    def get_negative_target_nodes(self, data, pos_target_embeds, curr_pos_target_embeds, all_edge_types):

        if self.hparams.hp_dict['negative_sampler_approach'] == 'all':

            # get negative targets by shuffling positive targets

            if 'index_to_node_features_pos' in data:

                rand_index = torch.randperm(data.index_to_node_features_pos.size(0))

            else:

                rand_index = torch.randperm(curr_pos_target_embeds.size(0))

        elif self.hparams.hp_dict['negative_sampler_approach'] == 'by_edge_type':

            # get negative targets by shuffling positive targets within each edge type

            et_ids, et_counts = all_edge_types.unique(return_counts=True)

            targets_dict = self.create_target_dict(all_edge_types, et_ids) # indices into all_edge_types for each edge type

            rand_index = torch.tensor(np.vectorize(sample_node_for_et)(all_edge_types.cpu(), targets_dict)).to(device)

        if 'index_to_node_features_pos' in data:

            index_to_node_features_neg = data.index_to_node_features_pos[rand_index] #NOTE: currently possible to get the same node as positive & negative target

            curr_neg_target_embeds = pos_target_embeds[index_to_node_features_neg,:]

        else:

            curr_neg_target_embeds = curr_pos_target_embeds[rand_index,:]

        return curr_neg_target_embeds

    def create_target_dict(self, all_edge_types, et_ids):

        targets_dict = {}

        for k in et_ids:

            indices = (all_edge_types == int(k)).nonzero().cpu()

            targets_dict[int(k)] = indices

        return targets_dict

    def decode(self, data, source_embeds, pos_target_embeds, all_edge_types): 

        curr_source_embeds = source_embeds[data.index_to_node_features_pos,:]

        curr_pos_target_embeds = pos_target_embeds[data.index_to_node_features_pos,:]

        ts = time.time()

        curr_neg_target_embeds = self.get_negative_target_nodes(data, pos_target_embeds, curr_pos_target_embeds, all_edge_types)

        te = time.time()

        if self.hparams.hp_dict['time']:

            print(f"Negative sampling took {te - ts:0.4f} seconds")

        # Get source & targets for pos & negative edges

        source = torch.cat([curr_source_embeds, curr_source_embeds])

        target = torch.cat([curr_pos_target_embeds, curr_neg_target_embeds])

        all_edge_types = torch.cat([all_edge_types, all_edge_types])

        data.all_edge_types = all_edge_types

        if self.decoder_type == "dot": 

            return data, self.decoder(source, target)

        else: 

            relation = self.relation_weights[all_edge_types]

            return data, self.decoder(source, relation, target)

    def get_predictions(self, data, embed):

        # Apply decoder

        source_embed, target_embed = embed.split(embed.size(0) // 2, dim=0)

        data, raw_pred = self.decode(data, source_embed, target_embed, data.pos_edge_types)

        # Apply activation

        if self.loss_type != "max-margin": 

            pred = torch.sigmoid(raw_pred)

        else: 

            pred = torch.tanh(raw_pred)

        return data, raw_pred, pred

    def get_link_labels(self, edge_types):

        num_links = edge_types.size(0) 

        link_labels = torch.zeros(num_links, dtype=torch.float, device=edge_types.device)

        link_labels[:(int(num_links/2))] = 1.

        link_labels[(int(num_links/2)):] = 0.

        return link_labels

    def _step(self, data, dataset_type):

        if not self.combined_training: 

            ts = time.time()

            data = get_batched_data(data, self.all_data) 

            tm = time.time()

            data = get_edges(data, self.all_data, dataset_type)

            te=time.time()

        data = data.to(device)

        # Get predictions

        t0 = time.time()

        out, gat_attn = self.forward(data.n_id, data.adjs) 

        t1 = time.time()

        data, raw_pred, pred = self.get_predictions(data, out)

        t2 = time.time()

        # Calculate loss

        link_labels = self.get_link_labels(data.all_edge_types)

        loss = self.calc_loss(pred, link_labels)

        t3 = time.time()

        # Calculate metrics

        if self.loss_type == "max-margin":

            metric_pred = torch.sigmoid(raw_pred)

            self.logger.experiment.log({f'{dataset_type}/node_predicted_probs': wandb.Histogram(metric_pred.cpu().detach().numpy())})

        else: metric_pred = pred

        roc_score, ap_score, acc, f1 = calc_metrics(metric_pred.cpu().detach().numpy(), link_labels.cpu().detach().numpy(), self.pred_threshold)

        self.logger.experiment.log({f'{dataset_type}/node_roc_curve': plot_roc_curve(metric_pred.cpu().detach().numpy(), link_labels.cpu().detach().numpy())})

        t4 = time.time()

        if self.hparams.hp_dict['time']:

            print(f'It took {tm-ts:0.2f}s to get batched data, {te-tm:0.2f}s to get edges, {t1-t0:0.2f}s to complete forward pass, {t2-t1:0.2f}s to decode, {t3-t2:0.2f}s to calc loss, and {t4-t3:0.2f}s to calc other metrics.')

        return data, loss, pred, link_labels, roc_score, ap_score, acc, f1

    def training_step(self, data, data_idx):

        data, loss, pred, link_labels, roc_score, ap_score, acc, f1 = self._step(data, 'train')

        logs = {"train/node_batch_loss": loss.detach(), 

                "train/node_roc": roc_score, 

                "train/node_ap": ap_score, 

                "train/node_acc": acc, 

                "train/node_f1": f1

               }

        rel_logs = metrics_per_rel(pred, link_labels, self.edge_attr_dict, data.all_edge_types, "train", self.pred_threshold)

        logs.update(rel_logs)

        self._logger(logs)

        return {'loss': loss, 'logs': logs}

    def training_epoch_end(self, outputs):     

        roc_train = []

        ap_train = []

        acc_train = []

        f1_train = []

        total_train_loss = []

        for batch_log in outputs:

            roc_train.append(batch_log['logs']["train/node_roc"])

            ap_train.append(batch_log['logs']["train/node_ap"])

            acc_train.append(batch_log['logs']["train/node_acc"])

            f1_train.append(batch_log['logs']["train/node_f1"])

            total_train_loss.append(batch_log['logs']["train/node_batch_loss"])

        self._logger({"train/node_total_loss": torch.mean(torch.Tensor(total_train_loss)), 

                      "train/node_total_roc": np.mean(roc_train), 

                      "train/node_total_ap": np.mean(ap_train), 

                      "train/node_total_acc": np.mean(acc_train), 

                      "train/node_total_f1": np.mean(f1_train)})

        self._logger({'node_curr_epoch': self.current_epoch})

    def validation_step(self, data, data_idx):

        data, loss, pred, link_labels, roc_score, ap_score, acc, f1 = self._step(data, 'val')

        logs = {"val/node_batch_loss": loss.detach().cpu(), 

                "val/node_roc": roc_score, 

                "val/node_ap": ap_score, 

                "val/node_acc": acc, 

                "val/node_f1": f1

               }

        rel_logs = metrics_per_rel(pred, link_labels, self.edge_attr_dict, data.all_edge_types, "val", self.pred_threshold)

        logs.update(rel_logs)

        self._logger(logs)

        return logs

    def validation_epoch_end(self, outputs):

        roc_val = []

        ap_val = []

        acc_val = []

        f1_val = []

        total_val_loss = []

        for batch_log in outputs:

            roc_val.append(batch_log["val/node_roc"])

            ap_val.append(batch_log["val/node_ap"])

            acc_val.append(batch_log["val/node_acc"])

            f1_val.append(batch_log["val/node_f1"])

            total_val_loss.append(batch_log["val/node_batch_loss"])

        self._logger({"val/node_total_loss": torch.mean(torch.Tensor(total_val_loss)), 

                      "val/node_total_roc": np.mean(roc_val), 

                      "val/node_total_ap": np.mean(ap_val), 

                      "val/node_total_acc": np.mean(acc_val), 

                      "val/node_total_f1": np.mean(f1_val)})

        self._logger({'node_curr_epoch': self.current_epoch})

    def test_step(self, data, data_idx):

        data, loss, pred, link_labels, roc_score, ap_score, acc, f1 = self._step(data, 'test')

        logs = {"test/node_batch_loss": loss.detach().cpu(), 

                "test/node_roc": roc_score, 

                "test/node_ap": ap_score, 

                "test/node_acc": acc, 

                "test/node_f1": f1

               }

        rel_logs = metrics_per_rel(pred, link_labels, self.edge_attr_dict, data.all_edge_types, "test", self.pred_threshold)

        logs.update(rel_logs)

        self._logger(logs)

        return logs

    def test_epoch_end(self, outputs):

        roc = []

        ap = []

        acc = []

        f1 = []

        for batch_log in outputs:

            roc.append(batch_log["test/node_roc"])

            ap.append(batch_log["test/node_ap"])

            acc.append(batch_log["test/node_acc"])

            f1.append(batch_log["test/node_f1"])

        self._logger({"test/node_total_roc": np.mean(roc), 

                      "test/node_total_ap": np.mean(ap), 

                      "test/node_total_acc": np.mean(acc), 

                      "test/node_total_f1": np.mean(f1)})

        self._logger({'node_curr_epoch': self.current_epoch})

    def predict(self, data):

        n_id = torch.arange(self.node_emb.weight.shape[0], device=self.device)

        x = self.node_emb(n_id)

        gat_attn = []

        for i in range(len(self.convs)):

            # Update node embeddings

            x, (edge_i, alpha) = self.convs[i](x, data.edge_index.to(self.device), return_attention_weights=True) #

            edge_i = edge_i.detach().cpu()

            alpha = alpha.detach().cpu()

            edge_i[0,:] = n_id[edge_i[0,:]]

            edge_i[1,:] = n_id[edge_i[1,:]]

            gat_attn.append((edge_i, alpha))

            # Normalize

            if i != self.n_layers - 1:

                if self.norm_method in ["batch", "layer"]:

                    x = self.norms[i](x)

                elif self.norm_method == "batch_layer":

                    x = self.layer_norms[i](x)

                x = F.leaky_relu(x)

                if self.norm_method == "batch_layer":

                    x = self.batch_norms[i](x)

        assert x.shape[0] == self.node_emb.weight.shape[0]

        return x, gat_attn

    def predict_step(self, data, data_idx):

        x, gat_attn = self.predict(data)

        return x, gat_attn

    def configure_optimizers(self):

        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr, weight_decay = self.wd)

        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=self.lr_factor, patience=self.lr_patience, threshold=self.lr_threshold, threshold_mode=self.lr_threshold_mode, cooldown=self.lr_cooldown, min_lr=self.min_lr, eps=self.eps)

        return {

                "optimizer": optimizer,

                "lr_scheduler": 

                    {

                    "scheduler": scheduler,

                    "monitor": "val/node_total_loss",

                    'name': 'curr_lr'

                    },

                }

    def _logger(self, logs):

        for k, v in logs.items():

            self.log(k, v)

    def calc_loss(self, pred, y):

        if self.loss_type == "BCE":

            loss = F.binary_cross_entropy(pred, y, reduction='none')

            norm_loss = torch.mean(loss)

        elif self.loss_type == "max-margin": 

            loss = ((1 - (pred[y == 1] - pred[y != 1])).clamp(min=0).mean())

            norm_loss = loss

        return norm_loss