import pandas as pd
import numpy as np
import datetime
from django.shortcuts import get_object_or_404
from bb.models import Throw, Player, DailyLoser, Gladiator

#bbdataset=pd.read_csv("C:\\Users\\yemmit\\Documents\\bb\\datenbasis\\BB-Mappe.csv", sep=';')
#bbdataset=pd.read_csv("H:\\BB-Mappe.csv", sep=';')
bbdataset=pd.read_csv("N:\\temp\\tim\\privat\\datenbasis\\BB-Mappe_ergz.csv", sep=';', dtype=str)
bbdataset_with0=bbdataset.replace(to_replace=np.nan, value='')

for x,date_str in enumerate(bbdataset.iloc[:,0]):
    player_i=0
    for player in bbdataset.columns[1:]:
        player_i+=1
        throw_list=bbdataset.ix[x,player]
        if type(throw_list) == type('str'):
            throw_list=throw_list.replace(',','')
            throw_i=0
            for throwround in range(len(throw_list)):
                throw_i+=1
                event_time=datetime.datetime(int(date_str[-4:]), int(date_str[3:5]), int(date_str[:2]),0,throw_i, player_i)
                input_player = get_object_or_404(Player, player_name=player)
                new_throw = Throw(result=int(throw_list[throwround]), event_time=event_time, player=input_player, round=throw_i)
                new_throw.save()
    throwround_lens = [len(z.replace(',','')) for z in bbdataset_with0.ix[x, 1:]]
    if max(throwround_lens) > 0:
        last_round_ix=[q+1 for q,z in enumerate(throwround_lens) if z==max(throwround_lens)]
        last_round_points=[int(z[-1]) for z in bbdataset_with0.iloc[x,last_round_ix]]
        loser = bbdataset.columns[[last_round_ix[q] for q, z in enumerate(last_round_points) if z == min(last_round_points)][0]]
        loser_da = DailyLoser(day=datetime.date(int(date_str[-4:]), int(date_str[3:5]), int(date_str[:2])),
                              loser=get_object_or_404(Player, player_name=loser), round=max(throwround_lens))
        loser_da.save()
        gladiator1_ix = [last_round_ix[q] for q, z in enumerate(last_round_points) if z > min(last_round_points) and z == 1]
        if len(gladiator1_ix) > 0:
            gladiator1 = bbdataset.columns[gladiator1_ix]
            for gladiator in gladiator1:
                gladiator_da = Gladiator(day=datetime.date(int(date_str[-4:]), int(date_str[3:5]), int(date_str[:2])),
                                     gladiator=get_object_or_404(Player, player_name=gladiator),
                                     points=1, round=max(throwround_lens))
                gladiator_da.save()
        gladiator15_ix=[last_round_ix[q] for q, z in enumerate(last_round_points) if z > min(last_round_points) and z == 2]
        if len(gladiator15_ix) > 0:
            gladiator15 = bbdataset.columns[gladiator15_ix]
            for gladiator in gladiator15:
                gladiator_da = Gladiator(day=datetime.date(int(date_str[-4:]), int(date_str[3:5]), int(date_str[:2])),
                                         gladiator=get_object_or_404(Player, player_name=gladiator),
                                         points=1.5, round=max(throwround_lens))
                gladiator_da.save()
        gladiator2_ix=[last_round_ix[q] for q, z in enumerate(last_round_points) if z > min(last_round_points) and z == 3]
        if len(gladiator2_ix) > 0:
            gladiator2 = bbdataset.columns[gladiator2_ix]
            for gladiator in gladiator2:
                gladiator_da = Gladiator(day=datetime.date(int(date_str[-4:]), int(date_str[3:5]), int(date_str[:2])),
                                         gladiator=get_object_or_404(Player, player_name=gladiator),
                                         points=2, round=max(throwround_lens))
                gladiator_da.save()


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

SEED = 0
import numpy as np

np.random.seed(SEED)
import tensorflow as tf

tf.set_random_seed(SEED)

import os, shutil, glob
from skimage import transform, io
from model import UNet
from utils import VIS, mean_IU
# configure args
from opts import *
import cv2

# save and compute metrics
vis = VIS(save_path=opt.checkpoint_path)
print('num_class = %d' % opt.num_class)

# configuration session
config = tf.ConfigProto(
    device_count={'GPU': 0}
)
# config.gpu_options.allow_growth = True
sess = tf.Session(config=config)

img_shape = [280, 500]  # Resized to be the same.

# define input holders
train_list = glob.glob(opt.data_path + '/train/img/0/*.png')
train_num = len(train_list)
test_list = glob.glob(opt.data_path + '/val/img/0/*.png')
test_num = len(test_list)
for idx, fname in enumerate(train_list):
    _, train_list[idx] = os.path.split(fname)
for idx, fname in enumerate(test_list):
    _, test_list[idx] = os.path.split(fname)

opt.iter_epoch = int(train_num / opt.batch_size)

label = tf.placeholder(tf.int32, shape=[None] + img_shape)
# define model
with tf.name_scope('unet'):
    model = UNet().create_model(img_shape=img_shape + [4], num_class=opt.num_class)
    img = model.input
    pred = model.output
# define loss
with tf.name_scope('weighted_cross_entropy'):
    cross_entropy_loss = tf.reduce_mean(
        tf.nn.sparse_softmax_cross_entropy_with_logits(labels=label, logits=pred))
# define optimizer
global_step = tf.Variable(0, name='global_step', trainable=False)
with tf.name_scope('learning_rate'):
    learning_rate = tf.train.exponential_decay(opt.learning_rate, global_step,
                                               opt.iter_epoch, opt.lr_decay, staircase=True)
train_step = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cross_entropy_loss,
                                                                          global_step=global_step)
''' Tensorboard visualization '''
# cleanup pervious info
if opt.load_from_checkpoint == '':
    cf = os.listdir(opt.checkpoint_path)
    for item in cf:
        if 'event' in item:
            os.remove(os.path.join(opt.checkpoint_path, item))
# define summary for tensorboard
tf.summary.scalar('cross_entropy_loss', cross_entropy_loss)
tf.summary.scalar('learning_rate', learning_rate)
summary_merged = tf.summary.merge_all()
# define saver
train_writer = tf.summary.FileWriter(opt.checkpoint_path, sess.graph)
saver = tf.train.Saver(max_to_keep=20)  # must be added in the end

''' Main '''
tot_iter = opt.iter_epoch * opt.epoch
init_op = tf.global_variables_initializer()

with tf.Session().as_default() as sess:  # .as_default():
    # restore from a checkpoint if exists
    # the name_scope can not change
    sess.run(init_op)
    if opt.load_from_checkpoint != '':
        try:
            # saver = tf.train.import_meta_graph('./checkpoints/unetxPascal/model-0.meta')
            saver.restore(sess, tf.train.latest_checkpoint(opt.load_from_checkpoint))
            print('--> load from checkpoint ' + opt.load_from_checkpoint)
        except:
            print('unable to load checkpoint ...')

    # debug
    start = global_step.eval()
    for it in range(start, tot_iter):
        # print('Iter = %d' % it)
        if (it + 1) % opt.iter_epoch == 0 or it == start:

            saver.save(sess, opt.checkpoint_path + 'model', global_step=global_step)
            print('save a checkpoint at ' + opt.checkpoint_path + 'model-' + str(it))
            print('start testing {} samples...'.format(test_num))
            for ti in range(50):
                ori_name, ext = os.path.splitext(test_list[ti])
                x = io.imread(opt.data_path + '/val/img/0/' + test_list[ti])
                x = transform.resize(x, img_shape)
                map_choice = np.random.randint(1, 16)
                selection_rect_img = io.imread(opt.data_path + '/my-VOC2012/rects/R' +
                                               ori_name[:-2] + '_' + str(map_choice) + '.png')
                ori_shape = selection_rect_img.shape
                selection_rect_img = transform.resize(selection_rect_img, output_shape=img_shape) > 0
                nonzero_list = np.nonzero(selection_rect_img)
                left, right, up, down = nonzero_list[1][0], nonzero_list[1][-1], \
                                        nonzero_list[0][0], nonzero_list[0][-1]
                inside_mask = np.zeros(selection_rect_img.shape, dtype=np.uint8)
                inside_mask[up:down + 1, left:right + 1] = 1

                selection_rect_img = (selection_rect_img == 0).astype(np.uint8)
                dis_map = cv2.distanceTransform(selection_rect_img, cv2.cv.CV_DIST_L2, cv2.cv.CV_DIST_MASK_PRECISE)
                dis_map = dis_map - 2 * inside_mask * dis_map + 128
                # cv2.imwrite('dmap.png', dis_map)

                # dmap = io.imread(opt.data_path+'/dis_map/R'+
                #                          ori_name[:-2]+'_'+str(map_choice)+'.png')
                dis_map = np.reshape(dis_map, dis_map.shape + (1,))
                x_batch = np.concatenate((x, dis_map), axis=2)
                x_batch = np.reshape(x_batch, (1,) + x_batch.shape)

                y = io.imread(opt.data_path + '/val/gt/0/' + ori_name[:-2] + '.png')
                y = transform.resize(y, img_shape) > 0
                y = y.astype(np.int32)
                y_batch = np.reshape(y, (1,) + y.shape)

                # tensorflow wants a different tensor order
                feed_dict = {
                    img: x_batch,
                    label: y_batch,
                }
                loss, pred_logits = sess.run([cross_entropy_loss, pred], feed_dict=feed_dict)
                pred_map_batch = np.argmax(pred_logits, axis=3)
                # import pdb; pdb.set_trace()
                for pred_map, y in zip(pred_map_batch, y_batch):
                    score = vis.add_sample(pred_map, y)
            vis.compute_scores(suffix=it)

        # x_batch, y_batch = next(train_generator)

        batch_index = np.random.randint(0, train_num, opt.batch_size)
        x_batch = np.zeros([opt.batch_size] + img_shape + [4])
        y_batch = np.zeros([opt.batch_size] + img_shape)
        for cnt, i in enumerate(batch_index):
            # print('batch_iter = %d' % cnt)
            ori_name, __ = os.path.splitext(train_list[i])
            x = io.imread(opt.data_path + '/train/img/0/' + train_list[i])
            x = transform.resize(x, img_shape) # Turned into [0, 1] float32 automatically
            map_choice = np.random.randint(1, 16)
            selection_rect_img = io.imread(opt.data_path + '/my-VOC2012/rects/R' +
                                           ori_name[:-2] + '_' + str(map_choice) + '.png')
            ori_shape = selection_rect_img.shape
            selection_rect_img = transform.resize(selection_rect_img, output_shape=img_shape) > 0
            nonzero_list = np.nonzero(selection_rect_img)
            left, right, up, down = nonzero_list[1][0], nonzero_list[1][-1], \
                                    nonzero_list[0][0], nonzero_list[0][-1]
            inside_mask = np.zeros(selection_rect_img.shape, dtype=np.uint8)
            inside_mask[up:down + 1, left:right + 1] = 1

            selection_rect_img = (selection_rect_img == 0).astype(np.uint8)
            dis_map = cv2.distanceTransform(selection_rect_img, cv2.cv.CV_DIST_L2, cv2.cv.CV_DIST_MASK_PRECISE)
            dis_map = dis_map - 2 * inside_mask * dis_map + 128

            dis_map = np.reshape(dis_map, dis_map.shape + (1,))
            x_batch[cnt] = np.concatenate((x, dis_map), axis=2)

            y = io.imread(opt.data_path + '/train/gt/0/' + ori_name[:-2] + '.png')
            y = transform.resize(y, img_shape) > 0
            y = y.astype(np.int32)
            y_batch[cnt] = y

        feed_dict = {img: x_batch,
                     label: y_batch
                     }
        _, loss, summary, lr, pred_logits = sess.run([train_step,
                                                      cross_entropy_loss,
                                                      summary_merged,
                                                      learning_rate,
                                                      pred
                                                      ], feed_dict=feed_dict)
        global_step.assign(it).eval()

        train_writer.add_summary(summary, it)
        if (it + 1) % 20 == 0:
            pred_map = np.argmax(pred_logits[0], axis=2)
            score, _ = mean_IU(pred_map, y_batch[0])
            print('max_pred = %d, max_gt = %d ' % (
                np.max(pred_map, (0, 1)), np.max(y_batch[0], (0, 1))))  # ('backup/'+str(it)+'.png',

            print(
                '[iter %d, epoch %.3f]: lr=%f loss=%f, mean_IU=%f' % (
                    it + 1, float(it + 1) / opt.iter_epoch, lr, loss, score))


"""The controllable constraint class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from robovat.math.pose import Pose
from robovat.simulation.constraint import Constraint
from robovat.utils.logging import logger

# These default constants.
NUM_STEPS_CHECK = 100
TIMEOUT = 1.0
POSITION_THRESHOLD = 0.01
EULER_THRESHOLD = np.pi / 36

class ControllableConstraint(Constraint):
    """Controllable constraint."""

    def __init__(self,
                 parent,
                 child,
                 joint_type='fixed',
                 joint_axis=[0, 0, 0],
                 parent_frame_pose=None,
                 child_frame_pose=None,
                 max_linear_velocity=None, max_angular_velocity=None,
                 max_force=None,
                 name=None):
        """Initialize.

        parent: The parent of the constraint.
        child: The child of the constraint.
        joint_type: Type of the joint.
        joint_axis: The rotation axis of the joint.
        parent_frame_pose: The joint pose in the parent frame.
        child_frame_pose: The joint pose in the child frame.
        max_linear_velocity: The maximum linear velocity of the joint.
        max_augular_velocity: The maximum angular velocity of the joint.
        max_force: The maximum force of the joint.
        name: Name of the joint.
        """
        super(ControllableConstraint, self).__init__(
            parent=parent,
            child=child,
            joint_type=joint_type,
            joint_axis=joint_axis,
            parent_frame_pose=parent_frame_pose,
            child_frame_pose=child_frame_pose,
            max_force=max_force,
            name=name)

        self._max_linear_velocity = max_linear_velocity
        self._max_angular_velocity = max_angular_velocity

        self.reset_targets()

    def reset_targets(self):
        """Reset the control variables.
        """
        self._target_pose = None
        self._target_linear_velocity = None
        self._target_angular_velocity = None
        self._start_time = None
        self._stop_time = None

    def is_ready(self):
        """Check if the constraint is ready.

        Returns:
            True if all control commands are done, False otherwise.
        """
        return self._target_pose is None

    def set_target_pose(self,
                        pose,
                        linear_velocity=None,
                        angular_velocity=None,
                        timeout=TIMEOUT):
        """Set target pose.
        """
        self._target_pose = pose
        self._target_linear_velocity = self.physics.time_step * (
            linear_velocity or self._max_linear_velocity)
        self._target_angular_velocity = self.physics.time_step * (
            angular_velocity or self._max_angular_velocity)
        self._start_time = self.physics.time()
        self._stop_time = self._start_time + timeout

        self._position_threshold = POSITION_THRESHOLD
        self._euler_threshold = EULER_THRESHOLD

    def update(self):
        """Update control and disturbances."""
        # Call the update function of the super class.
        super(ControllableConstraint, self).update()

        if self._target_pose is not None:
            self._update_pose_control()

            if self.physics.num_steps % NUM_STEPS_CHECK == 0:
                if self.check_reached() or self.check_timeout():
                    self.reset_targets()

    def _update_pose_control(self):
        """Update the pose control."""
        delta_position = self._target_pose.position - self.pose.position
        delta_position /= np.linalg.norm(delta_position)
        delta_position *= self._target_linear_velocity
        new_position = self.pose.position + delta_position

        delta_euler = (self._target_pose.euler - self.pose.euler + np.pi
                       ) % (2 * np.pi) - np.pi
        delta_euler = np.minimum(np.maximum(
            delta_euler,
            -self._target_angular_velocity),
            self._target_angular_velocity)
        new_euler = self.pose.euler + delta_euler
        new_euler[0] = (new_euler[0] + np.pi) % (2 * np.pi) - np.pi
        new_euler[1] = (new_euler[1] + 0.5 * np.pi) % np.pi - 0.5 * np.pi
        new_euler[2] = (new_euler[2] + np.pi) % (2 * np.pi) - np.pi

        new_pose = Pose((new_position, new_euler))

        self.pose = new_pose

    def check_reached(self):
        """Check if the specified joint positions are reached.

        Returns:
            True if the target has been reached, False otherwise.
        """
        delta_position = self._target_pose.position - self.pose.position
        position_reached = (
            abs(delta_position)[0] < self._position_threshold and
            abs(delta_position)[1] < self._position_threshold and
            abs(delta_position)[2] < self._position_threshold)

        delta_euler = self._target_pose.euler - self.pose.euler
        euler_reached = (
            abs(delta_euler)[0] % (2 * np.pi) < self._euler_threshold and
            abs(delta_euler)[1] % (2 * np.pi) < self._euler_threshold and
            abs(delta_euler)[2] % (2 * np.pi) < self._euler_threshold)

        if not (position_reached and euler_reached):
            return False
        else:
            return True

    def check_timeout(self):
        """Check if the joint is timeout.

        Returns:
            True if it is timeout, False otherwise.
        """
        is_timeout = self.physics.time() >= self._stop_time

        if is_timeout:
            logger.warning(
                'Time out (%.2f) with target_pose = %s, current_pose = %s.'
                % (self._stop_time - self._start_time,
                    self._target_pose,
                    self.pose)
            )

        return is_timeout


import cog
import random
import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image
from pathlib import Path
import tempfile
from torch.autograd import Variable
from nltk.tokenize import RegexpTokenizer
from AttnGAN_CL.code.miscc.config import cfg, cfg_from_file
from AttnGAN_CL.code.datasets import TextDataset
from AttnGAN_CL.code.trainer import condGANTrainer as trainer
from AttnGAN_CL.code.model import RNN_ENCODER
from AttnGAN_CL.code.model import G_DCGAN, G_NET
from AttnGAN_CL.code.miscc.utils import build_super_images, build_super_images2

class Predictor(cog.Predictor):
    def setup(self):
        config()

    @cog.input("sentence", type=str, help="text for image generation, describe an image with a bird")
    def predict(self, sentence):
        algo, data_dic = gen_image(sentence)
        output = gen_example(algo, data_dic)
        return output

def config():
    cfg_file = 'AttnGAN_CL/code/cfg/eval_bird.yml'
    cfg_from_file(cfg_file)
    if torch.cuda.is_available():
        cfg.GPU_ID = 0
    else:
        cfg.CUDA = False
        cfg.GPU_ID = -1

    manual_seed = 100
    random.seed(manual_seed)
    np.random.seed(manual_seed)
    torch.manual_seed(manual_seed)
    if cfg.CUDA:
        torch.cuda.manual_seed_all(manual_seed)

def gen_image(sentence):
    split_dir = 'test'
    bshuffle = True
    imsize = cfg.TREE.BASE_SIZE * (2 ** (cfg.TREE.BRANCH_NUM - 1))
    image_transform = transforms.Compose([
        transforms.Resize(int(imsize * 76 / 64)),
        transforms.RandomCrop(imsize),
        transforms.RandomHorizontalFlip()])

    dataset = TextDataset(cfg.DATA_DIR, split_dir,
                          base_size=cfg.TREE.BASE_SIZE,
                          transform=image_transform)

    assert dataset
    dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=cfg.TRAIN.BATCH_SIZE,
        drop_last=True, shuffle=bshuffle, num_workers=int(cfg.WORKERS))

    output_dir = './output'
    algo = trainer(output_dir, dataloader, dataset.n_words, dataset.ixtoword, dataset)
    data_dic = gen_data_dic(dataset.wordtoix, sentence)
    return algo, data_dic

def gen_data_dic(wordtoix, sentence):
    data_dic = {}
    sentences = ['b', sentence]
    # a list of indices for a sentence
    captions = []
    cap_lens = []
    for sent in sentences:
        if len(sent) == 0:
            continue
        sent = sent.replace("\ufffd\ufffd", " ")
        tokenizer = RegexpTokenizer(r'\w+')
        tokens = tokenizer.tokenize(sent.lower())
        if len(tokens) == 0:
            continue

        rev = []
        for t in tokens:
            t = t.encode('ascii', 'ignore').decode('ascii')
            if len(t) > 0 and t in wordtoix:
                rev.append(wordtoix[t])
        captions.append(rev)
        cap_lens.append(len(rev))

    max_len = np.max(cap_lens)

    sorted_indices = np.argsort(cap_lens)[::-1]
    cap_lens = np.asarray(cap_lens)
    cap_lens = cap_lens[sorted_indices]
    cap_array = np.zeros((len(captions), max_len), dtype='int64')
    for i in range(len(captions)):
        idx = sorted_indices[i]
        cap = captions[idx]
        c_len = len(cap)
        cap_array[i, :c_len] = cap
    data_dic['current'] = [cap_array, cap_lens, sorted_indices]
    return data_dic

def gen_example(algo, data_dic):

    if cfg.TRAIN.NET_G == '':
        print('Error: the path for morels is not found!')
    else:
        # Build and load the generator
        text_encoder = \
            RNN_ENCODER(algo.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
        state_dict = \
            torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
        text_encoder.load_state_dict(state_dict)
        print('Load text encoder from:', cfg.TRAIN.NET_E)
        if cfg.CUDA:
            text_encoder = text_encoder.cuda()
        text_encoder.eval()

        if cfg.GAN.B_DCGAN:
            netG = G_DCGAN()
        else:
            netG = G_NET()
        model_dir = cfg.TRAIN.NET_G
        state_dict = \
            torch.load(model_dir, map_location=lambda storage, loc: storage)
        netG.load_state_dict(state_dict)
        print('Load G from: ', model_dir)
        if cfg.CUDA:
            netG.cuda()
        netG.eval()
        for key in data_dic:
            captions, cap_lens, sorted_indices = data_dic[key]
            batch_size = captions.shape[0]
            nz = cfg.GAN.Z_DIM
            captions = Variable(torch.from_numpy(captions), volatile=True)
            cap_lens = Variable(torch.from_numpy(cap_lens), volatile=True)
            if cfg.CUDA:
                captions = captions.cuda()
                cap_lens = cap_lens.cuda()

            noise = Variable(torch.FloatTensor(batch_size, nz), volatile=True)
            if cfg.CUDA:
                noise = noise.cuda()
            #######################################################
            # (1) Extract text embeddings
            ######################################################
            hidden = text_encoder.init_hidden(batch_size)
            words_embs, sent_emb = text_encoder(captions, cap_lens, hidden)
            mask = (captions == 0)
            #######################################################
            # (2) Generate fake images
            ######################################################
            noise.data.normal_(0, 1)
            fake_imgs, attention_maps, _, _ = netG(noise, sent_emb, words_embs, mask)
            im = fake_imgs[-1][0].data.cpu().numpy()
            im = (im + 1.0) * 127.5
            im = im.astype(np.uint8)
            im = np.transpose(im, (1, 2, 0))
            im = Image.fromarray(im)
            out_path = Path(tempfile.mkdtemp()) / "out.png"
            im.save(str(out_path))

    return out_path


#!/usr/bin/env python3

# This work is licensed under the terms of the MIT license.
# For a copy, see <https://opensource.org/licenses/MIT>.
#
# Author: Konstantin Ivanchenko
# Date: August 03, 2019

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

from DatasetHandler import DatasetHandler
from ModelHandler import ModelHandler

from ImageHandler import ImageHandler

from sklearn.model_selection import train_test_split as tts

import cv2
from keras.preprocessing.image import array_to_img

im_handler = ImageHandler()

ds_handler = DatasetHandler()
ds_handler.init_apollo_data_source()
ds_handler.read_frame_big_batch('Apollo')

print(ds_handler.data_paths_apollo_leftcam)

model_train = False
new_model = False

# TODO: uncomment to see augmentations
image_paths_l = ds_handler.data_paths_apollo_leftcam#[100]
image_paths_r = ds_handler.data_paths_apollo_rightcam
output_paths_l = ds_handler.label_paths_apollo_leftcam#[100]
output_paths_r = ds_handler.label_paths_apollo_rightcam
#image_handler.augment_batch_visualize(image_paths, valid_paths)

if len(image_paths_l) > len(image_paths_r):
    del image_paths_l[len(image_paths_r):]
    del output_paths_l[len(image_paths_r):]
elif len(image_paths_l) < len(image_paths_r):
    del image_paths_r[len(image_paths_l):]
    del output_paths_r[len(image_paths_l):]

train_input_l, valid_input_l, train_input_r, valid_input_r, train_output_l, valid_output_l, \
train_output_r, valid_output_r = tts(image_paths_l, image_paths_r, output_paths_l, output_paths_r,
                                         test_size=0.01, random_state=5)

# plt.imshow(o)

if model_train:
    if new_model:
        model_handler = ModelHandler()
        model_handler.get_model_parameters()
        history = model_handler.model_fit(train_input_l, train_input_r, train_output_l, train_output_r,
                                valid_input_l, valid_input_r, valid_output_l, valid_output_r)

        model_handler.model_save()
        model_handler.model_plot_result(history)
    else:
        model_handler = ModelHandler(load_from='model_x.h5')
        model_handler.get_model_parameters()
        history = model_handler.model_fit(train_input_l, train_input_r, train_output_l, train_output_r,
                                          valid_input_l, valid_input_r, valid_output_l, valid_output_r)

        model_handler.model_save()
        model_handler.model_plot_result(history)

else:
    model_handler = ModelHandler(load_from='model_x.h5')
    model_handler.get_model_parameters()

    i_batch = []
    i_batch_size = 2
    i_batch_start_ind = 20

    for ind in range(i_batch_size):
        path_i = image_paths_r[i_batch_start_ind+ind]
        path_o = output_paths_r[i_batch_start_ind+ind]
        i, o = im_handler.load_resized(path_i, path_o)
        i = model_handler.input_image_to_array(i)
        i_batch.append(i)

    i_batch = np.array(i_batch)

    #i = np.reshape(i, (None, np.shape(i)))

    y_batch = model_handler.get_model_output(i_batch)

    y_batch = np.array(y_batch)

    fig_img, axs = plt.subplots(i_batch_size, 2, figsize=(25, 50))
    fig_img.tight_layout()

    # create overlay
    for ind in range(i_batch_size):
        i_gray = cv2.cvtColor(i_batch[ind], cv2.COLOR_BGR2GRAY)
        o_gray = np.copy(i_gray)

        print(i_gray.dtype)

        # inp_image = array_to_img(i_gray)
        # output = array_to_img(i_gray)
        # print(np.shape(i_batch[ind]))

        raise_rank = np.expand_dims(y_batch[ind], axis=2)
        #raise_rank = raise_rank[np.newaxis]
        print(np.shape(raise_rank))
        print(raise_rank.dtype)
        raise_rank = raise_rank.astype(np.float32)
        raise_rank = raise_rank #/ 36.0
        # out_overlay = array_to_img(raise_rank)
        cv2.addWeighted(raise_rank, 0.2, i_gray, 0.8, 0, o_gray)

        axs[ind][0].imshow(i_batch[ind])
        axs[ind][0].set_title('Input image')
        axs[ind][1].imshow(o_gray)
        axs[ind][1].set_title('Trained model output')

    plt.show(fig_img)

''' Use it for a single image test
    print(np.shape(i))  # this is a Numpy array with shape (320, 480, 3)

    i = i[np.newaxis, ...]  # dimension added to fit input size

    y = model_handler.get_model_output(i)

    print(np.shape(y))
    y = np.squeeze(y, axis=0)
    print(np.shape(y))

    plt.imshow(y)
    plt.show()
'''


from typing import Union

import numpy as np

from .exceptions import UnitError
from .types import Numeric, UnitValue
from .unit_base import UnitBase
from .unit_interface import UnitInterface
from .unit_type import UnitType

class UnitlessType(UnitType):
    """
    A UnitType that represents a unitless value. This may seem kind of silly,
    but it makes things easier in some regards, because it allows us to
    do things like guarantee that a division operation always returns a unit
    and not occasionally a raw value, and to represent constructs like
    s ^ -1.

    We make it so that UnitlessTypes are not compatible with anything. You
    might think it more logical to instead have them be compatible with
    everything, but this introduces a nasty reflexivity problem into
    compatibility checks, e.g. the result is different depending on
    whether you do a.is_compatible(b) or b.is_compatible(a). Furthermore,
    I think it's better for PyUnits to force users to initialize a new
    unit if they want to give a unit-less value units, instead of just
    doing it implicitly.
    """

@UnitlessType.decorate
class Unitless(UnitBase):
    """
    Represents a unit-less value. By design, these are pretty locked-down, so
    that generally the user has to explicitly extract the raw value in order
    to do anything with them.
    """

    def __init__(self, unit_type: UnitlessType,
                 value: Union[Numeric, 'Unitless']):
        """
        :param unit_type: The UnitlessType instance used to create this class.
        :param value: The unitless value to wrap.
        """
        super().__init__(unit_type)

        if isinstance(value, UnitInterface):
            if not value.type.is_compatible(self.type):
                # Implicit initialization from normal units are disallowed,
                # because my judgement is that this has too high of a propensity
                # for causing bugs.
                raise UnitError("To initialize a Unitless value from another "
                                "unit, you must explicitly take the raw value.")

            self.__value = value.raw
        else:
            # Numeric type.
            self.__value = np.asarray(value)

    def __mul__(self, other: UnitValue) -> UnitInterface:
        """
        See superclass for documentation.
        """
        if not isinstance(other, UnitInterface):
            # Multiplying by a raw numeric value.
            # Our only concern in this case is making sure the result is a
            # Unitless value.
            return self.type(self.raw * other)
        elif isinstance(other, type(self)):
            # In this case, we just multiply the raw values.
            return self.type(self.raw * other.raw)
        else:
            # We don't handle normal multiplication in this class, and instead
            # rely on the other unit's reflected multiplication operation.
            return NotImplemented

    def __truediv__(self, other: UnitValue) -> UnitInterface:
        """
        See superclass for documentation.
        """
        if not isinstance(other, UnitInterface):
            # Dividing by a raw numeric value.
            return self.type(self.raw / other)
        if isinstance(other, type(self)):
            # In this case, we just divide the raw values.
            return self.type(self.raw / other.raw)
        else:
            # We don't handle normal division in this class, and instead rely
            # on the other unit's reflected division operator.
            return NotImplemented

    def __rtruediv__(self, other: UnitValue) -> UnitInterface:
        # The only way that we should ever get here is if we are trying to
        # divide a raw numerical value.
        assert not isinstance(other, UnitInterface)

        return self.type(other / self.raw)

    def __add__(self, other: UnitValue) -> UnitInterface:
        if not isinstance(other, UnitInterface):
            # Our only concern in this case is making sure the result is a
            # Unitless value.
            return self.type(self.raw + other)
        elif isinstance(other, type(self)):
            # In this case, we just add the raw values.
            return self.type(self.raw + other.raw)
        else:
            # We don't handle normal addition in this class, and instead
            # rely on the other unit's reflected addition operation.
            return NotImplemented

    @classmethod
    def is_standard(cls) -> bool:
        """
        See superclass for documentation.
        """
        # This unit is already in standard form.
        return True

    def to_standard(self) -> 'Unitless':
        """
        See superclass for documentation.
        """
        # The standard version of a unitless value is just itself.
        return self

    @property
    def raw(self) -> np.ndarray:
        """
        See superclass for documentation.
        """
        return self.__value

    @property
    def name(self) -> str:
        """
        See superclass for documentation.
        """
        # The name is always empty.
        return ""

    def cast_to(self, out_type: UnitType) -> UnitInterface:
        """
        See superclass for documentation.
        """
        # If you're trying to cast a unitless value, something has gone
        # terribly wrong.
        raise NotImplementedError("A unitless value should not ever need to be "
                                  "casted.")


import re
from time import sleep
from pathlib import Path
from typing import Union
from multiprocessing import Pipe, Process, Queue

import faiss
import numpy as np

from .base_indexer import *
from .faiss_cache import faiss_cache

__all__ = ['DeployShards', 'RangeShards']

class DeployShards(BaseIndexer):
    def __init__(self, shard_dir, nprobe: int = 4):
        """
        For deploying multiple, pre-made IVF indexes as shards.
            (intended for on-disk indexes that do not fit in memory)

        Note: The index shards must be true partitions with no overlapping ids

        :param shard_dir: Dir containing faiss index shards
        :param nprobe: Number of clusters to visit during search
                       (speed accuracy trade-off)
        """
        super().__init__()
        self.paths_to_shards = self.get_index_paths(shard_dir)
        self.nprobe = nprobe

        # Load shards
        shards = list()
        for shard_path in self.paths_to_shards:
            shard = self.load_shard(path_to_shard=shard_path, nprobe=self.nprobe)
            shards.append(shard)

        # Merge shards
        self.index = faiss.IndexShards(512, threaded=True, successive_ids=False)
        for shard in shards:
            self.index.add_shard(shard)

    @staticmethod
    def load_shard(path_to_shard: Union[str, Path], nprobe: int = 4):
        try:
            shard = faiss.read_index(path_to_shard)
        except RuntimeError:
            shard = faiss.read_index(path_to_shard, faiss.IO_FLAG_ONDISK_SAME_DIR)
        shard.nprobe = nprobe
        return shard

    def add_shard(self, new_shard_path: Union[str, Path]):
        if new_shard_path in self.paths_to_shards:
            print('WARNING: This shard is already online \n'
                  '         Aborting...')
            return
        self.paths_to_shards.append(new_shard_path)
        shard = self.load_shard(path_to_shard=new_shard_path, nprobe=self.nprobe)
        self.index.add_shard(shard)

#### Parallelized Nearest Neighbor Search ####
class Shard(Process):
    def __init__(self, shard_name, shard_path: Union[str, Path],
                 input_pipe: Pipe, output_queue: Queue,
                 nprobe: int = 4, daemon: bool = False):
        """ RangeShards search worker """
        super().__init__(name=shard_name)
        self.daemon = daemon

        try:
            self.index = faiss.read_index(shard_path)
        except RuntimeError:
            self.index = faiss.read_index(shard_path, faiss.IO_FLAG_ONDISK_SAME_DIR)
        self.index.nprobe = nprobe
        self.input = input_pipe
        self.output = output_queue

    def run(self):

        @faiss_cache(64)
        def neighborhood(index, query, radius):
            _, ddd, iii = index.range_search(query, radius)
            return ddd, iii

        if self.input.poll():
            (query_vector, radius_limit) = self.input.recv()
            dd, ii = neighborhood(self.index, query_vector, radius_limit)
            self.output.put((dd, ii), block=False)

class RangeShards(BaseIndexer):
    def __init__(self, shard_dir: Union[str, Path], nprobe: int = 4,
                 get_nested: bool = False):
        """
        For deploying multiple, pre-made IVF indexes as shards.
            (intended for on-disk indexes that do not fit in memory)

        Note: The index shards must be true partitions with no overlapping ids

        :param shard_dir: Dir containing faiss index shards
        :param nprobe: Number of clusters to visit during search
                       (speed accuracy trade-off)
        :param get_nested: Load indexes in sub directories of shard_dir
        """
        super().__init__()
        self.paths_to_shards = self.get_index_paths(shard_dir, recursive=get_nested)
        self.nprobe = nprobe
        self.dynamic = True
        self.lock = False
        self.date_seed = str('\d{4}[-/]\d{2}[-/]\d{2}')

        self.results = Queue()
        self.shards = dict()
        self.n_shards = 0
        for shard_path in self.paths_to_shards:
            self.load_shard(shard_path)
        for shard_name, (handler_pipe, shard) in self.shards.items():
            shard.start()
            self.n_shards += 1

    @faiss_cache(256)
    def search(self, query_vector: np.array, k: int, radius: float = 0.65,
               start: str = '0000-00-00', end: str = '9999-99-99'
               ) -> FaissSearch:

        if len(query_vector.shape) < 2 or query_vector.shape[0] > 1:
            query_vector = np.reshape(query_vector, (1, query_vector.shape[0]))

        # Lock search while loading index or actively searching
        while self.lock:
            sleep(1)

        # Lock out other searches
        self.lock = True

        # Start parallel range search
        n_results = 0
        for shard_name, (hpipe, shard) in self.shards.items():
            shard_date = re.search(self.date_seed, shard_name).group()
            if start <= shard_date <= end:
                hpipe.send((query_vector, radius))
                shard.run()
                n_results += 1

        # Aggregate results
        D, I = list(), list()
        while n_results > 0 or not self.results.empty():
            dd, ii = self.results.get()
            D.extend(dd[:k]), I.extend(ii[:k])
            n_results -= 1

        self.lock = False
        return self.joint_sort([D], [I])

    def load_shard(self, shard_path: Union[str, Path]):
        shard_name = str(shard_path).replace('.index', '')
        shard_pipe, handler_pipe = Pipe(False)
        shard = Shard(shard_name, str(shard_path),
                      input_pipe=shard_pipe, output_queue=self.results,
                      nprobe=self.nprobe, daemon=False)
        self.shards[shard_name] = (handler_pipe, shard)

    def add_shard(self, new_shard_path: Union[str, Path]):
        # Lock search while deploying a new shard
        self.lock = True

        shard_name = str(new_shard_path).replace('.index', '')
        if new_shard_path in self.paths_to_shards or \
                shard_name in self.shards:
            print('WARNING: This shard is already online \n'
                  '         Aborting...')
        else:
            self.paths_to_shards.append(new_shard_path)
            self.load_shard(new_shard_path)

        # Release lock
        self.lock = False


# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Parameter provider for JAX networks."""
import io
from typing import Any, Dict, Optional, Tuple

from absl import logging
import dill
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
import tree

from deepmind.diplomacy.environment import action_utils
from deepmind.diplomacy.environment import tree_utils

def apply_unbatched(f, *args, **kwargs):
  batched = f(*tree_utils.tree_expand_dims(args),
              **tree_utils.tree_expand_dims(kwargs))
  return tree.map_structure(lambda arr: np.squeeze(arr, axis=0), batched)

def fix_waives(action_list):
  """Fix action_list so that there is at most one waive, with that at the end.

  This means that build lists are invariant to order, and can be a fixed length.

  Args:
    action_list: a list of actions.

  Returns:
    A copy of action_list, modified so that any waives are truncated to 1
    and moved to the end.
  """
  non_waive_actions = [a for a in action_list if not action_utils.is_waive(a)]
  waive_actions = [a for a in action_list if action_utils.is_waive(a)]
  if waive_actions:
    # Return non-waives, then one waive.
    return non_waive_actions + waive_actions[:1]
  return non_waive_actions

def fix_actions(actions_lists):
  """Fixes network action outputs to be compatible with game_runners.

  Args:
    actions_lists: Actions for all powers in a single board state (i.e. output
    of a single inference call). Note that these are shrunk actions
    (see action_utils.py).

  Returns:
    A sanitised actions_lists suitable for stepping the environment
  """
  non_zero_actions = []
  for single_power_actions in actions_lists:
    non_zero_actions.append([])
    for unit_action in single_power_actions:
      if unit_action != 0:
        non_zero_actions[-1].append(
            action_utils.POSSIBLE_ACTIONS[unit_action >> 16])

  # Fix waives.
  final_actions = [
      fix_waives(single_power_actions)
      for single_power_actions in non_zero_actions
  ]
  return final_actions

class ParameterProvider:
  """Loads and exposes network params that have been saved to disk."""

  def __init__(self, file_handle: io.IOBase):
    self._params, self._net_state, self._step = dill.load(file_handle)

  def params_for_actor(self) -> Tuple[hk.Params, hk.Params, jnp.ndarray]:
    """Provides parameters for a SequenceNetworkHandler."""
    return self._params, self._net_state, self._step

class SequenceNetworkHandler:
  """Plays Diplomacy with a Network as policy.

  This class turns a Network into a Diplomacy bot by forwarding observations,
  and receiving policy outputs. It handles the network parameters, batching and
  observation processing.
  """

  def __init__(self,
               network_cls,
               network_config: Dict[str, Any],
               rng_seed: Optional[int],
               parameter_provider: ParameterProvider):
    if rng_seed is None:
      rng_seed = np.random.randint(2**16)
      logging.info("RNG seed %s", rng_seed)
    self._rng_key = jax.random.PRNGKey(rng_seed)
    self._network_cls = network_cls
    self._network_config = network_config

    self._rng_key, subkey = jax.random.split(self._rng_key)
    self._observation_transformer = network_cls.get_observation_transformer(
        network_config, subkey
    )

    def transform(fn_name, static_argnums=()):

      def fwd(*args, **kwargs):
        net = network_cls(**network_config)
        fn = getattr(net, fn_name)
        return fn(*args, **kwargs)

      apply = hk.transform_with_state(fwd).apply
      return jax.jit(apply, static_argnums=static_argnums)

    self._parameter_provider = parameter_provider
    # The inference method of our Network does not modify its arguments, Jax can
    # exploit this information when jitting this method to make it more
    # efficient. Network.inference takes 4 arguments.
    self._network_inference = transform("inference", 4)
    self._network_shared_rep = transform("shared_rep")
    self._network_initial_inference = transform("initial_inference")
    self._network_step_inference = transform("step_inference")
    self._network_loss_info = transform("loss_info")
    self._params = None
    self._state = None
    self._step_counter = -1

  def reset(self):
    if self._parameter_provider:
      learner_state = self._parameter_provider.params_for_actor()
      (self._params, self._state, self._step_counter) = learner_state

  def _apply_transform(self, transform, *args, **kwargs):
    self._rng_key, subkey = jax.random.split(self._rng_key)
    output, unused_state = transform(
        self._params, self._state, subkey, *args, **kwargs)
    return tree.map_structure(np.asarray, output)

  def batch_inference(self, observation, num_copies_each_observation=None):
    """Do inference on unbatched observation and state."""
    if num_copies_each_observation is not None:
      # Cast this to a tuple so that static_argnums recognises it as unchanged
      # and doesn't recompile.
      num_copies_each_observation = tuple(num_copies_each_observation)
    initial_output, step_output = self._apply_transform(
        self._network_inference, observation, num_copies_each_observation)
    final_actions = [
        fix_actions(single_board_actions)
        for single_board_actions in step_output["actions"]
    ]

    return (initial_output, step_output), final_actions

  def compute_losses(self, *args, **kwargs):
    return self._apply_transform(self._network_loss_info, *args, **kwargs)

  def inference(self, *args, **kwargs):
    outputs, (final_actions,) = apply_unbatched(self.batch_inference, *args,
                                                **kwargs)
    return outputs, final_actions

  def batch_loss_info(self, step_types, rewards, discounts, observations,
                      step_outputs):
    return tree.map_structure(
        np.asarray,
        self._network_loss_info(step_types, rewards, discounts, observations,
                                step_outputs))

  @property
  def step_counter(self):
    return self._step_counter

  def observation_transform(self, *args, **kwargs):
    return self._observation_transformer.observation_transform(
        *args, **kwargs)

  def zero_observation(self, *args, **kwargs):
    return self._observation_transformer.zero_observation(*args, **kwargs)

  def observation_spec(self, num_players):
    return self._observation_transformer.observation_spec(num_players)

  def variables(self):
    return self._params


# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Ke Sun (sunk@mail.ustc.edu.cn)
# ------------------------------------------------------------------------------

import os

import cv2
import numpy as np

import torch
from torch.nn import functional as F

from .base_dataset import BaseDataset

class LIP(BaseDataset):
    def __init__(self,
                 root,
                 list_path,
                 num_samples=None,
                 num_classes=20,
                 multi_scale=True,
                 flip=True,
                 ignore_label=-1,
                 base_size=473,
                 crop_size=(473, 473),
                 downsample_rate=1,
                 scale_factor=11,
                 center_crop_test=False,
                 mean=[0.485, 0.456, 0.406],
                 std=[0.229, 0.224, 0.225]):

        super(LIP, self).__init__(ignore_label, base_size,
                crop_size, downsample_rate, scale_factor, mean, std)

        self.root = root
        self.num_classes = num_classes
        self.list_path = list_path
        self.class_weights = None

        self.multi_scale = multi_scale
        self.flip = flip
        self.img_list = [line.strip().split() for line in open(root+list_path)]

        self.files = self.read_files()
        if num_samples:
            self.files = self.files[:num_samples]

    def read_files(self):
        files = []
        for item in self.img_list:
            image_path, label_path = item[:2]
            name = os.path.splitext(os.path.basename(label_path))[0]
            sample = {"img": image_path,
                      "label": label_path,
                      "name": name,}
            files.append(sample)
        return files

    def resize_image(self, image, label, size):
        image = cv2.resize(image, size, interpolation = cv2.INTER_LINEAR)
        label = cv2.resize(label, size, interpolation=cv2.INTER_NEAREST)
        return image, label

    def __getitem__(self, index):
        item = self.files[index]
        name = item["name"]

        image = cv2.imread(os.path.join(
                    self.root, 'lip/TrainVal_images/', item["img"]),
                    cv2.IMREAD_COLOR)  # BGR
        label = cv2.imread(os.path.join(
                    self.root, 'lip/TrainVal_parsing_annotations/',
                    item["label"]),
                    cv2.IMREAD_GRAYSCALE)
        size = label.shape

        if 'testval' in self.list_path:
            image = cv2.resize(image, self.crop_size,
                               interpolation = cv2.INTER_LINEAR)
            image = self.input_transform(image)
            image = image.transpose((2, 0, 1))

            return image.copy(), label.copy(), np.array(size), name

        if self.flip:
            flip = np.random.choice(2) * 2 - 1
            image = image[:, ::flip, :]
            label = label[:, ::flip]

            if flip == -1:
                right_idx = [15, 17, 19]
                left_idx = [14, 16, 18]
                for i in range(0, 3):
                    right_pos = np.where(label == right_idx[i])
                    left_pos = np.where(label == left_idx[i])
                    label[right_pos[0], right_pos[1]] = left_idx[i]
                    label[left_pos[0], left_pos[1]] = right_idx[i]

        image, label = self.resize_image(image, label, self.crop_size)
        image, label = self.gen_sample(image, label,
                                self.multi_scale, False)

        return image.copy(), label.copy(), np.array(size), name

    def inference(self, model, image, flip):
        size = image.size()
        pred = model(image)
        pred = F.upsample(input=pred,
                          size=(size[-2], size[-1]),
                          mode='bilinear')
        if flip:
            flip_img = image.numpy()[:,:,:,::-1]
            flip_output = model(torch.from_numpy(flip_img.copy()))
            flip_output = F.upsample(input=flip_output,
                            size=(size[-2], size[-1]),
                            mode='bilinear')
            flip_output = flip_output.cpu().numpy()
            flip_pred = flip_output.copy()
            flip_pred[:,14,:,:] = flip_output[:,15,:,:]
            flip_pred[:,15,:,:] = flip_output[:,14,:,:]
            flip_pred[:,16,:,:] = flip_output[:,17,:,:]
            flip_pred[:,17,:,:] = flip_output[:,16,:,:]
            flip_pred[:,18,:,:] = flip_output[:,19,:,:]
            flip_pred[:,19,:,:] = flip_output[:,18,:,:]
            flip_pred = torch.from_numpy(flip_pred[:,:,:,::-1].copy()).cuda()
            pred += flip_pred
            pred = pred * 0.5
        return pred.exp()


import math
import random
import operator
import sys

import numpy as np

from codes.e_utils.experience import ExperienceSource, ExperienceSourceFirstLast
from codes.e_utils.experience_single import ExperienceSourceSingleEnvFirstLast

class TrajectoryBuffer:
    def __init__(self, experience_source):
        self.experience_source = experience_source
        self.experience_source_iter = None if experience_source is None else iter(experience_source)
        self.buffer = []

    def __len__(self):
        return len(self.buffer)

    def set_experience_source(self, experience_source):
        self.experience_source = experience_source
        self.experience_source_iter = None if experience_source is None else iter(experience_source)

    def populate(self, num_samples):
        entry = None
        for _ in range(num_samples):
            entry = next(self.experience_source_iter)
            self.buffer.append(entry)

        return entry

    def sample_all(self):
        return self.buffer

    def clear(self):
        self.buffer.clear()

class ExperienceSourceBuffer:
    """
    The same as ExperienceSource, but takes episodes from the simple buffer
    """

    def __init__(self, buffer, n_step=1):
        """
        Create buffered experience source
        :param buffer: list of episodes, each is a list of Experience object
        :param n_step: count of steps in every entry
        """
        self.update_buffer(buffer)
        self.n_step = n_step

    def update_buffer(self, buffer):
        self.buffer = buffer
        self.lens = list(map(len, buffer))

    def __iter__(self):
        """
        Infinitely sample episode from the buffer and then sample item offset
        """
        while True:
            episode = random.randrange(len(self.buffer))
            ofs = random.randrange(self.lens[episode] - self.n_step - 1)
            yield self.buffer[episode][ofs:ofs + self.n_step]

class ExperienceReplayBuffer:
    def __init__(self, experience_source, buffer_size):
        assert isinstance(
            experience_source,
            (ExperienceSource, ExperienceSourceFirstLast, ExperienceSourceSingleEnvFirstLast, type(None))
        )
        assert isinstance(buffer_size, int)
        self.experience_source = experience_source
        self.experience_source_iter = None if experience_source is None else iter(experience_source)
        self.buffer = []
        self.capacity = buffer_size
        self.pos = 0

    def __len__(self):
        return len(self.buffer)

    def __iter__(self):
        return iter(self.buffer)

    def set_experience_source(self, experience_source):
        self.experience_source = experience_source
        self.experience_source_iter = None if experience_source is None else iter(experience_source)

    def sample_all_for_on_policy(self, expected_model_version):
        return [sample for sample in self.buffer if sample.model_version == expected_model_version]

    def sample(self, batch_size):
        """
        Get one random batch from experience replay
        TODO: implement sampling order policy
        :param batch_size:
        :return:
        """
        if batch_size is None or len(self.buffer) <= batch_size:
            return self.buffer

        # Warning: replace=False makes random.choice O(n)
        keys = np.random.choice(len(self.buffer), batch_size, replace=True)
        return [self.buffer[key] for key in keys]

    def add_sample(self, sample):
        if len(self.buffer) < self.capacity:
            self.buffer.append(sample)
        else:
            self.buffer[self.pos] = sample
        self.pos = (self.pos + 1) % self.capacity

    def populate(self, num_samples):
        """
        Populates samples into the buffer
        :param samples: how many samples to populate
        """
        entry = None
        for _ in range(num_samples):
            entry = next(self.experience_source_iter)
            self.add_sample(entry)

        return entry

    def populate_with_action_count(self, num_samples, action_count):
        """
        Populates samples into the buffer
        :param samples: how many samples to populate
        """
        for _ in range(num_samples):
            entry = next(self.experience_source_iter)
            action_count[entry.action] += 1
            self.add_sample(entry)

    def update_priorities(self, batch_indices, batch_priorities):
        raise NotImplementedError()

    def update_beta(self, idx):
        raise NotImplementedError()

    def clear(self):
        self.buffer.clear()
        self.pos = 0

    def size(self):
        return len(self.buffer)

class PrioReplayBufferNaive:
    def __init__(self, experience_source, buffer_size, prob_alpha=0.6):
        self.experience_source_iter = iter(experience_source)
        self.prob_alpha = prob_alpha
        self.capacity = buffer_size
        self.pos = 0
        self.buffer = []
        self.priorities = np.zeros((buffer_size,), dtype=np.float32)

    def __len__(self):
        return len(self.buffer)

    def set_experience_source(self, experience_source):
        self.experience_source = experience_source
        self.experience_source_iter = None if experience_source is None else iter(experience_source)

    def populate(self, num_samples):
        max_prio = self.priorities.max() if self.buffer else 1.0
        for _ in range(num_samples):
            sample = next(self.experience_source_iter)
            if len(self.buffer) < self.capacity:
                self.buffer.append(sample)
            else:
                self.buffer[self.pos] = sample
            self.priorities[self.pos] = max_prio
            self.pos = (self.pos + 1) % self.capacity

    def sample(self, batch_size, beta=0.4):
        if len(self.buffer) == self.capacity:
            prios = self.priorities
        else:
            prios = self.priorities[:self.pos]
        probs = np.array(prios, dtype=np.float32) ** self.prob_alpha

        probs /= probs.sum()
        indices = np.random.choice(len(self.buffer), batch_size, p=probs, replace=True)
        samples = [self.buffer[idx] for idx in indices]
        total = len(self.buffer)
        weights = (total * probs[indices]) ** (-beta)
        weights /= weights.max()
        return samples, indices, np.array(weights, dtype=np.float32)

    def update_priorities(self, batch_indices, batch_priorities):
        for idx, prio in zip(batch_indices, batch_priorities):
            self.priorities[idx] = prio

    def size(self):
        return len(self.buffer)

# sumtree 사용 버전
class PrioritizedReplayBuffer(ExperienceReplayBuffer):
    def __init__(self, experience_source, buffer_size, alpha=0.6, n_step=1, beta_start=0.4, beta_frames=100000):
        super(PrioritizedReplayBuffer, self).__init__(experience_source, buffer_size)
        assert alpha > 0
        self.alpha = alpha
        self.beta = beta_start
        self.n_step = n_step
        self.beta_start = beta_start
        self.beta_frames = beta_frames
        self.buffer_size = buffer_size

        it_capacity = 1
        while it_capacity < self.buffer_size:
            it_capacity *= 2

        self._it_sum = SumSegmentTree(it_capacity)
        self._it_min = MinSegmentTree(it_capacity)
        self._max_priority = 1.0

    def clear(self):
        self.buffer.clear()

        it_capacity = 1
        while it_capacity < self.buffer_size:
            it_capacity *= 2

        self._it_sum = SumSegmentTree(it_capacity)
        self._it_min = MinSegmentTree(it_capacity)
        self._max_priority = 1.0

    def update_beta(self, idx):
        v = self.beta_start + idx * (1.0 - self.beta_start) / self.beta_frames
        self.beta = min(1.0, v)
        return self.beta

    def add_sample(self, *args, **kwargs):
        idx = self.pos
        super().add_sample(*args, **kwargs)
        self._it_sum[idx] = self._max_priority ** self.alpha
        self._it_min[idx] = self._max_priority ** self.alpha

    def _sample_proportional(self, batch_size):
        assert len(self) > self.n_step
        res = []
        for _ in range(batch_size):
            while True:
                mass = random.random() * self._it_sum.sum(0, len(self) - 1)
                idx = self._it_sum.find_prefixsum_idx(mass)

                upper = self.pos
                lower = (self.pos - self.n_step)
                if lower < 0:
                    lower = self.capacity + lower
                if lower < upper:
                    if not lower <= idx < upper:
                        res.append(idx)
                        break
                else:
                    if upper <= idx < lower:
                        res.append(idx)
                        break
        return res

    def sample(self, batch_size):
        assert self.beta > 0

        idxes = self._sample_proportional(batch_size)
        # print("#################")
        # print(idxes)
        # print(self.pos)
        # print("#################")

        weights = []
        p_min = self._it_min.min() / self._it_sum.sum()
        max_weight = (p_min * len(self)) ** (-self.beta)

        for idx in idxes:
            p_sample = self._it_sum[idx] / self._it_sum.sum()
            weight = (p_sample * len(self)) ** (-self.beta)
            weights.append(weight / max_weight)

        weights = np.array(weights, dtype=np.float32)
        samples = [self.buffer[idx] for idx in idxes]
        return samples, idxes, weights

    def update_priorities(self, idxes, priorities):
        # with torch.no_grad():
        assert len(idxes) == len(priorities)
        for idx, priority in zip(idxes, priorities):
            assert priority > 0.0, priority
            assert 0 <= idx < len(self), idx
            self._it_sum[idx] = priority ** self.alpha
            self._it_min[idx] = priority ** self.alpha

            self._max_priority = max(self._max_priority, priority)

class SegmentTree(object):
    def __init__(self, capacity, operation, neutral_element):
        """Build a Segment Tree data structure.

        https://en.wikipedia.org/wiki/Segment_tree

        Can be used as regular array, but with two
        important differences:

            a) setting item's value is slightly slower.
               It is O(lg capacity) instead of O(1).
            b) user has access to an efficient `reduce`
               operation which reduces `operation` over
               a contiguous subsequence of items in the
               array.

        Paramters
        ---------
        capacity: int
            Total size of the array - must be a power of two.
        operation: lambda obj, obj -> obj
            and operation for combining elements (eg. sum, max)
            must for a mathematical group together with the set of
            possible values for array elements.
        neutral_element: obj
            neutral element for the operation above. eg. float('-inf')
            for max and 0 for sum.
        """
        assert capacity > 0 and capacity & (capacity - 1) == 0, "capacity must be positive and a power of 2."
        self._capacity = capacity
        self._value = [neutral_element for _ in range(2 * capacity)]
        self._operation = operation

    def _reduce_helper(self, start, end, node, node_start, node_end):
        if start == node_start and end == node_end:
            return self._value[node]
        mid = (node_start + node_end) // 2
        if end <= mid:
            return self._reduce_helper(start, end, 2 * node, node_start, mid)
        else:
            if mid + 1 <= start:
                return self._reduce_helper(start, end, 2 * node + 1, mid + 1, node_end)
            else:
                return self._operation(
                    self._reduce_helper(start, mid, 2 * node, node_start, mid),
                    self._reduce_helper(mid + 1, end, 2 * node + 1, mid + 1, node_end)
                )

    def reduce(self, start=0, end=None):
        """Returns result of applying `self.operation`
        to a contiguous subsequence of the array.

            self.operation(arr[start], operation(arr[start+1], operation(... arr[end])))

        Parameters
        ----------
        start: int
            beginning of the subsequence
        end: int
            end of the subsequences

        Returns
        -------
        reduced: obj
            result of reducing self.operation over the specified range of array elements.
        """
        if end is None:
            end = self._capacity
        if end < 0:
            end += self._capacity
        end -= 1
        return self._reduce_helper(start, end, 1, 0, self._capacity - 1)

    def __setitem__(self, idx, val):
        # index of the leaf
        idx += self._capacity
        self._value[idx] = val
        idx //= 2
        while idx >= 1:
            self._value[idx] = self._operation(
                self._value[2 * idx],
                self._value[2 * idx + 1]
            )
            idx //= 2

    def __getitem__(self, idx):
        assert 0 <= idx < self._capacity
        return self._value[self._capacity + idx]

class SumSegmentTree(SegmentTree):
    def __init__(self, capacity):
        super(SumSegmentTree, self).__init__(
            capacity=capacity,
            operation=operator.add,
            neutral_element=0.0
        )

    def sum(self, start=0, end=None):
        """Returns arr[start] + ... + arr[end]"""
        return super(SumSegmentTree, self).reduce(start, end)

    def find_prefixsum_idx(self, prefixsum):
        """Find the highest index `i` in the array such that
            sum(arr[0] + arr[1] + ... + arr[i - i]) <= prefixsum

        if array values are probabilities, this function
        allows to sample indexes according to the discrete
        probability efficiently.

        Parameters
        ----------
        perfixsum: float
            upperbound on the sum of array prefix

        Returns
        -------
        idx: int
            highest index satisfying the prefixsum constraint
        """
        assert 0 <= prefixsum <= self.sum() + 1e-5
        idx = 1
        while idx < self._capacity:  # while non-leaf
            if self._value[2 * idx] > prefixsum:
                idx = 2 * idx
            else:
                prefixsum -= self._value[2 * idx]
                idx = 2 * idx + 1
        return idx - self._capacity

class MinSegmentTree(SegmentTree):
    def __init__(self, capacity):
        super(MinSegmentTree, self).__init__(
            capacity=capacity,
            operation=min,
            neutral_element=float('inf')
        )

    def min(self, start=0, end=None):
        """Returns min(arr[start], ...,  arr[end])"""

        return super(MinSegmentTree, self).reduce(start, end)

class RankBasedPrioritizedReplayBuffer(ExperienceReplayBuffer):
    def __init__(self, experience_source, buffer_size, params, alpha=0.7, beta_start=0.5, beta_frames=100000):
        super(RankBasedPrioritizedReplayBuffer, self).__init__(experience_source, buffer_size)
        assert alpha > 0
        self.alpha = alpha
        self.beta = beta_start
        self.beta_start = beta_start
        self.beta_frames = beta_frames
        self._max_priority = 1.0

        self.replace_flag = True
        self.learn_start = params.MIN_REPLAY_SIZE_FOR_TRAIN
        self.total_steps = params.MAX_GLOBAL_STEP
        # partition number N, split total size to N part
        self.partition_num = 100
        self.batch_size = params.BATCH_SIZE

        self.index = 0
        self.record_size = 0
        self.isFull = False

        self.buffer = {}
        # self._experience = {}
        self.priority_queue = BinaryHeap(self.capacity)
        self.distributions = self.build_distributions()

        # self.beta_grad = (1 - self.beta_start) / (self.total_steps - self.learn_start)

    def build_distributions(self):
        """
        preprocess pow of rank
        (rank i) ^ (-alpha) / sum ((rank i) ^ (-alpha))
        :return: distributions, dict
        """
        res = {}
        n_partitions = self.partition_num
        partition_num = 1
        # each part size
        partition_size = int(math.floor(1.0*self.capacity / n_partitions))

        for n in range(partition_size, self.capacity + 1, partition_size):
            if self.learn_start <= n <= self.capacity:
                distribution = {}
                # P(i) = (rank i) ^ (-alpha) / sum ((rank i) ^ (-alpha))
                pdf = list(
                    map(lambda x: math.pow(x, -self.alpha), range(1, n + 1))
                )
                pdf_sum = math.fsum(pdf)
                distribution['pdf'] = list(map(lambda x: x / pdf_sum, pdf))
                # split to k segment, and than uniform sample in each k
                # set k = batch_size, each segment has total probability is 1 / batch_size
                # strata_ends keep each segment start pos and end pos
                cdf = np.cumsum(distribution['pdf'])
                strata_ends = {1: 0, self.batch_size + 1: n}
                step = 1.0 / self.batch_size
                index = 1
                for s in range(2, self.batch_size + 1):
                    while cdf[index] < step:
                        index += 1
                    strata_ends[s] = index
                    step += 1.0 / self.batch_size

                distribution['strata_ends'] = strata_ends

                res[partition_num] = distribution

            partition_num += 1

        return res

    def fix_index(self):
        """
        get next insert index
        :return: index, int
        """
        if self.record_size < self.capacity:#self.record_size <= self.size:
            self.record_size += 1
        if self.index % self.capacity == 0:
            self.isFull = True if len(self.buffer) == self.capacity else False
            if self.replace_flag:
                self.index = 1
                return self.index
            else:
                sys.stderr.write('Experience replay buff is full and replace is set to FALSE!\n')
                return -1
        else:
            self.index += 1
            return self.index

    def store(self, sample):
        """
        store experience, suggest that experience is a tuple of (s1, a, r, s2, t)
        so each experience is valid
        :param sample: maybe a tuple, or list
        :return: bool, indicate insert status
        """
        insert_index = self.fix_index()
        if insert_index > 0:
            if insert_index in self.buffer:
                del self.buffer[insert_index]
            self.buffer[insert_index] = sample
            # add to priority queue
            priority = self.priority_queue.get_max_priority()
            self.priority_queue.update(priority, insert_index)
            return True
        else:
            sys.stderr.write('Insert failed\n')
            return False

    def retrieve(self, indices):
        """
        get experience from indices
        :param indices: list of experience id
        :return: experience replay sample
        """
        return [self.buffer[v] for v in indices]

    def rebalance(self):
        """
        rebalance priority queue
        :return: None
        """
        self.priority_queue.balance_tree()

    def update_priorities(self, idxes, priorities):
        """
        update priority according indices and deltas
        :param idxes: list of experience id
        :param priorities: list of delta, order correspond to indices
        :return: None
        """
        assert len(idxes) == len(priorities)
        for i in range(0, len(idxes)):
            self.priority_queue.update(math.fabs(priorities[i]), idxes[i])

    def update_beta(self, idx):
        v = self.beta_start + idx * (1.0 - self.beta_start) / self.beta_frames
        self.beta = min(1.0, v)
        return self.beta

    def add_sample(self, sample):
        self.store(sample)

    def sample(self, _):
        """
        sample a mini batch from experience replay
        :return: experience, list, samples
        :return: w, list, weights
        :return: rank_e_id, list, samples id, used for update priority
        """
        if self.record_size < self.learn_start:
            sys.stderr.write('Record size less than learn start! Sample failed\n')
            return False, False, False
        dist_index = int(math.floor(1.0 * self.record_size / self.capacity * self.partition_num))
        partition_size = int(math.floor(1.0 * self.capacity / self.partition_num))
        partition_max = dist_index * partition_size
        distribution = self.distributions[dist_index]
        rank_list = []
        # sample from k segments
        for n in range(1, self.batch_size + 1):
            if distribution['strata_ends'][n] + 1 <= distribution['strata_ends'][n + 1]:
                index = random.randint(distribution['strata_ends'][n] + 1,
                                       distribution['strata_ends'][n + 1])
            else:
                index = random.randint(distribution['strata_ends'][n + 1],
                                       distribution['strata_ends'][n] + 1)
            rank_list.append(index)

        # beta, increase by global_step, max 1
        # beta = min(self.beta_zero + (global_step - self.learn_start - 1) * self.beta_grad, 1)
        beta = self.beta
        # find all alpha pow, notice that pdf is a list, start from 0
        alpha_pow = [distribution['pdf'][v - 1] for v in rank_list]
        # w = (N * P(i)) ^ (-beta) / max w
        w = np.power(np.array(alpha_pow) * partition_max, -beta)
        w_max = max(w)
        w = np.divide(w, w_max)
        # rank list is priority id
        # convert to experience id
        rank_e_id = self.priority_queue.priority_to_experience(rank_list)
        # get experience id according rank_e_id
        experience = self.retrieve(rank_e_id)
        return experience, rank_e_id, w

class BinaryHeap(object):
    def __init__(self, priority_size=100, priority_init=None, replace=True):
        self.e2p = {}
        self.p2e = {}
        self.replace = replace

        if priority_init is None:
            self.priority_queue = {}
            self.size = 0
            self.max_size = priority_size
        else:
            # not yet test
            self.priority_queue = priority_init
            self.size = len(self.priority_queue)
            self.max_size = None or self.size

            experience_list = list(map(lambda x: self.priority_queue[x], self.priority_queue))
            self.p2e = list_to_dict(experience_list)
            self.e2p = exchange_key_value(self.p2e)
            for i in range(int(self.size / 2), -1, -1):
                self.down_heap(i)

    def __repr__(self):
        """
        :return: string of the priority queue, with level info
        """
        if self.size == 0:
            return 'No element in heap!'
        to_string = ''
        level = -1
        max_level = int(math.floor(math.log(self.size, 2)))

        for i in range(1, self.size + 1):
            now_level = int(math.floor(math.log(i, 2)))
            if level != now_level:
                to_string = to_string + ('\n' if level != -1 else '') \
                            + '    ' * (max_level - now_level)
                level = now_level

            to_string = to_string + '%.2f ' % self.priority_queue[i][1] + '    ' * (max_level - now_level)

        return to_string

    def check_full(self):
        return self.size > self.max_size

    def _insert(self, priority, e_id):
        """
        insert new experience id with priority
        (maybe don't need get_max_priority and implement it in this function)
        :param priority: priority value
        :param e_id: experience id
        :return: bool
        """
        self.size += 1

        if self.check_full() and not self.replace:
            sys.stderr.write('Error: no space left to add experience id %d with priority value %f\n' % (e_id, priority))
            return False
        else:
            self.size = min(self.size, self.max_size)

        self.priority_queue[self.size] = (priority, e_id)
        self.p2e[self.size] = e_id
        self.e2p[e_id] = self.size

        self.up_heap(self.size)
        return True

    def update(self, priority, e_id):
        """
        update priority value according its experience id
        :param priority: new priority value
        :param e_id: experience id
        :return: bool
        """
        if e_id in self.e2p:
            p_id = self.e2p[e_id]
            self.priority_queue[p_id] = (priority, e_id)
            self.p2e[p_id] = e_id

            self.down_heap(p_id)
            self.up_heap(p_id)
            return True
        else:
            # this e id is new, do insert
            return self._insert(priority, e_id)

    def get_max_priority(self):
        """
        get max priority, if no experience, return 1
        :return: max priority if size > 0 else 1
        """
        if self.size > 0:
            return self.priority_queue[1][0]
        else:
            return 1

    def pop(self):
        """
        pop out the max priority value with its experience id
        :return: priority value & experience id
        """
        if self.size == 0:
            sys.stderr.write('Error: no value in heap, pop failed\n')
            return False, False

        pop_priority, pop_e_id = self.priority_queue[1]
        self.e2p[pop_e_id] = -1
        # replace first
        last_priority, last_e_id = self.priority_queue[self.size]
        self.priority_queue[1] = (last_priority, last_e_id)
        self.size -= 1
        self.e2p[last_e_id] = 1
        self.p2e[1] = last_e_id

        self.down_heap(1)

        return pop_priority, pop_e_id

    def up_heap(self, i):
        """
        upward balance
        :param i: tree node i
        :return: None
        """
        if i > 1:
            parent = int(math.floor(i / 2))
            if self.priority_queue[parent][0] < self.priority_queue[i][0]:
                tmp = self.priority_queue[i]
                self.priority_queue[i] = self.priority_queue[parent]
                self.priority_queue[parent] = tmp
                # change e2p & p2e
                self.e2p[self.priority_queue[i][1]] = i
                self.e2p[self.priority_queue[parent][1]] = parent
                self.p2e[i] = self.priority_queue[i][1]
                self.p2e[parent] = self.priority_queue[parent][1]
                # up heap parent
                self.up_heap(parent)

    def down_heap(self, i):
        """
        downward balance
        :param i: tree node i
        :return: None
        """
        if i < self.size:
            greatest = i
            left, right = i * 2, i * 2 + 1
            if left < self.size and self.priority_queue[left][0] > self.priority_queue[greatest][0]:
                greatest = left
            if right < self.size and self.priority_queue[right][0] > self.priority_queue[greatest][0]:
                greatest = right

            if greatest != i:
                tmp = self.priority_queue[i]
                self.priority_queue[i] = self.priority_queue[greatest]
                self.priority_queue[greatest] = tmp
                # change e2p & p2e
                self.e2p[self.priority_queue[i][1]] = i
                self.e2p[self.priority_queue[greatest][1]] = greatest
                self.p2e[i] = self.priority_queue[i][1]
                self.p2e[greatest] = self.priority_queue[greatest][1]
                # down heap greatest
                self.down_heap(greatest)

    def get_priority(self):
        """
        get all priority value
        :return: list of priority
        """
        return list(map(lambda x: x[0], self.priority_queue.values()))[0:self.size]

    def get_e_id(self):
        """
        get all experience id in priority queue
        :return: list of experience ids order by their priority
        """
        return list(map(lambda x: x[1], self.priority_queue.values()))[0:self.size]

    def balance_tree(self):
        """
        rebalance priority queue
        :return: None
        """
        sort_array = sorted(self.priority_queue.values(), key=lambda x: x[0], reverse=True)
        # reconstruct priority_queue
        self.priority_queue.clear()
        self.p2e.clear()
        self.e2p.clear()
        cnt = 1
        while cnt <= self.size:
            priority, e_id = sort_array[cnt - 1]
            self.priority_queue[cnt] = (priority, e_id)
            self.p2e[cnt] = e_id
            self.e2p[e_id] = cnt
            cnt += 1
        # sort the heap
        for i in range(int(math.floor(self.size / 2)), 1, -1):
            self.down_heap(i)

    def priority_to_experience(self, priority_ids):
        """
        retrieve experience ids by priority ids
        :param priority_ids: list of priority id
        :return: list of experience id
        """
        return [self.p2e[i] for i in priority_ids]

def list_to_dict(in_list):
    return dict((i, in_list[i]) for i in range(0, len(in_list)))

def exchange_key_value(in_dict):
    return dict((in_dict[i], i) for i in in_dict)


#!/usr/bin/python3
# -*- coding: utf-8 -*-

import logging
import matplotlib
import multiprocessing as mp
import numpy as np
import os
import cv2
import sys
# Fix problem: no $DISPLAY environment variable
matplotlib.use('Agg')

from argparse import ArgumentParser
from datetime import datetime as dt
from pprint import pprint

from config import cfg

import json
import paddle
import utils.binvox_visualization
import utils.data_loaders
import utils.data_transforms
import utils.network_utils
import utils.voxel

from datetime import datetime as dt

from models.encoder import Encoder
from models.decoder import Decoder
from models.merger import Merger

def get_args_from_command_line():
    parser = ArgumentParser(description='Parser of Runner of Pix2Vox')
    parser.add_argument('--gpu',
                        dest='gpu_id',
                        help='GPU device id to use [cuda0]',
                        default=cfg.CONST.DEVICE,
                        type=str)

    parser.add_argument('--imgs', dest='imgs', help='demo dir', default=None)
    parser.add_argument('--weights', dest='weights', help='Initialize network from the weights file', default=None)
    parser.add_argument('--out', dest='out_path', help='Set output path', default=cfg.DIR.OUT_PATH)
    args = parser.parse_args()
    return args

def main():
    # Get args from command line
    args = get_args_from_command_line()

    if args.gpu_id is not None:
        cfg.CONST.DEVICE = args.gpu_id
    if args.out_path is not None:
        cfg.DIR.OUT_PATH = args.out_path
    if args.weights is not None:
        cfg.CONST.WEIGHTS = args.weights

    # Print config
    print('Demo Use Config:')
    pprint(cfg)

    # Set GPU to use
    if type(cfg.CONST.DEVICE) == str:
        os.environ["CUDA_VISIBLE_DEVICES"] = cfg.CONST.DEVICE

    # Start train/test process
    if args.imgs:
        demo_net(cfg, args.imgs)

def demo_net(cfg, imgs_path):
    encoder = Encoder(cfg)
    decoder = Decoder(cfg)
    merger = Merger(cfg)

    print('[INFO] %s Loading weights from %s ...' % (dt.now(), cfg.CONST.WEIGHTS))
    encoder_state_dict = paddle.load(os.path.join(cfg.CONST.WEIGHTS, "encoder.pdparams"))
    encoder.set_state_dict(encoder_state_dict)
    decoder_state_dict = paddle.load(os.path.join(cfg.CONST.WEIGHTS, "decoder.pdparams"))
    decoder.set_state_dict(decoder_state_dict)

    if cfg.NETWORK.USE_MERGER:
        merger_state_dict = paddle.load(os.path.join(cfg.CONST.WEIGHTS, "merger.pdparams"))
        merger.set_state_dict(merger_state_dict)

    # Switch models to evaluation mode
    encoder.eval()
    decoder.eval()
    merger.eval()

    rendering_images = []
    if os.path.isfile(imgs_path):
        print("demo img")
        rendering_image = cv2.imread(imgs_path, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255.
        rendering_image = np.asarray(rendering_image)[np.newaxis,:,:,:]
        # print(rendering_image.shape)
        IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W
        CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W
        test_transforms = utils.data_transforms.Compose([
            utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE),
            utils.data_transforms.RandomBackground(cfg.TEST.RANDOM_BG_COLOR_RANGE),
            utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD),
            utils.data_transforms.ToTensor(),
        ])

        rendering_image = test_transforms(rendering_image)
        # print(rendering_image)
        rendering_image = paddle.reshape(rendering_image, [1, 1, 3, 224, 224])
        with paddle.no_grad():
            # Get data from data loader
            rendering_image = utils.network_utils.var_or_cuda(rendering_image)

            # Test the encoder, decoder and merger
            image_features = encoder(rendering_image)
            raw_features, generated_volume = decoder(image_features)

            if cfg.NETWORK.USE_MERGER:
                generated_volume = merger(raw_features, generated_volume)
            else:
                generated_volume = paddle.mean(generated_volume, axis=1)

            for th in cfg.TEST.DEMO_VOXEL_THRESH:
                _volume = paddle.greater_equal(generated_volume, paddle.to_tensor(th)).astype("float32")
                _volume = paddle.reshape(_volume, [32, 32, 32])
                # print(_volume.shape)
                # print(_volume)
                # Append generated volumes to TensorBoard
                if cfg.DIR.OUT_PATH:
                    # Volume Visualization
                    pred_file_name = os.path.join(cfg.DIR.OUT_PATH, imgs_path.split('/')[-1].split('.')[0]+'.obj')
                    print("save ",pred_file_name)
                    utils.voxel.voxel2obj(pred_file_name, _volume.cpu().numpy())

    elif os.path.isdir(imgs_path):
        print("demo dir")
        rendering_files_path = os.listdir(imgs_path)
        for rendering_file_path in rendering_files_path:
            if '.png' not in rendering_file_path:
                continue
            print(os.path.join(imgs_path, rendering_file_path))
            rendering_image = cv2.imread(os.path.join(imgs_path, rendering_file_path), cv2.IMREAD_UNCHANGED).astype(np.float32) / 255.
            rendering_image = np.asarray(rendering_image)[np.newaxis,:,:,:]
            # print(rendering_image.shape)
            IMG_SIZE = cfg.CONST.IMG_H, cfg.CONST.IMG_W
            CROP_SIZE = cfg.CONST.CROP_IMG_H, cfg.CONST.CROP_IMG_W
            test_transforms = utils.data_transforms.Compose([
                utils.data_transforms.CenterCrop(IMG_SIZE, CROP_SIZE),
                utils.data_transforms.RandomBackground(cfg.TEST.RANDOM_BG_COLOR_RANGE),
                utils.data_transforms.Normalize(mean=cfg.DATASET.MEAN, std=cfg.DATASET.STD),
                utils.data_transforms.ToTensor(),
            ])

            rendering_image = test_transforms(rendering_image)
            # print(rendering_image)
            rendering_image = paddle.reshape(rendering_image, [1, 1, 3, 224, 224])
            with paddle.no_grad():
                # Get data from data loader
                rendering_image = utils.network_utils.var_or_cuda(rendering_image)

                # Test the encoder, decoder and merger
                image_features = encoder(rendering_image)
                raw_features, generated_volume = decoder(image_features)

                if cfg.NETWORK.USE_MERGER:
                    generated_volume = merger(raw_features, generated_volume)
                else:
                    generated_volume = paddle.mean(generated_volume, axis=1)

                # for th in cfg.TEST.VOXEL_THRESH:
                #     _volume = paddle.greater_equal(generated_volume, paddle.to_tensor(th)).astype("float32")
                #     print(_volume.shape)

                # Append generated volumes to TensorBoard
                if cfg.DIR.OUT_PATH:
                    # Volume Visualization
                    gv = generated_volume.detach().cpu().numpy()
                    pred_file_name = os.path.join(cfg.DIR.OUT_PATH, imgs_path, rendering_file_path.split('.')[0]+'.obj')
                    utils.voxel.voxel2obj(pred_file_name, gv[0, 1] > cfg.TEST.DEMO_VOXEL_THRESH)
    else:
        raise Exception("error input path")

if __name__ == '__main__':
    # Check python version
    if sys.version_info < (3, 0):
        raise Exception("Please Check python version")

    # Setup logger
    mp.log_to_stderr()
    logger = mp.get_logger()
    logger.setLevel(logging.INFO)

    main()