import os
import torch
import hydra
import random
import pytorch_lightning as pl
from pytorch_lightning import loggers as pl_loggers
from lib.utils.common import hydra_params_to_dotdict
import numpy as np

torch.manual_seed(3)
random.seed(3)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(3)

def set_cfg(cfg):
    mode = cfg['run_mode']
    assert mode == 'fit' or mode == 'test'
    if mode == 'fit':
        wb_logger = pl_loggers.WandbLogger(name=cfg['exp_name'], project=cfg['project_name'],
                                           save_dir=cfg['experiment_path'], sync_step=True, offline=False)
        wb_upload_path = wb_logger.experiment.dir
        cfg['preds_path'] = os.path.join(wb_upload_path, mode, 'overall/predictions')
        cfg['best_preds_path'] = os.path.join(wb_upload_path, mode, 'best/predictions')
        cfg['checkpoints_path'] = os.path.join(wb_upload_path, mode, 'best/checkpoints')

    elif mode == 'test':
        wb_logger = None
        if cfg['manual_ckpt']:
            print("[INFO.TEST.LOAD_MANUAL_CHECKPOINT]")
            pass # checkpoint will loaded inside on_test_start callback.
        else:
            print("[INFO.TEST.LOAD_AUTOMATIC_CHECKPOINT]")
            best_checkpoint_dir = os.path.join(cfg['experiment_path'], 'wandb', 'latest-run', 'files', 'fit', 'best',
                                            'checkpoints')
            best_checkpoint = [i for i in os.listdir(best_checkpoint_dir) if '.ckpt' in i]
            print(best_checkpoint)
            assert len(best_checkpoint) == 1
            cfg['best_checkpoint_path'] = os.path.join(best_checkpoint_dir, best_checkpoint[0])

        eval_path = os.path.join(cfg['experiment_path'], 'wandb', 'latest-run', 'files', 'test')
        cfg['preds_path'] = os.path.join(eval_path, 'overall/predictions')
        cfg['best_preds_path'] = os.path.join(eval_path, 'best/predictions')

    else:
        raise NotImplementedError('Mode {} not implemented.'.format(mode))

    return cfg, wb_logger

@hydra.main("lib/config/config.yaml")
def main(cfg):
    cfg, wb_logger = set_cfg(cfg=cfg)
    model = hydra.utils.instantiate(cfg.task_data_model, hydra_params_to_dotdict(cfg))
    gpus = list(range(torch.cuda.device_count()))
    trainer = pl.Trainer(
        gpus=gpus,
        min_epochs=cfg.epochs,
        max_epochs=cfg.epochs,
        check_val_every_n_epoch=1,
        num_sanity_val_steps=0,
        distributed_backend=cfg.distrib_backend,
        logger=wb_logger
    )

    if cfg['run_mode'] == 'fit':
        trainer.fit(model)
        trainer.test(ckpt_path='best')
        
    elif cfg['run_mode'] == 'test':
        if cfg['manual_ckpt']:
            pass
        else:
            print("________________________________________________________")
            print("\n[INFO.CHECKPOINT]: {}".format(cfg['best_checkpoint_path']))
            checkpoint = torch.load(cfg['best_checkpoint_path'], map_location=lambda storage, loc: storage)
            model.load_state_dict(checkpoint['state_dict'], strict=False)  # eg: embeddings will be loaded separately on model init
        trainer.test(model)
    else:
        raise NotImplementedError;


if __name__ == "__main__":
    main()
