# General
import numpy as np
import random
import argparse
import os
import sys
from pathlib import Path
from datetime import datetime
from collections import Counter
import pandas as pd
import pickle
import time

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

# Pytorch
import torch
import torch.nn as nn
from torch_geometric.utils.convert import to_networkx, to_scipy_sparse_matrix
from torch_geometric.data import Data, DataLoader, NeighborSampler
from torch_geometric.utils import negative_sampling
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split, SubsetRandomSampler

# Pytorch lightning
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import ModelCheckpoint

# W&B
import wandb

# multiprocessing
import torch.multiprocessing
torch.multiprocessing.set_sharing_strategy('file_system')

# Own code
import project_config
from shepherd.dataset import PatientDataset
from shepherd.gene_prioritization_model import CombinedGPAligner
from shepherd.patient_nca_model import CombinedPatientNCA 
from shepherd.samplers import PatientNeighborSampler

import preprocess
from hparams import get_pretrain_hparams, get_train_hparams


import os
os.environ['CUDA_LAUNCH_BLOCKING'] = "1" 
import faulthandler; faulthandler.enable()



def parse_args():
    parser = argparse.ArgumentParser(description="Learning node embeddings.")
    
    # Input files/parameters
    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('--saved_node_embeddings_path', type=str, default=None, help='Path within kg_embeddings folder to the saved KG embeddings')
    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("--aug_sim", type=str, default=None, help="File with the similarity dictionary")
    parser.add_argument("--aug_gene_by_deg", type=bool, default=False, help="Augment gene by degree")
    parser.add_argument("--aug_gene_w", type=float, default=0.7, help="Contribution of augmentation (gene)")
    parser.add_argument("--n_sim_genes", type=int, default=3, help="K similar genes for augmentation")
    parser.add_argument("--n_transformer_layers", type=int, default=3, help="Number of transformer layers")
    parser.add_argument("--n_transformer_heads", type=int, default=8, help="Number of transformer heads")

    # Tunable parameters
    parser.add_argument('--sparse_sample', default=200, type=int)
    parser.add_argument('--lr', default=0.0001, type=float)
    parser.add_argument('--upsample_cand', default=1, type=int)
    parser.add_argument('--neighbor_sampler_size', default=-1, type=int)
    parser.add_argument('--lmbda', type=float, default=0.5, help='Lambda')
    parser.add_argument('--alpha', type=float, default=0, help='Alpha')
    parser.add_argument('--kappa', type=float, default=0.3, help='Kappa (Only used for combined model with link prediction loss)')
    parser.add_argument('--seed', default=33, type=int)
    parser.add_argument('--batch_size', default=64, type=int) 
    
    # Resume / run inference with best checkpoint
    parser.add_argument('--resume', default="", type=str)
    parser.add_argument('--do_inference', action='store_true')
    parser.add_argument('--best_ckpt', type=str, default=None, help='Name of the best performing checkpoint')
    
    parser.add_argument('--use_wandb', type=bool, default=True)

    args = parser.parse_args()
    return args


def load_patient_datasets(hparams, inference=False):
    print('loading patient datasets')

    if inference:
        train_dataset = None
        val_dataset = None
    else:
        train_dataset = PatientDataset(project_config.PROJECT_DIR / 'patients' / hparams['train_data'],  time=hparams['time'])
        val_dataset = PatientDataset(project_config.PROJECT_DIR / 'patients' / hparams['validation_data'], time=hparams['time'])

    if inference:
        test_dataset = PatientDataset(project_config.PROJECT_DIR / 'patients' / hparams['test_data'], time=hparams['time'])
    else:
        test_dataset = None
    
    print('finished loading patient datasets')
    return train_dataset, val_dataset, test_dataset


def get_dataloaders(hparams, all_data, nid_to_spl_dict, n_nodes, gene_phen_dis_node_idx, train_dataset, val_dataset, test_dataset, inference=False):
    print('Get dataloaders', flush=True)
    shuffle = False if hparams['debug'] or inference else True
    if not hparams['sample_from_gpd']: gene_phen_dis_node_idx = None
    batch_sz = hparams['inference_batch_size'] if inference else hparams['batch_size']
    sparse_sample = 1 if inference else hparams['sparse_sample']

    #get phenotypes & genes found in train patients
    if hparams['sample_edges_from_train_patients']:
        phenotype_counter = Counter()
        gene_counter = Counter()
        for patient in train_dataset:
            phenotype_node_idx, candidate_gene_node_idx, correct_genes_node_idx, disease_node_idx, labels, additional_labels, patient_ids = patient

            phenotype_counter += Counter(list(phenotype_node_idx.numpy()))
            gene_counter += Counter(list(candidate_gene_node_idx.numpy()))
    else:
        phenotype_counter=None
        gene_counter=None

    print('Loading SPL...')
    if hparams['spl'] is not None:
        spl = np.load(project_config.PROJECT_DIR / 'patients' / hparams['spl'])  
    else: spl = None
    if hparams['spl_index'] is not None and (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 # TODO: short term fix for simulated patients, get rid once we create this dict
    
    print('Loaded SPL information')

    if args.aug_sim is not None:
        with open(str(project_config.PROJECT_DIR / 'knowledge_graph/8.9.21_kg' / ('top_10_similar_genes_sim=%s.pkl' % args.aug_sim)), "rb") as input_file:
            gene_similarity_dict = pickle.load(input_file)
        print("Using augment gene similarity: %s" % args.aug_sim)
    else: gene_similarity_dict=None

    with open("/home/ema30/zaklab/rare_disease_dx/formatted_patients/degree_dict_8.9.21_kg.pkl", "rb") as input_file:
        gene_deg_dict = pickle.load(input_file)

    if inference:
        train_dataloader = None
        val_dataloader = None
    else:
        print('setting up train dataloader')         
        train_dataloader = PatientNeighborSampler('train', all_data.edge_index[:,all_data.train_mask], all_data.edge_index[:,all_data.train_mask], 
                        sizes = hparams['neighbor_sampler_sizes'], patient_dataset=train_dataset, batch_size = batch_sz, 
                        sparse_sample = sparse_sample, do_filter_edges=hparams['filter_edges'], 
                        all_edge_attributes=all_data.edge_attr, n_nodes = n_nodes, relevant_node_idx=gene_phen_dis_node_idx,
                        shuffle = shuffle, train_phenotype_counter=phenotype_counter, train_gene_counter=gene_counter, sample_edges_from_train_patients=hparams['sample_edges_from_train_patients'], num_workers=hparams['num_workers'], 
                        upsample_cand=hparams['upsample_cand'], n_cand_diseases=hparams['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,
                        hparams=hparams, pin_memory=hparams['pin_memory'],
                        gene_similarity_dict = gene_similarity_dict,
                        gene_deg_dict = gene_deg_dict)
        print('finished setting up train dataloader')
        print('setting up val dataloader')
        val_dataloader = PatientNeighborSampler('val', all_data.edge_index, all_data.edge_index[:,all_data.val_mask], 
                        sizes = [-1,10,5], 
                        patient_dataset=val_dataset, batch_size = batch_sz, 
                        sparse_sample = sparse_sample, all_edge_attributes=all_data.edge_attr, n_nodes = n_nodes, 
                        relevant_node_idx=gene_phen_dis_node_idx, 
                        shuffle = False, train_phenotype_counter=phenotype_counter, train_gene_counter=gene_counter, sample_edges_from_train_patients=hparams['sample_edges_from_train_patients'], num_workers=hparams['num_workers'],
                        n_cand_diseases=hparams['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,
                        hparams=hparams,  pin_memory=hparams['pin_memory'],
                        gene_similarity_dict = gene_similarity_dict, 
                        gene_deg_dict = gene_deg_dict)
        print('finished setting up val dataloader')
    
    print('setting up test dataloader')
    if inference:
        sizes = [-1,10,5]
        print('SIZES: ', sizes)
        test_dataloader = PatientNeighborSampler('test', all_data.edge_index, all_data.edge_index[:,all_data.test_mask], 
                        sizes = sizes, patient_dataset=test_dataset, batch_size = len(test_dataset), 
                        sparse_sample = sparse_sample, all_edge_attributes=all_data.edge_attr, n_nodes = n_nodes, relevant_node_idx=gene_phen_dis_node_idx,
                        shuffle = False, num_workers=hparams['num_workers'],
                        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,
                        hparams=hparams, pin_memory=hparams['pin_memory'],
                        gene_similarity_dict = gene_similarity_dict, 
                        gene_deg_dict = gene_deg_dict) 
    else: test_dataloader = None
    print('finished setting up test dataloader')
    
    return train_dataloader, val_dataloader, test_dataloader


def get_model(args, hparams, node_hparams, all_data, edge_attr_dict, n_nodes, load_from_checkpoint=False):
    print("setting up model", hparams['model_type'])
    # get patient model 
    if hparams['model_type'] == 'aligner':
        if load_from_checkpoint: 
            comb_patient_model = CombinedGPAligner.load_from_checkpoint(checkpoint_path=str(Path(project_config.PROJECT_DIR /  args.best_ckpt)), 
                                    edge_attr_dict=edge_attr_dict, all_data=all_data, n_nodes=n_nodes, node_ckpt = hparams["saved_checkpoint_path"], node_hparams=node_hparams)
        else:
            comb_patient_model = CombinedGPAligner(edge_attr_dict=edge_attr_dict, all_data=all_data, n_nodes=n_nodes, hparams=hparams, node_ckpt = hparams["saved_checkpoint_path"], node_hparams=node_hparams)
    elif hparams['model_type'] == 'patient_NCA':
        if load_from_checkpoint:
            comb_patient_model = CombinedPatientNCA.load_from_checkpoint(checkpoint_path=str(Path(project_config.PROJECT_DIR) /  args.best_ckpt), 
                                    all_data=all_data, edge_attr_dict=edge_attr_dict, n_nodes=n_nodes, node_ckpt=hparams["saved_checkpoint_path"])
        else:
            comb_patient_model = CombinedPatientNCA(edge_attr_dict=edge_attr_dict, all_data=all_data, n_nodes=n_nodes, node_ckpt=hparams["saved_checkpoint_path"], hparams=hparams)
    else:
        raise NotImplementedError
    print('finished setting up model')
    return comb_patient_model


def train(args, hparams):
    print('Training Model', flush=True)

    # Hyperparameters
    node_hparams = get_pretrain_hparams(args, combined=True)
    print('Edge List: ', args.edgelist,  flush=True)
    print('Node Map: ', args.node_map, flush=True)

    # Set seed
    pl.seed_everything(hparams['seed'])

    # Read input data
    print('Read data', flush=True)
    all_data, edge_attr_dict, nodes = preprocess.preprocess_graph(args)
    n_nodes = len(nodes["node_idx"].unique())
    print(f'Number of nodes: {n_nodes}')
    gene_phen_dis_node_idx = torch.LongTensor(nodes.loc[nodes['node_type'].isin(['gene/protein', 'effect/phenotype', 'disease']), 'node_idx'].values)
    
    if args.resume != "":
        print('Resuming Run')
        # create Weights & Biases Logger
        if ":" in args.resume: # colons are not allowed in ID/resume name
            resume_id = "_".join(args.resume.split(":"))
        run_name = args.resume
        wandb_logger = WandbLogger(run_name, project=hparams['wandb_project_name'], entity='rare_disease_dx', save_dir=hparams['wandb_save_dir'], id=resume_id, resume=resume_id)
        
        #add run name to hparams dict
        hparams['run_name'] = run_name
        
        # get patient model 
        comb_patient_model = get_model(args, hparams, node_hparams, all_data, edge_attr_dict, n_nodes, load_from_checkpoint=True)
        
    else:
        print('Creating new W&B Logger')
        # create Weights & Biases Logger
        curr_time = datetime.now().strftime("%m_%d_%y:%H:%M:%S")
        lr = hparams['lr']   
        val_data = str(hparams['validation_data']).split('.txt')[0].replace('/', '.')
        run_name = "{}_val_{}".format(curr_time, val_data).replace('patients', 'pats') 
        run_name = run_name + f'_seed={args.seed}'
        run_name = run_name.replace('5_candidates_mapped_only', '5cand_map').replace('8.9.21_kgsolved_manual_baylor_nobgm_distractor_genes', 'manual').replace('patient_disease_NCA', 'pd_NCA').replace('_distractor', '')
        wandb_logger = WandbLogger(name=run_name, project=hparams['wandb_project_name'], entity='rare_disease_dx', save_dir=hparams['wandb_save_dir'],
                        id="_".join(run_name.split(":")), resume="allow") 
        
        #add run name to hparams dict
        print('Run name', run_name)
        hparams['run_name'] = run_name

        # get patient model 
        comb_patient_model = get_model(args, hparams, node_hparams, all_data, edge_attr_dict, n_nodes, load_from_checkpoint=False)

    # get model & dataloaders
    nid_to_spl_dict = {nid: idx for idx, nid in enumerate(nodes[nodes["node_type"] == "gene/protein"]["node_idx"].tolist())}
    train_dataset, val_dataset, test_dataset = load_patient_datasets(hparams)
    patient_train_dataloader, patient_val_dataloader, patient_test_dataloader = get_dataloaders(hparams, all_data, nid_to_spl_dict,
                                                                                                n_nodes, gene_phen_dis_node_idx, 
                                                                                                train_dataset, val_dataset, test_dataset)

    # callbacks
    print('Init callbacks')
    checkpoint_path = (project_config.PROJECT_DIR / 'checkpoints' / hparams['model_type'] / run_name) 
    hparams['checkpoint_path'] = checkpoint_path
    print('Checkpoint path: ', checkpoint_path)
    if not os.path.exists(project_config.PROJECT_DIR / 'checkpoints' / hparams['model_type']): (project_config.PROJECT_DIR / 'checkpoints' / hparams['model_type']).mkdir()
    if not os.path.exists(checkpoint_path): checkpoint_path.mkdir()
    monitor_type =  'val/mrr' if args.run_type == 'disease_characterization' or args.run_type == 'patients_like_me' else 'val/gp_val_epoch_mrr'
    fname = 'epoch={epoch:02d}-val_mrr={val/mrr:.2f}' if args.run_type == 'disease_characterization' or args.run_type == 'patients_like_me'  else 'epoch={epoch:02d}-val_mrr={val/gp_val_epoch_mrr:.2f}'
    patient_checkpoint_callback = ModelCheckpoint(
        monitor=monitor_type,
        dirpath=checkpoint_path,
        filename=fname,
        save_top_k=-1,
        mode='max',
        auto_insert_metric_name = False
    )

    # log gradients with logger
    print('wandb logger watch')
    wandb_logger.watch(comb_patient_model, log='all')

    #initialize trainer
    if hparams['debug']: 
        limit_train_batches = 1
        limit_val_batches = 1 
        hparams['max_epochs'] = 6
    else: 
        limit_train_batches=1.0
        limit_val_batches=1.0

    print('initialize trainer')
    patient_trainer = pl.Trainer(gpus=hparams['n_gpus'], 
                                logger=wandb_logger, 
                                max_epochs=hparams['max_epochs'], 
                                callbacks=[patient_checkpoint_callback],
                                profiler=hparams['profiler'],
                                log_gpu_memory=hparams['log_gpu_memory'],
                                limit_train_batches=limit_train_batches, 
                                limit_val_batches=limit_val_batches,
                                weights_summary="full",
                                gradient_clip_val=hparams['gradclip'])

    #  Train
    patient_trainer.fit(comb_patient_model, patient_train_dataloader, patient_val_dataloader)

@torch.no_grad()
def inference(args, hparams):
    print('Running inference')
    # Hyperparameters
    node_hparams = get_pretrain_hparams(args, combined=True)

    hparams.update({'add_similar_patients': False})

    # Seed
    pl.seed_everything(hparams['seed'])

    # Read data
    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 logger & trainer
    curr_time = datetime.now().strftime("%m_%d_%y:%H:%M:%S")
    lr = hparams['lr']   
    test_data = hparams['test_data'].split('.txt')[0].replace('/', '.')
    run_name = "{}_lr_{}_test_{}".format(curr_time, lr, test_data)
    wandb_logger = WandbLogger(run_name, project=hparams['wandb_project_name'], entity='rare_disease_dx', save_dir=hparams['wandb_save_dir'])
    print('Run name: ', run_name)
    hparams['run_name'] = run_name

    # Get datasets
    train_dataset, val_dataset, test_dataset = load_patient_datasets(hparams, inference=True)

    # Get dataloader
    nid_to_spl_dict = {nid: idx for idx, nid in enumerate(nodes[nodes["node_type"] == "gene/protein"]["node_idx"].tolist())}
    _, _, test_dataloader = get_dataloaders(hparams, all_data, nid_to_spl_dict,
                                                                        n_nodes, gene_phen_dis_node_idx, 
                                                                        train_dataset, val_dataset, test_dataset, inference=True)

    # Get patient model 
    model = get_model(args, hparams, node_hparams, all_data, edge_attr_dict, n_nodes, load_from_checkpoint=True)

    trainer = pl.Trainer(gpus=0, logger=wandb_logger)
    results = trainer.test(model, dataloaders=test_dataloader)
    print(results)
    print('---- RESULTS ----')


if __name__ == "__main__":
    
    # Get hyperparameters
    args = parse_args()
    hparams = get_train_hparams(args)

    # Run model
    if args.do_inference:
        inference(args, hparams)
    else:
        train(args, hparams)