import copy
import os
import random
import time

import core
import core.models.decoders as decoders
import core.models.backbones as backbones
import core.models.necks as necks
import core.data.datasets as datasets
from core.models.model_entry import model_entry
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist

from core.distributed_utils import DistributedGivenIterationSampler

from torch.utils.data import DataLoader

from .solver import Solver

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 SolverDeter(Solver):

    def __init__(self, C):
        super().__init__(C)

        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.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)
            torch.cuda.manual_seed_all(seed)
            os.environ['PYTHONHASHSEED'] = str(seed)
            if dist.is_initialized():
                dist.barrier()
            if self.C.rank == 0:
                self.logger.info(f'deterministic mode, seed: {seed}, worker_rank: {worker_rank},\
                                   cudnn_deterministic: {self.config.get("cudnn_deterministic", True)}')
            if dist.is_initialized():
                dist.barrier()
        else:
            self.worker_init_fn = None

    def create_dataloader(self):
        config = self.config
        ginfo = self.ginfo

        self.sampler = DistributedGivenIterationSampler(
                        self.dataset, config.max_iter, config.sampler.batch_size,
                        world_size=ginfo.task_size, rank=ginfo.task_rank,
                        last_iter=self.last_iter,
                        shuffle_strategy=config.sampler.shuffle_strategy,
                        random_seed=ginfo.task_random_seed,
                        ret_save_path=config.sampler.get('ret_save_path', None))


        self.loader = DataLoader(self.dataset, batch_size=config.sampler.batch_size,
                            shuffle=False, num_workers=config.workers,
                            pin_memory=False, sampler=self.sampler, worker_init_fn=self.worker_init_fn)