import torch
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler
import os

class Parallel_Trainer(object):
    def __init__(self, synch, config):
        self.synch = synch
        if synch:
            local_rank = int(os.environ['LOCAL_RANK'])
            torch.cuda.set_device(local_rank)
            self.device = torch.device("cuda", local_rank)
            self.local_rank = local_rank
            torch.distributed.init_process_group("nccl", world_size = config.n_gpus, 
                                                rank = local_rank, init_method = 'env://')
        else:
            self.device = config.device
    
    def get_device(self):
        return self.device

    def model_parallel(self, model):
        model = model.to(self.device)
        model.train()
        if self.synch:
            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)   
            model = DistributedDataParallel(model.cuda(self.local_rank), device_ids=[self.local_rank],
                                            find_unused_parameters=True)
            return model
        else:
            return model
    
    def dataset_parallel(self, dataset):
        if self.synch:
            train_sampler = DistributedSampler(dataset)
            return train_sampler
        else:
            return None
    
    def get_loader(self, dataset, epoch, bsz, num_workers, train_sampler = None):
        if self.synch:
            train_sampler.set_epoch(epoch)
            dataloader = DataLoader(dataset, bsz, pin_memory=True,
                                num_workers= num_workers, sampler=train_sampler)
            return dataloader
        else:
            dataloader = DataLoader(dataset, bsz, num_workers = num_workers, shuffle = True,
                                    drop_last = False)
        return dataloader
    
    def save_model(self, ckpt, model, path):
        if self.synch:
            if self.local_rank == 0:
                ckpt['state_dict'] = model.module.state_dict()
                torch.save(ckpt, path)
        
        else:
            ckpt['state_dict'] = model.state_dict()
            torch.save(ckpt, path)
        
    def write_log(self, log_file, content):
        if self.synch:
            if self.local_rank == 0:
                log_file.write(content + '\n')
                log_file.flush()
            
        else:
            log_file.write(content + '\n')
            log_file.flush()
    
    def call_func(self, func, *args, **kwargs):
        if self.synch:
            if self.local_rank == 0:
                func(*args, **kwargs)
        
        else:
            func(*args, **kwargs)
        
    def get_model(self, model):
        if self.synch:
            return model.module

        else:
            return model