import random

import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist

from .tester import Tester

class WorkerInit(object):
    def __init__(self, rank, num_workers):
        self.rank = rank
        self.num_workers = num_workers
    def func(self, pid):
        print(f'[rank{self.rank}] setting worker seed {self.rank*self.num_workers+pid}', flush=True)
        np.random.seed(self.rank*self.num_workers+pid)

class TesterDeter(Tester):

    def __init__(self, C_train, C_test):
        super().__init__(C_train, C_test)

        if self.config.get('deterministic', False):
            if self.config.get('cudnn_deterministic', True):
                cudnn.deterministic = True
                cudnn.benchmark = False
            else:
                cudnn.benchmark = True
            seed = self.config.get('random_seed', 0)
            worker_rank = self.config.get('worker_rank', False)
            if worker_rank:
                worker_init = WorkerInit(self.C_train.rank, self.config.workers)
            else:
                worker_init = WorkerInit(0, 0)
            self.worker_init_fn = worker_init.func
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            torch.cuda.manual_seed(seed)
            dist.barrier()
            if self.C_train.rank == 0:
                self.logger.info(f'deterministic mode, seed: {seed}, worker_rank: {worker_rank},\
                                   cudnn_deterministic: {self.config.get("cudnn_deterministic", True)}')
            dist.barrier()
        else:
            self.worker_init_fn = None