# General
import numpy as np
import random
import argparse
import os
from copy import deepcopy
from pathlib import Path
import sys
from datetime import datetime

# Pytorch
import torch
import torch.nn as nn

# Pytorch Lightning
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor

# Pytorch Geo
from torch_geometric.data.sampler import NeighborSampler as PyGeoNeighborSampler
from torch_geometric.data import Data, DataLoader

# W&B
import wandb

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

# Own code
import preprocess
from node_embedder_model import NodeEmbeder
import project_config
from hparams import get_pretrain_hparams
from samplers import NeighborSampler


def parse_args():
    parser = argparse.ArgumentParser(description="Learn 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('--save_dir', type=str, default=None, help='Directory for saving files')
    
    # Tunable parameters
    parser.add_argument('--nfeat', type=int, default=2048, help='Dimension of embedding layer')
    parser.add_argument('--hidden', default=256, type=int)
    parser.add_argument('--output', default=128, type=int)
    parser.add_argument('--n_heads', default=4, type=int)
    parser.add_argument('--wd', default=0.0, type=float)
    parser.add_argument('--dropout', type=float, default=0.3, help='Dropout')
    parser.add_argument('--lr', default=0.0001, type=float)
    parser.add_argument('--max_epochs', default=1000, type=int)
    
    # Resume with best checkpoint
    parser.add_argument('--resume', default="", type=str)
    parser.add_argument('--best_ckpt', type=str, default=None, help='Name of the best performing checkpoint')
    
    # Output
    parser.add_argument('--save_embeddings', action='store_true')

    args = parser.parse_args()
    return args


def get_dataloaders(hparams, all_data):
    print('get dataloaders')
    train_dataloader = NeighborSampler('train', all_data.edge_index[:,all_data.train_mask], all_data.edge_index[:,all_data.train_mask], sizes = hparams['neighbor_sampler_sizes'], batch_size = hparams['batch_size'], shuffle = True, num_workers=hparams['num_workers'], do_filter_edges=hparams['filter_edges'])
    val_dataloader = NeighborSampler('val', all_data.edge_index[:,all_data.train_mask], all_data.edge_index[:,all_data.val_mask], sizes = hparams['neighbor_sampler_sizes'], batch_size = hparams['batch_size'], shuffle = False, num_workers=hparams['num_workers'], do_filter_edges=hparams['filter_edges'])
    test_dataloader = NeighborSampler('test', all_data.edge_index[:,all_data.train_mask], all_data.edge_index[:,all_data.test_mask], sizes = hparams['neighbor_sampler_sizes'], batch_size = hparams['batch_size'], shuffle = False, num_workers=hparams['num_workers'], do_filter_edges=hparams['filter_edges'])
    return train_dataloader, val_dataloader, test_dataloader 


def train(args, hparams):

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

    # Read input data
    all_data, edge_attr_dict, nodes = preprocess.preprocess_graph(args)

    # Set up
    if args.resume != "":
        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='kg-train', entity='rare_disease_dx', save_dir=hparams['wandb_save_dir'], id=resume_id, resume=resume_id)
        model = NodeEmbeder.load_from_checkpoint(checkpoint_path=str(Path(args.save_dir) / 'checkpoints' /  args.best_ckpt), 
                                                 all_data=all_data, edge_attr_dict=edge_attr_dict, 
                                                 num_nodes=len(nodes["node_idx"].unique()), combined_training=False) 
    else:
        curr_time = datetime.now().strftime("%H:%M:%S")
        run_name = f"{curr_time}_run"
        wandb_logger = WandbLogger(run_name, project='kg-train', entity='rare_disease_dx', save_dir=hparams['wandb_save_dir'], id="_".join(run_name.split(":")), resume="allow")
        model = NodeEmbeder(all_data, edge_attr_dict, hp_dict=hparams, num_nodes=len(nodes["node_idx"].unique()), combined_training=False)

    checkpoint_callback = ModelCheckpoint(monitor='val/node_total_acc', dirpath=Path(args.save_dir) / 'checkpoints', filename=f'{run_name}' + '_{epoch}', save_top_k=1, mode='max')
    lr_monitor = LearningRateMonitor(logging_interval='step')
    wandb_logger.watch(model, log='all')

    if hparams['debug']:
        limit_train_batches = 1
        limit_val_batches = 1.0 
        hparams['max_epochs'] = 3
    else:
        limit_train_batches = 1.0
        limit_val_batches = 1.0 

    trainer = pl.Trainer(gpus=hparams['n_gpus'], logger=wandb_logger, 
                         max_epochs=hparams['max_epochs'], 
                         callbacks=[checkpoint_callback, lr_monitor], 
                         gradient_clip_val=hparams['gradclip'],
                         profiler=hparams['profiler'],
                         log_every_n_steps=hparams['log_every_n_steps'],
                         limit_train_batches=limit_train_batches, 
                         limit_val_batches=limit_val_batches,
                        ) 
    train_dataloader, val_dataloader, test_dataloader = get_dataloaders(hparams, all_data)

    # Train
    trainer.fit(model, train_dataloader, val_dataloader)
    
    # Test
    trainer.test(ckpt_path='best', test_dataloaders=test_dataloader)

@torch.no_grad()
def save_embeddings(args, hparams):
    print('Saving Embeddings')

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

    # Read input data
    all_data, edge_attr_dict, nodes = preprocess.preprocess_graph(args)
    all_data.num_nodes = len(nodes["node_idx"].unique())

    model = NodeEmbeder.load_from_checkpoint(checkpoint_path=str(Path(args.save_dir) / 'checkpoints' /  args.best_ckpt), 
                                            all_data=all_data, edge_attr_dict=edge_attr_dict, 
                                            num_nodes=len(nodes["node_idx"].unique()), combined_training=False) 
   
    dataloader = DataLoader([all_data], batch_size=1)
    trainer = pl.Trainer(gpus=0, 
                        gradient_clip_val=hparams['gradclip']
                    ) 
    embeddings = trainer.predict(model, dataloaders=dataloader)  
    embed_path = Path(args.save_dir) / (str(args.best_ckpt).split('.ckpt')[0] + '.embed')
    torch.save(embeddings[0], str(embed_path))
    print(embeddings[0].shape)



if __name__ == "__main__":
    
    # Get hyperparameters
    args = parse_args()
    hparams = get_pretrain_hparams(args, combined=False) 
    
    if args.save_embeddings:
        # save node embeddings from a trained model
        save_embeddings(args, hparams)
    else:
        # Train model
        train(args, hparams)