import os.path as osp
import gym
import time
import joblib
import logging
import numpy as np
import tensorflow as tf
import tensorflow.contrib as tc
from baselines import logger

from baselines.common import set_global_seeds, explained_variance
from baselines.common.FakeVecEnv import FakeVecEnvWrapper
from baselines.common.vec_env.subproc_vec_env import SubprocVecEnv
from baselines.common.misc_util import (
    parse_args,
    SimpleMonitor
)
import itertools

from baselines.utils.utils import discounted_returns, td_returns
from baselines.utils.utils import Scheduler, make_path, find_trainable_variables
from baselines.models.decmodels import CollectiveDecActorGrid
from baselines.models.centralized_critic_models import GridDeepVPN
from baselines.utils.utils import cat_entropy, mse
from baselines.environments.patrolling import PatrollingGame
from baselines.utils.utils import flatten

TINY = 1e-8

class Model(object):
    def __init__(self, actor, critic, observation_space, action_space, num_procs,
                 ent_coef=0.01, critic_l2_reg=0.01, actor_lr=1e-4, critic_lr=1e-3, population=0, rewardScaler=0.0,
                 critic_training_type=1):
        config = tf.ConfigProto(allow_soft_placement=True,
                                intra_op_parallelism_threads=num_procs,
                                inter_op_parallelism_threads=num_procs)
        config.gpu_options.allow_growth = True
        sess = tf.Session(config=config)
        # nact = ac_space.shape
        # nbatch = nenvs*nsteps

        obs0 = tf.placeholder(tf.float32, shape=(None, observation_space.shape[0]), name='obs')

        state_count = []  # tf.placeholder(tf.float32, shape=(None,) + action_space.shape, name='state_count')
        action_count = []

        # immediate revenue
        r = tf.placeholder(tf.float32, shape=(None,))

        # next state accumulated return
        next_R = tf.placeholder(tf.float32, shape=(None,))
        vf_mean_tf = tf.placeholder(tf.float32, shape=())
        #
        for component in range(len(action_space.shapes)):
            action_count.append(tf.placeholder(tf.float32, shape=(None, action_space.shapes[component]),
                                               name=('action' + str(component))))
            state_count.append(tf.placeholder(tf.float32, shape=(None, action_space.shapes[component]),
                                              name=('state_count' + str(component))))

        action_count = tf.stack(action_count, axis=1, name='joint_actions')
        state_count = tf.stack(state_count, axis=1, name='joint_state_counts')

        actor_tf = actor(obs0)
        # actor_funcs = actor(obs0)
        # actor_tf = actor_funcs['output']
        # local_obs_tf = actor_funcs['local_obs']
        critic_outputs = critic(obs0, action_count)
        critic_tf = critic_outputs['symbolic_val']
        immediate_reward_tf = critic_outputs['immediate_reward']
        next_return_tf = critic_outputs['next_return']
        next_state_count = critic_outputs['next_state_count']

        # critic_tf_with_actor_tf = critic(obs0, actor_tf*state_count, reuse=True)

        # entropy = -tf.reduce_mean(tf.reduce_sum(state_count * actor_tf * tf.log(actor_tf + TINY), axis=(1, 2)))
        # critic_tf_with_actor_tf = critic(obs0, actor_tf * state_count, reuse=True)['symbolic_val']
        ADV = tf.placeholder(tf.float32, shape=(None, len(action_space.shapes), action_space.shapes[0]))
        action_grads = tf.gradients(critic_tf, action_count)
        entropy = -tf.reduce_mean(tf.reduce_sum(state_count * actor_tf * tf.log(actor_tf + TINY), axis=(1, 2)))
        pg_loss = -tf.reduce_mean(
            tf.reduce_sum(ADV * action_count * tf.log(actor_tf + TINY), axis=(1, 2)))/population# - 0.01*entropy/population

        vf_loss = tf.reduce_mean(tf.square(immediate_reward_tf - r)) \
                  + tf.reduce_mean(tf.square(next_return_tf - next_R))

        # tf.reduce_mean(tf.reduce_sum(tf.square(r_tf - r), axis=1)) +
        vf_immediate_loss = tf.reduce_mean(tf.square(immediate_reward_tf - r))
        vf_next_return_loss = tf.reduce_mean(tf.square(next_return_tf - next_R))

        if critic_l2_reg > 0:
            critic_reg_vars = [var for var in critic_tf.trainable_vars if
                               'kernel' in var.name and 'output' not in var.name]
            for var in critic_reg_vars:
                logger.info('  regularizing: {}'.format(var.name))
            logger.info('  applying l2 regularization with {}'.format(critic_l2_reg))
            critic_reg = tc.layers.apply_regularization(
                tc.layers.l2_regularizer(critic_l2_reg),
                weights_list=critic_reg_vars
            )
            vf_loss += critic_reg

        # loss = pg_loss  + vf_loss * vf_coef #pg_loss + entropy*ent_coef + vf_loss * vf_coef

        actor_params = actor.trainable_vars
        actor_grads = list(zip(tf.gradients(pg_loss, actor_params), actor_params))
        actor_optimizer = tf.train.AdamOptimizer(learning_rate=actor_lr)
        actor_update = actor_optimizer.apply_gradients(actor_grads)
        critic_params = critic.trainable_vars
        critic_grads = list(zip(tf.gradients(vf_loss, critic_params), critic_params))
        critic_optimizer = tf.train.AdamOptimizer(learning_rate=critic_lr)
        critic_update = critic_optimizer.apply_gradients(critic_grads)

        pre_critic_optimizer = tf.train.AdamOptimizer(learning_rate=critic_lr * 10)
        pre_critic_update = pre_critic_optimizer.apply_gradients(critic_grads)

        def get_next_return(obs, action_count_samples):
            td_map = {obs0: obs, action_count: action_count_samples}
            next_state_revenue = sess.run(
                next_return_tf,
                td_map
            )
            return next_state_revenue

        # def get_local_obs(obs):
        #     return sess.run(local_obs_tf, {obs0: obs})
        #
        # self.get_local_obs = get_local_obs

        def get_immediate_reward(obs, action_count_samples):
            td_map = {obs0: obs, action_count: action_count_samples}
            next_state_revenue = sess.run(
                immediate_reward_tf,
                td_map
            )
            return next_state_revenue

        def get_next_state_count(obs, action_count_samples):
            td_map = {obs0: obs, action_count: action_count_samples}
            next_state_revenue = sess.run(
                next_state_count,
                td_map
            )
            return next_state_revenue

        self.get_next_return = get_next_return
        self.get_immediate_reward = get_immediate_reward
        self.get_next_state_count = get_next_state_count

        def train(obs, COMA_ADV, state_count_samples, action_count_samples, sample_r, sample_next_R, vf_mean):
            td_map = {obs0: obs, ADV: COMA_ADV, state_count: state_count_samples, action_count: action_count_samples, r: sample_r,
                      next_R: sample_next_R, vf_mean_tf: vf_mean}
            policy_loss, value_loss, _, _, immediate_val_loss, revenue_VF_loss = sess.run(
                [pg_loss, vf_loss, actor_update, critic_update, vf_immediate_loss,
                 vf_next_return_loss],#, actor_grads
                td_map
            )#, actor_grad_vals
            return policy_loss, value_loss, immediate_val_loss, revenue_VF_loss#, actor_grad_vals

        def trainCritic(obs, state_count_samples, action_count_samples, sample_r, sample_next_R):
            td_map = {obs0: obs, state_count: state_count_samples, action_count: action_count_samples, r: sample_r,
                      next_R: sample_next_R}
            value_loss, _, immediate_val_loss, revenue_VF_loss = sess.run(
                [vf_loss, pre_critic_update, vf_immediate_loss, vf_next_return_loss],
                td_map
            )

            return value_loss, immediate_val_loss, revenue_VF_loss

        self.trainCritic = trainCritic

        def get_action_grads(obs, action_count_samples):
            td_map = {obs0: obs, action_count: action_count_samples}
            return sess.run(action_grads, td_map)

        self.get_action_grads = get_action_grads

        def step(obs):
            return sess.run(actor_tf, {obs0: obs})

        def value(obs, action_count_input):
            return sess.run(critic_tf, {obs0: obs, action_count: action_count_input})

        def get_next_state_count(obs, action_count_input):
            return sess.run(next_state_count, {obs0: obs, action_count: action_count_input})

        self.get_next_state_count = get_next_state_count

        def cost(obs, action_count_input):
            return sess.run(critic_outputs['cost'], {obs0: obs, action_count: action_count_input})

        def costRaw(obs, action_count_input):
            return sess.run(critic_outputs['cost'], {obs0: obs, action_count: action_count_input})

        def reward(obs, action_count_input):
            return sess.run(critic_outputs['immediateRewards'], {obs0: obs, action_count: action_count_input})

        def next_vals(obs, action_count_input):
            return sess.run(critic_outputs['next_vals'], {obs0: obs, action_count: action_count_input})

        self.cost = cost
        self.costRaw = costRaw
        self.reward = reward
        self.next_vals = next_vals
        saver = tf.train.Saver()

        def save(save_path, fileName, checkPoint):
            # actor_param_values = sess.run(actor_params)
            # critic_param_values = sess.run(critic_params)
            make_path(save_path)
            save_file = saver.save(sess, save_path + fileName, global_step=checkPoint)
            print("Model saved in file: %s" % save_file)

        def load(load_path, ):
            saver.restore(sess, load_path)

        self.load = load

        def loadActor(load_path):
            loaded_params = joblib.load(load_path)
            restores = []
            for p, loaded_p in zip(actor_params, loaded_params):
                restores.append(p.assign(loaded_p))
            ps = sess.run(restores)

        self.loadActor = loadActor

        def get_actor_params():
            return sess.run(actor_params)

        def get_critic_params():
            return sess.run(critic_params)

        def get_d(obs):
            return sess.run(critic_outputs['d'], {obs0: obs})

        def get_future_trip_weights(obs):
            return sess.run(critic_outputs['future_trip_weights'], {obs0: obs})

        def get_V_norm(obs):
            return sess.run(critic_outputs['V_norm'], {obs0: obs})

        self.get_V_norm = get_V_norm

        def get_customer_flow(obs, action_count_input):
            return sess.run(critic_outputs['customer_flow'], {obs0: obs, action_count: action_count_input})

        self.get_customer_flow = get_customer_flow

        def get_trip_weights():
            return sess.run(critic_outputs['trip_weights'])

        self.get_d = get_d
        self.get_future_trip_weights = get_future_trip_weights
        self.get_trip_weights = get_trip_weights
        #
        # def load(load_path):
        #     loaded_params = joblib.load(load_path)
        #     restores = []
        #     for p, loaded_p in zip(params, loaded_params):
        #         restores.append(p.assign(loaded_p))
        #     ps = sess.run(restores)
        self.get_critic_params = get_critic_params
        self.get_actor_params = get_actor_params
        self.train = train
        self.actor = actor_tf
        self.critic = critic_tf
        self.step = step
        self.value = value
        self.save = save
        self.population = float(population)
        self.rewardScaler = rewardScaler
        # self.load = load
        tf.global_variables_initializer().run(session=sess)

class Runner(object):
    def __init__(self, env, model, nsteps=5, nstack=4, gamma=0.99, lam=0.8, vf_normalization=False):
        self.env = env
        self.model = model
        # nh, nw, nc = env.observation_space.shape
        nenv = env.num_envs
        self.nenv = nenv
        self.batch_ob_shape = (nenv * nsteps, env.observation_space.shape[0])
        self.batch_ac_shape = (nenv * nsteps, len(env.action_space.shapes), env.action_space.shapes[0])
        self.batch_inflow_shape = (nenv * nsteps, len(env.action_space.shapes))
        self.batch_ac_state_shape = (nenv * nsteps, len(env.action_space.shapes))
        self.batch_ac_unflatten_shape = (nenv, nsteps, len(env.action_space.shapes), env.action_space.shapes[0])
        self.obs = np.zeros((nenv, env.observation_space.shape[0]), dtype=np.float32)
        self.obs, self.stateCount = env.reset()
        # self.update_obs(obs)
        self.gamma = gamma
        self.lam = lam
        self.nsteps = nsteps
        self.dones = [False for _ in range(nenv)]
        self.actor_params = []
        self.critic_params = []
        self.vf_normalization = vf_normalization
        self.epoch = 0
        self.meanAds = 1

    # def update_obs(self, obs):
    #     self.obs = np.roll(self.obs, shift=-1, axis=3)
    #     self.obs[:, :, :, -1] = obs[:, :, :, 0]

    def run(self):
        mb_obs, mb_actions, mb_rewards, mb_action_probs, mb_dones, mb_values = [], [], [], [], [], []
        for n in range(self.nsteps):
            action_probs = self.model.step(self.obs)
            mb_obs.append(np.copy(self.obs))
            mb_action_probs.append(action_probs)
            mb_dones.append(self.dones)
            obs, action, rewards, dones, info = self.env.step(action_probs)
            # action /=20.0
            count = [info[i]['count'] for i in range(len(info))]
            mb_actions.append(action)
            self.dones = dones
            for n, done in enumerate(dones):
                if done:
                    self.obs[n] = self.obs[n] * 0
            self.obs = obs
            mb_rewards.append(rewards)
        mb_dones.append(self.dones)
        last_values = self.model.get_next_return(mb_obs[-1], np.array(mb_actions[-1]))# / self.model.population)

        # batch of steps to batch of rollouts

        # self.model.value(mb_obs, mb_actions[0] / self.model.population)
        mb_action_probs = np.asarray(mb_action_probs, dtype=np.float32).swapaxes(1, 0)
        mb_actions = np.asarray(mb_actions, dtype=np.float32).swapaxes(1, 0)

        average_reward = (np.sum(mb_rewards) / self.env.num_envs)

        mb_rewards = (np.asarray(mb_rewards, dtype=np.float32).swapaxes(1,
                                                                        0) / self.model.rewardScaler)# / self.model.population

        action_count = np.array(mb_actions).reshape(self.batch_ac_shape)
        # mb_values = np.asarray(mb_values, dtype=np.float32).swapaxes(1, 0)
        action_count[:, :, :] = np.sum(action_count, axis=2)[:, :, np.newaxis]  # HACK
        mb_states = action_count.reshape(self.batch_ac_shape)

        mb_action_probs = mb_action_probs.reshape(self.batch_ac_shape)

        mb_dones = np.asarray(mb_dones, dtype=np.bool).swapaxes(1, 0)
        mb_masks = mb_dones[:, :-1]
        self.obs, self.stateCount = self.env.reset()

        mb_actions = mb_actions# / self.model.population

        mb_obs = np.asarray(mb_obs, dtype=np.float32).swapaxes(1, 0)

        mb_obs = mb_obs.reshape(self.batch_ob_shape)  # / self.model.population
        # mb_masks = mb_masks.flatten()

        mb_states = mb_states #/ self.model.population
        mb_actions = mb_actions.reshape(
            self.batch_ac_shape)

        mb_returns = []
        mb_next_vals = []
        # values = self.model.value(mb_obs, mb_actions)
        # for n, atomic_rewards in enumerate(mb_rewards):
        #     rewards = td_returns(atomic_rewards, values, self.gamma, self.lam)#discounted_returns(atomic_rewards, last_values[n], self.gamma)
        #     mb_returns.append(rewards)
        #     rewards = np.roll(rewards, -1, axis=0)
        #     rewards[-1] = last_values[n]
        #     mb_next_vals.append(rewards)

        for n, atomic_rewards in enumerate(mb_rewards):
            rewards = discounted_returns(atomic_rewards, last_values[n], self.gamma)
            mb_returns.append(rewards)
            rewards = np.roll(rewards, -1, axis=0)
            rewards[-1] = last_values[n]
            mb_next_vals.append(rewards)

        mb_rewards = mb_rewards.flatten()
        mb_next_vals = np.array(mb_next_vals).flatten()

        action_grads = self.model.get_action_grads(mb_obs, mb_actions)[0]
        if self.vf_normalization:
            if self.epoch == 0:
                # Thien check dimension here
                self.meanAds = np.sum(action_grads * mb_actions) / (len(action_grads)*self.model.population)
                self.stdAds = np.sqrt(np.sum(np.square(action_grads - self.meanAds) * mb_actions/self.model.population) / (len(action_grads)))
                self.epoch = 1
            else:
                self.meanAds1 = np.sum(action_grads * mb_actions) / (len(action_grads)*self.model.population)
                self.stdAds1 = np.sqrt(np.sum(np.square(action_grads - self.meanAds) * mb_actions/self.model.population) / (len(action_grads)))
                self.meanAds = 0.9 * self.meanAds1 + 0.1 * self.meanAds
                self.stdAds = 0.9 * self.stdAds1 + 0.1 * self.stdAds
            action_grads = (action_grads - self.meanAds) / (self.stdAds + 1e-8)
        mb_baselines = np.sum(action_grads * mb_action_probs, axis=2)
        mb_advantages = action_grads - mb_baselines[:, :, np.newaxis]
        # print('max advantage, min advantage:' + str(np.max(mb_advantages)) + ',' + str(np.min(mb_advantages)))
        return mb_obs, mb_actions, mb_states, mb_rewards, mb_masks, mb_action_probs, mb_next_vals, mb_advantages, average_reward, self.meanAds

from os.path import dirname, abspath
def learn(actor, critic, env, seed, nsteps=5, nstack=4, total_timesteps=int(80e6), ent_coef=0.01, gamma=0.9,
          lam=0.7, log_interval=1, actor_lr=1e-4, critic_lr=1e-3, critic_l2_reg=0.01, population=0, rewardScaler=0.0,
          layer_norm=False, batch_norm=False, critic_training_type=1, vf_normalization=False, victimNum=1, N=1
          ):
    tf.reset_default_graph()
    set_global_seeds(seed)
    norm_type = ''
    if layer_norm:
        norm_type = 'LayerNorm'
    if batch_norm:
        norm_type = 'BatchNorm'
    if vf_normalization:
        norm_type += '.criticNormalization'
    nenvs = env.num_envs
    ob_space = env.observation_space
    ac_space = env.action_space
    num_procs = len(env.remotes)  # HACK
    model = Model(actor=actor, critic=critic, observation_space=ob_space, action_space=ac_space,
                  num_procs=num_procs, ent_coef=ent_coef, actor_lr=actor_lr, critic_lr=critic_lr,
                  critic_l2_reg=critic_l2_reg, population=population, rewardScaler=rewardScaler,
                  critic_training_type=critic_training_type)
    runner = Runner(env, model, nsteps=nsteps, nstack=nstack, gamma=gamma, lam=lam, vf_normalization=vf_normalization)

    nbatch = nenvs * nsteps
    tstart = time.time()
    average_running_reward = 0
    average_running_value_loss = 0
    average_running_immediate_val_loss = 0
    average_running_revenue_VF_loss = 0
    # model.load(
    #   './a2cGETaxiStatePenaltyandRewardTaxiCentralizedPenaltyandRewardDeepVPNLayerNorm.satisfied_percentage0.9/a2cGETaxiStatePenaltyandRewardPseudoFlowMinoutACWithCostVPNSequelLayerNorm.param-6000')
    print('Intializing with training critic function')
    for k in range(500):
        print('Iteration ' + str(k))
        obs, action_count_samples, state_count_samples, sample_r, mb_masks, mb_action_probs, sample_next_R, mb_advantages, average_reward, meanAds = runner.run()
        value_loss, immediate_val_loss, revenue_VF_loss = model.trainCritic(obs, state_count_samples,
                                                                                action_count_samples,
                                                                                sample_r, sample_next_R)

        print(str(k) + ',' + str(value_loss) + ',' + str(immediate_val_loss) + ',' + str(revenue_VF_loss))
    # model.load('./scaleda2cGEGridPatrollingDeepVPNApproxCOMA' + norm_type + '.criticTraining' + str(
    #                 critic_training_type) + '.lr' + str(actor_lr) + '-' + str(critic_lr) + '.victimNum' + str(
    #                 victimNum) + '.N' + str(N) + '/'+
    #                        'a2cGETaxiStatePenaltyandRewardPseudoFlowMinoutACWithCostVPNSequel' + norm_type + '.param-5000')
    # model.loadActor('./a2cFictitiousSamplingTaxiDecNormalizedObsRewardScaler/fileName')
    # log_name = 'ApproxCOMA' + '.lr' + str(actor_lr) + '-' + str(critic_lr) + '.victimNum' + str(victimNum) + '.N' + str(N) + '.txt'
    # log_file = open(osp.join(dirname(abspath(__file__)), log_name), 'wt')
    # logger.Logger.CURRENT = logger.Logger(output_formats=[logger.HumanOutputFormat(log_file)], dir=None)
    for update in range(1, 20205):
        obs, action_count_samples, state_count_samples, sample_r, mb_masks, mb_action_probs, sample_next_R, mb_advantages, average_reward, meanAds = runner.run()
        average_running_reward += average_reward
        # (obs, state_count_samples, action_count_samples, returns)
        policy_loss, value_loss, immediate_val_loss, revenue_VF_loss = model.train(obs, mb_advantages, state_count_samples, action_count_samples, sample_r, sample_next_R, meanAds)
        # actor_grads = np.array(list(flatten(actor_grads)))
        # max_actor_grads = np.max([np.max(i) for i in actor_grads])
        # min_actor_grads = np.min([np.min(i) for i in actor_grads])
        # print('max grad, min grad:' + str(max_actor_grads) + ',' + str(min_actor_grads))
        nseconds = time.time() - tstart
        fps = int((update * nbatch) / nseconds)
        average_running_value_loss += value_loss
        average_running_immediate_val_loss += immediate_val_loss
        average_running_revenue_VF_loss += revenue_VF_loss
        if update % log_interval == 0:
            # ev = explained_variance(values, returns)
            logger.record_tabular("nupdates", update)
            logger.record_tabular("total_timesteps", update * nbatch)
            logger.record_tabular("fps", fps)
            logger.record_tabular("value_loss", (float(average_running_value_loss)/ log_interval))
            logger.record_tabular("immediate_val_loss", (float(average_running_immediate_val_loss)/ log_interval))
            logger.record_tabular("revenue_VF_loss", (float(average_running_revenue_VF_loss)/ log_interval))
            logger.record_tabular("policy_loss", float(policy_loss))
            # logger.record_tabular("max_predict", np.max(abs(values)))            # logger.record_tabular("explained_variance", float(ev))
            logger.record_tabular("return", (float(average_running_reward) / log_interval))
            logger.dump_tabular()
            if update % 10000 == 0:
                # if (average_running_reward >1e6)|(update% 20 == 0):
                best_value = average_running_reward
                model.save('./GridPatrollingCOMA' + norm_type + '.criticTraining' + str(
                    critic_training_type) + '.lr' + str(actor_lr) + '-' + str(critic_lr) + '.victimNum' + str(
                    victimNum) + '.N' + str(N) + '/',
                           'FCOMA' + norm_type + '.param',
                           checkPoint=update)
            average_running_reward = 0
            average_running_value_loss = 0
            average_running_immediate_val_loss = 0
            average_running_revenue_VF_loss = 0
    env.close()

def main():
    seed = 42
    args = parse_args()
    nenvs = args['num_cpu']
    edge1 = args['grid_size']
    edge2 = args['grid_size']
    stateNum = edge1 * edge2
    initialDistribution = np.zeros(stateNum)
    initialDistribution[int(edge1/2)*edge2 + int(edge2/2)] = 1.0
    victimNum = args['victim_num']
    N = args['population_size']
    env = PatrollingGame(N, initialDistribution, victimNum, edge1=edge1, edge2=edge2)
    populationSize = env.N
    rewardScaler = args['reward_scaler']
    adjacentMatrix = env.neighbors
    adjacent_list = []
    for i in range(stateNum):
        adjacent_list.append(adjacentMatrix[i])
    print('population size' + str(populationSize))

    env.seed(seed)
    incomeFlows = env.incomeFlows
    env = FakeVecEnvWrapper(env)
    vf_normalization = args['vf_normalization']
    set_global_seeds(seed)
    hidden_layers = ()
    hidden_unit_nums = args['hidden_unit_num']
    if hidden_unit_nums > 0:
        hidden_layers = (hidden_unit_nums, hidden_unit_nums)
    critic = GridDeepVPN(env.action_space.shapes, incomeFlows, adjacent_list, layer_norm=args['layer_norm'], batch_norm=args['batch_norm'], stateNum=stateNum, N=N)
    actor = CollectiveDecActorGrid(env.action_space.shapes, layer_norm=args['layer_norm'],
                                   batch_norm=args['batch_norm'],
                                   hidden_sizes=hidden_layers, adjacent_list=adjacent_list, stateNum=stateNum, N=N)
    # critic = GridCollectiveCritic(env.action_space.shapes)
    # actor = GridCollectiveActor(env.action_space.shapes)
    learn(actor, critic, env, seed, nsteps=100, ent_coef=args['ent_coef'], actor_lr=args['actor_lr'],
          critic_lr=args['critic_lr'], critic_l2_reg=args['critic_l2_reg'], lam=args['lambda'],
          population=populationSize, rewardScaler=rewardScaler, layer_norm=args['layer_norm'],
          batch_norm=args['batch_norm'],
          critic_training_type=args['critic_training_type'], vf_normalization=vf_normalization, victimNum=victimNum,
          N=N)

if __name__ == '__main__':
    main()


import mxnet as mx
import numpy as np
import cv2
from tools.rand_sampler import RandSampler
import os, sys
import time
np.set_printoptions(formatter={"float":lambda x:"%.3f "%x},suppress=True)
import math

class DetRecordIter(mx.io.DataIter):
    """
    The new detection iterator wrapper for mx.io.ImageDetRecordIter which is
    written in C++, it takes record file as input and runs faster.
    Supports various augment operations for object detection.

    Parameters:
    -----------
    path_imgrec : str
        path to the record file
    path_imglist : str
        path to the list file to replace the labels in record
    batch_size : int
        batch size
    data_shape : tuple
        (3, height, width)
    label_width : int
        specify the label width, use -1 for variable length
    label_pad_width : int
        labels must have same shape in batches, use -1 for automatic estimation
        in each record, otherwise force padding to width in case you want t
        rain/validation to match the same width
    label_pad_value : float
        label padding value
    resize_mode : str
        force - resize to data_shape regardless of aspect ratio
        fit - try fit to data_shape preserving aspect ratio
        shrink - shrink to data_shape only, preserving aspect ratio
    mean_pixels : list or tuple
        mean values for red/green/blue
    kwargs : dict
        see mx.io.ImageDetRecordIter

    Returns:
    ----------

    """
    def __init__(self, path_imgrec, batch_size, data_shape, path_imglist="",
                 label_width=-1, label_pad_width=-1, label_pad_value=-1,
                 resize_mode='force',  mean_pixels=[123.68, 116.779, 103.939],
                 **kwargs):
        super(DetRecordIter, self).__init__()
        self.rec = mx.io.ImageDetRecordIter(
            path_imgrec     = path_imgrec,
            path_imglist    = path_imglist,
            label_width     = label_width,
            label_pad_width = label_pad_width,
            label_pad_value = label_pad_value,
            batch_size      = batch_size,
            data_shape      = data_shape,
            mean_r          = mean_pixels[0],
            mean_g          = mean_pixels[1],
            mean_b          = mean_pixels[2],
            resize_mode     = resize_mode,
            **kwargs)

        self.provide_label = None
        self._get_batch()
        if not self.provide_label:
            raise RuntimeError("Invalid ImageDetRecordIter: " + path_imgrec)
        self.reset()

    @property
    def provide_data(self):
        return self.rec.provide_data

    def reset(self):
        self.rec.reset()

    def iter_next(self):
        return self._get_batch()

    def next(self):
        if self.iter_next():
            return self._batch
        else:
            raise StopIteration

    def _get_batch(self):
        self._batch = self.rec.next()
        if not self._batch:
            return False

        if self.provide_label is None:
            # estimate the label shape for the first batch, always reshape to n*5
            first_label = self._batch.label[0][0].asnumpy()
            self.batch_size = self._batch.label[0].shape[0]
            self.label_header_width = int(first_label[4])
            self.label_object_width = int(first_label[5])
            assert self.label_object_width >= 5, "object width must >=5"
            self.label_start = 4 + self.label_header_width
            self.max_objects = (first_label.size - self.label_start) // self.label_object_width
            self.label_shape = (self.batch_size, self.max_objects, self.label_object_width)
            self.label_end = self.label_start + self.max_objects * self.label_object_width
            self.provide_label = [('label', self.label_shape)]

        # modify label
        label = self._batch.label[0].asnumpy()
        label = label[:, self.label_start:self.label_end].reshape(
            (self.batch_size, self.max_objects, self.label_object_width))
        self._batch.label = [mx.nd.array(label)]
        return True

class DetIter(mx.io.DataIter):
    """
    Detection Iterator, which will feed data and label to network
    Optional data augmentation is performed when providing batch

    Parameters:
    ----------
    imdb : Imdb
        image database
    batch_size : int
        batch size
    data_shape : int or (int, int)
        image shape to be resized
    mean_pixels : float or float list
        [R, G, B], mean pixel values
    rand_samplers : list
        random cropping sampler list, if not specified, will
        use original image only
    rand_mirror : bool
        whether to randomly mirror input images, default False
    shuffle : bool
        whether to shuffle initial image list, default False
    rand_seed : int or None
        whether to use fixed random seed, default None
    max_crop_trial : bool
        if random crop is enabled, defines the maximum trial time
        if trial exceed this number, will give up cropping
    is_train : bool
        whether in training phase, default True, if False, labels might
        be ignored
    """
    def __init__(self, imdb, batch_size, data_shape, \
                 mean_pixels=[128, 128, 128], rand_samplers=[], \
                 rand_mirror=False, shuffle=False, rand_seed=None, \
                 is_train=True, max_crop_trial=50):
        super(DetIter, self).__init__()

        self._imdb = imdb
        self.batch_size = batch_size
        if isinstance(data_shape, int):
            data_shape = (data_shape, data_shape)
        else:
            assert len(data_shape)==3 and data_shape[0]==3
            data_shape = (data_shape[1], data_shape[2])
        self._data_shape = data_shape
        self._mean_pixels = mx.nd.array(mean_pixels).reshape((3,1,1))
        if not rand_samplers:
            self._rand_samplers = []
        else:
            if not isinstance(rand_samplers, list):
                rand_samplers = [rand_samplers]
            assert isinstance(rand_samplers[0], RandSampler), "Invalid rand sampler"
            self._rand_samplers = rand_samplers
        self.is_train = is_train
        self._rand_mirror = rand_mirror
        self._shuffle = shuffle
        if rand_seed:
            np.random.seed(rand_seed) # fix random seed
        self._max_crop_trial = max_crop_trial

        self._current = 0
        self._size = imdb.num_images
        self._index = np.arange(self._size)

        self._data = None
        self._label = None
        self._get_batch()

    @property
    def provide_data(self):
        return [(k, v.shape) for k, v in self._data.items()]

    @property
    def provide_label(self):
        if self.is_train:
            return [(k, v.shape) for k, v in self._label.items()]
        else:
            return []

    def reset(self):
        self._current = 0
        if self._shuffle:
            np.random.shuffle(self._index)

    def iter_next(self):
        return self._current < self._size

    def next(self):
        if self.iter_next():
            self._get_batch()
            data_batch = mx.io.DataBatch(data=self._data.values(),
                                   label=self._label.values(),
                                   pad=self.getpad(), index=self.getindex())
            self._current += self.batch_size
            return data_batch
        else:
            raise StopIteration

    def getindex(self):
        return self._current // self.batch_size

    def getpad(self):
        pad = self._current + self.batch_size - self._size
        return 0 if pad < 0 else pad

    def _get_batch(self):
        """
        Load data/label from dataset
        """
        batch_data = mx.nd.zeros((self.batch_size, 3, self._data_shape[0], self._data_shape[1]))
        batch_label = []
        for i in range(self.batch_size):
            if (self._current + i) >= self._size:
                if not self.is_train:
                    continue
                # use padding from middle in each epoch
                idx = (self._current + i + self._size // 2) % self._size
                index = self._index[idx]
            else:
                index = self._index[self._current + i]
            # index = self.debug_index
            im_path = self._imdb.image_path_from_index(index)
            with open(im_path, 'rb') as fp:
                img_content = fp.read()
            img = mx.img.imdecode(img_content)
            gt = self._imdb.label_from_index(index).copy() if self.is_train else None
            data, label = self._data_augmentation(img, gt)
            batch_data[i] = data
            if self.is_train:
                batch_label.append(label)
        self._data = {'data': batch_data}
        if self.is_train:
            self._label = {'label': mx.nd.array(np.array(batch_label))}
        else:
            self._label = {'label': None}

    def _data_augmentation(self, data, label):
        """
        perform data augmentations: crop, mirror, resize, sub mean, swap channels...
        """
        if self.is_train and self._rand_samplers:
            rand_crops = []
            for rs in self._rand_samplers:
                rand_crops += rs.sample(label)
            num_rand_crops = len(rand_crops)
            # randomly pick up one as input data
            if num_rand_crops > 0:
                index = int(np.random.uniform(0, 1) * num_rand_crops)
                width = data.shape[1]
                height = data.shape[0]
                crop = rand_crops[index][0]
                xmin = int(crop[0] * width)
                ymin = int(crop[1] * height)
                xmax = int(crop[2] * width)
                ymax = int(crop[3] * height)
                if xmin >= 0 and ymin >= 0 and xmax <= width and ymax <= height:
                    data = mx.img.fixed_crop(data, xmin, ymin, xmax-xmin, ymax-ymin)
                else:
                    # padding mode
                    new_width = xmax - xmin
                    new_height = ymax - ymin
                    offset_x = 0 - xmin
                    offset_y = 0 - ymin
                    data_bak = data
                    data = mx.nd.full((new_height, new_width, 3), 128, dtype='uint8')
                    data[offset_y:offset_y+height, offset_x:offset_x + width, :] = data_bak
                label = rand_crops[index][1]
        if self.is_train:
            interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, \
                              cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
        else:
            interp_methods = [cv2.INTER_LINEAR]
        interp_method = interp_methods[int(np.random.uniform(0, 1) * len(interp_methods))]
        data = mx.img.imresize(data, self._data_shape[1], self._data_shape[0], interp_method)
        if self.is_train and self._rand_mirror:
            if np.random.uniform(0, 1) > 0.5:
                data = mx.nd.flip(data, axis=1)
                valid_mask = np.where(label[:, 0] > -1)[0]
                tmp = 1.0 - label[valid_mask, 1]
                label[valid_mask, 1] = 1.0 - label[valid_mask, 3]
                label[valid_mask, 3] = tmp
        data = mx.nd.transpose(data, (2,0,1))
        data = data.astype('float32')
        data = data - self._mean_pixels
        return data, label

class MultiTaskRecordIter(mx.io.DataIter):
    """
    The new detection iterator wrapper for mx.io.ImageDetRecordIter which is
    written in C++, it takes record file as input and runs faster.
    Supports various augment operations for object detection.

    Parameters:
    -----------
    path_imgrec : str
        path to the record file
    path_imglist : str
        path to the list file to replace the labels in record
    batch_size : int
        batch size
    data_shape : tuple
        (3, height, width)
    label_width : int
        specify the label width, use -1 for variable length
    label_pad_width : int
        labels must have same shape in batches, use -1 for automatic estimation
        in each record, otherwise force padding to width in case you want t
        rain/validation to match the same width
    label_pad_value : float
        label padding value
    resize_mode : str
        force - resize to data_shape regardless of aspect ratio
        fit - try fit to data_shape preserving aspect ratio
        shrink - shrink to data_shape only, preserving aspect ratio
    mean_pixels : list or tuple
        mean values for red/green/blue
    kwargs : dict
        see mx.io.ImageDetRecordIter

    Returns:
    ----------

    """
    def __init__(self, path_imgrec, batch_size, data_shape, path_imglist="",
                 label_width=-1, label_pad_width=-1, label_pad_value=-1,
                 resize_mode='force',  mean_pixels=[123.68, 116.779, 103.939],
                 enable_aug=True, **kwargs):
        super(MultiTaskRecordIter, self).__init__()
        path_imgidx = path_imgrec.replace(".rec",".idx")
        path_imglst = path_imgrec.replace(".rec",".lst")
        data_shape = tuple(data_shape)
        self.enable_aug = enable_aug
        self.batch_size = batch_size
        self.data_shape = data_shape
        self.mean_pixels = mean_pixels
        print(path_imgrec)
        print(path_imgidx)
        self.angle_range = (-5, 5)
        self.scale_range = (.5, 2.)
        self.ratio_range = (.8, 1.2)
        # self.trans_range = (.2*data_shape[2], .2*data_shape[1])

        # build look-up-table (LUT) for mapping segmentation label
        from dataset.cs_labels import labels
        lut = np.ones(256)*255
        for l in labels:
            if l.trainId>=0:
                lut[l.id]=l.id #trainId
        self.lut = lut

        # self.rec = mx.io.ImageDetRecordIter(
        #     path_imgrec     = path_imgrec,
        #     path_imglist    = path_imglist,
        #     label_width     = label_width,
        #     label_pad_width = label_pad_width,
        #     label_pad_value = label_pad_value,
        #     batch_size      = batch_size,
        #     data_shape      = data_shape,
        #     mean_r          = mean_pixels[0],
        #     mean_g          = mean_pixels[1],
        #     mean_b          = mean_pixels[2],
        #     resize_mode     = resize_mode,
        #     **kwargs)

        self.num_samples = sum(1 for line in open(path_imgidx,"r"))
        self.index_table = np.arange(self.num_samples)
        np.random.seed(233) # initialize random seed
        np.random.shuffle(self.index_table)
        self._reset_aug_params()
        self.curr_index = 0

        self.imglst = {}
        with open(path_imglst,"r") as fp:
            lines = fp.readlines()
            dirname = os.path.dirname(path_imglst)
            for line in lines:
                patch = line.rstrip("\n").split()
                segfile = patch[-1].replace("leftImg8bit.jpg","gtFine_labelTrainIds.png")
                segfile = segfile.replace("JPEGImages","SegmentationClass")
                self.imglst[patch[0]] = os.path.join(dirname,"cityscapes",segfile)

        self.rec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, "r")
        self.provide_data = [('data',[self.batch_size, self.data_shape[0], self.data_shape[1], self.data_shape[2]])]

        # self.rec = mx.io.ImageRecordIter(batch_size=batch_size, data_shape=data_shape, path_imgrec=path_imgrec, label_width=128)

        self.provide_label = None
        self._get_batch()
        if not self.provide_label:
            raise RuntimeError("Invalid ImageDetRecordIter: " + path_imgrec)
        self.reset()

    # @property
    # def provide_data(self):
    #     return self.rec.provide_data

    def reset(self):
        # self.rec.reset()
        np.random.shuffle(self.index_table)
        self.curr_index = 0
        self._reset_aug_params()

    def _reset_aug_params(self):
        self.aug_params = np.zeros((self.num_samples,6)) # flip, theta, sx, sy, tx, ty
        self.aug_params[:,0] = np.random.rand(self.num_samples)>.5 # flip
        self.aug_params[:,1] = np.radians(self.angle_range[0]+np.random.rand(self.num_samples)*(self.angle_range[1]-self.angle_range[0])) # rotate
        self.aug_params[:,2] = self.scale_range[0]+np.random.rand(self.num_samples)*(self.scale_range[1]-self.scale_range[0]) # x-scaling
        self.aug_params[:,3] = self.aug_params[:,2]*(self.ratio_range[0]+(np.random.rand(self.num_samples))*(self.ratio_range[1]-self.ratio_range[0])) # y-scaling
        self.aug_params[:,4] = -(np.random.rand(self.num_samples))*self.data_shape[2]*(self.aug_params[:,2]-1.) # tx
        self.aug_params[:,5] = -(np.random.rand(self.num_samples))*self.data_shape[1]*(self.aug_params[:,3]-1.) # ty

    def iter_next(self):
        # return self._get_batch()
        return (self.curr_index+self.batch_size) <= self.num_samples

    def next(self):
        if self.iter_next():
            # tic = time.time()
            self._get_batch()
            # print("elapsed: %.0fms" % ((time.time()-tic)*1000.,)),
            return self._batch, self._fnames
        else:
            raise StopIteration

    def _get_resized(self, img, hdr, seg, data_shape):
        hh, ww, ch = img.shape
        theta, sx, sy, tx, ty = 0., 1.*(data_shape[2]/float(ww)), 1.*(data_shape[1]/float(hh)), 0, 0
        M = np.array([[sx*math.cos(theta),-sy*math.sin(theta),tx],[sx*math.sin(theta),sy*math.cos(theta),ty]])
        img = cv2.warpAffine(img,M,(data_shape[2],data_shape[1]), flags=cv2.INTER_LINEAR)
        if seg is not None:
            seg = cv2.warpAffine(seg,M,(data_shape[2],data_shape[1]), flags=cv2.INTER_NEAREST, borderValue=(0,0,0))
        hdr_reshaped = hdr[3:].reshape((-1,6))

        ### bbox processing ###
        cls = hdr[3:].reshape((-1,6))[:,0]
        idx = np.where(cls>=0)
        pts = hdr_reshaped[:,1:5]
        dist = hdr_reshaped[idx,5]
        xop = np.ones((idx[0].shape[0],1))*self.data_shape[2]
        yop = np.ones((idx[0].shape[0],1))*self.data_shape[1]
        pts_new = pts[idx]*np.hstack((xop,yop,xop,yop))
        if pts_new.shape[0]<1:
            return img, hdr, seg
        #### remove invalid boxes ###
        xmin, ymin, xmax, ymax = hdr_reshaped[:,1],hdr_reshaped[:,2],hdr_reshaped[:,3],hdr_reshaped[:,4]
        areas = (xmax-xmin)*self.data_shape[2]*(ymax-ymin)*self.data_shape[1]
        idx_small = np.where(areas<100) # small bounding boxes
        hdr_reshaped[idx_small,:]=-1
        # idx_small = np.where(dist*255.>100) # remove objects too far away
        # hdr_reshaped[idx_small,:]=-1
        #### move annotations to top ####
        idx_valid = np.where(xmax>-.5) # removed boxes
        reshaped_top = np.squeeze(hdr_reshaped[idx_valid,:])
        hdr_reshaped.fill(-1)
        hdr_reshaped[:reshaped_top.shape[0],:] = reshaped_top

        #### visualize augmented results ####
        # cv2.imshow("seg",np.vstack((img,cv2.merge((seg,seg,seg)))))
        # [exit(0) if cv2.waitKey()&0xff==27 else None]
        return img, hdr, seg

    def _get_augmented(self, img, hdr, seg, data_shape):
        hh, ww, ch = img.shape
        aug_args = self.aug_params[self.curr_index,:]
        # print(aug_args)
        flip, theta, sx, sy, tx, ty = tuple(aug_args)
        sx2, sy2 = sx*(data_shape[2]/float(ww)), sy*(data_shape[1]/float(hh))
        M = np.array([[sx2*math.cos(theta),-sy2*math.sin(theta),tx],[sx2*math.sin(theta),sy2*math.cos(theta),ty]])
        img = cv2.warpAffine(img,M,(data_shape[2],data_shape[1]), flags=cv2.INTER_LINEAR, borderValue=(128,128,128))
        seg = cv2.warpAffine(seg,M,(data_shape[2],data_shape[1]), flags=cv2.INTER_NEAREST, borderValue=(255,255,255))
        cls = hdr[3:].reshape((-1,6))[:,0]
        idx = np.where(cls>=0)
        hdr_reshaped = hdr[3:].reshape((-1,6))
        pts = hdr_reshaped[:,1:5]
        dist = hdr_reshaped[idx,5]
        xop = np.ones((idx[0].shape[0],1))*self.data_shape[2]
        yop = np.ones((idx[0].shape[0],1))*self.data_shape[1]
        pts_new = pts[idx]*np.hstack((xop,yop,xop,yop))
        if pts_new.shape[0]<1:
            return img, hdr, seg

        ################## transform ##################
        pts_new = np.expand_dims(np.vstack((pts_new[:,:2],pts_new[:,2:])),axis=0)
        M = np.array([[sx*math.cos(theta),-sy*math.sin(theta),tx],[sx*math.sin(theta),sy*math.cos(theta),ty]])
        pts_new = np.squeeze(cv2.transform(pts_new,M))
        if flip>.5:
            pts_new[:,0] = self.data_shape[2]-pts_new[:,0]
        pts_new = pts_new*np.hstack((np.ones((pts_new.shape[0],1))/self.data_shape[2],
                                     np.ones((pts_new.shape[0],1))/self.data_shape[1]))
        pts_new = np.hstack((pts_new[:idx[0].shape[0],:],pts_new[idx[0].shape[0]:,:]))
        if flip>.5:
            tmp = pts_new[:,2].copy()
            pts_new[:,2] = pts_new[:,0]
            pts_new[:,0] = tmp
        #### remove invalid boxes
        xmin, ymin, xmax, ymax = pts_new[:,0],pts_new[:,1],pts_new[:,2],pts_new[:,3]
        idx_small = np.where(xmax>-.5) # removed boxes
        xmin[idx_small] = np.clip(xmin[idx_small], 0, 1)
        xmax[idx_small] = np.clip(xmax[idx_small], 0, 1)
        ymin[idx_small] = np.clip(ymin[idx_small], 0, 1)
        ymax[idx_small] = np.clip(ymax[idx_small], 0, 1)
        pts_new[:,0],pts_new[:,1],pts_new[:,2],pts_new[:,3] = xmin, ymin, xmax, ymax
        #### copy back to reshaped header
        hdr_reshaped[idx,1:5]=pts_new
        hdr_reshaped[idx,5]=dist/math.sqrt(sx*sy)
        ################## transform ##################

        #### remove invalid boxes
        xmin, ymin, xmax, ymax = hdr_reshaped[:,1],hdr_reshaped[:,2],hdr_reshaped[:,3],hdr_reshaped[:,4]
        areas = (xmax-xmin)*self.data_shape[2]*(ymax-ymin)*self.data_shape[1]
        idx_small = np.where(areas<100) # small bounding boxes
        hdr_reshaped[idx_small,:]=-1
        idx_small = np.where(xmax<.01) # out-of-image
        hdr_reshaped[idx_small,:]=-1
        idx_small = np.where(xmin>.99) # out-of-image
        hdr_reshaped[idx_small,:]=-1
        idx_small = np.where(ymax<.01) # out-of-image
        hdr_reshaped[idx_small,:]=-1
        idx_small = np.where(ymin>.99) # out-of-image
        hdr_reshaped[idx_small,:]=-1
        #### move annotations to top ####
        idx_valid = np.where(xmax>-.5) # removed boxes
        reshaped_top = np.squeeze(hdr_reshaped[idx_valid,:])
        hdr_reshaped.fill(-1)
        hdr_reshaped[:reshaped_top.shape[0],:] = reshaped_top

        if flip>.5:
            img = cv2.flip(img,1)
            seg = cv2.flip(seg,1)

        #### visualize augmented results ####
        # cv2.imshow("seg",np.vstack((img,cv2.merge((seg,seg,seg)))))
        # [exit(0) if cv2.waitKey()&0xff==27 else None]
        return img, hdr, seg

    def _get_batch(self):
        # self._batch = self.rec.next()

        data = np.zeros((self.batch_size, self.data_shape[0], self.data_shape[1], self.data_shape[2]))
        label = np.ones((self.batch_size, 1206))*-1
        label[:,:3] = np.array(list(self.data_shape))
        seg_out_label = mx.nd.zeros((self.batch_size,self.data_shape[1]/4,self.data_shape[2]/4))
        self._fnames = []
        for batch_idx in xrange(self.batch_size):
            item = self.rec.read_idx(self.index_table[self.curr_index])
            header, img = mx.recordio.unpack_img(item)
            hdr = np.array([header.label.shape[0]]+header.label.tolist())
            id0 = header.id
            seg = cv2.imread(self.imglst[str(id0)],-1)
            self._fnames.append(self.imglst[str(id0)])
            if self.enable_aug:
                assert seg is not None, self.imglst[str(id0)] + " not found."
                img, hdr, seg = self._get_augmented(img, hdr, seg, self.data_shape)
            else:
                img, hdr, seg = self._get_resized(img, hdr, seg, self.data_shape)
            for chidx in xrange(3):
                data[batch_idx,chidx,:,:] = img[:,:,2-chidx] - self.mean_pixels[chidx]
            if seg is not None:
                hh, ww = seg.shape
                seg = cv2.resize(seg, (ww/4,hh/4), interpolation=cv2.INTER_NEAREST)
                seg = cv2.LUT(seg,self.lut).astype(np.uint8) # apply colormap to segmentation label
                seg_out_label[batch_idx,:,:] = seg.reshape((1,self.data_shape[1]/4,self.data_shape[2]/4))
            label[batch_idx,3:3+hdr.shape[0]] = hdr
            self.curr_index += 1
        self._batch = mx.io.DataBatch(data=[mx.ndarray.array(data)], label=[mx.ndarray.array(label), seg_out_label])

        if self.provide_label is None:
            # estimate the label shape for the first batch, always reshape to n*5
            first_label = self._batch.label[0][0].asnumpy()
            print(map(int,first_label[:6]))
            print(", ".join(map(lambda x:"%.3f"%x,first_label[6:12])))
            print(", ".join(map(lambda x:"%.3f"%x,first_label[12:18])))
            self.batch_size = self._batch.label[0].shape[0]
            self.label_header_width = int(first_label[4])
            self.label_object_width = int(first_label[5])
            assert self.label_object_width >= 5, "object width must >=5"
            self.label_start = 4 + self.label_header_width
            self.max_objects = (first_label.size - self.label_start) // self.label_object_width
            self.label_shape = (self.batch_size, self.max_objects, self.label_object_width)
            self.label_end = self.label_start + self.max_objects * self.label_object_width
            self.provide_label = [('label_det', self.label_shape),('seg_out_label', seg_out_label.shape)]
            print(self.provide_label)

        # modify label
        label = self._batch.label[0].asnumpy()
        label = label[:, self.label_start:self.label_end].reshape(
            (self.batch_size, self.max_objects, self.label_object_width))
        self._batch.label = [mx.nd.array(label), seg_out_label]
        # return True


import numpy as np
import pandas as pd
import tensorflow as tf
import SimpleITK as sitk
import skimage.transform import resize

from keras.models import *
from keras import backend as K

from function import *

def predict(model_path, load_dcm_path, save_image_path):

    ################################################################################################################
    # Load Image
    ################################################################################################################

    image = sitk.ReadImage(load_dcm_path)

    image_array = sitk.GetArrayFromImage(image)

    if image_array.shape[0] == 1:

        image_array = image_array[0][:,:,0]

        image_array = np.expand_dims(image_array, axis=2)

    else:

        image_array = image_array.transpose(1,2,0)

    resize_image_array = resize(image_array, (512,512,1))

    test_image = resize_image_array / 255

    ################################################################################################################
    # Load Model
    ################################################################################################################

    model = tf.lite.Interpreter(model_path)

    model.allocate_tensors()

    input_details = model_1.get_input_details()

    output_details = model_1.get_output_details()

    model.set_tensor(input_details[0]['index'], test_image)

    model.invoke()

    ################################################################################################################
    # Predcit Image
    ################################################################################################################

    pred_probability = model.get_tensor(output_details[0]['index'])

    output_data = (pred_probability > 0.5) * 1

    return output_data