import os
from tqdm import tqdm
import torch.backends.cudnn as cudnn
import torch
from datasets import ImageNetInstance, ImageNetInstanceLMDB
from torchvision import transforms
import argparse
from BaseTaskModel.task_network import get_moco_network, get_swav_network, get_selfboost_network, get_minmaxent_network, get_simclr_network, get_sup_network, get_dino_network
from torch.utils.data import DataLoader
from PIL import ImageFile, Image
import torch.distributed as dist
from lars import *
ImageFile.LOAD_TRUNCATED_IMAGES = True
import warnings
warnings.filterwarnings('ignore')

def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    tensors_gather = [torch.ones_like(tensor)
        for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)

    output = torch.cat(tensors_gather, dim=0)
    return output

def main():
    parser = argparse.ArgumentParser("The first stage of BoostrapSelfSup")
    parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed parallel')
    parser.add_argument("--task", type=str, default="moco", help="the pretraining models")
    parser.add_argument("--pretrained_path", type=str, default="", help="the pretraining models")
    parser.add_argument("--save_path", type=str, default="", help="where to save the memory_bank")
    parser.add_argument("--backbone", type=str, default="resnet50")
    parser.add_argument("--data_path", type=str, default="~/ILSVRC2012/", help="the data path")
    parser.add_argument("--batch_size", type=int, default=32, help="batch size")
    parser.add_argument("--img_size", type=int, default=224, help="image size")
    parser.add_argument("--feat_dim", type=int, default=128, help="feat dimension")
    parser.add_argument("--feature_layer", type=str, default='lowdim', help="feature layer")
    parser.add_argument('--use-lmdb', action='store_true')
    args = parser.parse_args()
    
    pretrained_path = os.path.expanduser(args.pretrained_path)
    save_path = os.path.expanduser(args.save_path)
    data_path = os.path.expanduser(args.data_path)
    batch_size = args.batch_size
    feat_dim = args.feat_dim

    dist.init_process_group(backend='nccl')
    torch.cuda.set_device(args.local_rank)

    # network = ResNet(50, frozen_stages=4)
    if args.task == 'moco':
        network = get_moco_network(pretrained_path, feature_layer=args.feature_layer)
    elif args.task == 'swav':
        network = get_swav_network(pretrained_path, feature_layer=args.feature_layer)
    elif args.task == 'selfboost':
        network = get_selfboost_network(pretrained_path, feature_layer=args.feature_layer)
    elif args.task == 'minmaxent':
        network = get_minmaxent_network(args.backbone, pretrained_path, feature_layer=args.feature_layer)
    elif args.task == 'dino':
        network = get_dino_network(args.backbone, pretrained_path, feature_layer=args.feature_layer)
    elif args.task == 'simclr':
        network = get_simclr_network(args.backbone, pretrained_path, feature_layer=args.feature_layer)
    elif args.task == 'sup':
        network = get_sup_network(args.backbone, pretrained_path, feature_layer=args.feature_layer)
    else:
        raise NotImplementedError
    
    network.cuda(args.local_rank)
    network = torch.nn.parallel.DistributedDataParallel(network, device_ids=[args.local_rank])

    cudnn.benchmark = True


    augmentation = transforms.Compose([
            transforms.Resize(int(256*args.img_size/224), interpolation=Image.BICUBIC),
            transforms.CenterCrop(args.img_size),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    
    if args.use_lmdb:
        train_dataset = ImageNetInstanceLMDB(root=data_path, list_file='train.lmdb', transform=augmentation)
        val_dataset = ImageNetInstanceLMDB(root=data_path, list_file='val.lmdb', transform=augmentation)
    else:
        train_dataset = ImageNetInstance(root=os.path.join(data_path, 'train'), transform=augmentation)
        val_dataset = ImageNetInstance(root=os.path.join(data_path, 'val'), transform=augmentation)

    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=False, rank=args.local_rank)
    val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False, rank=args.local_rank)

    n_train_points = len(train_dataset)
    n_val_points = len(val_dataset)
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, pin_memory=True, num_workers=4)
    val_dataloader   = DataLoader(val_dataset,   batch_size=batch_size, sampler=val_sampler, pin_memory=True, num_workers=4)

    print("Initializing train memory bank: {} points.".format(n_train_points))
    train_memory_bank = torch.zeros(n_train_points, feat_dim).to("cpu").detach()

    print("Initializing val memory bank: {} points.".format(n_val_points))
    val_memory_bank = torch.zeros(n_val_points, feat_dim).to("cpu").detach()

    network.eval()
    train_sampler.set_epoch(0)
    val_sampler.set_epoch(0)
    for data in tqdm(train_dataloader):
        idx, img, _ = data
        idx = idx.cuda(args.local_rank, non_blocking=True)
        img = img.cuda(args.local_rank, non_blocking=True)
        if True: #args.backbone.startswith('resnet'):
            feature = network(img)
        else:
            feature = network.module.get_intermediate_layers(img, 4)
            feature = [x[:, 0] for x in feature]
            feature = torch.cat(feature, dim=-1)

        feature = concat_all_gather(feature.contiguous())
        idx = concat_all_gather(idx)

        with torch.no_grad():
            train_memory_bank[idx,:] = feature.detach().cpu()

    for data in tqdm(val_dataloader):
        idx, img, _ = data
        idx = idx.cuda(args.local_rank, non_blocking=True)
        img = img.cuda(args.local_rank, non_blocking=True)
        if True: #args.backbone.startswith('resnet'):
            feature = network(img)
        else:
            feature = network.module.get_intermediate_layers(img, 4)
            feature = [x[:, 0] for x in feature]
            feature = torch.cat(feature, dim=-1)

        feature = concat_all_gather(feature.contiguous())
        idx = concat_all_gather(idx)

        with torch.no_grad():
            val_memory_bank[idx,:] = feature.detach().cpu()

    if args.local_rank == 0:
        torch.save(
            {'train_memory_bank': train_memory_bank,
            'val_memory_bank':  val_memory_bank
            },
            args.save_path
        )

if __name__ == '__main__':
    main()

[((15, 0, 15, 33), 'warnings.filterwarnings', 'warnings.filterwarnings', ({(15, 24, 15, 32): '"""ignore"""'}, {}), "('ignore')", False, 'import warnings\n'), ((24, 4, 24, 72), 'torch.distributed.all_gather', 'torch.distributed.all_gather', (), '', False, 'import torch\n'), ((26, 13, 26, 45), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((30, 13, 30, 74), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({(30, 37, 30, 73): '"""The first stage of BoostrapSelfSup"""'}, {}), "('The first stage of BoostrapSelfSup')", False, 'import argparse\n'), ((44, 22, 44, 62), 'os.path.expanduser', 'os.path.expanduser', ({(44, 41, 44, 61): 'args.pretrained_path'}, {}), '(args.pretrained_path)', False, 'import os\n'), ((45, 16, 45, 50), 'os.path.expanduser', 'os.path.expanduser', ({(45, 35, 45, 49): 'args.save_path'}, {}), '(args.save_path)', False, 'import os\n'), ((46, 16, 46, 50), 'os.path.expanduser', 'os.path.expanduser', ({(46, 35, 46, 49): 'args.data_path'}, {}), '(args.data_path)', False, 'import os\n'), ((50, 4, 50, 43), 'torch.distributed.init_process_group', 'dist.init_process_group', (), '', True, 'import torch.distributed as dist\n'), ((51, 4, 51, 42), 'torch.cuda.set_device', 'torch.cuda.set_device', ({(51, 26, 51, 41): 'args.local_rank'}, {}), '(args.local_rank)', False, 'import torch\n'), ((72, 14, 72, 94), 'torch.nn.parallel.DistributedDataParallel', 'torch.nn.parallel.DistributedDataParallel', (), '', False, 'import torch\n'), ((91, 20, 91, 119), 'torch.utils.data.distributed.DistributedSampler', 'torch.utils.data.distributed.DistributedSampler', (), '', False, 'import torch\n'), ((92, 18, 92, 115), 'torch.utils.data.distributed.DistributedSampler', 'torch.utils.data.distributed.DistributedSampler', (), '', False, 'import torch\n'), ((96, 23, 96, 126), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((97, 23, 97, 124), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader\n'), ((108, 16, 108, 38), 'tqdm.tqdm', 'tqdm', ({(108, 21, 108, 37): 'train_dataloader'}, {}), '(train_dataloader)', False, 'from tqdm import tqdm\n'), ((125, 16, 125, 36), 'tqdm.tqdm', 'tqdm', ({(125, 21, 125, 35): 'val_dataloader'}, {}), '(val_dataloader)', False, 'from tqdm import tqdm\n'), ((22, 22, 22, 45), 'torch.ones_like', 'torch.ones_like', ({(22, 38, 22, 44): 'tensor'}, {}), '(tensor)', False, 'import torch\n'), ((55, 18, 55, 85), 'BaseTaskModel.task_network.get_moco_network', 'get_moco_network', (), '', False, 'from BaseTaskModel.task_network import get_moco_network, get_swav_network, get_selfboost_network, get_minmaxent_network, get_simclr_network, get_sup_network, get_dino_network\n'), ((85, 24, 85, 108), 'datasets.ImageNetInstanceLMDB', 'ImageNetInstanceLMDB', (), '', False, 'from datasets import ImageNetInstance, ImageNetInstanceLMDB\n'), ((86, 22, 86, 104), 'datasets.ImageNetInstanceLMDB', 'ImageNetInstanceLMDB', (), '', False, 'from datasets import ImageNetInstance, ImageNetInstanceLMDB\n'), ((143, 8, 148, 9), 'torch.save', 'torch.save', ({(144, 12, 146, 13): "{'train_memory_bank': train_memory_bank, 'val_memory_bank': val_memory_bank}", (147, 12, 147, 26): 'args.save_path'}, {}), "({'train_memory_bank': train_memory_bank, 'val_memory_bank':\n    val_memory_bank}, args.save_path)", False, 'import torch\n'), ((57, 18, 57, 85), 'BaseTaskModel.task_network.get_swav_network', 'get_swav_network', (), '', False, 'from BaseTaskModel.task_network import get_moco_network, get_swav_network, get_selfboost_network, get_minmaxent_network, get_simclr_network, get_sup_network, get_dino_network\n'), ((79, 12, 79, 48), 'torchvision.transforms.CenterCrop', 'transforms.CenterCrop', ({(79, 34, 79, 47): 'args.img_size'}, {}), '(args.img_size)', False, 'from torchvision import transforms\n'), ((80, 12, 80, 33), 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ({}, {}), '()', False, 'from torchvision import transforms\n'), ((81, 12, 81, 87), 'torchvision.transforms.Normalize', 'transforms.Normalize', (), '', False, 'from torchvision import transforms\n'), ((117, 22, 117, 48), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((122, 13, 122, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((134, 22, 134, 48), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((139, 13, 139, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((23, 23, 23, 57), 'torch.distributed.get_world_size', 'torch.distributed.get_world_size', ({}, {}), '()', False, 'import torch\n'), ((59, 18, 59, 90), 'BaseTaskModel.task_network.get_selfboost_network', 'get_selfboost_network', (), '', False, 'from BaseTaskModel.task_network import get_moco_network, get_swav_network, get_selfboost_network, get_minmaxent_network, get_simclr_network, get_sup_network, get_dino_network\n'), ((88, 46, 88, 78), 'os.path.join', 'os.path.join', ({(88, 59, 88, 68): 'data_path', (88, 70, 88, 77): '"""train"""'}, {}), "(data_path, 'train')", False, 'import os\n'), ((89, 44, 89, 74), 'os.path.join', 'os.path.join', ({(89, 57, 89, 66): 'data_path', (89, 68, 89, 73): '"""val"""'}, {}), "(data_path, 'val')", False, 'import os\n'), ((61, 18, 61, 105), 'BaseTaskModel.task_network.get_minmaxent_network', 'get_minmaxent_network', (), '', False, 'from BaseTaskModel.task_network import get_moco_network, get_swav_network, get_selfboost_network, get_minmaxent_network, get_simclr_network, get_sup_network, get_dino_network\n'), ((100, 24, 100, 61), 'torch.zeros', 'torch.zeros', ({(100, 36, 100, 50): 'n_train_points', (100, 52, 100, 60): 'feat_dim'}, {}), '(n_train_points, feat_dim)', False, 'import torch\n'), ((103, 22, 103, 57), 'torch.zeros', 'torch.zeros', ({(103, 34, 103, 46): 'n_val_points', (103, 48, 103, 56): 'feat_dim'}, {}), '(n_val_points, feat_dim)', False, 'import torch\n'), ((63, 18, 63, 100), 'BaseTaskModel.task_network.get_dino_network', 'get_dino_network', (), '', False, 'from BaseTaskModel.task_network import get_moco_network, get_swav_network, get_selfboost_network, get_minmaxent_network, get_simclr_network, get_sup_network, get_dino_network\n'), ((65, 18, 65, 102), 'BaseTaskModel.task_network.get_simclr_network', 'get_simclr_network', (), '', False, 'from BaseTaskModel.task_network import get_moco_network, get_swav_network, get_selfboost_network, get_minmaxent_network, get_simclr_network, get_sup_network, get_dino_network\n'), ((67, 18, 67, 99), 'BaseTaskModel.task_network.get_sup_network', 'get_sup_network', (), '', False, 'from BaseTaskModel.task_network import get_moco_network, get_swav_network, get_selfboost_network, get_minmaxent_network, get_simclr_network, get_sup_network, get_dino_network\n')]