from src.optim import Optimizer
from src.optim.lr_scheduler import ReduceOnPlateauScheduler, RsqrtScheduler, NoamScheduler
from src.modules.rnn import RNN
from src.utils.logging import *
from src.utils.common_utils import Collections, Saver
from reinforces.reinforce_utils import SquashedNormal
from reinforces.Annunciator import IntrinsicPredictor
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import src.utils.init as my_init
import numpy as np
import copy
import torch.multiprocessing as _mp
mp = _mp.get_context('spawn')

"""
there are actor (Rephraser) and Q estimator (CriticNet)
"""
class Rephraser(nn.Module):
    """
    which is basically MLP, the output vector is then projected to actor actions and critic value
    using actor critic.

    with double critic network (double q trick)
    states: current embedding of the sentence
    actions: the modification vector (means and deviations) on t step embedding
    rewards: BLEU values and discriminator values
    new_states: hidden representation of the perturbed sentence

    """
    def __init__(self, d_word_vec=512,
                 d_model=256, limit_dist=0.1,
                 dropout=0.0, reparam_noise=1e-6):
        super(Rephraser, self).__init__()
        self.input_size = d_word_vec
        self.action_dim = d_word_vec  # modification on embeddings
        self.hidden_size = d_model
        self.action_range = limit_dist  # action range
        self.reparam_noise = reparam_noise
        self.dropout_rate = dropout
        self.dropout = nn.Dropout(dropout)
        self.log_std_bound = [-5, 4]  # default log std bound

        # current sequence as ctx features
        self.src_gru = RNN(type="gru", batch_first=True, input_size=self.input_size,
                           hidden_size=self.hidden_size, bidirectional=True)

        # Linears for input step features: current embeddings, avg_seqs as ctx
        self.ctx_linear = nn.Linear(
            in_features=2*self.hidden_size,
            out_features=self.hidden_size)
        self.input_linear = nn.Linear(
            in_features=self.input_size,
            out_features=self.hidden_size)

        # layer norm for inputs feature
        self.rephrase_LN = nn.LayerNorm(self.hidden_size, elementwise_affine=True)
        # outputs: actor policy distribution
        # Gaussian policy: mean and std;
        self.rephraser_linear_base_mu = nn.Linear(
            in_features=self.hidden_size,
            out_features=self.hidden_size)
        self.rephraser_linear_mu = nn.Linear(
            in_features=self.hidden_size,
            out_features=self.action_dim)
        self.rephraser_linear_base_log_sig = nn.Linear(
            in_features=self.hidden_size,
            out_features=self.hidden_size)
        self.rephraser_linear_log_sig = nn.Linear(
            in_features=self.hidden_size,
            out_features=self.action_dim)
        # # intrinsic curiosity module
        # self.icm = IntrinsicPredictor(
        #     d_model=self.hidden_size, action_dim=self.action_dim, dropout=dropout)

        # initialize parameter
        self._reset_parameters()

    def _reset_parameters(self):
        my_init.default_init(self.ctx_linear.weight)
        my_init.default_init(self.input_linear.weight)
        self.rephrase_LN.reset_parameters()
        # RNN params are initialized by the RNN module
        my_init.default_init(self.rephraser_linear_base_mu.weight)
        my_init.default_init(self.rephraser_linear_base_log_sig.weight)
        my_init.default_init(self.rephraser_linear_mu.weight)
        my_init.default_init(self.rephraser_linear_log_sig.weight)

    def preprocess(self, x_emb, x_pad_indicator, label_emb):
        """
        compute the s_t representation
        :param label_emb: batched labels [batch, 3, dim], including previous\current\next label
        :paran x_emb: the current sentence embedding. [batch, len, emb_dim]. type=float
        :param x_pad_indicator: PAD indicator for the sequence x.detach().eq(PAD). [batch, len]. type=long
        :return: batched actions in shape [batch, emb_dim]. type=float
        """
        label_emb = label_emb.sum(dim=1)  # [batch, dim]
        ctx_x, _ = self.src_gru(x_emb, x_pad_indicator)
        x_mask = 1.0 - x_pad_indicator.float()
        x_ctx_mean = (ctx_x * x_mask.unsqueeze(2)).sum(1) / x_mask.unsqueeze(2).sum(1)

        rephrase_feature = (self.input_linear(label_emb)+self.ctx_linear(x_ctx_mean))
        rephrase_feature = self.rephrase_LN(rephrase_feature)
        return rephrase_feature

    def forward(self, x_emb, x_pad_indicator, label_emb):
        """
        launch and return the attacked output vec (distribution on the actionspace)
        :param label_emb: batched labels [batch, 3, embeding]
        :param x_emb: batched sentences [batch, max_seq_len]
        :param x_pad_indicator: pad indicator for input sequences: x.detach().eq(PAD).long()
        :return: batched actions in shape [batch, 2*emb_dim] and
                  value estimation for the state averaged value over batch
        """
        # returns the attack actions (vec of expectation \mu and vairation \sigma )
        rephrase_feature = self.preprocess(x_emb, x_pad_indicator, label_emb)
        mu_out = self.rephraser_linear_base_mu(self.dropout(rephrase_feature))
        mu_out = self.dropout(F.relu(mu_out))
        mu = self.rephraser_linear_mu(mu_out)

        log_sig_out = self.rephraser_linear_base_log_sig(self.dropout(rephrase_feature))
        log_sig_out = self.dropout(F.relu(log_sig_out))
        log_sig_out = F.relu(self.rephraser_linear_log_sig(log_sig_out))
        # the output of the Gaussian policy output must be clipped
        log_sig = torch.tanh(log_sig_out)   # clamp to the 0-1
        # constraint the log sig in the bound and convert to sig for distribution
        bound_min, bound_max = self.log_std_bound
        log_sig = bound_min + 0.5 * (bound_max-bound_min) * (log_sig+1)
        sig = log_sig.exp()

        return mu, sig

    def sample_normal(self, x_emb, x_pad_indicator, label_emb, reparamization=True):
        """
        yield action for the policy
        :param x_emb: state
        :param x_pad_indicator: pad indicator for input sequences
        :param label_emb: current step for rephrasing
        :param reparamization: reparamization for sampling, when learning actor
        :return: the action to add on the current embedding
        """
        mu, std = self.forward(x_emb, x_pad_indicator, label_emb)
        p = SquashedNormal(mu, std)
        if reparamization:  # for learning
            raw_action = p.rsample()
        else:
            raw_action = p.sample()
        log_probs = p.log_prob(raw_action)
        # squash the action within certain range, expand to the action range:
        action = raw_action * self.action_range
        return action, log_probs

    def random_step_perturb(self,
                            x_emb, x_pad_indicator,
                            half_mode=False, rand_act=False):
        """
        provide perturbed embedding for training environment annunciator.
        this provides a single step perturb on random position of sequences.
        device is determined by the inputs: x_emb
        1 sample a random timestep(position) in the sequence,;
        2 mask all the perturbation after that timestep(the src)
        :param x_emb: origin source embedding
        :param x_pad_indicator: mask indicating PAD
        :param half_mode: perturb half of the batch if True
        :param rand_act: sample random caps given rephraser output.
        :return: origin_emb, perturbed_emb, flags.
        """
        # randomly choose one position for every sequence in batch
        batch_size, max_seq_len = x_pad_indicator.shape

        # generate "lengthed" mask for the random step perturb
        x_mask = 1 - x_pad_indicator.int()
        x_len = x_mask.sum(dim=1).detach().cpu().numpy()
        _x_np = np.full((batch_size, max_seq_len), fill_value=0, dtype="int64")  # mask for former
        x_np_ = np.full((batch_size, max_seq_len), fill_value=0, dtype="int64")  # mask for perturbed
        for ii in range(batch_size):
            # choose a random time step, range excludes SOS(0) and EOS(len)
            perturb_step = np.random.choice(x_len[ii]-2, 1)+1
            x_np_[ii, :perturb_step[0]+1] = np.ones(perturb_step+1)
            _x_np[ii, :perturb_step[0]] = np.ones(perturb_step)
        _x_np = torch.tensor(_x_np).float().to(x_emb.device)
        x_np_ = torch.tensor(x_np_).float().to(x_emb.device)

        # perturb the entire batch
        with torch.no_grad():
            # former_emb = x_emb.detach().clone()
            perturbed_emb = x_emb.detach().clone()
            self.eval()   # set the agent to evaluation mode
            for t in range(1, max_seq_len-1):
                former_emb = perturbed_emb
                input_emb = former_emb[:, t-1:t+2, :]
                if rand_act:
                    action, _ = self.sample_normal(
                        x_emb=former_emb, x_pad_indicator=x_pad_indicator, label_emb=input_emb,
                        reparamization=False)
                else:
                    mu, _ = self.forward(
                        x_emb=former_emb, x_pad_indicator=x_pad_indicator, label_emb=input_emb)
                    action = mu * self.action_range
                perturbed_emb[:, t, :] += action.float()

            # apply the mask for outputs
            former_emb = perturbed_emb * _x_np.unsqueeze(dim=2) + \
                         x_emb * (1. - _x_np).unsqueeze(dim=2)
            perturbed_emb = perturbed_emb * x_np_.unsqueeze(dim=2) + \
                            x_emb * (1. - x_np_).unsqueeze(dim=2)

            # randomly choose the attacking sequence
            flags = torch.ones(batch_size).to(x_emb.device)
            if half_mode:  # randomly choose half of the batch
                flags = torch.bernoulli(0.5 * flags).to(x_emb.device)
                former_emb = former_emb * flags.unsqueeze(dim=1).unsqueeze(dim=2) \
                             + x_emb * (1. - flags).unsqueeze(dim=1).unsqueeze(dim=2)
                perturbed_emb = perturbed_emb * flags.unsqueeze(dim=1).unsqueeze(dim=2) \
                                + x_emb * (1. - flags).unsqueeze(dim=1).unsqueeze(dim=2)
        return former_emb, perturbed_emb, flags.long()

    def random_seq_perturb(self, x_emb, x_pad_indicator,
                           half_mode=False, rand_act=False,
                           enable_UNK=False):
        """
        sequence perturb for the embedding, half of the batch is perturbed
        x_emb, x_pad_indicator, for the source.
        :param half_mode: randomly perturb half the embedding
        :param rand_act: randomly sample by rephraser's Gaussian policy
        :param enable_UNK: randomly perturb to UNK by the possibility of action capsule length
        :return: 1.perturbed batch of embedding (for self-improvement learning)
                 2.the indicating mask for UNK replacement.
        """
        self.eval()
        with torch.no_grad():
            perturbed_emb = x_emb.detach().clone()
            batch_size, seq_len = x_pad_indicator.shape
            # # UNK mask for the embeddings(additional one for BOS and EOS row)
            # UNK_position = []
            # UNK_position.append(perturbed_emb.new_zeros(batch_size, 1))
            for t in range(1, seq_len-1):
                input_emb = perturbed_emb[:, t-1:t+2, :]
                # action len is also the possibility of the perturbation
                if rand_act:
                    action, _ = self.sample_normal(
                        x_emb=perturbed_emb, x_pad_indicator=x_pad_indicator, label_emb=input_emb,
                        reparamization=False)
                else:
                    mu, _ = self.forward(
                        x_emb=perturbed_emb, x_pad_indicator=x_pad_indicator, label_emb=input_emb)
                    action = mu * self.action_range
                # # randomly inject some UNK. probability is the dropout rate
                # UNK_position.append(torch.bernoulli(
                #     self.dropout_rate * torch.ones(batch_size, 1).to(x_emb.device)))
                perturbed_emb[:, t, :] += action.float()
            # UNK_position.append(perturbed_emb.new_zeros(batch_size, 1))

            # randomly choose the attacking sequence
            flags = torch.ones(batch_size).to(x_emb.device)
            if half_mode:  # randomly choose half of the batch
                flags = torch.bernoulli(0.5 * flags).to(x_emb.device)
                # perturbed embedding shape: [batch_size, sen_len, embedding]
                perturbed_emb = perturbed_emb * flags.unsqueeze(dim=1).unsqueeze(dim=2) \
                                + x_emb * (1. - flags).unsqueeze(dim=1).unsqueeze(dim=2)
            # # indicating mask of perturbing the embedding to UNK
            # mask_to_UNK = torch.cat(UNK_position, dim=1)  # [batch_size, sen_len]
            # if not enable_UNK:  # reset all indicator to 0
            #     mask_to_UNK = mask_to_UNK.new_zeros(mask_to_UNK.size())

        return x_emb, perturbed_emb, flags.long()

    def sync_from(self, rephrase_model):
        # synchronize model parameteres from other attacker. might be on different device
        self.load_state_dict(rephrase_model.state_dict())

    def ensure_shared_grads(self, shared_model):
        """
        for multi-thread(agent) learning updates
        :param shared_model: a global model to share grad parameters
        :return:
        """
        for param, shared_param in zip(self.parameters(),
                                       shared_model.parameters()):
            # only copy the gradients to the shared when shared grad is None
            # the sync in the training process means that the update is processed
            # thread by thread
            if shared_param.grad is not None:
                return
            shared_param._grad = param.grad

class CriticNet(nn.Module):
    """
    models the V(s_t), s_t is the current embedding and the current step inputs
    value net is updated by target net (which is a moving-average updated value nets to avoid over-estimation)
    there are two critic net inside of the class
    """
    def __init__(self, d_word_vec=512,
                 d_model=256, limit_dist=0.1,
                 dropout=0.0, reparam_noise=1e-6):
        super(CriticNet, self).__init__()
        self.input_size = d_word_vec
        self.hidden_size = d_model
        self.limit_dist = limit_dist
        self.reparam_noise = reparam_noise
        self.dropout = nn.Dropout(dropout)

        # current sequences as ctx
        self.src_gru = RNN(type="gru", batch_first=True, input_size=self.input_size,
                           hidden_size=self.hidden_size, bidirectional=True)

        # Linear for input step features
        self.ctx_linear = nn.Linear(
            in_features=2*self.hidden_size,
            out_features=self.hidden_size)
        self.input_linear = nn.Linear(
            in_features=self.input_size,
            out_features=self.hidden_size)
        self.action_linear = nn.Linear(
            in_features=self.input_size,
            out_features=self.hidden_size)
        self.critic_LN = nn.LayerNorm(self.hidden_size, elementwise_affine=True)

        # double-q trick with 2 critics (sharing input features),
        # Q(s_t, a_t) is the smaller one
        self.critic1_linear_base = nn.Linear(
            in_features=self.hidden_size,
            out_features=self.hidden_size)
        self.critic1_linear = nn.Linear(
            in_features=self.hidden_size,
            out_features=1)
        self.critic2_linear_base = nn.Linear(
            in_features=self.hidden_size,
            out_features=self.hidden_size)
        self.critic2_linear = nn.Linear(
            in_features=self.hidden_size,
            out_features=1)
        # initialize the parameter
        self._reset_parameters()

    def _reset_parameters(self):
        my_init.default_init(self.ctx_linear.weight)
        my_init.default_init(self.input_linear.weight)
        my_init.default_init(self.action_linear.weight)
        self.critic_LN.reset_parameters()
        my_init.default_init(self.critic1_linear_base.weight)
        my_init.default_init(self.critic2_linear_base.weight)
        my_init.default_init(self.critic1_linear.weight)
        my_init.default_init(self.critic2_linear.weight)

    def forward(self, x_emb, x_pad_indicator, label_emb, action):
        """
        takes in current embedding and step inputs
        :param label_emb: batched labels [batch, 3, dim], including previous\current\next label
        :param x_emb: batched sentences embedding [batch, max_seq_len, dim]
        :param x_pad_indicator: indicating PAD in the sequence x.detach().eq(PAD)
        :param action: batched actual actions that was add to the states [batch, dim]
        :return: batched value estimation for the state [batch, 1]
        """
        label_emb = label_emb.sum(dim=1)  # [batch, dim]
        ctx_x, _ = self.src_gru(x_emb, x_pad_indicator)
        x_mask = 1.0 - x_pad_indicator.float()
        x_ctx_mean = (ctx_x * x_mask.unsqueeze(2)).sum(1) / x_mask.unsqueeze(2).sum(1)

        critic_feature = self.input_linear(label_emb) + \
                         self.ctx_linear(x_ctx_mean) + self.action_linear(action)
        critic_feature = self.critic_LN(critic_feature)

        critic1_feature = self.critic1_linear_base(self.dropout(critic_feature))
        critic1 = self.critic1_linear(self.dropout(critic1_feature))
        critic2_feature = self.critic2_linear_base(self.dropout(critic_feature))
        critic2 = self.critic2_linear(self.dropout(critic2_feature))
        critic = torch.min(critic1, critic2)

        return critic

    def sync_from(self, critic_model):
        # synchronize model parameteres from other attacker
        self.load_state_dict(critic_model.state_dict())

    def ensure_shared_grads(self, shared_model):
        """
        for multi-thread(agent) learning updates
        :param shared_model: a global model to share grad parameters
        :return:
        """
        for param, shared_param in zip(self.parameters(),
                                       shared_model.parameters()):
            # only copy the gradients to the shared when shared grad is None
            # the sync in the training process means that the update is processed
            # thread by thread
            if shared_param.grad is not None:
                return
            shared_param._grad = param.grad

class SACAgent(object):
    def __init__(self, device="cpu", d_word_vec=512,
                 d_model=256, limit_dist=0.1,
                 dropout=0.0, reparam_noise=1e-6,
                 **kwargs):
        self.device = device
        self.actor = Rephraser(d_word_vec=d_word_vec,
                               d_model=d_model, limit_dist=limit_dist,
                               dropout=dropout, reparam_noise=reparam_noise).to(device)
        self.critic = CriticNet(d_word_vec=d_word_vec,
                                d_model=d_model, limit_dist=limit_dist,
                                dropout=dropout, reparam_noise=reparam_noise).to(device)
        self.saver = Saver(save_prefix="{0}.ckpt".format(os.path.join(kwargs["save_to"], "ACmodel")),
                           num_max_keeping=kwargs["num_kept_checkpoints"])
        self.soft_update_lock = mp.Lock()
        # the entropy regularization weight for SAC learning
        self.learnable_temperature = kwargs["learnable_temperature"]
        self.target_entropy = - d_word_vec  # act_dim (d_word_vec) as the expected entropy base
        self.log_alpha = torch.tensor(np.log(kwargs["init_temperature"])).to(self.device)
        self.log_alpha.requires_grad = True
        # initialize the training mode for the Agent
        self.train()
        self._init_local_optims(kwargs["rephraser_optimizer_configs"])
        # self.load_model()  # always reload model if there is any in the path

    def to(self, device):
        self.actor.to(device)
        self.critic.to(device)
        self.log_alpha.to(device)
        return self

    def share_memory(self):
        # global model needs to share memory with other threads
        self.actor.share_memory()
        self.critic.share_memory()

    @property
    def alpha(self):
        return self.log_alpha.exp()

    def load_model(self, load_final_path:str=None):
        """
        load from path by self.saver
        :param load_final_path: final model path dir, final model doesn't have optim_params
        :return: training step count int
        """
        step = 0
        model_collections = Collections()
        if load_final_path:
            # self.saver.load_latest(
            #     actor_model=self.actor, critic_model=self.critic
            # )  # load from the latest ckpt model
            state_dict = torch.load(os.path.join(load_final_path))
            self.actor.load_state_dict(state_dict["actor_model"])
            self.critic.load_state_dict(state_dict["critic_model"])
        else:
            self.saver.load_latest(
                collections=model_collections,
                actor_model=self.actor, critic_model=self.critic,
                actor_optim=self.actor_optimizer, critic_optim=self.critic_optimizer,
                actor_scheduler=self.actor_scheduler, critic_scheduler=self.critic_scheduler)
            step = model_collections.get_collection("step", [0])[-1]
        return step

    def save_model(self, step=None, save_to_final=None):  # save model parameters, optims, lr_steps
        model_collections = Collections()
        if step is not None:
            model_collections.add_to_collection("step", step)
            self.saver.save(
                global_step=step, collections=model_collections,
                actor_model=self.actor, critic_model=self.critic,
                actor_optim=self.actor_optimizer, critic_optim=self.critic_optimizer,
                actor_scheduler=self.actor_scheduler, critic_scheduler=self.critic_scheduler)
        else:  # only save the model parameters
            assert save_to_final is not None, "final model saving dir must be provided"
            collection = dict()
            collection["actor_model"] = self.actor.state_dict()
            collection["critic_model"] = self.critic.state_dict()
            torch.save(collection,
                       os.path.join(save_to_final, "ACmodel.final"))
        return

    def _init_local_optims(self, rephraser_optimizer_configs):
        """ actor, critic, alpha optimizers and lr scheduler if necessary
        rephraser_optimizer_configs:
            optimizer: "adafactor"
            learning_rate: 0.01
            grad_clip: -1.0
            optimizer_params: ~
            schedule_method: rsqrt
            scheduler_configs:
                d_model: *dim
                warmup_steps: 100
        """
        # initiate local optimizer
        if rephraser_optimizer_configs is None:
            self.actor_optimizer = None
            self.critic_optimizer = None
            self.log_alpha_optimizer = None
            # self.actor_icm_optimizer = None
            self.actor_scheduler = None
            self.critic_scheduler = None
        else:
            self.actor_optimizer = Optimizer(name=rephraser_optimizer_configs["optimizer"],
                                             model=self.actor,
                                             lr=rephraser_optimizer_configs["learning_rate"],
                                             grad_clip=rephraser_optimizer_configs["grad_clip"],
                                             optim_args=rephraser_optimizer_configs["optimizer_params"])
            self.critic_optimizer = Optimizer(name=rephraser_optimizer_configs["optimizer"],
                                              model = self.critic,
                                              lr=rephraser_optimizer_configs["learning_rate"],
                                              grad_clip=rephraser_optimizer_configs["grad_clip"],
                                              optim_args=rephraser_optimizer_configs["optimizer_params"])
            # hardcoded entropy weight updates and icm updates
            self.log_alpha_optimizer = torch.optim.Adam([self.log_alpha], lr=1e-4, betas=(0.9, 0.999))
            # self.actor_icm_optimizer = torch.optim.Adam(self.actor.icm.parameters(), lr=1e-3, )

            # Build scheduler for optimizer if needed
            if rephraser_optimizer_configs['schedule_method'] is not None:
                if rephraser_optimizer_configs['schedule_method'] == "loss":
                    self.actor_scheduler = ReduceOnPlateauScheduler(optimizer=self.actor_optimizer,
                                                                    **rephraser_optimizer_configs["scheduler_configs"])
                    self.critic_scheduler = ReduceOnPlateauScheduler(optimizer=self.critic_optimizer,
                                                                     **rephraser_optimizer_configs["scheduler_configs"])
                elif rephraser_optimizer_configs['schedule_method'] == "noam":
                    self.actor_scheduler = NoamScheduler(optimizer=self.actor_optimizer,
                                                         **rephraser_optimizer_configs["scheduler_configs"])
                    self.critic_scheduler = NoamScheduler(optimizer=self.critic_optimizer,
                                                          **rephraser_optimizer_configs["scheduler_configs"])
                elif rephraser_optimizer_configs["schedule_method"] == "rsqrt":
                    self.actor_scheduler = RsqrtScheduler(optimizer=self.actor_optimizer,
                                                          **rephraser_optimizer_configs["scheduler_configs"])
                    self.critic_scheduler = RsqrtScheduler(optimizer=self.critic_optimizer,
                                                           **rephraser_optimizer_configs["scheduler_configs"])
                else:
                    WARN("Unknown scheduler name {0}. Do not use lr_scheduling.".format(
                        rephraser_optimizer_configs['schedule_method']))
                    self.actor_scheduler = None
                    self.critic_scheduler = None
            else:
                self.actor_scheduler = None
                self.critic_scheduler = None

    def sync_from(self, sac_agent):
        with sac_agent.soft_update_lock, self.soft_update_lock:
            self.actor.sync_from(sac_agent.actor)
            self.critic.sync_from(sac_agent.critic)

    def train(self, training=True):
        # default training is true
        self.training=training
        self.actor.train(training)
        self.critic.train(training)
        return self

    def update_critic(self,
                      states, masks,
                      rephrase_positions, actions,
                      rewards, survive_and_no_maxs,
                      target_critic,
                      update_step, discount_factor, summary_writer,
                      update_trust_region=0.8):
        """
        update critic by using a target_critic net (usually a global critic model) and a buffer
        SARSA for TD learning
        :param states:
        :param masks:
        :param rephrase_positions:
        :param actions: actions
        :param rewards: the rewards
        :param survive_and_no_maxs: able to rollout next step for TD learning

        :param target_critic: provides target value estimation(global model is usually on cpu)
        :param discount_factor: for discounted rewards update
        :param update_step: learning steps
        :param update_trust_region: discount for loss updates
        """

        label_emb = slice_by_indices(states, rephrase_positions, device=self.device)  # next_rephrase_positions to label emb
        next_action, log_probs = self.actor.sample_normal(states, 1.-masks, label_emb, reparamization=True)
        log_probs = log_probs.sum(dim=-1, keepdims=True)
        next_states = transition(states, masks, actions, rephrase_positions)
        next_rephrase_positions = rephrase_positions + 1
        next_label_emb = slice_by_indices(next_states, next_rephrase_positions, device=self.device)

        # # note that with intrinsic curiosity module, the rewards will add curiosity bonus
        # self.actor.icm.eval()
        # rephrase_feature = self.actor.preprocess(states, 1.-masks, label_emb)
        # next_rephrase_feature = self.actor.preprocess(next_states, 1.-masks, next_label_emb)
        # bonus = self.actor.icm.get_surprise_bonus(rephrase_feature, next_rephrase_feature, actions).detach()
        # bonus = 0.01 * bonus * survive_and_no_maxs
        # print("bonus:", bonus.sum())
        # rewards += bonus
        # print("rewards:", rewards.squeeze())

        # note that log_probs has the same dimension with the action. thus the log_prob of a whole action is the sum along dimensions.
        target_critic.eval()
        target_V = target_critic(next_states, 1.-masks, next_label_emb, next_action) - log_probs * self.alpha.detach()
        target_Q = rewards + (survive_and_no_maxs) * discount_factor * target_V  #  we have next states for TD learning rollout
        target_Q = target_Q.detach()

        # get current Q estimates
        current_Q = self.critic(states, 1.-masks, label_emb, actions)
        critic_loss = F.mse_loss(current_Q, target_Q)
        critic_loss *= update_trust_region
        print("critic_loss", critic_loss.sum())

        # Optimize the critic
        self.critic_optimizer.zero_grad()
        if self.critic_scheduler is not None:
            self.critic_scheduler.step(global_step=update_step)
        critic_loss.backward()
        self.critic_optimizer.step()

        # logging: entropy/target_entropy ratio, critic_loss,
        summary_writer.add_scalar("critic_loss", scalar_value=critic_loss, global_step=update_step)

    def update_actor_and_alpha(self,
                               states, masks, rephrase_positions,
                               target_critic,
                               update_step, summary_writer,
                               update_trust_region=0.5):
        """
        :param states: tensor states from the buffer samples
        :param masks: indicats the valid token positions
        :param rephrase_positions: induce the next states by the given states
        :param update_trust_region: current annnunciator trust_acc (valid).
        served as a trust region for RL updates;
        also the weight of rewind or reinforce
        trust_acc * rewind_loss + (1-trust_acc) * policy_loss
        """
        self.actor.train()
        label_emb = slice_by_indices(states, rephrase_positions, device=self.device)
        actions, log_probs = self.actor.sample_normal(states, 1.-masks, label_emb, reparamization=True)
        log_probs = log_probs.sum(dim=-1, keepdims=True)
        actor_Q = self.critic(states, 1.-masks, label_emb, actions)
        policy_loss = (self.alpha.detach() * log_probs - actor_Q).mean()

        summary_writer.add_scalar('policy_loss', policy_loss, update_step)
        summary_writer.add_scalar('entropy_ratio', -log_probs.mean()/self.target_entropy, update_step)

        # the policy rewind loss, the rewind is determined by target value estimates (estimated survival + improvements)
        # negative means rewind needed.
        target_Q = target_critic(states, 1.-masks, label_emb, actions).detach()
        rewind_mask = target_Q.lt(0.).detach().float()  # [batch, 1]
        next_states = transition(states, masks, actions, rephrase_positions).detach()
        next_label_emb = slice_by_indices(next_states, rephrase_positions, device=self.device).detach()
        rewind_action, _ = self.actor.forward(next_states, 1.-masks, next_label_emb)
        target_actions = -actions.detach()
        rewind_loss = F.mse_loss(
            rewind_action*rewind_mask*self.actor.action_range,
            target_actions*rewind_mask)
        summary_writer.add_scalar('rewind_loss', rewind_loss, update_step)

        # the higher trust region means less indicative the perturbations are, policy should focus more on the rewind.
        actor_loss = (update_trust_region) * rewind_loss + (1.- update_trust_region) * policy_loss
        ## update the intrinsic reward module: action reconstruction and feature prediction mse
        # self.actor.icm.train()
        # next_states = transition(states, masks, actions, rephrase_positions)
        # next_label_emb = slice_by_indices(next_states, rephrase_positions, device=self.device)
        # rephrase_feature = self.actor.preprocess(states, 1.0-masks, label_emb)
        # next_rephrase_feature = self.actor.preprocess(next_states, 1.0-masks, next_label_emb)
        # if update_step<3000:
        #     # icm updates does not propagate to the policy on the early stage
        #     icm_loss = self.actor.icm(rephrase_feature.detach(), next_rephrase_feature.detach(), actions)
        # else:
        #     icm_loss = self.actor.icm(rephrase_feature, next_rephrase_feature, actions)
        # summary_writer.add_scalar("intrinsic_curiosity_loss", icm_loss, update_step)
        # # the 0.1 is the setting by Intrinsic curiosity learning
        # actor_loss = actor_loss + icm_loss

        # optimize the actor
        self.actor_optimizer.zero_grad()
        # self.actor_icm_optimizer.zero_grad()
        if self.actor_scheduler is not None:
            self.actor_scheduler.step(global_step=update_step)
        actor_loss.backward()
        self.actor_optimizer.step()
        # self.actor_icm_optimizer.step()

        if self.learnable_temperature:
            self.log_alpha_optimizer.zero_grad()
            alpha_loss = (self.alpha *
                          (-log_probs - self.target_entropy).detach()).mean()
            summary_writer.add_scalar('alpha_loss', alpha_loss, update_step)
            summary_writer.add_scalar('alpha', self.alpha, update_step)
            alpha_loss.backward()
            self.log_alpha_optimizer.step()

    def update_local_net(self, local_agent_configs, replay_buffer, target_critic,
                         update_step, discount_factor, summary_writer,
                         update_trust_region=1.0):
        """
        :param local_agent_configs: provides agent update freq
        :param replay_buffer: provides the SARSA listed below
            [states, masks, actions, rephrase_positions, rewards, terminal_signals]
            states: the embedding as states. [batch, len, emb_dim] float
            masks: the indicator of valid token for embedding. [batch, len] float
                actions: the action embedding on the position. [batch, emb_dim] float
            rephrase_positions: the position to rephrase. [batch, 1]  long
            rewards: the rewards for the transition. [batch, 1] float
            terminal_signals: the terminal signals for the transition. [batch, 1] float
        :param target_critic: provides the global critic
        :param update_step: for lr scheduler and logging
        :param discount_factor: rollout-rewards discount
        :param summary_writer: logging
        :param update_trust_region: discount for loss updates
        """
        learn_batch_size = local_agent_configs["rephraser_learning_batch"]

        states, masks, \
        actions, rephrase_positions, \
        rewards, _, survive_and_no_maxs = replay_buffer.sample(learn_batch_size, device=self.device)
        INFO("update local agent critics on device: %s"%self.device)
        self.update_critic(
            states, masks, rephrase_positions, actions,
            rewards, survive_and_no_maxs,
            target_critic,
            update_step, discount_factor, summary_writer, update_trust_region)
        if update_step % local_agent_configs["actor_update_freq"] == 0:
            INFO("update local agent policy on device: %s"%self.device)
            self.update_actor_and_alpha(
                states, masks, rephrase_positions,
                target_critic,
                update_step, summary_writer, update_trust_region)

    def soft_update_target_net(self, target_SACAgent, tau):
        # soft update the target network. first move to CPU, than move back to local
        # mind not to update global model while reading and synch local models.
        self.to(target_SACAgent.device)
        with target_SACAgent.soft_update_lock:
            for param, target_param in zip(self.critic.parameters(), target_SACAgent.critic.parameters()):
                target_param.data.copy_(tau * param.data + (1-tau) * target_param.data)

            for param, target_param in zip(self.actor.parameters(), target_SACAgent.actor.parameters()):
                target_param.data.copy_(tau * param.data + (1-tau) * target_param.data)
        self.to(self.device)

def transition(states, masks, actions, rephrase_positions):
    """
    given s_t, a_t, t, returns torch.tensor s_{t+1}
    :param states: s_t torch.tensor [batch, len, dim], float
    :param masks: masks to indicate valid token for src, torch.tensor [batch, len], float
    :param actions: a_t, torch.tensor [batch, dim], float
    :param rephrase_positions: t, torch.tensor torch.tensor [batch, 1], long
    """
    batch, _, _ = states.shape
    new_states = copy.deepcopy(states)
    for i in range(batch):
        new_states[i][rephrase_positions[i], :] += actions[i] * masks[i][rephrase_positions[i]]
    return new_states.detach()

def slice_by_indices(states, rephrase_positions, device="cpu"):
    """
    extract the exact positions of states for label_emb, [index-1:index+1]
    return sliced states and move to device: torch.tensor [batch, 3, dim]
    """
    batch, sent_len, _ = states.shape
    label_emb = []
    for i in range(batch):
        stuff = states[i, rephrase_positions[i]-1:rephrase_positions[i]+2, :].unsqueeze(dim=0)
        label_emb.append(stuff)
    label_emb = torch.cat(label_emb, dim=0).to(device)
    return label_emb


import time
import os

import numpy as np
from numpy.linalg import inv, norm

import abcmesh
import meshrender
import meshsmooth
import cythonmagick
import av

F_MAX = float("+inf")
F_MIN = float("-inf")
from concurrent.futures import ThreadPoolExecutor,as_completed

def read_image_data16(path, width=None, height=None):
    container = av.open(path)
    for i, frame in enumerate(container.decode(video=0)):
        if not width:
            width = frame.width
        if not height:
            height = frame.height

        rgba = frame.reformat(width, height, format="rgba64le")
        return meshrender.MeshTexture(rgba.planes[0], width, height, depth=16)

def save_exr(texture, path):
    import OpenEXR
    import Imath
    start = time.time()
    hdr = OpenEXR.Header(texture.width, texture.height)

    pixels = {}
    channels = {}
    chan_format = Imath.Channel(Imath.PixelType(Imath.PixelType.HALF))
    # chan_format = Imath.Channel(Imath.PixelType(Imath.PixelType.FLOAT))
    c_start = time.time()
    for c, data in [("R", texture.r_f16), ("G",texture.g_f16), ("B", texture.b_f16), ("A", texture.a_f16)]:
        # pixels[c] = np.array(data).astype(np.float16)
        pixels[c] = data
        channels[c]= chan_format
    print "channel convert in %f secs" % (time.time() - c_start)

    hdr['channels'] = channels
    hdr['compression'] = Imath.Compression(Imath.Compression.ZIPS_COMPRESSION)
    exr = OpenEXR.OutputFile(path, hdr)
    exr.writePixels(pixels)
    exr.close()

    print "wrote exr in %f secs" % ( time.time() - start)

def save_image16(texture, path, resize=None):

    # with ThreadPoolExecutor(max_workers=16) as e:
    #     for i in xrange(1):
    #         start = time.time()
    #         # texture = texture.grow()
    #
    #         texture = mulithread_texture_grow(e, texture)
    #         print "texture grow in %f secs" % (time.time() - start)

    start = time.time()
    # texture = texture.resample(texture.width/2, texture.height/2)
    # print "resize in %f secs" % (time.time() - start)

    print texture.width, texture.height

    # save_exr(texture, path)
    # return
    #print texture.rgba
    start = time.time()
    i = cythonmagick.Image()
    i.verbose=True
    i.from_rawbuffer(bytearray(texture.rgba16), texture.width, texture.height, 'RGBA', 'short')
    print "imagemagick read data in %f secs" % (time.time()- start)
    # start = time.time()

    i.attributes['colorspace'] = "RGB"

    name, ext = os.path.splitext(path)
    if resize:
        i.filter = "Catrom"
        # start = time.time()

        #i.distort("resize", [4096,4096])
        # print i.size()
        # i.scale(resize)
        #i.resize(resize)
        # print "resize in %f secs" % (time.time() - start)

    if ext.lower() == ".tif":
        i.compress = "lzw"

    elif ext.lower() == '.exr':
        # start = time.time()
        # alpha = i.channel("alpha")
        # alpha.negate()
        # i.type = "TrueColor"
        # i.composite(alpha, "multiply")
        # i.composite(alpha, "copyopacity")
        # print "premult in %f secs" % (time.time() - start)
        i.compress = "zips"

    start = time.time()
    i.write(path)
    print "imagemagick save in %f secs" % (time.time()- start)

def autoclipping(scene_bbox, camera_transform):
    mat4 = np.matrix(camera_transform.T)
    scene_bbox = scene_bbox.dot(mat4)
    scene_bounds =  np.array(abcmesh.bbox_to_bounds(scene_bbox))
    znear = -1 * scene_bounds[1][2]
    zfar =  -1 * scene_bounds[0][2]

    znear = max(znear, 0.01)
    zfar = max(zfar, 1)
    return znear, zfar

def process_mesh(mesh, persp_matrix, camera_transform, viewport_matrix, seconds, level):
    start  = time.time()
    mesh_name = mesh.full_name
    mesh.time = seconds
    bbox = np.array(mesh.bbox)

    if level == 0:
        rmesh = meshrender.Mesh()
        rmesh.face_counts = mesh.face_counts
        rmesh.face_indices = mesh.face_indices
        rmesh.uvs = mesh.uv_values
        rmesh.uv_indices = np.array(mesh.uv_indices, dtype=np.int32)

        normal_indices = mesh.normal_indices
        rmesh.normal_indices = None
        if normal_indices:
            rmesh.normal_indices = np.array(mesh.normal_indices, dtype=np.int32)
            rmesh.normals = mesh.normal_values
        rmesh.vertices = mesh.vertices_vec4

    else:
        channels = []
        uv_indices = mesh.uv_indices
        if uv_indices:
            channels.append(meshsmooth.FVarChannel("uvs", np.array(uv_indices,  dtype=np.int32),
                                                  mesh.uv_values))

        normal_indices = mesh.normal_indices
        if normal_indices:
            channels.append(meshsmooth.FVarChannel("normals", np.array(mesh.normal_indices, dtype=np.int32),
                                                      mesh.normal_values))

        vert_channel = meshsmooth.FVarChannel("verts", mesh.face_indices, mesh.vertices_vec4)
        src_mesh =  meshsmooth.Mesh(mesh.face_counts, vert_channel, channels)

        # "LINEAR_NONE"
        # "LINEAR_CORNERS_ONLY"
        # "LINEAR_CORNERS_PLUS1"
        # "LINEAR_CORNERS_PLUS2"
        # "LINEAR_BOUNDARIES"
        # "LINEAR_ALL"

        boundry = "NONE"
        # boundry = "EDGE_ONLY"
        boundry = "EDGE_AND_CORNER"

        fvar_interp = "LINEAR_NONE"
        fvar_interp = "LINEAR_CORNERS_ONLY"
        # fvar_interp = "LINEAR_CORNERS_PLUS2"
        fvar_interp = "LINEAR_BOUNDARIES"
        # fvar_interp = "LINEAR_ALL"

        refiner = meshsmooth.TopologyRefiner(src_mesh, BoundaryInterpolation = boundry,
                                                       FVarInterpolation = fvar_interp)
        smooth_mesh = refiner.refine_uniform(level)

        rmesh = meshrender.Mesh()
        normals = None
        if normal_indices:
            normals = smooth_mesh.fvchannels[1]
        rmesh.from_meshsmooth(smooth_mesh, smooth_mesh.fvchannels[0], normals)

    model_transform = np.matrix(mesh.world_matrix, dtype=np.float32)

    mat4 = viewport_matrix * persp_matrix * camera_transform * model_transform
    mat4 = np.matrix(mat4.T,  dtype=np.float32)

    v  = np.array(rmesh.vertices)

    rmesh.objspace_vertices = np.empty(v.shape, dtype=np.float32)
    rmesh.objspace_vertices[:] = v

    rmesh.vertices = v.dot(mat4)
    meshrender.perpsective_divide(rmesh.vertices)

    rmesh.bbox = bbox.dot(mat4)
    meshrender.perpsective_divide(rmesh.bbox)

    return rmesh

def mulithread_render_mesh(executor, renderer, mesh, x_tiles=4, y_tiles=4, texture=None):

    x_tile_size = renderer.width / x_tiles
    y_tile_size = renderer.height / y_tiles

    miny = 0
    maxy = y_tile_size
    render_futures = []

    for y in range(y_tiles):
        minx = 0
        maxx = x_tile_size

        for x in range(x_tiles):
            rect = [(minx, miny), (maxx-1, maxy-1)]
            render_futures.append(executor.submit(renderer.render_mesh, mesh, rect=rect, texture=texture))
            # renderer.render_mesh(mesh, rect=rect, texture=texture)

            minx += x_tile_size
            maxx += x_tile_size

        miny += y_tile_size
        maxy += y_tile_size

    for r_future in as_completed(render_futures):
        r_future.result()

def mulithread_texture_grow(executor, texture, x_tiles=4, y_tiles=4):

    x_tile_size = texture.width / x_tiles
    y_tile_size = texture.height / y_tiles

    miny = 0
    maxy = y_tile_size
    render_futures = []

    dest_texture = meshrender.MeshTexture(None, texture.width, texture.height)

    for y in range(y_tiles):
        minx = 0
        maxx = x_tile_size

        for x in range(x_tiles):
            rect = [(minx, miny), (maxx-1, maxy-1)]

            # print texture.width, texture.height
            # print rect
            # texture.grow(16, rect=rect, dst_texture=dest_texture)
            # if x % 2 == 0:
            render_futures.append(executor.submit(texture.grow, 16, rect=rect, dst_texture=dest_texture))
            # render_futures.append(executor.submit(renderer.render_mesh, mesh, rect=rect, texture=texture))
            # renderer.render_mesh(mesh, rect=rect, texture=texture)

            minx += x_tile_size
            maxx += x_tile_size

            # return dest_texture

        miny += y_tile_size
        maxy += y_tile_size

    for r_future in as_completed(render_futures):
        r_future.result()

    return dest_texture

def project3d(abc_path, image_path, dest_path, frame=None, subdiv_level=0, size=1024):
    start = time.time()
    abc = abcmesh.open(abc_path)
    if not abc.camera_list:
        raise Exception("abc files has no cameras")

    camera = abc.camera_list[0]

    if frame is None:
        frame = 1

    seconds = frame / 24.0

    image_texture = read_image_data16(image_path)

    width = image_texture.width
    height = image_texture.height

    scale = 4
    renderer = meshrender.MeshRenderer(1024 * scale, 1024 * scale)
    renderer.bake_projection = True
    renderer.uvspace = True
    renderer.wireframe = False
    scene_bounds = np.array(((F_MAX,F_MAX,F_MAX),
                             (F_MIN,F_MIN,F_MIN)), dtype=np.float32)

    for mesh in abc.mesh_list:
        mesh.time = seconds

        model_transform = np.matrix(mesh.world_matrix)

        mat4 = np.matrix(model_transform.T)
        bbox = np.array(mesh.bbox)
        bbox = bbox.dot(mat4)

        mesh_bounds = abcmesh.bbox_to_bounds(bbox)
        # m_bounds = mesh.bounds
        for i in range(3):
            scene_bounds[0][i] = min(scene_bounds[0][i], mesh_bounds[0][i])
        for i in range(3):
            scene_bounds[1][i] = max(scene_bounds[1][i], mesh_bounds[1][i])

    scene_bbox = np.array(abcmesh.bounds_to_bbox(scene_bounds))

    camera.time = seconds
    camera_transform = inv(np.matrix(camera.world_matrix, dtype=np.float32))

    znear, zfar = autoclipping(scene_bbox, camera_transform)

    persp_matrix = np.matrix(camera.perspective_matrix(width/float(height),
                             znear=znear, zfar=zfar), dtype=np.float32)

    viewport_matrix = np.matrix(meshrender.viewport_matrix(width-1, height-1), dtype=np.float32)

    result = meshrender.MeshTexture(None, renderer.width, renderer.height)
    result.clear()
    sample_coords = []
    r = 4
    if r <= 1:
        sample_coords=[(0,0)]
    else:
        for y in range(r):
            for x in range(r):
                y_p = float(y + 1)/float(r) - 0.5
                x_p = float(x + 1)/float(r) - 0.5
                sample_coords.append((x_p, y_p))

    with ThreadPoolExecutor(max_workers=16) as e:
        mesh_list = []
        for mesh in abc.mesh_list[:1]:
            rmesh = process_mesh(mesh, persp_matrix,
                                 camera_transform, viewport_matrix, seconds, subdiv_level)
            mesh_list.append(rmesh)
        blend = 1.0 / float(len(sample_coords))

        for samp_x, samp_y in sample_coords:
            renderer.clear()
            renderer.sample_offset_x = samp_x
            renderer.sample_offset_y = samp_y
            for rmesh in mesh_list:
                # renderer.render_mesh(rmesh, texture=image_texture)
                mulithread_render_mesh(e, renderer, rmesh,  texture=image_texture)

            texture = mulithread_texture_grow(e, renderer)
            result.add(texture, blend)
        # result.blend(texture, )

    print "projection rendered in %f secs"% ( time.time() - start)

    save_image16(result, dest_path, None)

if __name__ == "__main__":

    from optparse import OptionParser

    default_size = 1024
    # default_size = 1024
    parser = OptionParser()
    parser.add_option("--abc", dest="abc", metavar="FILE")
    parser.add_option("--image", dest="projected_image", metavar="FILE")
    parser.add_option("--overlay", dest="texture_image", metavar="FILE")
    parser.add_option("--frame", dest="frame", type="int", default=1)
    parser.add_option("--level", dest="level", type="int", help="subdiv level", default=1)
    parser.add_option("--size", dest="size", type="int", help="dest image size", default=default_size)

    (options, args) = parser.parse_args()

    if len(args) < 1:
        parser.error("not enought args")

    start = time.time()
    project3d(options.abc, options.projected_image, args[0],
              frame = options.frame,
              subdiv_level = options.level,
              size = options.size
    )
    print "projection in %f secs" % (time.time() - start)


# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import cv2
import csv
import copy
import os, sys
import itertools
import numpy as np
from tqdm import tqdm

from House3D.objrender import Vec3

import pdb

class House3DUtils():
    def __init__(
            self,
            env,
            rotation_sensitivity=9,
            move_sensitivity=0.5,
            build_graph=False,
            graph_dir='/path/to/3d-graphs',
            target_obj_conn_map_dir='/path/to/target_obj_connmaps',
            debug=True,
            load_semantic_classes=True,
            collision_reward=0.0,
            success_reward=1.0,
            dist_reward_scale=0.005,
            seeing_rwd=False):
        self.env = env
        self.debug = debug

        self.rotation_sensitivity = rotation_sensitivity
        self.move_sensitivity = move_sensitivity

        self.angles = [x for x in range(-180, 180, self.rotation_sensitivity)]
        self.angle_strings = {1: 'right', -1: 'left'}

        self.dirs, self.angle_map = self.calibrate_steps(reset=True)
        self.move_multiplier = self.move_sensitivity / np.array([np.abs(x).sum() for x in self.dirs]).mean()

        self.graph_dir = graph_dir
        self.graph = None

        self.target_obj_conn_map_dir = target_obj_conn_map_dir

        if build_graph == True:
            if os.path.exists(
                    os.path.join(graph_dir,
                                 self.env.house.house['id'] + '.pkl')):
                self.load_graph(
                    os.path.join(graph_dir,
                                 self.env.house.house['id'] + '.pkl'))
            else:
                self.build_graph(
                    save_path=os.path.join(
                        graph_dir, self.env.house.house['id'] + '.pkl'))

        self.rooms, self.objects = self._parse()

        self.collision_reward = collision_reward
        self.success_reward = success_reward
        self.dist_reward_scale = dist_reward_scale
        self.seeing_rwd = seeing_rwd

        if load_semantic_classes == True:
            self._load_semantic_classes()

    # Shortest paths are computed in 1000 x 1000 grid coordinates.
    # One step in the SUNCG continuous coordinate system however, can be
    # multiple grids in the grid coordinate system (since turns aren't 90 deg).
    # So even though the grid shortest path is fine-grained,
    # an equivalent best-fit path in SUNCG continuous coordinates
    # has to be computed by simulating steps. Sucks, but yeah.
    #
    # For now, we first explicitly calibrate how many steps in the gridworld
    # correspond to one step in continuous world, across all directions
    def calibrate_steps(self, reset=True):
        mults, angle_map = [], {}

        cx, cy = self.env.house.to_coor(50, 50)
        if reset == True:
            self.env.reset(x=cx, y=cy)

        for i in range(len(self.angles)):
            yaw = self.angles[i]

            self.env.cam.yaw = yaw
            self.env.cam.updateDirection()

            x1, y1 = self.env.house.to_grid(self.env.cam.pos.x,
                                            self.env.cam.pos.z)

            pos = self.env.cam.pos
            pos = pos + self.env.cam.front * self.move_sensitivity

            x2, y2 = self.env.house.to_grid(pos.x, pos.z)

            mult = np.array([x2, y2]) - np.array([x1, y1])
            mult = (mult[0], mult[1])

            angle_map[mult] = yaw
            mults.append(mult)

        return mults, angle_map

    # 0: forward
    # 1: left
    # 2: right
    # 3: stop
    #
    # returns observation, reward, done, info
    def step(self, action, step_reward=False):
        if action not in [0, 1, 2, 3]:
            raise IndexError

        if step_reward == True:
            pos = self.env.cam.pos
            x1, y1 = self.env.house.to_grid(self.env.cam.pos.x, self.env.cam.pos.z)
            init_target_dist = self.env.house.connMap[x1, y1]

        reward = 0
        done = False

        if action == 0:
            mv = self.env.move_forward(
                dist_fwd=self.move_sensitivity, dist_hor=0)
            obs = self.env.render()
            if mv == False:  # collision
                reward -= self.collision_reward
            elif mv != False and step_reward == True:
                # evaluate connMap dist here
                x2, y2 = self.env.house.to_grid(self.env.cam.pos.x,
                                                self.env.cam.pos.z)
                final_target_dist = self.env.house.connMap[x2, y2]
                reward += self.dist_reward_scale * ((init_target_dist - final_target_dist) / np.abs(
                    self.dirs[self.angles.index(self.env.cam.yaw % 180)]).sum())

        elif action == 1:
            self.env.rotate(-self.rotation_sensitivity)
            obs = self.env.render()

        elif action == 2:
            self.env.rotate(self.rotation_sensitivity)
            obs = self.env.render()

        elif action == 3:
            done = True
            obs = self.env.render()

        return obs, reward, done

    # pos: [x, y, z, yaw], or objrender.Vec3
    def get_dist_to_target(self, pos):
        if isinstance(pos, Vec3) == True:
            x, y = self.env.house.to_grid(pos.x, pos.z)
        else:
            x, y = self.env.house.to_grid(pos[0], pos[2])
        dist = self.env.house.connMap[x, y]
        return self.move_multiplier * dist

    def is_inside_room(self, pos, room):
        if isinstance(pos, Vec3) == True:
            x = pos.x
            y = pos.z
        else:
            x = pos[0]
            y = pos[2]
        if x >= room['bbox']['min'][0] and x <= room['bbox']['max'][0] and \
            y >= room['bbox']['min'][2] and y <= room['bbox']['max'][2]:
            return True
        return False

    # takes 200-300 seconds(!) when rotation_sensitivity == 9
    def build_graph(self, save_path=None):
        import time
        start_time = time.time()

        collide_res = self.env.house.n_row

        from dijkstar import Graph

        visit = dict()
        self.graph = Graph()

        self.mock_obs_map = np.zeros(
            (collide_res + 1, collide_res + 1), dtype=np.uint8)
        self.mock_obs_map[np.where(self.env.house.connMap == -1)] = 1

        for x in range(collide_res + 1):
            for y in range(collide_res + 1):
                pos = (x, y)
                if self.env.house.canMove(x, y) and pos not in visit:
                    que = [pos]
                    visit[pos] = True
                    ptr = 0
                    while ptr < len(que):
                        cx, cy = que[ptr]
                        ptr += 1

                        # add all angles for (cx, cy) here
                        # connect first and last
                        for ang in range(len(self.angles) - 1):
                            self.graph.add_edge((cx, cy, self.angles[ang]),
                                                (cx, cy, self.angles[ang + 1]),
                                                {
                                                    'cost': 1
                                                })
                            self.graph.add_edge((cx, cy, self.angles[ang + 1]),
                                                (cx, cy, self.angles[ang]), {
                                                    'cost': 1
                                                })
                        self.graph.add_edge((cx, cy, self.angles[-1]),
                                            (cx, cy, self.angles[0]), {
                                                'cost': 1
                                            })
                        self.graph.add_edge((cx, cy, self.angles[0]),
                                            (cx, cy, self.angles[-1]), {
                                                'cost': 1
                                            })

                        for deti in range(len(self.dirs)):
                            det = self.dirs[deti]
                            tx, ty = cx + det[0], cy + det[1]
                            if (self.env.house.inside(tx, ty) and
                                    self.mock_obs_map[min(cx, tx):max(cx, tx)+1,
                                                      min(cy, ty):max(cy, ty)+1].sum() == 0):
                                # make changes here to add edges for angle increments as well
                                #
                                # cost = 1 from one angle to the next,
                                # and connect first and last
                                # this would be for different angles for same tx, ty
                                #
                                # then there would be connections for same angle
                                # and from (cx, cy) to (tx, ty)
                                self.graph.add_edge(
                                    (cx, cy, self.angle_map[self.dirs[deti]]),
                                    (tx, ty, self.angle_map[self.dirs[deti]]),
                                    {
                                        'cost': 1
                                    })
                                tp = (tx, ty)
                                if tp not in visit:
                                    visit[tp] = True
                                    que.append(tp)

        if self.debug == True:
            print("--- %s seconds to build the graph ---" %
                  (time.time() - start_time))

        if save_path != None:
            start_time = time.time()

            print("saving graph to %s" % (save_path))
            self.graph.dump(save_path)

            if self.debug == True:
                print("--- %s seconds to save the graph ---" %
                      (time.time() - start_time))

    def load_graph(self, path):
        import time
        start_time = time.time()

        from dijkstar import Graph

        self.graph = Graph()
        self.graph.load(path)

        if self.debug == True:
            print("--- %s seconds to load the graph ---" %
                  (time.time() - start_time))

    # takes 1-5 seconds when rotation_sensitivity == 9
    def compute_shortest_path(self, source, target, graph=None):
        from dijkstar import find_path

        if graph == None:
            if self.graph == None:
                if os.path.exists(
                        os.path.join(self.graph_dir,
                                     self.env.house.house['id'] + '.pkl')):
                    self.load_graph(
                        os.path.join(self.graph_dir,
                                     self.env.house.house['id'] + '.pkl'))
                else:
                    self.build_graph(
                        save_path=os.path.join(
                            graph_dir, self.env.house.house['id'] + '.pkl'))
            graph = self.graph

        cost_func = lambda u, v, e, prev_e: e['cost']
        shortest_path = find_path(graph, source, target, cost_func=cost_func)

        return shortest_path

    def fit_grid_path_to_suncg(self, nodes, init_yaw=None, back_skip=2):

        # don't mess with the originals
        nodes = copy.deepcopy(nodes)

        # set initial position
        x, y = self.env.house.to_coor(nodes[0][0], nodes[0][1], True)
        x, y = x.astype(np.float32).item(), y.astype(np.float32).item()

        self.env.cam.pos.x, self.env.cam.pos.y, self.env.cam.pos.z = x, self.env.house.robotHei, y
        if init_yaw == None:
            self.env.cam.yaw = np.random.choice(self.angles)
        else:
            self.env.cam.yaw = init_yaw
        self.env.cam.updateDirection()

        pos_queue, action_queue = [], []

        current_pos = self._vec_to_array(self.env.cam.pos, self.env.cam.yaw)
        pos_queue = pos_queue + [current_pos]

        ptr = 0

        while ptr < len(nodes) - 1:
            turned = False

            # target rotation
            target_yaw = self.angle_map[tuple(
                np.array(nodes[ptr]) - np.array(nodes[ptr + 1]))]

            # turn
            if target_yaw != current_pos[3]:
                p_q, a_q = self.get_rotate_steps(current_pos, target_yaw)

                pos_queue = pos_queue + p_q
                action_queue = action_queue + a_q

                self.env.cam.yaw = target_yaw
                self.env.cam.updateDirection()

                turned = True
                current_pos = self._vec_to_array(self.env.cam.pos,
                                                 self.env.cam.yaw)

            # move
            cx, cz = self.env.house.to_coor(nodes[ptr + 1][0],
                                            nodes[ptr + 1][1], True)

            # if collision, find another sub-path, and delete that edge
            if self.env.move(cx, cz) == False:
                if nodes[ptr + 1] in self.graph[nodes[ptr]]:
                    del self.graph[nodes[ptr]][nodes[ptr + 1]]
                    print('deleted', nodes[ptr], nodes[ptr + 1])

                # delete the turns
                if turned == True:
                    pos_queue = pos_queue[:-len(p_q)]
                    action_queue = action_queue[:-len(a_q)]

                if back_skip != 0:
                    pos_queue = pos_queue[:-back_skip]
                    action_queue = action_queue[:-back_skip]

                dest_ptr = ptr + 1
                ptr = ptr - back_skip

                sub_shortest_path = self.compute_shortest_path(
                    nodes[ptr], nodes[dest_ptr])
                nodes = nodes[:ptr] + sub_shortest_path.nodes + nodes[dest_ptr
                                                                      + 1:]

                current_pos = pos_queue[-1]
            else:
                # this is the new position the agent moved to
                current_pos = self._vec_to_array(self.env.cam.pos,
                                                 self.env.cam.yaw)

                assert current_pos[3] == pos_queue[-1][3] and (
                    current_pos[0] != pos_queue[-1][0]
                    or current_pos[2] != pos_queue[-1][2])

                pos_queue = pos_queue + [current_pos]
                action_queue = action_queue + ['fwd']

                ptr = ptr + 1

        action_queue.append('stop')

        return pos_queue, action_queue

    # pos contains [x, y, z, yaw]
    # given a position and target yaw, this function
    # computes actions needed to turn there
    def get_rotate_steps(self, pos, target_yaw):

        direction = np.random.choice([1, -1])

        cur_yaw = pos[-1]
        ptr = self.angles.index(cur_yaw)
        pos_queue, action_queue = [], []

        while cur_yaw != target_yaw:
            if len(pos_queue) == len(self.angles) // 2:
                # reset
                direction = direction * -1
                cur_yaw = pos[-1]
                ptr = self.angles.index(cur_yaw)
                pos_queue, action_queue = [], []

            ptr = (ptr + direction) % len(self.angles)
            cur_yaw = self.angles[ptr]

            pos_queue.append([pos[0], pos[1], pos[2], self.angles[ptr]])
            action_queue.append(self.angle_strings[direction])

        return pos_queue, action_queue

    def _vec_to_array(self, pos, yaw):
        return [pos.x, pos.y, pos.z, yaw]

    # render images from camera position queue
    def render_images_from_pos_queue(self,
                                     pos_queue=[],
                                     img_dir='tmp/images',
                                     actions=None,
                                     values=None,
                                     rewards=None):
        if len(pos_queue) == 0:
            return False

        action_map = {0: 'FRWD', 1: 'LEFT', 2: 'RGHT', 3: 'STOP'}

        import scipy.misc

        sgx, sgy = self.env.house.to_grid(pos_queue[0][0], pos_queue[0][2])
        tgx, tgy = self.env.house.to_grid(pos_queue[-1][0], pos_queue[-1][2])

        for i in range(len(pos_queue)):
            # set position
            p = pos_queue[i]
            self.env.reset(x=p[0], y=p[2], yaw=p[3])

            # save image
            image = np.array(self.env.render(), copy=False)

            # put some text
            text = "[%02d]" % (i + 1)

            if actions != None and i < len(actions):
                text += "[%s]" % action_map[actions[i]]

            if values != None and i < len(values):
                text += "[V%.03f]" % values[i]

            if rewards != None and i > 0 and i <= len(rewards):
                text += "[R%.03f]" % rewards[i - 1]

            image = cv2.putText(
                img=np.copy(image),
                text=text,
                org=(20, 30),
                fontFace=3,
                fontScale=0.4,
                color=(255, 255, 255),
                thickness=1)

            scipy.misc.toimage(image).save(
                '%s/%s_%04d_%04d_%04d_%04d_%05d_%05d.jpg' %
                (img_dir, self.env.house.house['id'], sgx, sgy, tgx, tgy,
                 i + 1, len(pos_queue)))

        return True

    # render video from camera position queue
    #
    # NOTE: call `render_images_from_pos_queue` before calling this
    def render_video_from_pos_queue(self,
                                    pos_queue=[],
                                    img_dir='tmp/images',
                                    vid_dir='tmp/videos',
                                    fps=[5],
                                    tag_name='piano'):
        if len(pos_queue) == 0:
            return False

        import subprocess

        sgx, sgy = self.env.house.to_grid(pos_queue[0][0], pos_queue[0][2])
        tgx, tgy = self.env.house.to_grid(pos_queue[-1][0], pos_queue[-1][2])

        for fp in fps:
            subprocess.Popen([
                '/srv/share/abhshkdz/local/bin/ffmpeg', '-f', 'image2', '-r',
                str(fp), '-i',
                '%s/%s_%04d_%04d_%04d_%04d' %
                (img_dir, self.env.house.house['id'], sgx, sgy, tgx, tgy) +
                '_%05d_' + '%05d.jpg' % (len(pos_queue)), '-vcodec', 'libx264',
                '-crf', '25', '-y',
                '%s/%s_%04d_%04d_%s_%04d_%04d_%d.mp4' %
                (vid_dir, self.env.house.house['id'], sgx, sgy, tag_name, tgx,
                 tgy, fp)
            ])

            if self.debug == True:
                print('Rendered video to ' +
                      '%s/%s_%04d_%04d_%s_%04d_%04d_%d.mp4' %
                      (vid_dir, self.env.house.house['id'], sgx, sgy, tag_name,
                       tgx, tgy, fp))

        return True

    # Go over all nodes of house environment and accumulate objects room-wise.
    def _parse(self, levelsToExplore=[0]):
        rooms, objects = [], {}
        data = self.env.house.house

        modelCategoryMapping = {}

        import csv
        csvFile = csv.reader(open(self.env.house.metaDataFile, 'r'))
        headers = next(csvFile)

        for row in csvFile:
            modelCategoryMapping[row[headers.index('model_id')]] = {
                headers[x]: row[x]
                for x in range(2, len(headers))  # 0 is index, 1 is model_id
            }

        for i in levelsToExplore:
            for j in range(len(data['levels'][i]['nodes'])):
                assert data['levels'][i]['nodes'][j]['type'] != 'Box'

                if 'valid' in data['levels'][i]['nodes'][j]:
                    assert data['levels'][i]['nodes'][j]['valid'] == 1

                # Rooms
                if data['levels'][i]['nodes'][j]['type'] == 'Room':
                    if 'roomTypes' not in data['levels'][i]['nodes'][j]:
                        continue

                    # Can rooms have more than one type?
                    # Yes, they can; just found ['Living_Room', 'Dining_Room', 'Kitchen']
                    # assert len(data['levels'][i]['nodes'][j]['roomTypes']) <= 3

                    roomType = [
                        # ' '.join(x.lower().split('_'))
                        x.lower()
                        for x in data['levels'][i]['nodes'][j]['roomTypes']
                    ]

                    nodes = data['levels'][i]['nodes'][j][
                        'nodeIndices'] if 'nodeIndices' in data['levels'][i][
                            'nodes'][j] else []
                    rooms.append({
                        'type':
                        roomType,
                        'bbox':
                        data['levels'][i]['nodes'][j]['bbox'],
                        'nodes':
                        nodes,
                        'model_id':
                        data['levels'][i]['nodes'][j]['modelId']
                    })

                # Objects
                elif data['levels'][i]['nodes'][j]['type'] == 'Object':
                    if 'materials' not in data['levels'][i]['nodes'][j]:
                        material = []
                    else:
                        material = data['levels'][i]['nodes'][j]['materials']
                    objects[data['levels'][i]['nodes'][j]['id']] = {
                        'id':
                        data['levels'][i]['nodes'][j]['id'],
                        'model_id':
                        data['levels'][i]['nodes'][j]['modelId'],
                        'fine_class':
                        modelCategoryMapping[data['levels'][i]['nodes'][j][
                            'modelId']]['fine_grained_class'],
                        'coarse_class':
                        modelCategoryMapping[data['levels'][i]['nodes'][j][
                            'modelId']]['coarse_grained_class'],
                        'bbox':
                        data['levels'][i]['nodes'][j]['bbox'],
                        'mat':
                        material
                    }

        return rooms, objects

    # Spawn at a randomly selected point in a particular room
    def spawn_room(self, room=None):
        if room == None:
            return False, None

        target_room = '_'.join(room.lower().split(' '))

        if self.env.house.hasRoomType(target_room) == False:
            return False, None

        rooms = self.env.house._getRooms(target_room)
        room = np.random.choice(rooms)

        gx1, gy1, gx2, gy2 = self.env.house._getRoomBounds(room)

        available_coords = []
        for x in range(gx1, gx2 + 1):
            for y in range(gy1, gy2 + 1):
                if self.env.house.moveMap[x, y] > 0:
                    available_coords.append((x, y))

        # print(available_coords)
        spawn_coord_idx = np.random.choice(len(available_coords))
        spawn_coord = available_coords[spawn_coord_idx]

        return spawn_coord, room

    # Spawn close to an object
    # If room given, look for object within room
    def spawn_object(self, obj=None, room=None):
        if object == None:
            return False, None

        if isinstance(obj, list) == False:
            obj = [obj]

        is_door = False
        if 'door' in obj:
            is_door = True

        target_obj = ['_'.join(x.lower().split(' ')) for x in obj]

        if room != None:
            if 'nodeIndices' in room:
                objs = [
                    self.objects['0_' + str(x)] for x in room['nodeIndices']
                    if self.objects['0_' + str(x)]['fine_class'] in target_obj
                ]
            else:
                objs = [
                    self.objects['0_' + str(x)] for x in room['nodes']
                    if self.objects['0_' + str(x)]['fine_class'] in target_obj
                ]
        else:
            obj_id_list = list(
                itertools.chain.from_iterable(
                    [x['nodes'] for x in self.rooms if x['type'] != []]))
            objs = [
                self.objects['0_' + str(x)] for x in obj_id_list
                if self.objects['0_' + str(x)]['fine_class'] in target_obj
            ]

        if len(objs) == 0:
            return False, None, None

        obj_idx = np.random.choice(len(objs))
        obj = objs[obj_idx]

        self.target_obj_class = obj['fine_class'].lower()

        gx1, gy1, gx2, gy2 = self.env.house._getRoomBounds(obj)

        if room == None:
            obj_node_idx = int(obj['id'][2:])
            room = [
                x for x in self.env.house.all_rooms
                if 'nodeIndices' in x and obj_node_idx in x['nodeIndices']
            ][0]

        self.set_target_object(obj, room)

        available_x, available_y = np.where(self.env.house.connMap == 0)

        if len(available_x) == 0:
            return False, None, None

        spawn_coords = []
        for i in range(len(available_x)):
            spawn_coords.append((available_x[i], available_y[i]))

        return spawn_coords, obj, room

    # analogous to `setTargetRoom` in the House3D API
    def set_target_object(self, obj, room):
        object_tp = room['id'] + '_' + obj['id'] + '_' + obj['fine_class'].lower(
        )
        # Caching
        if object_tp in self.env.house.connMapDict:
            self.env.house.connMap, self.env.house.connectedCoors, self.env.house.inroomDist, self.env.house.maxConnDist = self.env.house.connMapDict[
                object_tp]
            return True  # object changed!
        elif os.path.exists(
                os.path.join(
                    self.target_obj_conn_map_dir,
                    self.env.house.house['id'] + '_' + object_tp + '.npy')):
            self.env.house.connMap = np.load(
                os.path.join(
                    self.target_obj_conn_map_dir,
                    self.env.house.house['id'] + '_' + object_tp + '.npy'))

            if self.env.house.connMap.shape[0] == self.env.house.n_row+1:
                self.env.house.connectedCoors, self.env.house.inroomDist, self.env.house.maxConnDist = None, None, None
                return True

        self.env.house.connMap = connMap = np.ones(
            (self.env.house.n_row + 1, self.env.house.n_row + 1),
            dtype=np.int32) * -1
        self.env.house.inroomDist = inroomDist = np.ones(
            (self.env.house.n_row + 1, self.env.house.n_row + 1),
            dtype=np.float32) * -1
        dirs = [[0, 1], [1, 0], [-1, 0], [0, -1]]
        que = []
        flag_find_open_components = True

        _ox1, _, _oy1 = obj['bbox']['min']
        _ox2, _, _oy2 = obj['bbox']['max']
        ocx, ocy = (_ox1 + _ox2) / 2, (_oy1 + _oy2) / 2
        ox1, oy1, ox2, oy2 = self.env.house.rescale(_ox1, _oy1, _ox2, _oy2)

        for _ in range(2):
            _x1, _, _y1 = room['bbox']['min']
            _x2, _, _y2 = room['bbox']['max']
            cx, cy = (_x1 + _x2) / 2, (_y1 + _y2) / 2
            x1, y1, x2, y2 = self.env.house.rescale(_x1, _y1, _x2, _y2)

            curr_components = self.env.house._find_components(
                x1,
                y1,
                x2,
                y2,
                dirs=dirs,
                return_open=flag_find_open_components
            )  # find all the open components
            if len(curr_components) == 0:
                print('No space found! =(')
                print('Found Target') #Himi changes #Target too close
                return True
                #raise ValueError('no space')
            if isinstance(curr_components[0],
                          list):  # join all the coors in the open components
                curr_major_coors = list(itertools.chain(*curr_components))
            else:
                curr_major_coors = curr_components
            min_dist_to_center, min_dist_to_edge = 1e50, 1e50
            for x, y in curr_major_coors:
                ###
                # Compute minimum dist to edge here
                if x in range(ox1, ox2):
                    dx = 0
                elif x < ox1:
                    dx = ox1 - x
                else:
                    dx = x - ox2

                if y in range(oy1, oy2):
                    dy = 0
                elif y < oy1:
                    dy = oy1 - y
                else:
                    dy = y - oy2

                assert dx >= 0 and dy >= 0

                if dx != 0 or dy != 0:
                    dd = np.sqrt(dx**2 + dy**2)
                elif dx == 0:
                    dd = dy
                else:
                    dd = dx

                if dd < min_dist_to_edge:
                    min_dist_to_edge = int(np.ceil(dd))
                ###
                tx, ty = self.env.house.to_coor(x, y)
                tdist = np.sqrt((tx - ocx)**2 + (ty - ocy)**2)
                if tdist < min_dist_to_center:
                    min_dist_to_center = tdist
                inroomDist[x, y] = tdist
            margin = min_dist_to_edge + 1
            for x, y in curr_major_coors:
                inroomDist[x, y] -= min_dist_to_center
            for x, y in curr_major_coors:
                if x in range(ox1 - margin, ox2 + margin) and y in range(
                        oy1 - margin, oy2 + margin):
                    connMap[x, y] = 0
                    que.append((x, y))
            if len(que) > 0: break
            if flag_find_open_components:
                flag_find_open_components = False
            else:
                break
            raise ValueError

        ptr = 0
        self.env.house.maxConnDist = 1
        while ptr < len(que):
            x, y = que[ptr]
            cur_dist = connMap[x, y]
            ptr += 1
            for dx, dy in dirs:
                tx, ty = x + dx, y + dy
                if self.env.house.inside(tx, ty) and self.env.house.canMove(
                        tx, ty) and not self.env.house.isConnect(tx, ty):
                    que.append((tx, ty))
                    connMap[tx, ty] = cur_dist + 1
                    if cur_dist + 1 > self.env.house.maxConnDist:
                        self.env.house.maxConnDist = cur_dist + 1
        self.env.house.connMapDict[object_tp] = (connMap, que, inroomDist,
                                                 self.env.house.maxConnDist)
        np.save(
            os.path.join(
                self.target_obj_conn_map_dir,
                self.env.house.house['id'] + '_' + object_tp + '.npy'),
            connMap)
        self.connectedCoors = que
        #print(' >>>> ConnMap Cached!')
        return True  # room changed!

    def _load_semantic_classes(self, color_file=None):
        if color_file == None:
            color_file = self.env.config['colorFile']

        self.semantic_classes = {}

        with open(color_file) as csv_file:
            reader = csv.DictReader(csv_file)
            for row in reader:
                c = np.array((row['r'], row['g'], row['b']), dtype=np.uint8)
                fine_cat = row['name'].lower()
                self.semantic_classes[fine_cat] = c

        return self.semantic_classes

    def _get_best_yaw_obj_from_pos(self, obj_id, grid_pos, height=1.0):
        obj = self.objects[obj_id]
        obj_fine_class = obj['fine_class']

        cx, cy = self.env.house.to_coor(grid_pos[0], grid_pos[1])

        self.env.cam.pos.x = cx
        self.env.cam.pos.y = height
        self.env.cam.pos.z = cy

        best_yaw, best_coverage = None, 0

        for yaw in self.angles:
            self.env.cam.yaw = yaw
            self.env.cam.updateDirection()

            seg = self.env.render(mode='semantic')
            c = self.semantic_classes[obj_fine_class.lower()]
            mask = np.all(seg == c, axis=2)
            coverage = np.sum(mask) / (seg.shape[0] * seg.shape[1])

            if best_yaw == None:
                best_yaw = yaw
                best_coverage = coverage
            else:
                if coverage > best_coverage:
                    best_yaw = yaw
                    best_coverage = coverage

        return best_yaw, best_coverage

    def _get_best_view_obj(self,
                           obj,
                           coverage_thres=0.5,
                           dist_add=0.5,
                           robot_height=False):
        bbox = obj['bbox']
        obj_fine_class = obj['fine_class']

        obj_max = np.asarray(bbox['max'])
        obj_min = np.asarray(bbox['min'])
        obj_center = (obj_min + obj_max) / 2

        c_x, c_y, c_z = obj_center
        max_radius = np.sqrt(
            (obj_max[0] - obj_min[0]) * (obj_max[0] - obj_min[0]) +
            (obj_max[2] - obj_min[2]) * (obj_max[2] - obj_min[2])) / 2.0
        max_radius += dist_add

        best_pos = None
        best_coverage = 0

        returned_pos_cov = []

        for yaw in self.angles:
            pos = [
                c_x - max_radius * np.cos(yaw * (2 * np.pi) / 360.0), c_y,
                c_z - max_radius * np.sin(yaw * (2 * np.pi) / 360.0), yaw
            ]

            if robot_height == True:
                pos[1] = min(max(0.75, c_y), 2.00)

            self.env.cam.pos.x = pos[0]
            self.env.cam.pos.y = pos[1]
            self.env.cam.pos.z = pos[2]
            self.env.cam.yaw = pos[3]

            self.env.cam.updateDirection()

            seg = self.env.render(mode='semantic')
            c = self.semantic_classes[obj_fine_class.lower()]
            mask = np.all(seg == c, axis=2)
            coverage = np.sum(mask) / (seg.shape[0] * seg.shape[1])

            returned_pos_cov.append([pos, coverage])

            if coverage > coverage_thres:
                return pos, coverage, returned_pos_cov
            elif coverage > best_coverage:
                best_coverage = coverage
                best_pos = pos

        return best_pos, best_coverage, returned_pos_cov


"""
Tests optimizing a unit in AlexNet/CaffeNet with DeePSiM generators 1-8
Running time (approx.; steps == 200; n_units == 1):
    5 mins (GTX 2080); 10 mins (GTX 1060)
"""

from time import time
from pathlib import Path
import h5py as h5
import numpy as np
from Experiments import CNNExperiment

save_root = Path('temp')
engine = 'pytorch'    # or caffe; will switch engine in generator and target
optimizer_name = 'genetic'
generator_names = (
    'deepsim-norm1', 'deepsim-norm2', 'deepsim-conv3', 'deepsim-conv4',
    'deepsim-pool5', 'deepsim-fc6', 'deepsim-fc7', 'deepsim-fc8')
steps = 200
target_unit = {    # goldfish neuron
    'caffe': ('caffenet', 'fc8', 1),
    'pytorch': ('alexnet', 'classifier.6', 1)}[engine]
init_rand_seed = 0
exp_settings = {
    'optimizer_name': optimizer_name,
    'optimizer_parameters': {'generator_parameters': {'engine': engine}},
    'scorer_parameters': {'engine': engine},
    'image_size': 85,
    'with_write': False,
    'max_optimize_steps': steps,
    'random_seed': 0,
    'stochastic': False,
    'config_file_path': __file__}
optimizer_setting_names = (
    'population_size', 'mutation_rate', 'mutation_size', 'selectivity',
    'heritability', 'n_conserve')
by_generator_optimizer_settings = {
    'deepsim-norm1': (15, 1, 1.5, 2, 0.5, 0),
    'deepsim-norm2': (10, 0.5, 0.7, 4, 0.5, 0),
    'deepsim-conv3': (12, 0.65, 0.75, 2.25, 0.5, 0),
    'deepsim-conv4': (10, 0.9, 0.75, 2.5, 0.5, 0),
    'deepsim-pool5': (10, 0.6, 1, 2.5, 0.75, 0),
    'deepsim-fc6': (20, 0.5, 0.5, 2, 0.5, 0),
    'deepsim-fc7': (45, 0.6, 0.3, 1.25, 0.5, 0),
    'deepsim-fc8': (20, 0.2, 0.6, 2, 0.5, 0)}
t0 = time()

# main loop
unit = np.atleast_1d(target_unit[2:])
target_s = '_'.join(map(str, unit))
init_randgen = np.random.RandomState(init_rand_seed)
for generator_name in generator_names:
    # set generator-specific parameters
    optimizer_parameters = exp_settings['optimizer_parameters']
    optimizer_parameters['generator_name'] = generator_name
    for p, v in zip(
            optimizer_setting_names,
            by_generator_optimizer_settings[generator_name]):
        optimizer_parameters[p] = v
    exp_settings['optimizer_parameters'] = optimizer_parameters

    # specify initialization (for reproducibility)
    gen_layer_name = generator_name.split('-')[1]
    with h5.File('test_data/stats_caffenet.h5', 'r') as f:
        mu = f[gen_layer_name]['mu'][()]
        sig = f[gen_layer_name]['sig'][()]
    pop_size = exp_settings['optimizer_parameters']['population_size']
    init_codes = init_randgen.normal(mu, sig, size=(pop_size, *mu.shape))
    if generator_name != 'deepsim-fc8':
        # generator was trained on post-ReLU input
        init_codes = np.clip(init_codes, 0, None)

    project_dir = (
        save_root / engine / optimizer_name / generator_name
        / 'random_dist'    # init name
        / target_unit[0] / target_unit[1] / target_s)
    try:
        project_dir.mkdir(parents=True)
    except FileExistsError:
        # skips finished, running, or aborted experiments
        # this allows running test1.py in parallel with (almost) no
        # interference (unless timing is so precise---for example in batch
        # jobs---that two processes attempted to create the same folder
        # at the same time. Downside is, for aborted experiments,
        # project_dir needs to be manually removed
        continue
    exp_settings['project_dir'] = project_dir
    exp_settings['target_neuron'] = target_unit
    experiment = CNNExperiment(**exp_settings)
    experiment.optimizer.set_initialization(init_codes)    # manual init
    experiment.run()

# load and print results
print('target unit:\t', target_unit)
for generator_name in generator_names:
    project_dir = (
            save_root / engine / optimizer_name / generator_name
            / 'random_dist'  # init name
            / target_unit[0] / target_unit[1] / target_s)
    acts = []
    for i in (0, steps-1):
        score_f = project_dir / f'scores_step{i:03d}.npz'
        try:
            acts.append(np.load(score_f)['scores'].max())
        except FileNotFoundError:
            acts.append(np.nan)
    print(f'{generator_name}\t{acts[0]:>4.1f} -> {acts[1]:>5.1f}')
print(f'took time: {time() - t0:.0f} s')

"""
Note: for unclear reasons, torch results are slightly different
  on different runs; caffe results are consistent.
results:
>>> cuda 10.1.243, python 3.7.5, Ubuntu 19.10
    >>> torch 1.3.1
        deepsim-norm1	-1.6 ->   8.5
        deepsim-norm2	 0.8 ->   6.7
        deepsim-conv3	 1.6 ->  10.1
        deepsim-conv4	 1.8 ->  12.8
        deepsim-pool5	 2.0 ->  18.2
        deepsim-fc6	 5.4 ->  31.6
        deepsim-fc7	 5.8 ->  58.5
        deepsim-fc8	 6.2 ->  38.7
    >>> caffe
        deepsim-norm1	 0.0 ->  16.4
        deepsim-norm2	 3.5 ->  10.8
        deepsim-conv3	 4.8 ->  15.3
        deepsim-conv4	 5.4 ->  21.8
        deepsim-pool5	 3.5 ->  27.3
        deepsim-fc6	 6.3 ->  51.1
        deepsim-fc7	 6.9 ->  74.1
        deepsim-fc8	 9.1 ->  52.7
>>> cuda 9.0, python 3.7.1 (anaconda), Windows 10
    >>> torch 1.0.1
        deepsim-norm1	-1.6 ->   8.5
        deepsim-norm2	 0.8 ->   6.3
        deepsim-conv3	 1.6 ->   9.5
        deepsim-conv4	 1.8 ->  13.7
        deepsim-pool5	 2.0 ->  19.2
        deepsim-fc6	 5.4 ->  26.6
        deepsim-fc7	 5.8 ->  45.7
        deepsim-fc8	 6.2 ->  35.9
    >>> caffe
        deepsim-norm1	 0.0 ->  15.2
        deepsim-norm2	 3.5 ->  11.4
        deepsim-conv3	 4.8 ->  15.7
        deepsim-conv4	 5.4 ->  21.1
        deepsim-pool5	 3.5 ->  27.6
        deepsim-fc6	 6.3 ->  49.8
        deepsim-fc7	 6.9 ->  64.3
        deepsim-fc8	 9.1 ->  49.1
"""


import os
import numpy as np
import pydoc
from omegaconf import OmegaConf
import yaml
import munch

import torch
from torch import nn

import utils

from utils.common import dict2device, tti, itt, to_sigm, to_tanh
from utils.common import get_rotation_matrix, segment_img, load_module
from utils import uv_renderer as uvr
from torch.utils.data import DataLoader
from models.efficientnet import EfficientNetLevelEncoder

def get_config(config_path, divide_n_channels=1):
    with open(config_path, 'r') as f:
        config_ = yaml.load(f, Loader=yaml.FullLoader)

    config = dict()
    for k, v in config_.items():
        if type(v) == dict and 'value' in v:
            config[k] = v['value']
        else:
            config[k] = v

    if 'fix_renderer' not in config:
        config['fix_renderer'] = False
    if 'renderer_checkpoint' not in config:
        config['renderer_checkpoint'] = ""
    if 'ainp_path' not in config:
        config['ainp_path'] = None
        config['ainp_scales'] = []

    config = munch.munchify(config)

    # config hacks
    config.segm_channels = 1
    config.ainp_path='data/spectral_texture16.pth'
    config.checkpoint512_path = None

    config.bidis_channel_multiplier = int(config.bidis_channel_multiplier / divide_n_channels)
    config.udis_channel_multiplier = int(config.udis_channel_multiplier / divide_n_channels)

    if 'alternating' not in config:
        config.alternating = False

    return config

class Inferer(nn.Module):
    def __init__(self, generator, renderer, generator_config, config):
        super().__init__()

        self.config = config
        self.generator_config = generator_config
        self.device = config.device

        # self.renderer = renderer.eval().to(self.device)
        self.renderer = renderer.eval().to(self.device)
        self.generator = generator.eval().to(self.device)

        self.uv_renderer = uvr.UVRenderer(self.generator_config.image_size, self.generator_config.image_size)
        self.uv_renderer.to(self.device)

        self.batch_size = 1
        self.n_rotsamples = 64
        # self.eval()

        # load encoder
        encoder_config_path = os.path.join(config.encoder.experiment_dir, "encoder_config.yaml")
        with open(encoder_config_path) as f:
            encoder_config = OmegaConf.load(f)
        self.encoder = EfficientNetLevelEncoder(
            **utils.common.return_empty_dict_if_none(encoder_config.model.encoder.args)
        )
        state_dict = torch.load(os.path.join(config.encoder.experiment_dir, "checkpoints", config.encoder.checkpoint_name))
        self.encoder.load_state_dict(state_dict['inferer']['encoder'])
        self.encoder.eval()
        for p in self.encoder.parameters():
            p.requires_grad = False

        # trainable parameters
        ## latent
        latent_mean, _ = self.calc_latent_stats(
            n_mean_latent=config.inferer.n_mean_latent
        )

        latent_init = latent_mean \
                     .detach() \
                     .clone() \
                     .view(1, 1, self.generator.style_dim) \
                     .repeat(1, self.generator.n_latent, 1)

        self.latent = nn.Parameter(latent_init, requires_grad=True)

        ## noise
        noise_init = self.generator.make_noise()
        self.noises = nn.ParameterList([nn.Parameter(x, requires_grad=True) for x in noise_init])

        ## ntexture
        self.ntexture = None

        # diff renderer
        self.diff_uv_renderer = utils.uv_renderer.NVDiffRastUVRenderer(config.inferer.faces_path, config.inferer.uv_vert_values_path)

    def sample_ntexture(self, ntexture, uv):
        return torch.nn.functional.grid_sample(ntexture, uv.permute(0, 2, 3, 1), align_corners=True)

    def infer_pass(self, noise, verts, ntexture=None, uv=None, sampled_ntexture=None):

        if verts is not None:
            B = verts.shape[0]
        else:
            B = uv.shape[0]

        if uv is None and sampled_ntexture is None:
            uv = self.uv_renderer(verts)

        if type(noise) == torch.Tensor and noise.ndim == 3:
            noise = torch.split(noise, 1, dim=1)
            noise = [x[0] for x in noise]

        if ntexture is None and sampled_ntexture is None:
            fake_ntexture = self.generator(noise, noise=self.noises, input_is_latent=True)
            fake_ntexture = torch.cat([fake_ntexture]*B, dim=0)
        else:
            fake_ntexture = ntexture

        if sampled_ntexture is None:
            sampled_ntexture = self.sample_ntexture(fake_ntexture, uv)

        renderer_dict_out = self.renderer(dict(uv=uv, sampled_ntextures=sampled_ntexture, nrender=sampled_ntexture))
        fake_img = renderer_dict_out['fake_rgb']
        fake_segm = renderer_dict_out['fake_segm']
        fake_img = segment_img(fake_img, fake_segm)

        out = dict(
            fake_img=fake_img,
            fake_segm=fake_segm,
            fake_ntexture=fake_ntexture,
            uv=uv,
            sampled_ntexture=sampled_ntexture
        )

        return out

    def calc_latent_stats(self, n_mean_latent=10000):
        with torch.no_grad():
            noise_sample = torch.randn(n_mean_latent, self.generator.style_dim).to(self.device)
            latent_out = self.generator.style(noise_sample)
            latent_mean = latent_out.mean(0)
            latent_std = ((latent_out - latent_mean).pow(2).sum() / n_mean_latent) ** 0.5

        return latent_mean, latent_std

    def make_noise_multiple(self, batch_size, latent_dim, n_noises):
        noises = torch.randn(n_noises, batch_size, latent_dim, device=self.device).unbind(0)
        return noises

    def get_state_dict(self):
        # collect generator parameters
        generator_params = []
        generator_params.extend(self.generator.conv1.parameters())
        for l in self.generator.convs:
            generator_params.extend(l.parameters())
        generator_params.extend(self.generator.to_rgb1.parameters())
        for l in self.generator.to_rgbs:
            generator_params.extend(l.parameters())

        state_dict = {
            'latent': self.latent,
            'noise': self.noise,
            'generator_params': generator_params,
            'ntexture': self.ntexture
        }

        return state_dict


import _init_paths
from tools.train_net import train
import os
import numpy as np
import caffe
from datafactory.imdb import IMDB
from datafactory.load import load_data_with_boxes

from tools.train_net_with_boxes import get_training_roidb, train_net
import fast_rcnn.config as fconfig
from fast_rcnn.config import cfg_from_file
from configuration.config import def_cfg

import sys
import pprint

gpu_id = 0
solver = '/home/zyb/VirtualDisk500/exhdd/Recognition-master/models/UniqueDET/solver_debug.prototxt'
max_iters = 100000
size = 224
imdb_name = 'UniqueDET'
out = 'out'
cfg = '/home/zyb/VirtualDisk500/exhdd/Recognition-master/experiments/cfgs/faster_rcnn_end2end.yml'
pretrained_model = '/home/zyb/VirtualDisk500/exhdd/Recognition-master/pretrained_models/VGG_CNN_M_1024.v2.caffemodel'

if __name__ == '__main__':

    def_cfg('UDET')

    cfg_from_file(cfg)
    pprint.pprint(fconfig.cfg)

    caffe.set_mode_gpu()
    caffe.set_device(gpu_id)
    # setup the dataset's path
    dataset = os.path.join('..', 'data', imdb_name)
    # load pixel mean
    pixel_means = None
    if os.path.exists(os.path.join(dataset, 'mean.npy')):
        pixel_means = np.load(os.path.join(dataset, 'mean.npy'))
        fconfig.cfg.PIXEL_MEANS = pixel_means
        print 'Loaded mean.npy: {}'.format(pixel_means)
    else:
        print 'Cannot find mean.npy and we will use default mean.'

    imdb = IMDB()
    imdb.get_roidb(load_data_with_boxes, dataset=dataset)
    roidb = get_training_roidb(imdb)

    np.random.seed(fconfig.cfg.RNG_SEED)
    caffe.set_random_seed(fconfig.cfg.RNG_SEED)

    train_net(solver, roidb, out,
              pretrained_model=pretrained_model, max_iters=max_iters)