import torch
import lib.modeling.resnet as resnet
import lib.modeling.semseg_heads as snet
import torch.nn as nn
import torch.optim as optim
import utils.resnet_weights_helper as resnet_utils
from torch.autograd import Variable
from roi_data.loader import RoiDataLoader, MinibatchSampler, collate_minibatch, collate_minibatch_semseg
from datasets.roidb import combined_roidb_for_training, combined_roidb_for_training_semseg
import os
import numpy as np
import nn as mynn
import cv2
from modeling.model_builder_3DSD import Generalized_3DSD
from modeling.model_builder_PSP3D import DispSeg
from core.config import cfg, cfg_from_file, cfg_from_list, assert_and_infer_cfg

#load net
class load_net(nn.Module):
    def __init__(self):
        super(load_net, self).__init__()

        build=snet.ModelBuilder()
        fc_dim = 2048
        self.encoder = build.build_encoder(
                arch= 'resnet50_dilated8',
                fc_dim=fc_dim)
        self.decoder = build.build_decoder(
                arch = 'ppm_bilinear',
                num_class=19,
                fc_dim=fc_dim,
                use_softmax=False)
    def _init_modules(self):
        resnet_utils.load_pretrained_imagenet_weights(self)

    def forward(self, data):
        pred=self.decoder(self.encoder(data, return_feature_maps=True))
        
        pred = nn.functional.interpolate(
                pred, size=[128,128],
                mode='bilinear', align_corners=False)
        pred = nn.functional.log_softmax(pred, dim=1)
        return pred
def dataloader(bs, gpus):
    inputs = {}
    inputs['data'] = Variable(torch.randn(2*bs, 3, 128, 128)).to('cuda')
    inputs['semseg_label_0'] = Variable(torch.LongTensor(
        np.random.randint(0, 19, (bs, 128//8, 128//8), dtype=np.long))).to('cuda')
    inputs['disp_label_0'] = Variable(torch.rand(bs,  128//8, 128//8)).to('cuda')
    inputs['disp_scans'] = Variable(torch.arange(0,
            cfg.DISP.MAX_DISPLACEMENT).float().view(1,cfg.DISP.MAX_DISPLACEMENT,1,1).repeat(bs,1,1,1)).to('cuda')
    inputs['semseg_scans'] = Variable(torch.arange(0,
            cfg.MODEL.NUM_CLASSES).float().view(1, cfg.MODEL.NUM_CLASSES, 1, 1).repeat(bs,1,1,1)).to('cuda')
    return inputs

cfg_file = 'e2e_segdisp-R-50_3Dpool_1x.yaml'
cfg_from_file(cfg_file)
print (cfg.SEM)
print (cfg.DISP)
#cfg_from_list(cfg_file)
#assert_and_infer_cfg()
devices_ids=[5]
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(ids) for ids in devices_ids])
torch.backends.cudnn.benchmark=True
#torch.cuda.set_device(3)
len_gpus = len(devices_ids)
batch_size = 2 * len_gpus
#net = mynn.DataParallel(load_net().to('cuda'), minibatch=True)
net = mynn.DataParallel(DispSeg().to('cuda'), minibatch=True)
optimizer = optim.SGD(net.parameters(), lr=0.000875, momentum=0.9)
criterion = nn.NLLLoss(ignore_index=255)
#dataloader= dataloader(batch_size, len_gpus)
for i in range(10):
#for i, inputs in zip(range(1000), dataloader):
    inputs = dataloader(batch_size, len_gpus)
    for key in inputs:
        inputs[key] = torch.chunk(inputs[key], chunks=len_gpus, dim=0)
    optimizer.zero_grad()
    loss=net(**inputs)
    optimizer.step()
    for k in loss['losses'].keys():
        print (loss['losses'][k].item())


[((57, 0, 57, 23), 'core.config.cfg_from_file', 'cfg_from_file', ({(57, 14, 57, 22): 'cfg_file'}, {}), '(cfg_file)', False, 'from core.config import cfg, cfg_from_file, cfg_from_list, assert_and_infer_cfg\n'), ((71, 12, 71, 40), 'torch.nn.NLLLoss', 'nn.NLLLoss', (), '', True, 'import torch.nn as nn\n'), ((23, 14, 23, 33), 'lib.modeling.semseg_heads.ModelBuilder', 'snet.ModelBuilder', ({}, {}), '()', True, 'import lib.modeling.semseg_heads as snet\n'), ((34, 8, 34, 59), 'utils.resnet_weights_helper.load_pretrained_imagenet_weights', 'resnet_utils.load_pretrained_imagenet_weights', ({(34, 54, 34, 58): 'self'}, {}), '(self)', True, 'import utils.resnet_weights_helper as resnet_utils\n'), ((39, 15, 41, 53), 'torch.nn.functional.interpolate', 'nn.functional.interpolate', (), '', True, 'import torch.nn as nn\n'), ((42, 15, 42, 53), 'torch.nn.functional.log_softmax', 'nn.functional.log_softmax', (), '', True, 'import torch.nn as nn\n'), ((77, 22, 77, 70), 'torch.chunk', 'torch.chunk', (), '', False, 'import torch\n'), ((69, 24, 69, 33), 'modeling.model_builder_PSP3D.DispSeg', 'DispSeg', ({}, {}), '()', False, 'from modeling.model_builder_PSP3D import DispSeg\n'), ((46, 30, 46, 60), 'torch.randn', 'torch.randn', ({(46, 42, 46, 46): '2 * bs', (46, 48, 46, 49): '3', (46, 51, 46, 54): '128', (46, 56, 46, 59): '128'}, {}), '(2 * bs, 3, 128, 128)', False, 'import torch\n'), ((49, 38, 49, 69), 'torch.rand', 'torch.rand', ({(49, 49, 49, 51): 'bs', (49, 54, 49, 60): '128 // 8', (49, 62, 49, 68): '128 // 8'}, {}), '(bs, 128 // 8, 128 // 8)', False, 'import torch\n'), ((48, 8, 48, 69), 'numpy.random.randint', 'np.random.randint', (), '', True, 'import numpy as np\n'), ((50, 36, 51, 38), 'torch.arange', 'torch.arange', ({(50, 49, 50, 50): '0', (51, 12, 51, 37): 'cfg.DISP.MAX_DISPLACEMENT'}, {}), '(0, cfg.DISP.MAX_DISPLACEMENT)', False, 'import torch\n'), ((52, 38, 53, 34), 'torch.arange', 'torch.arange', ({(52, 51, 52, 52): '0', (53, 12, 53, 33): 'cfg.MODEL.NUM_CLASSES'}, {}), '(0, cfg.MODEL.NUM_CLASSES)', False, 'import torch\n')]