from __future__ import division
from __future__ import print_function

import sys
import os
sys.path.append( '%s/gcn' % os.path.dirname(os.path.realpath(__file__)) )
# add the libary path for graph reduction and local search
# sys.path.append( '%s/kernel' % os.path.dirname(os.path.realpath(__file__)) )

import time
import scipy.io as sio
import numpy as np
import scipy.sparse as sp
import queue
from copy import deepcopy

# import the libary for graph reduction and local search
# from reduce_lib import reducelib

import tensorflow as tf
from utils import *
from models import GCN_DEEP_DIVER

N_bd = 32

# Settings
flags = tf.app.flags
FLAGS = flags.FLAGS
flags.DEFINE_string('model', 'gcn_cheby', 'Model string.')  # 'gcn', 'gcn_cheby', 'dense'
flags.DEFINE_float('learning_rate', 0.001, 'Initial learning rate.')
flags.DEFINE_integer('epochs', 201, 'Number of epochs to train.')
flags.DEFINE_integer('hidden1', 32, 'Number of units in hidden layer 1.')
flags.DEFINE_integer('diver_num', 32, 'Number of outputs.')
flags.DEFINE_float('dropout', 0, 'Dropout rate (1 - keep probaNUmbility).')
flags.DEFINE_float('weight_decay', 5e-4, 'Weight for L2 loss on embedding matrix.')
flags.DEFINE_integer('early_stopping', 1000, 'Tolerance for early stopping (# of epochs).')
flags.DEFINE_integer('max_degree', 1, 'Maximum Chebyshev polynomial degree.')
flags.DEFINE_integer('num_layer', 20, 'number of layers.')

# test data path
data_path = "./data"
val_mat_names = os.listdir(data_path)

# Some preprocessing

num_supports = 1 + FLAGS.max_degree
model_func = GCN_DEEP_DIVER

# Define placeholders
placeholders = {
    'support': [tf.sparse_placeholder(tf.float32) for _ in range(num_supports)],
    'features': tf.sparse_placeholder(tf.float32, shape=(None, N_bd)), # featureless: #points
    'labels': tf.placeholder(tf.float32, shape=(None, 2)), # 0: not linked, 1:linked
    'labels_mask': tf.placeholder(tf.int32),
    'dropout': tf.placeholder_with_default(0., shape=()),
    'num_features_nonzero': tf.placeholder(tf.int32)  # helper variable for sparse dropout
}

# Create model
model = model_func(placeholders, input_dim=N_bd, logging=True)

# use gpu 0
os.environ['CUDA_VISIBLE_DEVICES']=str(0)

# Initialize session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)

# Define model evaluation function
def evaluate(features, support, placeholders):
    t_test = time.time()
    feed_dict_val = construct_feed_dict4pred(features, support, placeholders)
    outs_val = sess.run([model.outputs_softmax], feed_dict=feed_dict_val)
    return (time.time() - t_test), outs_val[0]

def findNodeEdges(adj):
    nn = adj.shape[0]
    edges = []
    for i in range(nn):
        edges.append(adj.indices[adj.indptr[i]:adj.indptr[i+1]])
    return edges

def isis_v2(edges, nIS_vec_local, cn):
    return np.sum(nIS_vec_local[edges[cn]] == 1) > 0

def isis(edges, nIS_vec_local):
    tmp = (nIS_vec_local==1)
    return np.sum(tmp[edges[0]]*tmp[edges[1]]) > 0

def add_rnd_q(cns, nIS_vec_local):
    global adj_0

    nIS_vec_local[cns] = 1
    tmp = sp.find(adj_0[cns, :] == 1)
    nIS_vec_local[tmp[1]] = 0
    remain_vec_tmp = (nIS_vec_local == -1)
    adj = adj_0
    adj = adj[remain_vec_tmp, :]
    adj = adj[:, remain_vec_tmp]
    if reduce_graph(adj, nIS_vec_local):
        return True
    return False

def fake_reduce_graph(adj):
    reduced_node = -np.ones(adj.shape[0])
    reduced_adj = adj
    mapping = np.arange(adj.shape[0])
    reverse_mapping = np.arange(adj.shape[0])
    crt_is_size = 0
    return reduced_node, reduced_adj, mapping, reverse_mapping, crt_is_size

def fake_local_search(adj, nIS_vec):
    return nIS_vec.astype(int)

def reduce_graph(adj, nIS_vec_local):
    global best_IS_num
    global best_IS_vec
    global bsf_q
    global adj_0
    global q_ct
    global id
    global out_id
    global res_ct

    remain_vec = (nIS_vec_local == -1)

    # reduce graph
    # reduced_node, reduced_adj, mapping, reverse_mapping, crt_is_size = api.reduce_graph(adj)
    reduced_node, reduced_adj, mapping, reverse_mapping, crt_is_size = fake_reduce_graph(adj)
    nIS_vec_sub = reduced_node.copy()
    nIS_vec_sub_tmp = reduced_node.copy()
    nIS_vec_sub[nIS_vec_sub_tmp == 0] = 1
    nIS_vec_sub[nIS_vec_sub_tmp == 1] = 0
    reduced_nn = reduced_adj.shape[0]

    # update MIS after reduction
    tmp = sp.find(adj[nIS_vec_sub == 1, :] == 1)
    nIS_vec_sub[tmp[1]] = 0
    nIS_vec_local[remain_vec] = nIS_vec_sub
    nIS_vec_local[nIS_vec_local == 2] = -1

    # if the whole graph is reduced, we find a candidate
    if reduced_nn == 0:
        remain_vec_tmp = (nIS_vec_local == -1)
        if np.sum(remain_vec_tmp) == 0:
            # get a solution
            res_ct += 1
            # nIS_vec_local = api.local_search(adj_0, nIS_vec_local)
            nIS_vec_local = fake_local_search(adj_0, nIS_vec_local)
            if np.sum(nIS_vec_local) > best_IS_num:
                best_IS_num = np.sum(nIS_vec_local)
                best_IS_vec = deepcopy(nIS_vec_local)
                sio.savemat('./res_%04d/%s' % (
                    time_limit, val_mat_names[id]), {'er_graph': adj_0, 'nIS_vec': best_IS_vec})
            print("ID: %03d" % id, "QItem: %03d" % q_ct, "Res#: %03d" % res_ct,
                  "Current: %d" % (np.sum(nIS_vec_local)), "Best: %d" % best_IS_num, "Reduction")
            return True
        adj = adj_0
        adj = adj[remain_vec_tmp, :]
        adj = adj[:, remain_vec_tmp]
        bsf_q.append([adj, nIS_vec_local.copy(), remain_vec.copy(), reduced_adj, reverse_mapping.copy()])
    else:
        bsf_q.append([adj, nIS_vec_local.copy(), remain_vec.copy(), reduced_adj, reverse_mapping.copy()])

    return False

# Init variables
saver=tf.train.Saver(max_to_keep=1000)
sess.run(tf.global_variables_initializer())

ckpt=tf.train.get_checkpoint_state("./model")
print('loaded '+ckpt.model_checkpoint_path)
saver.restore(sess,ckpt.model_checkpoint_path)

noout = FLAGS.diver_num # number of outputs
time_limit = 600  # time limit for searching

if not os.path.isdir("./res_%04d"%time_limit):
    os.makedirs("./res_%04d"%time_limit)

# for graph reduction and local search
# api = reducelib()

for id in range(len(val_mat_names)):
    best_IS_num = -1
    mat_contents = sio.loadmat(data_path + '/' + val_mat_names[id])
    adj_0 = mat_contents['adj']
    # yy = mat_contents['indset_label']
    # opt_num = np.sum(yy[:,0])
    # edges_0 = sp.find(adj_0) # for isis version 1
    edges_0 = findNodeEdges(adj_0)
    nn = adj_0.shape[0]
    bsf_q = []
    q_ct = 0
    res_ct = 0
    out_id = -1

    start_time = time.time()
    while time.time()-start_time < time_limit:

        # if best_IS_num == opt_num:
        #     break

        if len(bsf_q) == 0:
            if reduce_graph(adj_0, -np.ones(nn)):
                break

        q_item = bsf_q.pop(np.random.randint(0,len(bsf_q)))
        q_ct += 1

        adj = q_item[0]
        remain_vec = deepcopy(q_item[2])
        reduced_adj = q_item[3]
        reverse_mapping = deepcopy(q_item[4])
        remain_nn = adj.shape[0]
        reduced_nn = reduced_adj.shape[0]

        if reduced_nn != 0:
            # GCN
            features = np.ones([reduced_nn, N_bd])
            features = sp.lil_matrix(features)
            features = preprocess_features(features)
            support = simple_polynomials(reduced_adj, FLAGS.max_degree)

            _, z_out = evaluate(features, support, placeholders)

            for out_id in range(noout):
                # if best_IS_num == opt_num:
                #     break

                nIS_vec = deepcopy(q_item[1])
                nIS_Prob_sub_t = z_out[:, 2 * out_id + 1]
                nIS_Prob_sub = np.zeros(remain_nn)
                nIS_Prob_sub[reverse_mapping] = nIS_Prob_sub_t
                nIS_Prob = np.zeros(nn)
                nIS_Prob[remain_vec] = nIS_Prob_sub

                # chosen nodes
                cns_sorted = np.argsort(1 - nIS_Prob)

                # tt = time.time()
                nIS_vec_tmp = deepcopy(nIS_vec)
                for cid in range(nn):
                    cn = cns_sorted[cid]
                    # check graph
                    if isis_v2(edges_0, nIS_vec_tmp, cn):
                        break
                    nIS_vec_tmp[cn] = 1
                    if np.random.random_sample() > 0.7:
                        add_rnd_q(cns_sorted[:(cid+1)], deepcopy(nIS_vec))

                # print("time=", "{:.5f}".format((time.time() - tt)))

                cns = cns_sorted[:cid]
                nIS_vec[cns] = 1
                tmp = sp.find(adj_0[cns, :] == 1)
                nIS_vec[tmp[1]] = 0
                remain_vec_tmp = (nIS_vec == -1)
                if np.sum(remain_vec_tmp) == 0:
                    # get a solution
                    res_ct += 1
                    # nIS_vec = api.local_search(adj_0, nIS_vec)
                    nIS_vec = fake_local_search(adj_0, nIS_vec)
                    if np.sum(nIS_vec) > best_IS_num:
                        best_IS_num = np.sum(nIS_vec)
                        best_IS_vec = deepcopy(nIS_vec)
                        sio.savemat('./res_%04d/%s' % (
                        time_limit, val_mat_names[id]), {'er_graph': adj_0, 'nIS_vec': best_IS_vec})
                    print("ID: %03d" % id, "QItem: %03d" % q_ct, "Res#: %03d" % res_ct,
                          "Current: %d" % (np.sum(nIS_vec)), "Best: %d" % best_IS_num, "Network")
                    continue
                adj = adj_0
                adj = adj[remain_vec_tmp, :]
                adj = adj[:, remain_vec_tmp]

                if reduce_graph(adj, nIS_vec):
                    continue
        else:
            nIS_vec = deepcopy(q_item[1])
            if reduce_graph(adj, nIS_vec):
                continue

    sio.savemat('./res_%04d/%s' % (time_limit, val_mat_names[id]), {'er_graph': adj_0, 'nIS_vec': best_IS_vec})


from abc import ABC, abstractmethod

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import backend as K

from ..base.errors import MissingModelError
from ..base.mixins import RandomStateMixin, NumActionsMixin, LoggerMixin
from ..policies.base import BasePolicy
from ..caching import NStepCache
from ..losses import SoftmaxPolicyLossWithLogits, ClippedSurrogateLoss
from ..utils import (
    project_onto_actions_np, softmax, argmax, check_numpy_array)

__all__ = (
    'BaseV',
    'BaseQTypeI',
    'BaseQTypeII',
    'BaseSoftmaxPolicy',
)

class BaseFunctionApproximator(ABC, LoggerMixin):
    @abstractmethod
    def __call__(self, *args, **kwargs):
        pass

    @abstractmethod
    def update(self, *args, **kwargs):
        pass

    @abstractmethod
    def batch_eval(self, *args, **kwargs):
        pass

    @abstractmethod
    def batch_update(self, *args, **kwargs):
        pass

    def _check_attrs(self):
        required_attrs = [
            'env', 'gamma', 'bootstrap_n', 'train_model', 'predict_model',
            'target_model', '_cache']

        if isinstance(self, BaseSoftmaxPolicy):
            required_attrs.remove('bootstrap_n')
            required_attrs.remove('gamma')
            required_attrs.remove('_cache')

        missing_attrs = ", ".join(
            attr for attr in required_attrs if not hasattr(self, attr))

        if missing_attrs:
            raise AttributeError(
                "missing attributes: {}".format(missing_attrs))

    def _train_on_batch(self, inputs, outputs):
        """
        Run self.train_model.train_on_batch(inputs, outputs) and return the
        losses as a dict of type: {loss_name <str>: loss_value <float>}.

        """
        losses = self.train_model.train_on_batch(inputs, outputs)

        # add metric names
        if len(self.train_model.metrics_names) > 1:
            assert len(self.train_model.metrics_names) == len(losses)
            losses = dict(zip(self.train_model.metrics_names, losses))
        else:
            assert isinstance(losses, (float, np.float32, np.float64))
            assert len(self.train_model.metrics_names) == 1
            losses = {self.train_model.metrics_names[0]: losses}

        if hasattr(self.env, 'record_losses'):
            self.env.record_losses(losses)

        return losses

    def sync_target_model(self, tau=1.0):
        """
        Synchronize the target model with the primary model.

        Parameters
        ----------
        tau : float between 0 and 1, optional

            The amount of exponential smoothing to apply in the target update:

            .. math::

                w_\\text{target}\\ \\leftarrow\\ (1 - \\tau)\\,w_\\text{target}
                + \\tau\\,w_\\text{primary}

        """
        if tf.__version__ >= '2.0':
            target_weights = self.target_model.trainable_variables
            primary_weights = self.predict_model.trainable_variables
            tf.group(*(
                K.update(wt, wt + tau * (wp - wt))
                for wt, wp in zip(target_weights, primary_weights)))
            self.logger.debug(
                "updated target_mode with tau = {:.3g}".format(tau))

        else:
            if not hasattr(self, '_target_model_sync_op'):
                target_weights = tf.get_collection(
                    tf.GraphKeys.GLOBAL_VARIABLES, scope='target')  # list
                primary_weights = tf.get_collection(
                    tf.GraphKeys.GLOBAL_VARIABLES, scope='primary')  # list

                if not target_weights:
                    raise MissingModelError(
                        "no model weights found in variable scope: 'target'")
                if not primary_weights:
                    raise MissingModelError(
                        "no model weights found in variable scope: 'primary'")
                assert len(primary_weights) == len(target_weights)

                self._target_model_sync_tau = tf.placeholder(
                    tf.float32, shape=())
                self._target_model_sync_op = tf.group(*(
                    K.update(wt, wt + self._target_model_sync_tau * (wp - wt))
                    for wt, wp in zip(target_weights, primary_weights)))

            K.get_session().run(
                self._target_model_sync_op,
                feed_dict={self._target_model_sync_tau: tau})
            self.logger.debug(
                "updated target_mode with tau = {:.3g}".format(tau))

class BaseV(BaseFunctionApproximator):
    """
    Base class for modeling a :term:`state value function`.

    A :term:`state value function` is implemented by mapping :math:`s\\mapsto
    v(s)`.

    Parameters
    ----------
    env : gym environment

        A gym environment.

    train_model : keras.Model(:term:`S`, :term:`V`)

        Used for training.

    predict_model : keras.Model(:term:`S`, :term:`V`)

        Used for predicting. For a :term:`state value function` the
        :term:`target_model` and :term:`predict_model` are the same.

    target_model : keras.Model(:term:`S`, :term:`V`)

        A :term:`target_model` is used to make predictions on a bootstrapping
        scenario. It can be advantageous to use a point-in-time copy of the
        :term:`predict_model` to construct a bootstrapped target.

    gamma : float, optional

        The discount factor for discounting future rewards.

    bootstrap_n : positive int, optional

        The number of steps in n-step bootstrapping. It specifies the number of
        steps over which we're willing to delay bootstrapping. Large :math:`n`
        corresponds to Monte Carlo updates and :math:`n=1` corresponds to
        TD(0).

    bootstrap_with_target_model : bool, optional

        Whether to use the :term:`target_model` when constructing a
        bootstrapped target. If False (default), the primary
        :term:`predict_model` is used.

    """
    def __init__(
            self, env, train_model, predict_model, target_model,
            gamma=0.9,
            bootstrap_n=1,
            bootstrap_with_target_model=False):

        self.env = env
        self.train_model = train_model
        self.predict_model = predict_model
        self.target_model = target_model
        self.gamma = float(gamma)
        self.bootstrap_n = int(bootstrap_n)
        self.bootstrap_with_target_model = bool(bootstrap_with_target_model)

        self._cache = NStepCache(self.env, self.bootstrap_n, self.gamma)

    def __call__(self, s, use_target_model=False):
        """
        Evaluate the Q-function.

        Parameters
        ----------
        s : state observation

            A single state observation.

        use_target_model : bool, optional

            Whether to use the :term:`target_model` internally. If False
            (default), the :term:`predict_model` is used.

        Returns
        -------
        V : float or array of floats

            The estimated value of the state :math:`v(s)`.

        """
        assert self.env.observation_space.contains(s)
        S = np.expand_dims(s, axis=0)
        V = self.batch_eval(S, use_target_model=use_target_model)
        check_numpy_array(V, shape=(1,))
        V = np.asscalar(V)
        return V

    def update(self, s, r, done):
        """
        Update the Q-function.

        Parameters
        ----------
        s : state observation

            A single state observation..

        r : float

            A single observed reward.

        done : bool

            Whether the episode has finished.

        """
        assert self.env.observation_space.contains(s)
        self._cache.add(s, 0, r, done)

        # eager updates
        while self._cache:
            S, _, Rn, In, S_next, _ = self._cache.pop()
            self.batch_update(S, Rn, In, S_next)

    def batch_update(self, S, Rn, In, S_next):
        """
        Update the value function on a batch of transitions.

        Parameters
        ----------
        S : nd array, shape: [batch_size, ...]

            A batch of state observations.

        Rn : 1d array, dtype: float, shape: [batch_size]

            A batch of partial returns. For example, in n-step bootstrapping
            this is given by:

            .. math::

                R^{(n)}_t\\ =\\ R_t + \\gamma\\,R_{t+1} + \\dots
                    \\gamma^{n-1}\\,R_{t+n-1}

            In other words, it's the non-bootstrapped part of the n-step
            return.

        In : 1d array, dtype: float, shape: [batch_size]

            A batch bootstrapping factor. For instance, in n-step bootstrapping
            this is given by :math:`I^{(n)}_t=\\gamma^n` if the episode is
            ongoing and :math:`I^{(n)}_t=0` otherwise. This allows us to write
            the bootstrapped target as:

            .. math::

                G^{(n)}_t=R^{(n)}_t+I^{(n)}_tQ(S_{t+n}, A_{t+n})

        S_next : nd array, shape: [batch_size, ...]

            A batch of next-state observations.

        Returns
        -------
        losses : dict

            A dict of losses/metrics, of type ``{name <str>: value <float>}``.

        """
        V_next = self.batch_eval(
            S_next, use_target_model=self.bootstrap_with_target_model)
        Gn = Rn + In * V_next
        losses = self._train_on_batch(S, Gn)
        return losses

    def batch_eval(self, S, use_target_model=False):
        """
        Evaluate the state value function on a batch of state observations.

        Parameters
        ----------
        S : nd array, shape: [batch_size, ...]

            A batch of state observations.

        use_target_model : bool, optional

            Whether to use the :term:`target_model` internally. If False
            (default), the :term:`predict_model` is used.

        Returns
        -------
        V : 1d array, dtype: float, shape: [batch_size]

            The predicted state values.

        """
        model = self.target_model if use_target_model else self.predict_model

        V = model.predict_on_batch(S)
        check_numpy_array(V, ndim=2, axis_size=1, axis=1)
        V = np.squeeze(V, axis=1)  # shape: [batch_size]
        return V

class BaseGenericQ(BaseFunctionApproximator, NumActionsMixin):
    UPDATE_STRATEGIES = ('sarsa', 'q_learning', 'double_q_learning')

    def __init__(
            self, env, train_model, predict_model, target_model,
            gamma=0.9,
            bootstrap_n=1,
            bootstrap_with_target_model=False,
            update_strategy='sarsa'):

        self.env = env
        self.train_model = train_model
        self.predict_model = predict_model
        self.target_model = target_model
        self.gamma = float(gamma)
        self.bootstrap_n = int(bootstrap_n)
        self.bootstrap_with_target_model = bool(bootstrap_with_target_model)
        self.update_strategy = update_strategy

        self._cache = NStepCache(self.env, self.bootstrap_n, self.gamma)

    def __call__(self, s, a=None, use_target_model=False):
        """
        Evaluate the Q-function.

        Parameters
        ----------
        s : state observation

            A single state observation.

        a : action, optional

            A single action.

        use_target_model : bool, optional

            Whether to use the :term:`target_model` internally. If False
            (default), the :term:`predict_model` is used.

        Returns
        -------
        Q : float or array of floats

            If action ``a`` is provided, a single float representing
            :math:`q(s,a)` is returned. If, on the other hand, ``a`` is left
            unspecified, a vector representing :math:`q(s,.)` is returned
            instead. The shape of the latter return value is ``[num_actions]``,
            which is only well-defined for discrete action spaces.

        """
        assert self.env.observation_space.contains(s)
        S = np.expand_dims(s, axis=0)
        if a is not None:
            assert self.env.action_space.contains(a)
            A = np.expand_dims(a, axis=0)
            Q = self.batch_eval(S, A, use_target_model=use_target_model)
            check_numpy_array(Q, shape=(1,))
            Q = np.asscalar(Q)
        else:
            Q = self.batch_eval(S, use_target_model=use_target_model)
            check_numpy_array(Q, shape=(1, self.num_actions))
            Q = np.squeeze(Q, axis=0)
        return Q

    def update(self, s, pi, r, done):
        """
        Update the Q-function.

        Parameters
        ----------
        s : state observation

            A single state observation.

        pi : int or 1d array, shape: [num_actions]

            Vector of action propensities under the behavior policy. This may
            be just an indicator if the action propensities are inferred
            through sampling. For instance, let's say our action space is
            :class:`Discrete(4)`, then passing ``pi = 2`` is equivalent to
            passing ``pi = [0, 0, 1, 0]``. Both would indicate that the action
            :math:`a=2` was drawn from the behavior policy.

        r : float

            A single observed reward.

        done : bool

            Whether the episode has finished.

        """
        assert self.env.observation_space.contains(s)
        pi = self.check_pi(pi)
        self._cache.add(s, pi, r, done)

        # eager updates
        while self._cache:
            self.batch_update(*self._cache.pop())  # pop with batch_size=1

    def batch_update(self, S, P, Rn, In, S_next, P_next=None):
        """
        Update the value function on a batch of transitions.

        Parameters
        ----------
        S : nd array, shape: [batch_size, ...]

            A batch of state observations.

        P : 2d Tensor, dtype: int, shape: [batch_size]

            A batch of action propensities. :term:`P` is typically just an
            indicator for which action was chosen by the behavior policy. In
            this sense, :term:`P` acts as a projector more than a prediction
            target. That is, :term:`P` is used to project our predicted values
            down to those for which we actually received the feedback signal.

        Rn : 1d array, dtype: float, shape: [batch_size]

            A batch of partial returns. For example, in n-step bootstrapping
            this is given by:

            .. math::

                R^{(n)}_t\\ =\\ R_t + \\gamma\\,R_{t+1} + \\dots
                    \\gamma^{n-1}\\,R_{t+n-1}

            In other words, it's the non-bootstrapped part of the n-step
            return.

        In : 1d array, dtype: float, shape: [batch_size]

            A batch bootstrapping factor. For instance, in n-step bootstrapping
            this is given by :math:`I^{(n)}_t=\\gamma^n` if the episode is
            ongoing and :math:`I^{(n)}_t=0` otherwise. This allows us to write
            the bootstrapped target as:

            .. math::

                G^{(n)}_t=R^{(n)}_t+I^{(n)}_tQ(S_{t+n}, A_{t+n})

        S_next : nd array, shape: [batch_size, ...]

            A batch of next-state observations.

        P_next : 2d Tensor, dtype: int, shape: [batch_size, num_actions]

            Action propensities :term:`P_next` for the (potential) next action.
            This argument is only used if ``update_strategy='sarsa'``.

        Returns
        -------
        losses : dict

            A dict of losses/metrics, of type ``{name <str>: value <float>}``.

        """
        G = self.bootstrap_target(Rn, In, S_next, P_next)
        losses = self._train_on_batch([S, G], P)
        return losses

    def bootstrap_target(self, Rn, In, S_next, P_next=None):
        """
        Get the bootstrapped target
        :math:`G^{(n)}_t=R^{(n)}_t+\\gamma^nQ(S_{t+n}, A_{t+n})`.

        Parameters
        ----------
        Rn : 1d array, dtype: float, shape: [batch_size]

            A batch of partial returns. For example, in n-step bootstrapping
            this is given by:

            .. math::

                R^{(n)}_t\\ =\\ R_t + \\gamma\\,R_{t+1} + \\dots
                    \\gamma^{n-1}\\,R_{t+n-1}

            In other words, it's the non-bootstrapped part of the n-step
            return.

        In : 1d array, dtype: float, shape: [batch_size]

            A batch bootstrapping factor. For instance, in n-step bootstrapping
            this is given by :math:`I^{(n)}_t=\\gamma^n` if the episode is
            ongoing and :math:`I^{(n)}_t=0` otherwise. This allows us to write
            the bootstrapped target as:

            .. math::

                G^{(n)}_t=R^{(n)}_t+I^{(n)}_tQ(S_{t+n},A_{t+n})

        S_next : nd array, shape: [batch_size, ...]

            A batch of next-state observations.

        P_next : 2d Tensor, dtype: int, shape: [batch_size, num_actions]

            Action propensities :term:`P_next` for the (potential) next action.
            This argument is only used if ``update_strategy='sarsa'``.

        Returns
        -------
        Gn : 1d array, dtype: int, shape: [batch_size]

            A batch of bootstrap-estimated returns
            :math:`G^{(n)}_t=R^{(n)}_t+I^{(n)}_tQ(S_{t+n},A_{t+n})` computed
            according to given ``update_strategy``.

        """
        if self.update_strategy == 'sarsa':
            assert P_next is not None
            Q_next = self.batch_eval(
                S_next, use_target_model=self.bootstrap_with_target_model)
            Q_next = np.einsum('ij,ij->i', Q_next, P_next)
        elif self.update_strategy == 'q_learning':
            Q_next = np.max(
                self.batch_eval(
                    S_next, use_target_model=self.bootstrap_with_target_model),
                axis=1)
        elif self.update_strategy == 'double_q_learning':
            if not self.bootstrap_with_target_model:
                raise ValueError(
                    "incompatible settings: "
                    "update_strategy='double_q_learning' requires that "
                    "bootstrap_with_target_model=True")
            A_next = np.argmax(
                self.batch_eval(S_next, use_target_model=False), axis=1)
            Q_next = self.batch_eval(S_next, use_target_model=True)
            Q_next = project_onto_actions_np(Q_next, A_next)
        else:
            raise ValueError("unknown update_strategy")

        Gn = Rn + In * Q_next
        return Gn

    @abstractmethod
    def batch_eval(self, S, A=None, use_target_model=False):
        """
        Evaluate the Q-function on a batch of state (or state-action)
        observations.

        Parameters
        ----------
        S : nd array, shape: [batch_size, ...]

            A batch of state observations.

        A : 1d array, dtype: int, shape: [batch_size], optional

            A batch of actions that were taken.

        use_target_model : bool, optional

            Whether to use the :term:`target_model` internally. If False
            (default), the :term:`predict_model` is used.

        Returns
        -------
        Q : 1d or 2d array of floats

            If action ``A`` is provided, a 1d array representing a batch of
            :math:`q(s,a)` is returned. If, on the other hand, ``A`` is left
            unspecified, a vector representing a batch of :math:`q(s,.)` is
            returned instead. The shape of the latter return value is
            ``[batch_size, num_actions]``, which is only well-defined for
            discrete action
            spaces.

        """
        pass

class BaseQTypeI(BaseGenericQ):
    """
    Base class for modeling :term:`type-I <type-I state-action value function>`
    Q-function.

    A :term:`type-I <type-I state-action value function>` Q-function is
    implemented by mapping :math:`(s, a)\\mapsto q(s,a)`.

    Parameters
    ----------
    env : gym environment

        A gym environment.

    train_model : keras.Model([:term:`S`, :term:`A`], :term:`Q_sa`)

        Used for training.

    predict_model : keras.Model([:term:`S`, :term:`A`], :term:`Q_sa`)

        Used for predicting. For :term:`type-I <type-I state-action value
        function>` Q-functions, the :term:`target_model` and
        :term:`predict_model` are the same.

    target_model : keras.Model(:term:`S`, :term:`Q_sa`)

        A :term:`target_model` is used to make predictions on a bootstrapping
        scenario. It can be advantageous to use a point-in-time copy of the
        :term:`predict_model` to construct a bootstrapped target.

    gamma : float, optional

        The discount factor for discounting future rewards.

    bootstrap_n : positive int, optional

        The number of steps in n-step bootstrapping. It specifies the number of
        steps over which we're willing to delay bootstrapping. Large :math:`n`
        corresponds to Monte Carlo updates and :math:`n=1` corresponds to
        TD(0).

    bootstrap_with_target_model : bool, optional

        Whether to use the :term:`target_model` when constructing a
        bootstrapped target. If False (default), the primary
        :term:`predict_model` is used.

    update_strategy : str, optional

        The update strategy that we use to select the (would-be) next-action
        :math:`A_{t+n}` in the bootstrapped target:

        .. math::

            G^{(n)}_t\\ =\\ R^{(n)}_t + \\gamma^n Q(S_{t+n}, A_{t+n})

        Options are:

            'sarsa'
                Sample the next action, i.e. use the action that was actually
                taken.

            'q_learning'
                Take the action with highest Q-value under the current
                estimate, i.e. :math:`A_{t+n} = \\arg\\max_aQ(S_{t+n}, a)`.
                This is an off-policy method.

            'double_q_learning'
                Same as 'q_learning', :math:`A_{t+n} = \\arg\\max_aQ(S_{t+n},
                a)`, except that the value itself is computed using the
                :term:`target_model` rather than the primary model, i.e.

                .. math::

                    A_{t+n}\\ &=\\
                        \\arg\\max_aQ_\\text{primary}(S_{t+n}, a)\\\\
                    G^{(n)}_t\\ &=\\ R^{(n)}_t
                        + \\gamma^n Q_\\text{target}(S_{t+n}, A_{t+n})

    """
    def batch_eval(self, S, A=None, use_target_model=False):
        model = self.target_model if use_target_model else self.predict_model

        if A is not None:
            Q = model.predict_on_batch([S, A])
            check_numpy_array(Q, ndim=2, axis_size=1, axis=1)
            Q = np.squeeze(Q, axis=1)
            return Q  # shape: [batch_size]
        else:
            Q = []
            for a in range(self.num_actions):
                A = a * np.ones(len(S), dtype='int')
                Q.append(self.batch_eval(S, A))
            Q = np.stack(Q, axis=1)
            check_numpy_array(Q, ndim=2, axis_size=self.num_actions, axis=1)
            return Q  # shape: [batch_size, num_actions]

class BaseQTypeII(BaseGenericQ):
    """
    Base class for modeling :term:`type-II <type-II state-action value
    function>` Q-function.

    A :term:`type-II <type-II state-action value function>` Q-function is
    implemented by mapping :math:`s\\mapsto q(s,.)`.

    Parameters
    ----------
    env : gym environment

        A gym environment.

    train_model : keras.Model([:term:`S`, :term:`G`], :term:`Q_s`)

        Used for training.

    predict_model : keras.Model(:term:`S`, :term:`Q_s`)

        Used for predicting.

    target_model : keras.Model(:term:`S`, :term:`Q_s`)

        A :term:`target_model` is used to make predictions on a bootstrapping
        scenario. It can be advantageous to use a point-in-time copy of the
        :term:`predict_model` to construct a bootstrapped target.

    gamma : float, optional

        The discount factor for discounting future rewards.

    bootstrap_n : positive int, optional

        The number of steps in n-step bootstrapping. It specifies the number of
        steps over which we're willing to delay bootstrapping. Large :math:`n`
        corresponds to Monte Carlo updates and :math:`n=1` corresponds to
        TD(0).

    bootstrap_with_target_model : bool, optional

        Whether to use the :term:`target_model` when constructing a
        bootstrapped target. If False (default), the primary
        :term:`predict_model` is used.

    update_strategy : str, optional

        The update strategy that we use to select the (would-be) next-action
        :math:`A_{t+n}` in the bootsrapped target:

        .. math::

            G^{(n)}_t\\ =\\ R^{(n)}_t + \\gamma^n Q(S_{t+n}, A_{t+n})

        Options are:

            'sarsa'
                Sample the next action, i.e. use the action that was actually
                taken.

            'q_learning'
                Take the action with highest Q-value under the current
                estimate, i.e. :math:`A_{t+n} = \\arg\\max_aQ(S_{t+n}, a)`.
                This is an off-policy method.

            'double_q_learning'
                Same as 'q_learning', :math:`A_{t+n} = \\arg\\max_aQ(S_{t+n},
                a)`, except that the value itself is computed using the
                :term:`target_model` rather than the primary model, i.e.

                .. math::

                    A_{t+n}\\ &=\\
                        \\arg\\max_aQ_\\text{primary}(S_{t+n}, a)\\\\
                    G^{(n)}_t\\ &=\\ R^{(n)}_t
                        + \\gamma^n Q_\\text{target}(S_{t+n}, A_{t+n})

    """
    def batch_eval(self, S, A=None, use_target_model=False):
        model = self.target_model if use_target_model else self.predict_model

        if A is not None:
            Q = model.predict_on_batch(S)  # shape: [batch_size, num_actions]
            check_numpy_array(Q, ndim=2, axis_size=self.num_actions, axis=1)
            check_numpy_array(
                A, ndim=1, dtype='int', axis_size=Q.shape[0], axis=0)
            Q = project_onto_actions_np(Q, A)
            return Q  # shape: [batch_size]
        else:
            Q = model.predict_on_batch(S)
            check_numpy_array(Q, ndim=2, axis_size=self.num_actions, axis=1)
            return Q  # shape: [batch_size, num_actions]

class BaseSoftmaxPolicy(BasePolicy, BaseFunctionApproximator, NumActionsMixin, RandomStateMixin):  # noqa: E501
    """
    Base class for modeling :term:`updateable policies <updateable policy>` for
    discrete action spaces.

    Parameters
    ----------
    env : gym environment

        A gym environment.

    train_model : keras.Model([:term:`S`, :term:`Adv`], :term:`Z`)

        Used for training.

    predict_model : keras.Model(:term:`S`, :term:`Z`)

        Used for predicting.

    target_model : keras.Model(:term:`S`, :term:`Z`)

        A :term:`target_model` is used to make predictions on a bootstrapping
        scenario. It can be advantageous to use a point-in-time copy of the
        :term:`predict_model` to construct a bootstrapped target.

    update_strategy : str, optional

        The strategy for updating our policy. This typically determines the
        loss function that we use for our policy function approximator.

        Options are:

            'vanilla'
                Plain vanilla policy gradient. The corresponding (surrogate)
                loss function that we use is:

                .. math::

                    J(\\theta)\\ =\\ -\\mathcal{A}(s,a)\\,\\ln\\pi(a|s,\\theta)

            'ppo'
                `Proximal policy optimization
                <https://arxiv.org/abs/1707.06347>`_ uses a clipped proximal
                loss:

                .. math::

                    J(\\theta)\\ =\\ \\min\\Big(
                        r(\\theta)\\,\\mathcal{A}(s,a)\\,,\\
                        \\text{clip}\\big(
                            r(\\theta), 1-\\epsilon, 1+\\epsilon\\big)
                                \\,\\mathcal{A}(s,a)\\Big)

                where :math:`r(\\theta)` is the probability ratio:

                .. math::

                    r(\\theta)\\ =\\ \\frac
                        {\\pi(a|s,\\theta)}
                        {\\pi(a|s,\\theta_\\text{old})}

            'cross_entropy'
                Straightforward categorical cross-entropy (from logits). This
                loss function does *not* make use of the advantages
                :term:`Adv`. Instead, it minimizes the cross entropy between
                the behavior policy :math:`\\pi_b(a|s)` and the learned policy
                :math:`\\pi_\\theta(a|s)`:

                .. math::

                    J(\\theta)\\ =\\ \\hat{\\mathbb{E}}_t\\left\\{
                        -\\sum_a \\pi_b(a|S_t)\\, \\log \\pi_\\theta(a|S_t)
                    \\right\\}

    ppo_clipping : float, optional

        The clipping parameter :math:`\\epsilon` in the PPO clipped surrogate
        loss. This option is only applicable if ``update_strategy='ppo'``.

    entropy_bonus : float, optional

        The coefficient of the entropy bonus term in the policy objective.

    random_seed : int, optional

        Sets the random state to get reproducible results.

    """
    UPDATE_STRATEGIES = ('vanilla', 'ppo', 'cross_entropy')

    def __init__(
            self, env, train_model, predict_model, target_model,
            update_strategy='vanilla',
            ppo_clipping=0.2,
            entropy_bonus=0.01,
            random_seed=None):

        self.env = env
        self.train_model = train_model
        self.predict_model = predict_model
        self.target_model = target_model
        self.update_strategy = update_strategy
        self.ppo_clipping = float(ppo_clipping)
        self.entropy_bonus = float(entropy_bonus)
        self.random_seed = random_seed  # sets self.random in RandomStateMixin

        # TODO: allow for non-discrete action spaces
        self._actions = np.arange(self.num_actions)

    def __call__(self, s, use_target_model=False):
        """
        Draw an action from the current policy :math:`\\pi(a|s)`.

        Parameters
        ----------
        s : state observation

            A single state observation.

        use_target_model : bool, optional

            Whether to use the :term:`target_model` internally. If False
            (default), the :term:`predict_model` is used.

        Returns
        -------
        a : action

            A single action proposed under the current policy.

        """
        proba = self.proba(s, use_target_model=use_target_model)
        a = self.random.choice(self._actions, p=proba)
        return a

    def proba(self, s, use_target_model=False):
        """
        Get the probabilities over all actions :math:`\\pi(a|s)`.

        Parameters
        ----------
        s : state observation

            A single state observation.

        use_target_model : bool, optional

            Whether to use the :term:`target_model` internally. If False
            (default), the :term:`predict_model` is used.

        Returns
        -------
        pi : 1d array, shape: [num_actions]

            Probabilities over all actions.

            **Note.** This hasn't yet been implemented for non-discrete action
            spaces.

        """
        assert self.env.observation_space.contains(s)
        S = np.expand_dims(s, axis=0)
        P = self.batch_eval(S, use_target_model=use_target_model)
        check_numpy_array(P, shape=(1, self.num_actions))
        pi = np.squeeze(P, axis=0)
        return pi

    def greedy(self, s, use_target_model=False):
        """
        Draw the greedy action, i.e. :math:`\\arg\\max_a\\pi(a|s)`.

        Parameters
        ----------
        s : state observation

            A single state observation.

        use_target_model : bool, optional

            Whether to use the :term:`target_model` internally. If False
            (default), the :term:`predict_model` is used.

        Returns
        -------
        a : action

            A single action proposed under the current policy.

        """
        a = argmax(self.proba(s, use_target_model=use_target_model))
        return a

    def update(self, s, pi, advantage):
        """
        Update the policy.

        Parameters
        ----------
        s : state observation

            A single state observation.

        pi : 1d array, shape: [num_actions]

            Vector of action propensities under the behavior policy. This may
            be just an indicator if the action propensities are inferred
            through sampling, e.g. if we have four possible actions ``pi = [0,
            0, 1, 0]`` would indicate that the action :math:`a=2` was drawn
            from the behavior policy.

        advantage : float

            A value for the advantage :math:`\\mathcal{A}(s,a) = q(s,a) -
            v(s)`. This might be sampled and/or estimated version of the true
            advantage.

        """
        assert self.env.observation_space.contains(s)
        check_numpy_array(pi, ndim=1, axis_size=self.num_actions, axis=0)

        S = np.expand_dims(s, axis=0)
        P = np.expand_dims(pi, axis=0)
        Adv = np.expand_dims(advantage, axis=0)
        self.batch_update(S, P, Adv)

    def batch_eval(self, S, use_target_model=False):
        """
        Evaluate the policy on a batch of state observations.

        Parameters
        ----------
        S : nd array, shape: [batch_size, ...]

            A batch of state observations.

        use_target_model : bool, optional

            Whether to use the :term:`target_model` internally. If False
            (default), the :term:`predict_model` is used.

        Returns
        -------
        P : 2d array, shape: [batch_size, num_actions]

            A batch of predicted action probabilities :math:`\\pi(a|s)`.

        """
        model = self.target_model if use_target_model else self.predict_model

        Z = model.predict_on_batch(S)
        check_numpy_array(Z, ndim=2, axis_size=self.num_actions, axis=1)
        P = softmax(Z, axis=1)
        return P  # shape: [batch_size, num_actions]

    def batch_update(self, S, P, Adv):
        """
        Update the policy on a batch of transitions.

        Parameters
        ----------
        S : nd array, shape: [batch_size, ...]

            A batch of state observations.

        P : 2d Tensor, dtype: int, shape: [batch_size]

            A batch of action propensities. :term:`P` is typically just an
            indicator for which action was chosen by the behavior policy. In
            this sense, :term:`P` acts as a projector more than a prediction
            target. That is, :term:`P` is used to project our predicted values
            down to those for which we actually received the feedback signal:
            :term:`Adv`.

        Adv : 1d array, dtype: float, shape: [batch_size]

            A value for the :term:`advantage <Adv>` :math:`\\mathcal{A}(s,a) = q(s,a) -
            v(s)`. This might be sampled and/or estimated version of the true
            advantage.

        Returns
        -------
        losses : dict

            A dict of losses/metrics, of type ``{name <str>: value <float>}``.

        """
        check_numpy_array(P, ndim=2, axis_size=self.num_actions, axis=1)
        losses = self._train_on_batch([S, Adv], P)
        return losses

    def _policy_loss(self, Adv, Z_target=None):
        if self.update_strategy == 'vanilla':
            return SoftmaxPolicyLossWithLogits(
                Adv, entropy_bonus=self.entropy_bonus)

        if self.update_strategy == 'ppo':
            assert Z_target is not None
            return ClippedSurrogateLoss(
                Adv, Z_target, entropy_bonus=self.entropy_bonus,
                epsilon=self.ppo_clipping)

        if self.update_strategy == 'cross_entropy':
            return keras.losses.CategoricalCrossentropy(from_logits=True)

        raise ValueError(
            "unknown update_strategy '{}'".format(self.update_strategy))


import numpy as np
import cv2

def py_cpu_nms_in_numpy(dets, thresh):
    """Pure Python NMS baseline."""
    # 所有图片的坐标信息，字典形式储存
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]##bbox打分
    # 计算出所有图片的面积
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    # 打分从大到小排列，取index
    order = scores.argsort()[::-1]

    keep = [] # 用来存放最后保留的图片的相应评分
    while order.size > 0:
        # order[0]是当前分数最大的窗口，肯定保留
        i = order[0]
        keep.append(i)# 保留改图片的值
        # 计算窗口i与其他所有窗口的交叠部分的面积
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        # 计算出各个相交矩形的面积
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        # 交/并得到iou值
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        # inds为所有与窗口i的iou值小于threshold值的窗口的index，其他窗口此次都被窗口i吸收
        inds = np.where(ovr <= thresh)[0]
        # order里面只保留与窗口i交叠面积小于threshold的那些窗口，由于ovr长度比order长度少1(不包含i)，所以inds+1对应到保留的窗口
        order = order[inds + 1]

    return keep

def point_form(boxes):
    """ Convert prior_boxes to (xmin, ymin, xmax, ymax)
    representation for comparison to point form ground truth data.
    Args:
        boxes: (tensor) center-size default boxes from priorbox layers.
    Return:
        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
    """
    """
    将以前的方框转换为（xmin、ymin、xmax、ymax）
	用于与点格式地面真实数据进行比较的表示。
	参数：
	box:（tensor）中心大小priorbox层的默认框。
	返回：
	盒子：（张量）转换成盒子的xmin，ymin，xmax，ymax形式。
    """
    return np.concatenate((boxes[:, :2] - boxes[:, 2:]/2,     # xmin, ymin
                     boxes[:, :2] + boxes[:, 2:]/2), 1)  # xmax, ymax

def center_size(boxes):
    """ Convert prior_boxes to (cx, cy, w, h)
    representation for comparison to center-size form ground truth data.
    Args:
        boxes: (tensor) point_form boxes
    Return:
        boxes: (tensor) Converted xmin, ymin, xmax, ymax form of boxes.
    """
    return np.concatenate(((boxes[:, 2:] + boxes[:, :2])/2, boxes[:, 2:] - boxes[:, :2]), 1) # cx, cy # w, h

def intersect(box_a, box_b):
    """ We resize both tensors to [A,B,2] without new malloc:
    [A,2] -> [A,1,2] -> [A,B,2]
    [B,2] -> [1,B,2] -> [A,B,2]
    Then we compute the area of intersect between box_a and box_b.
    Args:
      box_a: (tensor) bounding boxes, Shape: [A,4].
      box_b: (tensor) bounding boxes, Shape: [B,4].
    Return:
      (tensor) intersection area, Shape: [A,B].
      将以前的方框转换为（cx，cy，w，h）
	用于与地面真值数据中心大小进行比较的表示。
	参数：
	方块：（张量）点式方块
	返回：
	box:（张量）转换成xmin，ymin，xmax，ymax形式的box
    """
    box_a = box_a.cpu()
    box_b = box_b.cpu()
    A = box_a.size(0)
    B = box_b.size(0)
    max_xy = np.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
                       box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
    min_xy = np.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
                       box_b[:, :2].unsqueeze(0).expand(A, B, 2))
    inter = np.clip((max_xy - min_xy), min=0)
    return inter[:, :, 0] * inter[:, :, 1]

def jaccard(box_a, box_b):
    """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap
    is simply the intersection over union of two boxes.  Here we operate on
    ground truth boxes and default boxes.
    E.g.:
        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
    Args:
        box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]
        box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]
    Return:
        jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
 	计算两组盒子的重叠相似系数。

重叠相似只是两个框的并集上的交集。

在这里，我们操作地面真相框和默认框。

例如。：

A∩B/A∪B=A∩B/（面积（A）+面积（B）-A∩B）

参数：

box_a:（张量）地面真值边界框，形状：[num_aobjects，4]

box_b:（张量）来自priorbox层的先前盒，形状：[num_priors，4]

返回：

重叠相似：（张量）形状：[框a.尺寸（0），框b.尺寸（0）]
    """
    inter = intersect(box_a, box_b)
    area_a = ((box_a[:, 2]-box_a[:, 0]) *
              (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
    area_b = ((box_b[:, 2]-box_b[:, 0]) *
              (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
    union = area_a + area_b - inter
    return inter / union  # [A,B]

def matrix_iou(a, b):
    """
    return iou of a and b, numpy version for data augenmentation
    返回a和b的iou，numpy版本进行数据补充
    """
    lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
    rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])

    area_i = (np.prod(rb - lt, axis=2) * (lt < rb)).all(axis=2)
    area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
    area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)
    return area_i / (area_a[:, np.newaxis] + area_b - area_i)

def matrix_iof(a, b):
    """
    matrix_iof 的意思是裁剪后的roi除以boxes,有全覆盖的则留下。
    return iof of a and b, numpy version for data augenmentation
    返回a和b的iof，numpy版本进行数据补充
    """
    lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])#x1 y1 最大值
    rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])# x2 y2 最小值

    area_i = (np.prod(rb - lt, axis=2) * (lt < rb)).all(axis=2)#连乘 即面积
    area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
    return area_i / np.maximum(area_a[:, np.newaxis], 1)

def match(threshold, truths, priors, variances, labels, landms, loc_t, conf_t, landm_t, idx):
    """Match each prior box with the ground truth box of the highest jaccard
    overlap, encode the bounding boxes, then return the matched indices
    corresponding to both confidence and location preds.
    Args:
        threshold: (float) The overlap threshold used when mathing boxes.
        truths: (tensor) Ground truth boxes, Shape: [num_obj, 4].
        priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4].
        variances: (tensor) Variances corresponding to each prior coord,
            Shape: [num_priors, 4].
        labels: (tensor) All the class labels for the image, Shape: [num_obj].
        landms: (tensor) Ground truth landms, Shape [num_obj, 10].
        loc_t: (tensor) Tensor to be filled w/ endcoded location targets.
        conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds.
        landm_t: (tensor) Tensor to be filled w/ endcoded landm targets.
        idx: (int) current batch index
    Return:
        The matched indices corresponding to 1)location 2)confidence 3)landm preds.
    """
    # jaccard index
    # # 第1步,计算IOU
    overlaps = jaccard(
        truths,
        point_form(priors)
    )
    # (Bipartite Matching)
    # [1,num_objects] best prior for each ground truth
    # 第2步,为每个真实框匹配一个IOU最大的锚点框,GT框->锚点框
    # best_prior_overlap为每个真实框的最大IOU值,shape[num_objects,1]
    # best_prior_idx为对应的最大IOU的先验锚点框的Index,其元素值的范围为[0,num_priors]
    best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)

    # ignore hard gt
    valid_gt_idx = best_prior_overlap[:, 0] >= 0.2
    best_prior_idx_filter = best_prior_idx[valid_gt_idx, :]
    if best_prior_idx_filter.shape[0] <= 0:
        loc_t[idx] = 0
        conf_t[idx] = 0
        return

    # [1,num_priors] best ground truth for each prior
    # 第3步,若先验锚点框与GT框的IOU>阈值,也将这些锚点框匹配上,锚点框->GT框
    # best_truth_overlap为每个先验锚点框对应其中一个真实框的最大IOU,shape[1,num_priors]
    # best_truth_idx为每个先验锚点框对应的真实框的index,其元素值的范围为[0,num_objects]

    best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True)
    best_truth_idx.squeeze_(0)
    best_truth_overlap.squeeze_(0)
    best_prior_idx.squeeze_(1)
    best_prior_idx_filter.squeeze_(1)
    best_prior_overlap.squeeze_(1)
    # 第4步
    # index_fill_(self, dim: _int, index: Tensor, value: Number)对第dim行的index使用value进行填充
    # best_truth_overlap为第一步匹配的结果,需要使用到,使用best_prior_idx是第二步的结果,也是需要使用上的
    # 所以在best_truth_overlap上进行填充,表明选出来的正例
    # 使用2进行填充,是因为,IOU值的范围是[0,1],只要使用大于1的值填充,就表明肯定能被选出来
    best_truth_overlap.index_fill_(0, best_prior_idx_filter, 2)  # ensure best prior
    # TODO refactor: index  best_prior_idx with long tensor
    # ensure every gt matches with its prior of max overlap
    # 确保每个GT框都能匹配上最大IOU的先验锚点框
    # 得到每个先验锚点框都能有一个匹配上的数字
    # best_prior_idx的元素值的范围是[0,num_priors],长度为num_objects
    for j in range(best_prior_idx.size(0)):     # 判别此anchor是预测哪一个boxes
        best_truth_idx[best_prior_idx[j]] = j
    #box
    matches = truths[best_truth_idx]            # 取出最佳匹配的GT框,Shape: [num_priors,4]
    conf = labels[best_truth_idx]                # Shape: [num_priors],0为背景,
    conf[best_truth_overlap < threshold] = 0     # 置信度小于阈值的label设置为0
    loc = encode(matches, priors, variances)    # 进行位置编码

    #landm
    matches_landm = landms[best_truth_idx]      # 取出最佳匹配的landm
    landm = encode_landm(matches_landm, priors, variances)# 进行位置编码
    loc_t[idx] = loc    # [num_priors,4] encoded offsets to learn应该学习的编码偏差
    conf_t[idx] = conf  # [num_priors] top class label for each prior每个锚点框的label
    landm_t[idx] = landm #应该学习的编码偏差

def encode(matched, priors, variances):
    """Encode the variances from the priorbox layers into the ground truth boxes
    we have matched (based on jaccard overlap) with the prior boxes.
    Args:
        matched: (tensor) Coords of ground truth for each prior in point-form
            Shape: [num_priors, 4].
        priors: (tensor) Prior boxes in center-offset form
            Shape: [num_priors,4].
        variances: (list[float]) Variances of priorboxes
    Return:
        encoded boxes (tensor), Shape: [num_priors, 4]
        对坐标进行编码
        利用GT框和先验锚点框,计算偏差,用于回归
    """

    # dist b/t match center and prior's center
    g_cxcy = (matched[:, :2] + matched[:, 2:])/2 - priors[:, :2] # 计算GT框与锚点框中心点的距离
    # encode variance
    g_cxcy /= (variances[0] * priors[:, 2:])
    # match wh / prior wh
    g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:]
    g_wh = np.log(g_wh) / variances[1]
    # return target for smooth_l1_loss
    return np.concatenate([g_cxcy, g_wh], 1)  # [num_priors,4]

def encode_landm(matched, priors, variances):

    """Encode the variances from the priorbox layers into the ground truth boxes
    we have matched (based on jaccard overlap) with the prior boxes.
    Args:
        matched: (tensor) Coords of ground truth for each prior in point-form
            Shape: [num_priors, 10].
        priors: (tensor) Prior boxes in center-offset form
            Shape: [num_priors,4].
        variances: (list[float]) Variances of priorboxes
    Return:
        encoded landm (tensor), Shape: [num_priors, 10]
    对landm坐标进行编码
    """

    # dist b/t match center and prior's center
    matched = np.reshape(matched, (matched.size(0), 5, 2))
    priors_cx = priors[:, 0].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
    priors_cy = priors[:, 1].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
    priors_w = priors[:, 2].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
    priors_h = priors[:, 3].unsqueeze(1).expand(matched.size(0), 5).unsqueeze(2)
    priors = np.concatenate([priors_cx, priors_cy, priors_w, priors_h], 2)
    g_cxcy = matched[:, :, :2] - priors[:, :, :2]
    # encode variance
    g_cxcy /= (variances[0] * priors[:, :, 2:])
    # g_cxcy /= priors[:, :, 2:]
    g_cxcy = g_cxcy.reshape(g_cxcy.size(0), -1)
    # return target for smooth_l1_loss
    return g_cxcy

# Adapted from https://github.com/Hakuyume/chainer-ssd
def decode_in_numpy(loc, priors, variances):
    """Decode locations from predictions using priors to undo
    the encoding we did for offset regression at train time.
    Args:
        loc (tensor): location predictions for loc layers,
            Shape: [num_priors,4]网络预测的锚点框偏差信息
        priors (tensor): Prior boxes in center-offset form.
            Shape: [num_priors,4].先验锚点框
        variances: (list[float]) Variances of priorboxes预测框的坐标
    Return:
        decoded bounding box predictions
    对编码的坐标进行解码,返回预测框的坐标
    """
    # [中心点x,中心点y,宽,高]
    boxes = np.concatenate((
        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
        priors[:, 2:] * np.exp(loc[:, 2:] * variances[1])), 1)

    boxes[:, :2] -= boxes[:, 2:] / 2 # xmin,ymin
    boxes[:, 2:] += boxes[:, :2] # xmax,ymax
    return boxes

def decode_landm_in_numpy(pre, priors, variances):
    """Decode landm from predictions using priors to undo
    the encoding we did for offset regression at train time.
    Args:
        pre (tensor): landm predictions for loc layers,
            Shape: [num_priors,10]
        priors (tensor): Prior boxes in center-offset form.
            Shape: [num_priors,4].
        variances: (list[float]) Variances of priorboxes
    Return:
        decoded landm predictions
    """
    landms = np.concatenate((priors[:, :2] + pre[:, :2] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 2:4] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 4:6] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 6:8] * variances[0] * priors[:, 2:],
                        priors[:, :2] + pre[:, 8:10] * variances[0] * priors[:, 2:],
                        ), axis=1)#五个特征点
    return landms
def log_sum_exp(x):
    """Utility function for computing log_sum_exp while determining
    This will be used to determine unaveraged confidence loss across
    all examples in a batch.
    Args:
        x (Variable(tensor)): conf_preds from conf layers
    确定时用于计算的实用函数（log_sum_exp）

这将用于确定批处理中所有示例的不可用信心损失。

参数：

x（变量（张量）：来自conf层的conf preds
    """
    x_max = x.data.max()
    return np.log(np.sum(np.exp(x-x_max), 1, keepdims=True)) + x_max

def shift_img(img,dx,dy,fill_with =0):
    '''
    平移图片
    :param img: 需要操作的图片
    :param dx: 向右平移多少
    :param dy: 向下平移多少
    :param fill_with: 多的部分用什么填充
    :return:img 操作之后的图片
    '''
    M = np.array([[1,0,dx],[0,1,dy]])
    shifted = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))
    return shifted

def rotate_img(img,angle,circle):
    '''
    旋转图片
    :param img:需要操作的图片
    :param angle: 旋转角度
    :param circle: 旋转圆心
    :return: 操作之后的图片
    '''
    M = cv2.getRotationMatrix2D((circle[0], circle[1]), angle, 1)
    rotate = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))
    return rotate