# General

import numpy as np

import pickle

import random

import argparse

import os

import sys

from pathlib import Path

from datetime import datetime

import time

from collections import Counter

import pandas as pd

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

# Pytorch

import torch

import torch.nn as nn

import torch.nn.functional as F

# Pytorch lightning

import pytorch_lightning as pl

from pytorch_lightning.loggers import WandbLogger

# W&B

import wandb

# Own code

import project_config

from shepherd.dataset import PatientDataset

from shepherd.samplers import PatientNeighborSampler

import preprocess

from hparams import get_predict_hparams

from train import get_model, load_patient_datasets, get_dataloaders

import os

os.environ['CUDA_LAUNCH_BLOCKING'] = "1" 

'''

Example Command:

python predict.py \

--run_type causal_gene_discovery \

--patient_data test_predict \

--edgelist KG_edgelist_mask.txt \

--node_map KG_node_map.txt \

--saved_node_embeddings_path checkpoints/pretrain.ckpt \

--best_ckpt checkpoints/causal_gene_discovery.ckpt 

python predict.py \

--run_type patients_like_me \

--patient_data test_predict \

--edgelist KG_edgelist_mask.txt \

--node_map KG_node_map.txt \

--saved_node_embeddings_path checkpoints/pretrain.ckpt \

--best_ckpt checkpoints/patients_like_me.ckpt 

python predict.py \

--run_type disease_characterization \

--patient_data test_predict \

--edgelist KG_edgelist_mask.txt \

--node_map KG_node_map.txt \

--saved_node_embeddings_path checkpoints/pretrain.ckpt \

--best_ckpt checkpoints/disease_characterization.ckpt 

'''

def parse_args():

    parser = argparse.ArgumentParser(description="Predict using SHEPHERD")

    parser.add_argument("--edgelist", type=str, default=None, help="File with edge list")

    parser.add_argument("--node_map", type=str, default=None, help="File with node list")

    parser.add_argument('--patient_data', default="disease_simulated", type=str)

    parser.add_argument('--run_type', choices=["causal_gene_discovery", "disease_characterization", "patients_like_me"], type=str)

    parser.add_argument('--saved_node_embeddings_path', type=str, default=None, help='Path to pretrained model checkpoint')

    parser.add_argument('--best_ckpt', type=str, default=None, help='Name of the best performing checkpoint')

    args = parser.parse_args()

    return args

@torch.no_grad()

def predict(args):

    # Hyperparameters

    hparams = get_predict_hparams(args)

    # Seed

    pl.seed_everything(hparams['seed'])

    # Read KG

    all_data, edge_attr_dict, nodes = preprocess.preprocess_graph(args)

    n_nodes = len(nodes["node_idx"].unique())

    gene_phen_dis_node_idx = torch.LongTensor(nodes.loc[nodes['node_type'].isin(['gene/protein', 'effect/phenotype', 'disease']), 'node_idx'].values)

    # Get dataset

    print('Loading SPL...')

    spl = np.load(project_config.PROJECT_DIR / 'patients' / hparams['spl'])  

    if (project_config.PROJECT_DIR / 'patients' / hparams['spl_index']).exists():

        with open(str(project_config.PROJECT_DIR / 'patients' / hparams['spl_index']), "rb") as input_file:

            spl_indexing_dict = pickle.load(input_file)

    else: spl_indexing_dict = None 

    print('Loaded SPL information')

    dataset = PatientDataset(project_config.PROJECT_DIR / 'patients' / hparams['test_data'], time=hparams['time'])

    print(f'There are {len(dataset)} patients in the test dataset')

    hparams.update({'inference_batch_size': len(dataset)})

    print('batch size: ', hparams['inference_batch_size'])

    # Get dataloader

    nid_to_spl_dict = {nid: idx for idx, nid in enumerate(nodes[nodes["node_type"] == "gene/protein"]["node_idx"].tolist())}

    dataloader = PatientNeighborSampler('predict', all_data.edge_index, all_data.edge_index[:,all_data.test_mask], 

                                        sizes = [-1,10,5], patient_dataset=dataset, batch_size = hparams['inference_batch_size'], sparse_sample = 0, 

                                        all_edge_attributes=all_data.edge_attr, n_nodes = n_nodes, relevant_node_idx=gene_phen_dis_node_idx,

                                        n_cand_diseases=hparams['test_n_cand_diseases'],  use_diseases=hparams['use_diseases'], 

                                        nid_to_spl_dict=nid_to_spl_dict, gp_spl=spl, spl_indexing_dict=spl_indexing_dict,

                                        shuffle = False, num_workers=hparams['num_workers'],

                                        hparams=hparams, pin_memory=hparams['pin_memory']) # 'test'

    # Create Weights & Biases Logger

    run_name = 'test'

    wandb_logger = WandbLogger(name=run_name, project='rare_disease_dx_combined', entity='rare_disease_dx', save_dir=hparams['wandb_save_dir'],

                    id="_".join(run_name.split(":")), resume="allow") 

    # Get patient model 

    model = get_model(args, hparams, None, all_data, edge_attr_dict,  n_nodes,load_from_checkpoint=True)

    trainer = pl.Trainer(gpus=hparams['n_gpus'])

    t1 = time.time()

    results = trainer.predict(model, dataloaders=dataloader)

    t2 = time.time()

    print(f"Predicting took {t2 - t1:0.4f} seconds", len(dataset), "patients")

    scores_dfs, attn_dfs, gat_attn_df_1, gat_attn_df_2, gat_attn_df_3, phenotype_embeddings, disease_embeddings = zip(*results)

    print('---- RESULTS ----')

    if not os.path.exists(project_config.PROJECT_DIR / 'results'):

        os.mkdir(project_config.PROJECT_DIR / 'results')

    output_base = project_config.PROJECT_DIR / 'results' /  (str(args.best_ckpt).replace('/', '.').split('.ckpt')[0])     

    # Save scores

    scores_df = pd.concat(scores_dfs).reset_index(drop=True)

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

    print(scores_df)

    # Save patient phenotype attention

    attn_df = pd.concat(attn_dfs).reset_index(drop=True)

    attn_df.to_csv(str(output_base) + '_phenotype_attn.csv', index=False)

    print(attn_df)

    # Save patient phenotype embeddings

    if type(phenotype_embeddings) == tuple: phenotype_embeddings = phenotype_embeddings[0]

    torch.save(phenotype_embeddings, str(output_base) + '_phenotype_embeddings.pth')

    print("Phenotype embeddings", phenotype_embeddings)

    # Save disease embeddings

    if args.run_type == "disease_characterization":

        if type(disease_embeddings) == tuple: disease_embeddings = disease_embeddings[0]

        torch.save(disease_embeddings, str(output_base) + '_disease_embeddings.pth')

        print("Disease embeddings", disease_embeddings)

if __name__ == "__main__":

    # Get hyperparameters

    args = parse_args()

    # perform prediction

    predict(args)