from cgl.eval.evaluator import GraphEvaluator
from pathlib import Path
import argparse
import imp
import matplotlib.pyplot as plt
import numpy as np
import wandb
import torch
from tqdm import tqdm

import pytorch_lightning as pl
from pytorch_lightning import loggers as pl_loggers
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.callbacks import LearningRateMonitor

from cgl.utils.params import ParamDict
from cgl.utils.general import listdict2dictlist

class Trainer:
    def __init__(self, pargs) -> None:
        pl.seed_everything(pargs.seed)
        
        conf = self.get_conf(pargs)
        model = conf.model

        pl.seed_everything(pargs.seed)
        model.reset_parameters()

        ckpt_path = pargs.ckpt
        node_emb_ckpt = conf.get('node_emb_ckpt', pargs.node_emb_ckpt)

        wandb_run = wandb.init(
            project='cgl',
            name=pargs.run_name,
            dir=pargs.output,
            id=pargs.wandb_id,
            resume='allow',
            config=dict(seed=pargs.seed)
        )
        wandb_run.log_code(root=Path(pargs.path).parent)

        Path(pargs.output).mkdir(parents=True, exist_ok=True)
        callbacks = [LearningRateMonitor()]
        if conf.valid_dloader:
            ckpt_callback = ModelCheckpoint(
                # monitor='valid_loss_total_ema',
                # monitor='valid_loss_total',
                # filename='cgl-{epoch:02d}-{valid_loss_total:.4f}',
                # filename='cgl-{epoch:02d}-{valid_loss_total_ema:.4f}',
                # save_last=True,
                every_n_train_steps=pargs.max_steps,
            )
            callbacks.append(ckpt_callback)
        else:
            ckpt_callback = None

        wandb_logger = pl_loggers.WandbLogger(experiment=wandb_run, save_dir=pargs.output)
        trainer = pl.Trainer(
            # max_epochs=pargs.max_epochs, 
            max_steps=pargs.max_steps,
            check_val_every_n_epoch=1, # every 1 epochs
            gpus=pargs.gpus,
            logger=wandb_logger,
            callbacks=callbacks,
            resume_from_checkpoint=ckpt_path,
            profiler='simple' if pargs.profile else None,
            terminate_on_nan=True,
            log_every_n_steps=pargs.log_freq,
            track_grad_norm=1,
        )

        if pargs.train:
            # pre-load the node_emb_ckpt if given
            if node_emb_ckpt:
                print(f'Loading weights for the Node Embedding module from checkpoint {node_emb_ckpt} ...')
                state_dict = torch.load(node_emb_ckpt, map_location=model.device)['state_dict']
                model.node_emb_module.load_state_dict(state_dict)
                print('Checkpoint Loaded.')

            trainer.fit(model, train_dataloader=conf.train_dloader, val_dataloaders=conf.valid_dloader)
            if ckpt_callback:
                ckpt_path = ckpt_callback.best_model_path
                print(f'Best model is stored in {ckpt_path}.')


        # Test with the model stored on ckpt_path based on validation set
        if pargs.test:
            if not ckpt_path:
                raise ValueError('Checkpoint path is not given.')
            else:
                print(f'Loading the checkpoint {str(Path(ckpt_path).absolute())} ...')
                state_dict = torch.load(ckpt_path, map_location=model.device)['state_dict']
                model.load_state_dict(state_dict)
                print('Checkpoint Loaded.')

            mse_evaluator = GraphEvaluator('mse')
            bc_evaluator = GraphEvaluator('rocauc')
            loaders = dict(train=conf.train_dloader, valid=conf.valid_dloader, test=conf.test_dloader)
            for mode, loader in loaders.items():
                prediction_list = []
                true_list = []
                for batch in tqdm(loader):
                    device = torch.device('cuda') if pargs.gpus != 0 else torch.device('cpu')
                    model = model.to(device)
                    if isinstance(batch, dict):
                        batch = ParamDict({k: v.to(device) for k, v in batch.items()})
                    else:
                        batch = batch.to(device)
                    batch_results = model.predict(batch, compute_loss=True)
                    prediction_list.append({k:batch_results.output[k] for k in model.output_labels_dict})
                    true_list.append({k:batch[k] for k in model.output_labels_dict})
                    

                prediction_dict = listdict2dictlist(prediction_list)
                prediction_dict = {k: torch.cat(v, 0) for k, v in prediction_dict.items()}
                true_dict = listdict2dictlist(true_list)
                true_dict = {k: torch.cat(v, 0) for k, v in true_dict.items()}

                # Regression tasks
                if model.reg_keys:
                    # TODO: clean up the data so that we don't have to deal with non-sense in here
                    if 'ugbw' in model.reg_keys:
                        key = 'ugbw'
                        cond = (true_dict[key] < 5).squeeze(-1)
                        prediction_dict[key] = prediction_dict[key][cond]
                        true_dict[key] = true_dict[key][cond]

                    # if 'tset' in model.reg_keys:
                    #     key = 'tset'
                    #     cond = (true_dict[key] != true_dict[key].max()).squeeze(-1)
                    #     prediction_dict[key] = prediction_dict[key][cond]
                    #     true_dict[key] = true_dict[key][cond]


                    y_pred = torch.stack([prediction_dict[k] for k in model.reg_keys], -1)
                    # in mlp dataset the true output has shape of (batch, 1)
                    y_true = torch.stack([true_dict[k].squeeze(-1) for k in model.reg_keys], -1)
                    mse_dict = mse_evaluator.eval(dict(y_pred=y_pred, y_true=y_true))

                    for label in model.reg_keys:
                        y_pred = prediction_dict[label].detach().cpu().numpy()
                        y_true = true_dict[label].squeeze(-1).detach().cpu().numpy()
                        plt.close()
                        _, ax = plt.subplots(1, 1)
                        ax.scatter(y_true, y_pred, s=1)
                        min_true, max_true = y_true.min(), y_true.max()
                        x = np.linspace(min_true, max_true, 1000)
                        ax.plot(x, x, linestyle='dashed', color='k', linewidth=1)
                        ax.set_xlabel('True Value')
                        ax.set_ylabel('Pred Value')
                        ax.set_title(label)
                        wandb_run.log({f'{mode}_chart_{label}': wandb.Image(ax)})
                    
                    wandb_run.summary[f'{mode}_mse'] = mse_dict['mse']

                # classification tasks
                if model.class_keys:
                    y_pred = torch.stack([prediction_dict[k] for k in model.class_keys], -1)
                    y_true = torch.stack([true_dict[k].squeeze(-1) for k in model.class_keys], -1)
                    bc_dict = bc_evaluator.eval(dict(y_pred=y_pred, y_true=y_true))

                    wandb_run.summary[f'{mode}_rocauc'] = bc_dict['rocauc']


    def get_conf(self, pargs):
        conf_path = pargs.path
        print(f'Loading from the config file {conf_path}')
        conf_module = imp.load_source('conf', conf_path)

        conf = ParamDict(
            exp_name=conf_module.exp_name, 
            model=conf_module.model,
            train_dloader=conf_module.train_dloader,
            valid_dloader=conf_module.valid_dloader,
            test_dloader=conf_module.test_dloader,
            output_labels_dict=conf_module.output_labels_dict,
        )

        # add ckpt if it exists
        if hasattr(conf_module, 'node_emb_ckpt'):
            conf.update(node_emb_ckpt=conf_module.node_emb_ckpt)

        return conf


def parse_args():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--path', type=str, help='config file path')
    parser.add_argument('--train', type=int, default=1, help='set to zero to skip training')
    parser.add_argument('--gpus', type=int, default=0, help='set the training on gpu')
    parser.add_argument('--output', type=str, default='/store/nosnap/results/cgl', 
                        help='The output directory')
    parser.add_argument('--max_epochs', type=int, default=100, 
                        help='The maximum number of training epochs')
    parser.add_argument('--max_steps', type=int, default=None, 
                        help='The maximum number of training steps (if earlier than max_epochs)')
    parser.add_argument('--ckpt', type=str, help='Resume from this checkpoint if valid.')
    parser.add_argument('--node_emb_ckpt', type=str, help='Load the node embedding module from this checkpoint, priority is with specification in the config file.')
    parser.add_argument('--seed', type=int, default=0, help='Random seed')
    parser.add_argument('--wandb_id', type=str, help='Wandb id to allow for resuming')
    parser.add_argument('--run_name', type=str, default=None, help='Wandb run name, if not provided the deafult of wandb is used')
    parser.add_argument('--profile', type=int, default=0, help='Set to 1 to profile individual steps during training')
    parser.add_argument('--test', type=int, default=1, help='Set to 0 to skip running the test scripts')
    parser.add_argument('--log_freq', type=int, default=10, help='Wandb log every n steps')
    
    return parser.parse_args()


if __name__ == '__main__':
    Trainer(parse_args())