#pytorch lightning

import pytorch_lightning as pl

from pytorch_lightning.loggers import WandbLogger

# torch

from torch import nn

import torch

import torch.nn.functional as F

import numpy as np

from scipy.stats import rankdata

import pandas as pd

from pathlib import Path

import time

import wandb

import sys

sys.path.insert(0, '..') # add project_config to path

from node_embedder_model import NodeEmbeder

from task_heads.gp_aligner import GPAligner

import project_config

# import utils

from utils.pretrain_utils import get_edges, calc_metrics

from utils.loss_utils import MultisimilarityCriterion

from utils.train_utils import mean_reciprocal_rank, top_k_acc, average_rank

from utils.train_utils import fit_umap, mrr_vs_percent_overlap, plot_gene_rank_vs_x_intrain, plot_gene_rank_vs_hops, plot_degree_vs_attention, plot_nhops_to_gene_vs_attention, plot_gene_rank_vs_fraction_phenotype, plot_gene_rank_vs_numtrain, plot_gene_rank_vs_trainset

from utils.train_utils import weighted_sum

class CombinedGPAligner(pl.LightningModule):

    def __init__(self, edge_attr_dict, all_data, n_nodes=None, node_ckpt=None, hparams=None, node_hparams=None,  spl_pca=[], spl_gate=[]):

        super().__init__()

        print('Initializing Model')

        self.save_hyperparameters('hparams', ignore=["spl_pca", "spl_gate"]) # spl_pca and spl_gate never get used

        #print('Saved combined model hyperparameters: ', self.hparams)

        self.all_data = all_data

        self.all_train_nodes = {}

        self.train_patient_nodes = {}

        self.train_sparse_nodes = {}

        self.train_target_batch = {}

        self.train_corr_gene_nid = {}

        print(f"Loading Node Embedder from {node_ckpt}")

        # NOTE: loads in saved hyperparameters

        self.node_model = NodeEmbeder.load_from_checkpoint(checkpoint_path=node_ckpt, 

                                                           all_data=all_data,

                                                           edge_attr_dict=edge_attr_dict, 

                                                           num_nodes=n_nodes)

        self.patient_model = self.get_patient_model()

        print('End Patient Model Initialization')

    def get_patient_model(self):

        # NOTE: this will only work with GATv2Conv

        model = GPAligner(self.hparams.hparams, embed_dim=self.node_model.hparams.hp_dict['output']*self.node_model.hparams.hp_dict['n_heads'])

        return model

    def forward(self, batch, step_type):

        # Node Embedder

        t0 = time.time()

        print(len(batch.adjs))

        outputs, gat_attn = self.node_model.forward(batch.n_id, batch.adjs)

        pad_outputs = torch.cat([torch.zeros(1, outputs.size(1), device=outputs.device), outputs])

        t1 = time.time()

        # get masks

        phenotype_mask = (batch.batch_pheno_nid != 0)

        gene_mask = (batch.batch_cand_gene_nid != 0)

        # index into outputs using phenotype & gene batch node idx

        batch_sz, max_n_phen = batch.batch_pheno_nid.shape

        phenotype_embeddings = torch.index_select(pad_outputs, 0, batch.batch_pheno_nid.view(-1)).view(batch_sz, max_n_phen, -1)

        batch_sz, max_n_cand_genes = batch.batch_cand_gene_nid.shape

        cand_gene_embeddings = torch.index_select(pad_outputs, 0, batch.batch_cand_gene_nid.view(-1)).view(batch_sz, max_n_cand_genes, -1)

        if self.hparams.hparams['augment_genes']:            

            print("Augmenting genes...", self.hparams.hparams['aug_gene_w'])

            _, max_n_sim_cand_genes, k_sim_genes = batch.batch_sim_gene_nid.shape

            sim_gene_embeddings = torch.index_select(pad_outputs, 0, batch.batch_sim_gene_nid.view(-1)).view(batch_sz, max_n_sim_cand_genes, self.hparams.hparams['n_sim_genes'], -1)

            agg_sim_gene_embedding = weighted_sum(sim_gene_embeddings, batch.batch_sim_gene_sims)

            if self.hparams.hparams['aug_gene_by_deg']:

                print("Augmenting gene by degree...")

                aug_gene_w = self.hparams.hparams['aug_gene_w'] * torch.exp(-self.hparams.hparams['aug_gene_w'] * batch.batch_cand_gene_degs) + (1 - self.hparams.hparams['aug_gene_w'] - 0.1)

                aug_gene_w = (aug_gene_w * (torch.sum(batch.batch_sim_gene_sims, dim = -1) > 0)).unsqueeze(-1)

            else:

                aug_gene_w = (self.hparams.hparams['aug_gene_w'] * (torch.sum(batch.batch_sim_gene_sims, dim = -1) > 0)).unsqueeze(-1)

            cand_gene_embeddings = torch.mul(1 - aug_gene_w, cand_gene_embeddings) + torch.mul(aug_gene_w, agg_sim_gene_embedding)

        # Patient Embedder with or without disease information

        if self.hparams.hparams['use_diseases']: 

            disease_mask = (batch.batch_disease_nid != 0)

            batch_sz, max_n_dx = batch.batch_disease_nid.shape

            disease_embeddings = torch.index_select(pad_outputs, 0, batch.batch_disease_nid.view(-1)).view(batch_sz, max_n_dx, -1)

            t2 = time.time()

            phenotype_embedding, candidate_gene_embeddings, disease_embeddings, gene_mask, phenotype_mask, disease_mask, attn_weights = self.patient_model.forward(phenotype_embeddings, cand_gene_embeddings, disease_embeddings, phenotype_mask, gene_mask, disease_mask)

            t3 = time.time()

        else:

            t2 = time.time()

            phenotype_embedding, candidate_gene_embeddings, disease_embeddings, gene_mask, phenotype_mask, disease_mask, attn_weights = self.patient_model.forward(phenotype_embeddings, cand_gene_embeddings, phenotype_mask=phenotype_mask, gene_mask=gene_mask)

            t3 = time.time()

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

            print(f'It takes {t1-t0:0.4f}s for the node model, {t2-t1:0.4f}s for indexing into the output, and {t3-t2:0.4f}s for the patient model forward.')

        return outputs, gat_attn, phenotype_embedding, candidate_gene_embeddings, disease_embeddings, gene_mask, phenotype_mask, disease_mask, attn_weights

    def _step(self, batch, step_type):

        t0 = time.time()

        if step_type != 'test':

            batch = get_edges(batch, self.all_data, step_type)

        t1 = time.time()

        # Forward pass

        node_embeddings, gat_attn, phenotype_embedding, candidate_gene_embeddings, disease_embeddings, gene_mask, phenotype_mask, disease_mask, attn_weights = self.forward(batch, step_type)

        t2 = time.time()

        # Calculate similarities between patient phenotypes & candidate genes/diseases

        alpha = self.hparams.hparams['alpha']

        use_candidate_list = True if step_type != 'train' else False

        cand_gene_to_phenotypes_spl = batch.batch_cand_gene_to_phenotypes_spl if use_candidate_list else batch.batch_concat_cand_gene_to_phenotypes_spl

        disease_nid = batch.batch_disease_nid if self.hparams.hparams['use_diseases'] else None

        # calculate similarity between phen & genes for all genes in manual candidate list

        phen_gene_sims, raw_phen_gene_sims, phen_gene_mask, phen_gene_one_hot_labels = self.patient_model._calc_similarity(phenotype_embedding, candidate_gene_embeddings, None, batch.batch_cand_gene_nid, batch.batch_corr_gene_nid, disease_nid, batch.one_hot_labels, gene_mask, phenotype_mask, disease_mask, True, batch.batch_cand_gene_to_phenotypes_spl, alpha) 

        # calculate similarity for loss function  

        if self.hparams.hparams['loss'] == 'gene_multisimilarity' and use_candidate_list: # in this case, the similarities are the same

            sims = phen_gene_sims

            mask = phen_gene_mask

            one_hot_labels = phen_gene_one_hot_labels

        else:

            if self.hparams.hparams['loss'] == 'disease_multisimilarity': candidate_gene_embeddings = None

            elif self.hparams.hparams['loss'] == 'gene_multisimilarity': disease_embeddings = None

            sims, raw_sims, mask, one_hot_labels = self.patient_model._calc_similarity(phenotype_embedding, candidate_gene_embeddings, disease_embeddings, batch.batch_cand_gene_nid, batch.batch_corr_gene_nid, disease_nid, batch.one_hot_labels, gene_mask, phenotype_mask, disease_mask, use_candidate_list, cand_gene_to_phenotypes_spl, alpha)

        ## Rank genes

        correct_gene_ranks, phen_gene_sims = self.patient_model._rank_genes(phen_gene_sims, phen_gene_mask, phen_gene_one_hot_labels)

        t3 = time.time()

        ## Calculate patient embedding loss

        loss = self.patient_model.calc_loss(sims, mask, one_hot_labels)

        t4 = time.time()

        ## Calculate node embedding loss

        if step_type == 'test':

            node_embedder_loss = 0

            roc_score, ap_score, acc, f1 = 0,0,0,0

        else:

            # Get link predictions

            batch, raw_pred, pred = self.node_model.get_predictions(batch, node_embeddings)

            link_labels = self.node_model.get_link_labels(batch.all_edge_types)

            node_embedder_loss = self.node_model.calc_loss(pred, link_labels)

            # Calculate metrics

            metric_pred = torch.sigmoid(raw_pred)

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

        t5 = time.time()

        ## calc time

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

            print(f'It takes {t1-t0:0.4f}s to get edges, {t2-t1:0.4f}s for the forward pass,  {t3-t2:0.4f}s to rank genes, {t4-t3:0.4f}s to calc patient loss, and {t5-t4:0.4f}s to calc the node loss.')

        ## Plot gradients

        if self.hparams.hparams['plot_gradients']:

            for k, v in self.patient_model.state_dict().items():

                self.logger.experiment.log({f'gradients/{step_type}.gradients.%s' % k: wandb.Histogram(v.detach().cpu())})

        return node_embedder_loss, loss, correct_gene_ranks, roc_score, ap_score, acc, f1, node_embeddings, gat_attn, phenotype_embedding, candidate_gene_embeddings, attn_weights, phen_gene_sims, raw_phen_gene_sims, gene_mask, phenotype_mask

    def training_step(self, batch, batch_idx):

        print('training step')

        node_embedder_loss, patient_loss, correct_gene_ranks, roc_score, ap_score, acc, f1, node_embeddings, gat_attn, phenotype_embedding, candidate_gene_embeddings, attn_weights, phen_gene_sims, raw_phen_gene_sims, gene_mask, phenotype_mask = self._step(batch, 'train')

        loss = (self.hparams.hparams['lambda'] * node_embedder_loss) + ((1 - self.hparams.hparams['lambda']) *  patient_loss)

        self.log('train_loss/patient.train_overall_loss', loss, prog_bar=True, on_epoch=True)

        self.log('train_loss/patient.train_patient_loss', patient_loss, prog_bar=True, on_epoch=True)

        self.log('train_loss/patient.train_node_embedder_loss', node_embedder_loss, prog_bar=True, on_epoch=True)

        batch_sz, n_candidates, embed_dim = candidate_gene_embeddings.shape

        candidate_gene_embeddings_flattened = candidate_gene_embeddings.view(batch_sz*n_candidates, embed_dim)

        one_hot_labels_flattened = batch.one_hot_labels.view(batch_sz*n_candidates)

        return {'loss': loss, 

                'train/train_correct_gene_ranks': correct_gene_ranks, 

                "train/node.train_roc": roc_score, 

                "train/node.train_ap": ap_score, 

                "train/node.train_acc": acc, 

                "train/node.train_f1": f1, 

                'train/one_hot_labels': batch.one_hot_labels.detach().cpu(),

                'train/attention_weights': attn_weights.detach().cpu() if attn_weights != None else None,

                'train/phen_gene_sims': phen_gene_sims.detach().cpu(),

                'train/phenotype_names_degrees': batch.phenotype_names,

                }

    def validation_step(self, batch, batch_idx):

        node_embedder_loss, patient_loss, correct_gene_ranks, roc_score, ap_score, acc, f1, node_embeddings, gat_attn, phenotype_embedding, candidate_gene_embeddings, attn_weights, phen_gene_sims, raw_phen_gene_sims, gene_mask, phenotype_mask = self._step(batch, 'val')

        loss = (self.hparams.hparams['lambda'] * node_embedder_loss) + ((1 - self.hparams.hparams['lambda']) * patient_loss)

        self.log('val_loss/patient.val_overall_loss', loss, prog_bar=True, on_epoch=True)

        self.log('val_loss/patient.val_patient_loss', patient_loss, prog_bar=True)

        self.log('val_loss/patient.val_node_embedder_loss', node_embedder_loss, prog_bar=True)

        return {'loss/val_loss': loss, 

                'val/val_correct_gene_ranks': correct_gene_ranks, 

                "val/node.val_roc": roc_score, 

                "val/node.val_ap": ap_score, 

                "val/node.val_acc": acc, 

                "val/node.val_f1": f1, 

                'val/one_hot_labels': batch.one_hot_labels.detach().cpu(),

                'val/attention_weights': attn_weights.detach().cpu() if attn_weights != None else None,

                'val/phen_gene_sims': phen_gene_sims.detach().cpu(),

                'val/phenotype_names_degrees': batch.phenotype_names,

                }

    def write_results_to_file(self, batch_info, phen_gene_sims, gene_mask, phenotype_mask, attn_weights, correct_gene_ranks, gat_attn, node_embeddings, phenotype_embeddings, save=True, loop_type='predict'):

        # NOTE: only saves a single batch - to run at inference time, make sure batch size includes all patients

        # Save GAT attention weights

        #NOTE: assumes 3 layers to model

        attn_dfs = []

        layer = 0

        for edge_attn in gat_attn:

            edge_index, attn = edge_attn

            edge_index = edge_index.cpu()

            attn = attn.cpu()

            gat_attn_df = pd.DataFrame({'source': edge_index[0,:], 'target': edge_index[1,:]})

            for head in range(attn.shape[1]):

                gat_attn_df[f'attn_{head}'] =  attn[:,head]

            attn_dfs.append(gat_attn_df)

            layer += 1

        # Save scores

        all_sims, all_genes, all_patient_ids, all_labels = [], [], [], []

        for patient_id, sims, genes, g_mask in zip(batch_info["patient_ids"], phen_gene_sims, batch_info["cand_gene_names"], gene_mask):

            nonpadded_sims = sims[g_mask].tolist()

            all_sims.extend(nonpadded_sims)

            all_genes.extend(genes)

            all_patient_ids.extend([patient_id] * len(genes))

        results_df = pd.DataFrame({'patient_id': all_patient_ids, 'genes': all_genes, 'similarities': all_sims})

        # Save phenotype information

        if attn_weights is None:

            phen_df = None

        else:

            all_patient_ids, all_phens, all_attn_weights, all_degrees = [], [], [], []

            for patient_id, attn_w, phen_names, p_mask in zip(batch_info["patient_ids"], attn_weights, batch_info["phenotype_names"], phenotype_mask):

                p_names, degrees = zip(*phen_names)

                all_patient_ids.extend([patient_id] * len(phen_names))

                all_degrees.extend(degrees)

                all_phens.extend(p_names)

                all_attn_weights.extend(attn_w[p_mask].tolist())

            phen_df = pd.DataFrame({'patient_id': all_patient_ids, 'phenotypes': all_phens, 'degrees': all_degrees, 'attention':all_attn_weights})

            print(phen_df.head())

        return results_df, phen_df, attn_dfs, phenotype_embeddings.cpu(), None

    def test_step(self, batch, batch_idx):

        node_embedder_loss, patient_loss, correct_gene_ranks, roc_score, ap_score, acc, f1, node_embeddings, gat_attn, phenotype_embedding, candidate_gene_embeddings, attn_weights, phen_gene_sims, raw_phen_gene_sims, gene_mask, phenotype_mask = self._step(batch, 'test')

        return {'test/test_correct_gene_ranks': correct_gene_ranks, 

                'test/node.embed': node_embeddings.detach().cpu(), 

                'test/patient.phenotype_embed': phenotype_embedding.detach().cpu(),

                'test/one_hot_labels': batch.one_hot_labels.detach().cpu(), #one_hot_labels_flattened.detach().cpu(),

                'test/attention_weights': attn_weights.detach().cpu() if attn_weights != None else None,

                'test/phen_gene_sims': phen_gene_sims.detach().cpu(),

                'test/phenotype_names_degrees': batch.phenotype_names, # type = list

                'test/gene_mask': gene_mask.detach().cpu(),

                'test/phenotype_mask': phenotype_mask.detach().cpu(),

                "test/patient_ids": batch.patient_ids, # type = list

                "test/cand_gene_names": batch.cand_gene_names, # type = list

                'test/gat_attn': gat_attn, # type = list

                "test/n_id": batch.n_id[:batch.batch_size].detach().cpu(),

                }

    def inference(self, batch, batch_idx):

        outputs, gat_attn = self.node_model.predict(self.all_data)

        pad_outputs = torch.cat([torch.zeros(1, outputs.size(1), device=outputs.device), outputs])

        t1 = time.time()

        # get masks

        phenotype_mask = (batch.batch_pheno_nid != 0)

        gene_mask = (batch.batch_cand_gene_nid != 0)

        # index into outputs using phenotype & gene batch node idx

        batch_sz, max_n_phen = batch.batch_pheno_nid.shape

        phenotype_embeddings = torch.index_select(pad_outputs, 0, batch.batch_pheno_nid.view(-1)).view(batch_sz, max_n_phen, -1)

        batch_sz, max_n_cand_genes = batch.batch_cand_gene_nid.shape

        cand_gene_embeddings = torch.index_select(pad_outputs, 0, batch.batch_cand_gene_nid.view(-1)).view(batch_sz, max_n_cand_genes, -1)

        if self.hparams.hparams['augment_genes']:            

            print("Augmenting genes at inference...", self.hparams.hparams['aug_gene_w'])

            _, max_n_sim_cand_genes, k_sim_genes = batch.batch_sim_gene_nid.shape

            sim_gene_embeddings = torch.index_select(pad_outputs, 0, batch.batch_sim_gene_nid.view(-1)).view(batch_sz, max_n_sim_cand_genes, self.hparams.hparams['n_sim_genes'], -1)

            agg_sim_gene_embedding = weighted_sum(sim_gene_embeddings, batch.batch_sim_gene_sims)        

            aug_gene_w = (self.hparams.hparams['aug_gene_w'] * (torch.sum(batch.batch_sim_gene_sims, dim = -1) > 0)).unsqueeze(-1)

            cand_gene_embeddings = torch.mul(1 - aug_gene_w, cand_gene_embeddings) + torch.mul(aug_gene_w, agg_sim_gene_embedding)

        # Patient Embedder with or without disease information

        if self.hparams.hparams['use_diseases']: 

            disease_mask = (batch.batch_disease_nid != 0)

            batch_sz, max_n_dx = batch.batch_disease_nid.shape

            disease_embeddings = torch.index_select(pad_outputs, 0, batch.batch_disease_nid.view(-1)).view(batch_sz, max_n_dx, -1)

            t2 = time.time()

            phenotype_embedding, candidate_gene_embeddings, disease_embeddings, gene_mask, phenotype_mask, disease_mask, attn_weights = self.patient_model.forward(phenotype_embeddings, cand_gene_embeddings, disease_embeddings, phenotype_mask, gene_mask, disease_mask)

            t3 = time.time()

        else:

            t2 = time.time()

            phenotype_embedding, candidate_gene_embeddings, disease_embeddings, gene_mask, phenotype_mask, disease_mask, attn_weights = self.patient_model.forward(phenotype_embeddings, cand_gene_embeddings, phenotype_mask=phenotype_mask, gene_mask=gene_mask)

            t3 = time.time()

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

            print(f'It takes {t1-t0:0.4f}s for the node model, {t2-t1:0.4f}s for indexing into the output, and {t3-t2:0.4f}s for the patient model forward.')

        return outputs, gat_attn, phenotype_embedding, candidate_gene_embeddings, disease_embeddings, gene_mask, phenotype_mask, disease_mask, attn_weights

    def predict_step(self, batch, batch_idx):

        node_embeddings, gat_attn, phenotype_embedding, candidate_gene_embeddings, disease_embeddings, gene_mask, phenotype_mask, disease_mask, attn_weights = self.inference(batch, batch_idx)

        # Calculate similarities between patient phenotypes & candidate genes/diseases

        alpha = self.hparams.hparams['alpha']

        use_candidate_list = True

        disease_nid = batch.batch_disease_nid if self.hparams.hparams['use_diseases'] else None

        # calculate similarity between phen & genes for all genes in manual candidate list

        phen_gene_sims, raw_phen_gene_sims, phen_gene_mask, phen_gene_one_hot_labels = self.patient_model._calc_similarity(phenotype_embedding, candidate_gene_embeddings, None, batch.batch_cand_gene_nid, batch.batch_corr_gene_nid, disease_nid, batch.one_hot_labels, gene_mask, phenotype_mask, disease_mask, True, batch.batch_cand_gene_to_phenotypes_spl, alpha) 

        # Rank genes

        correct_gene_ranks, phen_gene_sims = self.patient_model._rank_genes(phen_gene_sims, phen_gene_mask, phen_gene_one_hot_labels)

        results_df, phen_df, attn_dfs, phenotype_embeddings, disease_embeddings = self.write_results_to_file(batch, phen_gene_sims, gene_mask, phenotype_mask, attn_weights, correct_gene_ranks, gat_attn, node_embeddings, phenotype_embedding, save=True)

        return results_df, phen_df, *attn_dfs, phenotype_embeddings, disease_embeddings

    def _epoch_end(self, outputs, loop_type):

        correct_gene_ranks = torch.cat([x[f'{loop_type}/{loop_type}_correct_gene_ranks'] for x in outputs], dim=0)

        if loop_type == "test":

            batch_info = {"n_id": torch.cat([x[f'{loop_type}/n_id'] for x in outputs], dim=0),

                          "patient_ids": [pat for x in outputs for pat in x[f'{loop_type}/patient_ids']],

                          "phenotype_names": [pat for x in outputs for pat in x[f'{loop_type}/phenotype_names_degrees']],

                          "cand_gene_names": [pat for x in outputs for pat in x[f'{loop_type}/cand_gene_names']],

                          "one_hot_labels": [pat for x in outputs for pat in x[f'{loop_type}/one_hot_labels']],

                          }

            phen_gene_sims = [pat for x in outputs for pat in x[f'{loop_type}/phen_gene_sims']] 

            gene_mask = [pat for x in outputs for pat in x[f'{loop_type}/gene_mask']] 

            phenotype_mask = [pat for x in outputs for pat in x[f'{loop_type}/phenotype_mask']] 

            attn_weights = [pat for x in outputs for pat in x[f'{loop_type}/attention_weights']] 

            gat_attn = [pat for x in outputs for pat in x[f'{loop_type}/gat_attn']] 

            node_embeddings = torch.cat([x[f'{loop_type}/node.embed'] for x in outputs], dim=0)

            phenotype_embedding = torch.cat([x[f'{loop_type}/patient.phenotype_embed'] for x in outputs], dim=0)

            results_df, phen_df, attn_dfs, phenotype_embeddings, disease_embeddings = self.write_results_to_file(batch_info, phen_gene_sims, gene_mask, phenotype_mask, attn_weights, correct_gene_ranks, gat_attn, node_embeddings, phenotype_embedding, loop_type=loop_type)

            print("Writing results for test...")

            output_base = "/home/ml499/public_repos/SHEPHERD/shepherd/results/gp"

            results_df.to_csv(str(output_base) + '_scores.csv', index=False)

            print(results_df)

            phen_df.to_csv(str(output_base) + '_phenotype_attention.csv', sep = ',', index=False)

            print(phen_df)

        # Plot embeddings

        if loop_type != "train" and len(self.train_patient_nodes) > 0 and self.hparams.hparams['plot_intrain']:

            correct_gene_nid = torch.cat([x[f'{loop_type}/corr_gene_nid_orig'] for x in outputs], dim=0)

            assert correct_gene_ranks.shape[0] == correct_gene_nid.shape[0]

            # Rank of gene vs. number of train patients with causal gene

            gene_rank_corr_gene_fig, gene_rank_corr_gene_counts = plot_gene_rank_vs_numtrain(correct_gene_ranks, correct_gene_nid, self.train_corr_gene_nid)

            gene_rank_cand_gene_fig, gene_rank_cand_gene_counts = plot_gene_rank_vs_numtrain(correct_gene_ranks, correct_gene_nid, self.train_patient_nodes)

            gene_rank_sparse_fig, gene_rank_sparse_counts = plot_gene_rank_vs_numtrain(correct_gene_ranks, correct_gene_nid, self.train_sparse_nodes)

            gene_rank_target_fig, gene_rank_target_counts = plot_gene_rank_vs_numtrain(correct_gene_ranks, correct_gene_nid, self.train_target_batch)

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_num_train_corr_genes': gene_rank_corr_gene_fig})

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_num_train_cand_genes': gene_rank_cand_gene_fig})

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_num_train_sparse': gene_rank_sparse_fig})

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_num_train_target': gene_rank_target_fig})

            gene_nid_trainset = torch.stack([torch.tensor(gene_rank_corr_gene_counts),

                                             torch.tensor(gene_rank_cand_gene_counts),

                                             torch.tensor(gene_rank_sparse_counts),

                                             torch.tensor(gene_rank_target_counts)], dim=1)

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_trainset': plot_gene_rank_vs_trainset(correct_gene_ranks, correct_gene_nid, gene_nid_trainset)})

        if self.hparams.hparams['plot_PG_embed']:

            self.logger.experiment.log({f'{loop_type}/patient_embed': fit_umap(patient_emb, patient_label)})

        # plot % overlap with train patients

        if loop_type != 'train' and self.hparams.hparams['mrr_vs_percent_overlap']:

            max_percent_overlap_train = torch.cat([torch.tensor(x[f'val/max_percent_phen_overlap_train']) for x in outputs], dim=0)

            self.logger.experiment.log({f'{loop_type}/mrr_vs_percent_overlap': mrr_vs_percent_overlap(correct_gene_ranks.detach().cpu(), max_percent_overlap_train.detach().cpu())})

        if self.hparams.hparams['plot_frac_rank']:

            # Rank of gene vs. fraction of phenotypes to disease 

            frac_p_with_direct_edge_to_dx = [pat[0][0] for x in outputs for pat in x[f'{loop_type}/frac_p_with_direct_edge_to_dx']] # NOTE: Currently ony select first disease. 

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_frac_p_with_direct_edge_to_dx': plot_gene_rank_vs_fraction_phenotype(correct_gene_ranks.cpu(), frac_p_with_direct_edge_to_dx)})

            # Rank of gene vs. fraction of phenotypes to gene 

            frac_p_with_direct_edge_to_g = [pat[0][0] for x in outputs for pat in x[f'{loop_type}/frac_p_with_direct_edge_to_g']] # NOTE: Currently ony select first gene. 

            self.logger.experiment.log({f'{loop_type}/frac_p_with_direct_edge_to_g': plot_gene_rank_vs_fraction_phenotype(correct_gene_ranks.cpu(), frac_p_with_direct_edge_to_g)})

        if self.hparams.hparams['plot_nhops_rank']:

            # Rank of gene vs. hops from disease

            nhops_g_d = [pat[0] for x in outputs for pat in x[f'{loop_type}/n_hops_g_d']] # NOTE Currently ony select first disease.

            fig_mean, fig_min = plot_gene_rank_vs_hops(correct_gene_ranks.cpu(), nhops_g_d)

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_mean_n_hops_g_d': fig_mean})

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_min_n_hops_g_d': fig_min})

            # Rank of gene vs. mean/min hops from phenotypes

            nhops_g_p = [pat[0] for x in outputs for pat in x[f'{loop_type}/n_hops_g_p']] # NOTE Currently ony select first gene. 

            fig_mean, fig_min = plot_gene_rank_vs_hops(correct_gene_ranks.cpu(), nhops_g_p)

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_mean_n_hops_g_p': fig_mean})

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_min_n_hops_g_p': fig_min})

            # # Rank of gene vs. distance between phenotypes

            nhops_p_p = [torch.tensor(pat) for x in outputs for pat in x[f'{loop_type}/n_hops_p_p']] 

            fig_mean, fig_min = plot_gene_rank_vs_hops(correct_gene_ranks.cpu(), nhops_p_p)

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_mean_n_hops_p_p': fig_mean})

            self.logger.experiment.log({f'{loop_type}/gene_rank_vs_min_n_hops_p_p': fig_min})

        if self.hparams.hparams['plot_attn_nhops']:

            # plot phenotype attention vs n_hops to gene and degree

            attn_weights = [torch.split(x[f'{loop_type}/attention_weights'],1) for x in outputs]

            attn_weights = [w[w > 0] for batch_w in attn_weights for w in batch_w]

            phenotype_names = [pat for x in outputs for pat in x[f'{loop_type}/phenotype_names_degrees']]

            attn_weights_cpu_reshaped = torch.cat(attn_weights, dim=0)

            self.logger.experiment.log({f"{loop_type}_attn/attention weights": wandb.Histogram(attn_weights_cpu_reshaped[attn_weights_cpu_reshaped != 0])})

            self.logger.experiment.log({f"{loop_type}_attn/single patient attention weights": wandb.Histogram(attn_weights[0])})

            self.logger.experiment.log({f"{loop_type}_attn/n hops to gene vs attention weights" : plot_nhops_to_gene_vs_attention(attn_weights, phenotype_names, nhops_g_p)})

            self.logger.experiment.log({f"{loop_type}_attn/single patient n hops to gene vs attention weights" : plot_nhops_to_gene_vs_attention(attn_weights, phenotype_names, nhops_g_p, single_patient=True)})

            self.logger.experiment.log({f"{loop_type}_attn/degree vs attention weights" : plot_degree_vs_attention(attn_weights, phenotype_names)})

            self.logger.experiment.log({f"{loop_type}_attn/single patient degree vs attention weights" : plot_degree_vs_attention(attn_weights, phenotype_names, single_patient=True)})

            data = [[p_name[0], w.item(), p_name[1], n_hops_to_g] for w, p_name, n_hops_to_g in zip(attn_weights[0], phenotype_names[0], nhops_g_p[0])]

            self.logger.experiment.log({f"{loop_type}_attn/phenotypes": wandb.Table(data=data, columns=["HPO Code", "Attention Weight", "Degree", "Num Hops to Gene" ])}) 

        if self.hparams.hparams['plot_phen_gene_sims']:

            all_phen_gene_sims, all_raw_phen_gene_sims, all_pg_spl, all_correct_sims, all_incorrect_sims = [], [], [], [], []

            for x in outputs:

                phen_gene_sims = x[f'{loop_type}/phen_gene_sims']

                one_hot_labels = x[f'{loop_type}/one_hot_labels']

                correct_phen_squeuegene_sims = all_correct_sims.append(phen_gene_sims[one_hot_labels == 1])

                incorrect_phen_gene_sims = all_incorrect_sims.append(phen_gene_sims[one_hot_labels != 1])

                phen_gene_sims_reshaped = all_phen_gene_sims.append(phen_gene_sims.view(-1))

            phen_gene_sims_reshaped = torch.cat(all_phen_gene_sims)

            correct_phen_gene_sims = torch.cat(all_correct_sims)

            incorrect_phen_gene_sims = torch.cat(all_incorrect_sims)

            if len(all_pg_spl) > 0: pg_spl_reshaped = torch.cat(all_pg_spl)

            else: pg_spl_reshaped = []

            self.logger.experiment.log({f"{loop_type}_pg_similarities/phenotype-gene similarities": wandb.Histogram(phen_gene_sims_reshaped[phen_gene_sims_reshaped != -100000])})

            self.logger.experiment.log({f"{loop_type}_pg_similarities/phenotype-correct gene similarities": wandb.Histogram(correct_phen_gene_sims[correct_phen_gene_sims != -100000])})

            self.logger.experiment.log({f"{loop_type}_pg_similarities/phenotype-incorrect gene similarities": wandb.Histogram(incorrect_phen_gene_sims[incorrect_phen_gene_sims != -100000])})

            if len(pg_spl_reshaped) > 0: self.logger.experiment.log({f"{loop_type}_pg_similarities/pg spl": wandb.Histogram(pg_spl_reshaped[pg_spl_reshaped != 0])})

            phen_gene_sims_patient = outputs[0][f'{loop_type}/phen_gene_sims'][0,:]

            one_hot_labels_patient = outputs[0][f'{loop_type}/one_hot_labels'][0,:]

            correct_phen_gene_sims_patient = phen_gene_sims_patient[one_hot_labels_patient == 1]

            assert len(correct_phen_gene_sims_patient) == 1

            incorrect_phen_gene_sims_patient = phen_gene_sims_patient[one_hot_labels_patient != 1]

            self.logger.experiment.log({f"{loop_type}_pg_similarities/single patient phenotype-gene similarities": wandb.Histogram(phen_gene_sims_patient[phen_gene_sims_patient != -100000])})

            self.logger.experiment.log({f"{loop_type}_pg_similarities/single patient phenotype-correct gene similarities": wandb.Histogram(correct_phen_gene_sims_patient[correct_phen_gene_sims_patient != -100000])})

            self.logger.experiment.log({f"{loop_type}_pg_similarities/single patient phenotype-incorrect gene similarities": wandb.Histogram(incorrect_phen_gene_sims_patient[incorrect_phen_gene_sims_patient != -100000])})

            if len(pg_spl_reshaped) > 0: self.logger.experiment.log({f"{loop_type}_pg_similarities/single patient pg spl": wandb.Histogram(pg_spl_reshaped[pg_spl_reshaped != 0])})

        # top k accuracy

        top_1_acc = top_k_acc(correct_gene_ranks, k=1)

        top_3_acc = top_k_acc(correct_gene_ranks, k=3)

        top_5_acc = top_k_acc(correct_gene_ranks, k=5)

        top_10_acc = top_k_acc(correct_gene_ranks, k=10)

        #mean reciprocal rank

        mrr = mean_reciprocal_rank(correct_gene_ranks)

        avg_rank = average_rank(correct_gene_ranks)

        self.log(f'{loop_type}/gp_{loop_type}_epoch_top1_acc', top_1_acc, prog_bar=False)

        self.log(f'{loop_type}/gp_{loop_type}_epoch_top3_acc', top_3_acc, prog_bar=False)

        self.log(f'{loop_type}/gp_{loop_type}_epoch_top5_acc', top_5_acc, prog_bar=False)

        self.log(f'{loop_type}/gp_{loop_type}_epoch_top10_acc', top_10_acc, prog_bar=False)

        self.log(f'{loop_type}/gp_{loop_type}_epoch_mrr', mrr, prog_bar=False)

        self.log(f'{loop_type}/gp_{loop_type}_epoch_avg_rank', avg_rank, prog_bar=False)

        if loop_type == 'val':

            self.log(f'curr_epoch', self.current_epoch, prog_bar=False)

    def training_epoch_end(self, outputs):

        if self.hparams.hparams['plot_intrain']:

            all_train_nodes, counts = torch.unique(torch.cat([x['train/n_id'] for x in outputs], dim=0), return_counts=True)

            curr_all_train_nodes = {n.item(): c.item() if n not in self.all_train_nodes else c.item() + self.all_train_nodes[n].item() for n, c in zip(all_train_nodes, counts)}

            self.all_train_nodes.update(curr_all_train_nodes)

            train_sparse_nodes, counts = torch.unique(torch.cat([x['train/sparse_idx'] for x in outputs], dim=0), return_counts=True)

            curr_train_sparse_nodes = {n.item(): c.item() if n not in self.train_sparse_nodes else c.item() + self.train_sparse_nodes[n].item() for n, c in zip(train_sparse_nodes, counts)}

            self.train_sparse_nodes.update(curr_train_sparse_nodes)

            train_target_batch, counts = torch.unique(torch.cat([x['train/target_batch'] for x in outputs], dim=0), return_counts=True)

            curr_train_target_batch = {n.item(): c.item() if n not in self.train_target_batch else c.item() + self.train_target_batch[n].item() for n, c in zip(train_target_batch, counts)}

            self.train_target_batch.update(curr_train_target_batch)

            train_patient_nodes, counts = torch.unique(torch.cat([x['train/cand_gene_nid_orig'] for x in outputs], dim=0), return_counts=True)

            self.train_patient_nodes = {n.item(): c.item() for n, c in zip(train_patient_nodes, counts)}

            train_corr_gene_nids, counts = torch.unique(torch.cat([x['train/corr_gene_nid_orig'] for x in outputs], dim=0), return_counts=True)

            self.train_corr_gene_nid = {g.item(): c.item() for g, c in zip(train_corr_gene_nids, counts)}

        self._epoch_end(outputs, 'train')

    def validation_epoch_end(self, outputs):

        self._epoch_end(outputs, 'val')

    def test_epoch_end(self, outputs):

        self._epoch_end(outputs, 'test')

    def configure_optimizers(self):

        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.hparams['lr'])

        return optimizer