import project_config

####################################################################

#

# NODE EMBEDDER MODEL HYPERPARAMETERS

#

####################################################################

def get_pretrain_hparams(args, combined=False):    

    print('node embedder args: ', args)

    # Default

    hparams = {

               # Tunable parameters

               'nfeat': args.nfeat if not combined else 4096,

               'hidden': args.hidden if not combined else 256,

               'output': args.output if not combined else 128,

               'n_heads': args.n_heads if not combined else 2,

               'wd': args.wd if not combined else 5e-4,

               'dropout': args.dropout if not combined else 0.2,

               'lr': args.lr if not combined else 0.0001,

               # Fixed parameters

               'decoder_type': 'bilinear',

               'norm_method': "batch_layer",

               'loss': 'max-margin',

               'pred_threshold': 0.5,

               'negative_sampler_approach': 'by_edge_type',

               'filter_edges': True,

               'n_gpus': 1,

               'num_workers': 4,

               'batch_size': 512,

               'inference_batch_size': 64,

               'neighbor_sampler_sizes': [15, 10, 5],

               'max_epochs': 200,

               'gradclip': 1.0,

               'lr_factor': 0.01,

               'lr_patience': 1000,

               'lr_threshold': 1e-4,

               'lr_threshold_mode': 'rel',

               'lr_cooldown': 0,

               'min_lr': 0,

               'eps': 1e-8,

               'seed': 1,

               'profiler': None,

               'wandb_save_dir': project_config.PROJECT_DIR / 'wandb' / 'preprocess',

               'log_every_n_steps': 10,

               'time': False,

               'debug': False

        }

    print('Pretrain hparams: ', hparams)

    return hparams

####################################################################

#

# TRAIN MODEL HYPERPARAMETERS

#

####################################################################

def get_train_hparams(args):

    print('Train model args: ', args)

    # Default

    hparams = {

               # Tunable parameters

               'sparse_sample': args.sparse_sample, # Randomly sample N nodes from KG

               'lr': args.lr,

               'upsample_cand': args.upsample_cand, 

               'neighbor_sampler_sizes': [args.neighbor_sampler_size, 10, 5],

               'lambda': args.lmbda, # Contribution of two loss functions

               'alpha': args.alpha, # Contribution of GP gate. NOTE: This is not used for patients-like-me or novel disease characterization

               'kappa': (1 - args.lmbda) * args.kappa,

               'seed': args.seed,

               'batch_size': args.batch_size,

               'augment_genes': True if args.aug_gene_w > 0 else False,

               'n_sim_genes': args.n_sim_genes,

               'aug_gene_w': args.aug_gene_w,

               'aug_gene_by_deg': args.aug_gene_by_deg,

               'n_transformer_layers': args.n_transformer_layers,

               'n_transformer_heads': args.n_transformer_heads,

               # Fixed parameters

               'pos_weight': 1,

               'neg_weight': 20,

               'margin': 0.4,

               'thresh': 1,

               'filter_edges': False,

               'softmax_scale': 1,

               'leaky_relu': 0.1,

               'decoder_type': 'bilinear',

               'combined_training': True,

               'sample_from_gpd': True,

               'attention_type': 'bilinear',

               'n_cand_diseases': 1000,

               'test_n_cand_diseases': -1, 

               'candidate_disease_type': 'all_kg_nodes',

               'patient_similarity_type': 'gene', # How we determine labels for similar patients in "Patients Like Me"

               'n_similar_patients': 2, # Number of patients with the same gene/disease that we add to the batch

               'only_hard_distractors': False, # Flag when true only uses the curated hard distractors at train time

               'sample_edges_from_train_patients': False, # Preferentially sample edges connected to training patients

               'gradclip': 1.0,

               'inference_batch_size': 64,

               'max_epochs': 100, 

               'n_gpus': 1, 

               'num_workers': 4,

               'wandb_save_dir' : project_config.PROJECT_DIR / 'wandb',

               'precision': 16, 

               'reload_dataloaders_every_n_epochs': 0,

               'profiler': 'simple',

               'pin_memory': False,

               'time': False,

               'log_gpu_memory': True,

               'debug': False, 

               'plot_softmax': False,

               'plot_intrain': False, # Flag to plot gene rank vs. in train sets

               'plot_PG_embed': False, # Flag to plot embeddings with phenotype and gene labels

               'plot_disease_embed': False, # Flag to plot embeddings with disease labels

               'plot_patient_embed': False, # Flag to plot embeddings for patients

               'plot_degree_rank': False, # Flag to plot degree vs. gene rank

               'plot_nhops_rank': False, # Flag to plot nhops vs. gene rank

               'plot_frac_rank': False, # Flag to plot fraction of ___ vs. gene rank

               'plot_gradients': False, # Flag to plot gradients

               'plot_attn_nhops': False, # Flag to plot attn weights vs. nhops

               'plot_phen_gene_sims': False, # Flag to plot phenotype-gene similarities

               'mrr_vs_percent_overlap': False, # Flag to plot MRR vs. percent overlap of phenotypes

               'saved_checkpoint_path': project_config.PROJECT_DIR  / f'{args.saved_node_embeddings_path}', 

    }

    # Get hyperparameters based on run type arguments

    hparams = get_run_type_args(args, hparams)

    # Get hyperparameters based on patient data arguments

    hparams = get_patient_data_args(args, hparams)

    print('Train hparams: ', hparams)

    return hparams

def get_run_type_args(args, hparams):

    if args.run_type == 'causal_gene_discovery':

        hparams.update({

                        'model_type': 'aligner', 

                        'loss': 'gene_multisimilarity', 

                        'use_diseases': False,

                        'add_cand_diseases': False,

                        'add_similar_patients': False,

                        'wandb_project_name': 'causal-gene-discovery'

                       })

    elif args.run_type == 'disease_characterization':

        hparams.update({

                        'model_type': 'patient_NCA',

                        'loss': 'patient_disease_NCA',

                        'use_diseases': True,

                        'add_cand_diseases': True ,

                        'add_similar_patients': False,

                        'wandb_project_name': 'disease-heterogeneity',

                       })

    elif args.run_type == 'patients_like_me':

        hparams.update({

                        'model_type': 'patient_NCA',

                        'loss': 'patient_patient_NCA',

                        'use_diseases': False,

                        'add_cand_diseases': False,

                        'add_similar_patients': True,

                        'wandb_project_name': 'patients-like-me',

                       })

    else:

        raise Exception('You must specify run type.')

    return hparams

def get_patient_data_args(args, hparams):

    if args.patient_data == "disease_simulated":

        hparams.update({'train_data': f'simulated_patients/disease_split_train_sim_patients_{project_config.CURR_KG}.txt',

                        'validation_data': f'simulated_patients/disease_split_val_sim_patients_{project_config.CURR_KG}.txt', 

                        'test_data': f'simulated_patients/disease_split_all_sim_patients_{project_config.CURR_KG}.txt',

                        'spl': f'simulated_patients/disease_split_all_sim_patients_{project_config.CURR_KG}_spl_matrix.npy',

                        'spl_index': f'simulated_patients/disease_split_all_sim_patients_{project_config.CURR_KG}_spl_index_dict.pkl'

                        })

    elif args.patient_data == "my_data":

        hparams.update({'train_data': project_config.MY_TRAIN_DATA,

                        'validation_data': project_config.MY_VAL_DATA,

                        'test_data': project_config.MY_TEST_DATA,

                        'spl': project_config.MY_SPL_DATA, # Result of add_spl_to_patients.py (suffix: _spl_matrix.npy)

                        'spl_index': project_config.MY_SPL_INDEX_DATA, # Result of add_spl_to_patients.py (suffix: _spl_index_dict.pkl)

                        })

    else:

        raise Exception('You must specify patient data.')

    return hparams

####################################################################

#

# PREDICTION HYPERPARAMETERS

#

####################################################################

def get_predict_hparams(args):

    hparams = {

               'seed': 33,

               'n_gpus': 0, # NOTE: currently predict scripts only work with CPU

               'num_workers': 4, 

               'profiler': 'simple',

               'pin_memory': False,

               'time': False,

               'log_gpu_memory': False,

               'debug': False,

               'augment_genes': True,

               'n_sim_genes': 3,

               'aug_gene_w': 0.5,

               'wandb_save_dir' : project_config.PROJECT_DIR / 'wandb',

               'saved_checkpoint_path': project_config.PROJECT_DIR  / f'{args.saved_node_embeddings_path}',

               'test_n_cand_diseases': -1, 

               'candidate_disease_type': 'all_kg_nodes', 

               'only_hard_distractors': False, # Flag when true only uses the curated hard distractors at train time

               'patient_similarity_type': 'gene', # How we determine labels for similar patients in "Patients Like Me"

               'n_similar_patients': 2, # (Patients Like Me only) Number of patients with the same gene/disease that we add to the batch

    }

    # Get hyperparameters based on run type arguments

    hparams = get_run_type_args(args, hparams)    

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

    hparams = get_patient_data_args(args, hparams)

    print('Predict hparams: ', hparams)

    return hparams