import multiprocessing as mp
import sys

import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from torch.utils.data import DataLoader
from pytorch_lightning.plugins import DDPPlugin

from feeder.ntu_feeder import Feeder_triple
from tools import load_config, PeriodicCheckpoint
from net.train_learners import TrainingLearner

pl.seed_everything(123)
# torch.use_deterministic_algorithms(True)


if __name__ == '__main__':

    cfg = load_config(name='Training')

    # data loading
    train_feeder = Feeder_triple(cfg.train_feeder_args.data_path,
                                 cfg.train_feeder_args.label_path)
    train_loader = DataLoader(
        dataset=train_feeder,
        batch_size=cfg.batch_size // len(cfg.device),
        shuffle=True,
        pin_memory=True,    # set True when memory is abundant
        num_workers=mp.cpu_count()//3,
        persistent_workers=True,
        drop_last=True)

    wandb_logger = None
    if not cfg.disable_wandb:
        # init wandb logger
        wandb_logger = WandbLogger(project='aimclr', group='dev')

    # init self-supervised learner
    learner = TrainingLearner(cfg)

    if wandb_logger is not None:
        wandb_logger.watch(learner, log_freq=10)

    checkpoint_callback = PeriodicCheckpoint(dirpath=cfg.work_dir, every=cfg.save_interval)

    # init trainer
    trainer = pl.Trainer(
        gpus=cfg.device,
        max_epochs=cfg.num_epoch,
        accumulate_grad_batches=1,
        sync_batchnorm=True,
        strategy=DDPPlugin(find_unused_parameters=False),
        num_nodes=1,
        logger=wandb_logger,
        callbacks=[checkpoint_callback])

    # start training
    ckpt_path = cfg.resume_from if cfg.resume_from != 'None' else None
    trainer.fit(learner, train_loader, ckpt_path=ckpt_path)
