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

def init_seeds(seed = 0, cuda_deterministic = True):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if cuda_deterministic:
        cudnn.deterministic = True
        cudnn.benchmark = False
    else:
        cudnn.determinisstic = False
        cudnn.benchmark = True
        
def reduce_mean(tensor, nprocs):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
    rt /= nprocs
    return rt

def process_input(inputs,local_rank,non_blocking):
    inputs['img_0'] = inputs['img_0'].cuda(local_rank, non_blocking)
    inputs['img_1'] = inputs['img_1'].cuda(local_rank, non_blocking)
    inputs['img_2'] = inputs['img_2'].cuda(local_rank, non_blocking)
    inputs['img_3'] = inputs['img_3'].cuda(local_rank, non_blocking)
    # inputs['img_4'] = inputs['img_4'].cuda(local_rank, non_blocking)
    # inputs['img_5'] = inputs['img_5'].cuda(local_rank, non_blocking)
    # inputs['img_6'] = inputs['img_6'].cuda(local_rank, non_blocking)
    # inputs['img_7'] = inputs['img_7'].cuda(local_rank, non_blocking)
    # inputs['img_8'] = inputs['img_8'].cuda(local_rank, non_blocking)
    # inputs['img_9'] = inputs['img_9'].cuda(local_rank, non_blocking)
    
    return inputs
def process_output(targets, local_rank, non_blocking):
    targets['annotation_laban_0'] = targets['annotation_laban_0'].cuda(local_rank, non_blocking)
    # targets['annotation_laban_1'] = targets['annotation_laban_1'].cuda(local_rank, non_blocking)
    # targets['annotation_laban_2'] = targets['annotation_laban_2'].cuda(local_rank, non_blocking)
    # targets['annotation_laban_3'] = targets['annotation_laban_3'].cuda(local_rank, non_blocking)
    # targets['annotation_laban_4'] = targets['annotation_laban_4'].cuda(local_rank, non_blocking)
    # targets['annotation_laban_5'] = targets['annotation_laban_5'].cuda(local_rank, non_blocking)
    # targets['annotation_laban_6'] = targets['annotation_laban_6'].cuda(local_rank, non_blocking)
    # targets['annotation_laban_7'] = targets['annotation_laban_7'].cuda(local_rank, non_blocking)
    # targets['annotation_laban_8'] = targets['annotation_laban_8'].cuda(local_rank, non_blocking)
    # targets['annotation_laban_9'] = targets['annotation_laban_9'].cuda(local_rank, non_blocking)
    
    targets['laban_0_mask'] = targets['laban_0_mask'].cuda(local_rank, non_blocking)
    # targets['laban_1_mask'] = targets['laban_1_mask'].cuda(local_rank, non_blocking)
    # targets['laban_2_mask'] = targets['laban_2_mask'].cuda(local_rank, non_blocking)
    # targets['laban_3_mask'] = targets['laban_3_mask'].cuda(local_rank, non_blocking)
    # targets['laban_4_mask'] = targets['laban_4_mask'].cuda(local_rank, non_blocking)
    # targets['laban_5_mask'] = targets['laban_5_mask'].cuda(local_rank, non_blocking)
    # targets['laban_6_mask'] = targets['laban_6_mask'].cuda(local_rank, non_blocking)
    # targets['laban_7_mask'] = targets['laban_7_mask'].cuda(local_rank, non_blocking)
    # targets['laban_8_mask'] = targets['laban_8_mask'].cuda(local_rank, non_blocking)
    # targets['laban_9_mask'] = targets['laban_9_mask'].cuda(local_rank, non_blocking)
    return targets
def process_input2device(inputs, device):
    inputs['img_0'] = inputs['img_0'].to(device)
    inputs['img_1'] = inputs['img_1'].to(device)
    inputs['img_2'] = inputs['img_2'].to(device)
    inputs['img_3'] = inputs['img_3'].to(device)
    # inputs['img_4'] = inputs['img_4'].to(device)
    # inputs['img_5'] = inputs['img_5'].to(device)
    # inputs['img_6'] = inputs['img_6'].to(device)
    # inputs['img_7'] = inputs['img_7'].to(device)
    # inputs['img_8'] = inputs['img_8'].to(device)
    # inputs['img_9'] = inputs['img_9'].to(device)
    
    return inputs
def process_output2device(targets, device):
    targets['annotation_laban_0'] = targets['annotation_laban_0'].to(device)
    # targets['annotation_laban_1'] = targets['annotation_laban_1'].to(device)
    # targets['annotation_laban_2'] = targets['annotation_laban_2'].to(device)
    # targets['annotation_laban_3'] = targets['annotation_laban_3'].to(device)
    # targets['annotation_laban_4'] = targets['annotation_laban_4'].to(device)
    # targets['annotation_laban_5'] = targets['annotation_laban_5'].to(device)
    # targets['annotation_laban_6'] = targets['annotation_laban_6'].to(device)
    # targets['annotation_laban_7'] = targets['annotation_laban_7'].to(device)
    # targets['annotation_laban_8'] = targets['annotation_laban_8'].to(device)
    # targets['annotation_laban_9'] = targets['annotation_laban_9'].to(device)
    
    targets['laban_0_mask'] = targets['laban_0_mask'].to(device)
    # targets['laban_1_mask'] = targets['laban_1_mask'].to(device)
    # targets['laban_2_mask'] = targets['laban_2_mask'].to(device)
    # targets['laban_3_mask'] = targets['laban_3_mask'].to(device)
    # targets['laban_4_mask'] = targets['laban_4_mask'].to(device)
    # targets['laban_5_mask'] = targets['laban_5_mask'].to(device)
    # targets['laban_6_mask'] = targets['laban_6_mask'].to(device)
    # targets['laban_7_mask'] = targets['laban_7_mask'].to(device)
    # targets['laban_8_mask'] = targets['laban_8_mask'].to(device)
    # targets['laban_9_mask'] = targets['laban_9_mask'].to(device)
    return targets
    

