# -*- coding: utf-8 -*-
"""

Created on 08/09/2020

Author : Carlos Eduardo Barbosa

"""
from __future__ import print_function, division

import os
import copy
import yaml
import re

import numpy as np
import astropy.units as u
from astropy.table import Table, hstack, vstack
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize, LinearSegmentedColormap
from matplotlib import cm
from astropy.io import fits
from scipy.ndimage.filters import gaussian_filter1d
from spectres import spectres
from tqdm import tqdm
try:
    import ppxf_util as util
except:
    import ppxf.ppxf_util as util
import emcee
from scipy.interpolate import interp1d

import context
from lick import Lick
import paintbox as pb
from run_paintbox import build_sed_model

import context

def get_colors(R, cmapname="Blues_r"):
    norm = Normalize(vmin=0, vmax=np.ceil(np.max(R)),
                                       clip=True)
    cmap = plt.get_cmap(cmapname)
    new_cmap = truncate_colormap(cmap, 0.0, 0.8)
    mapper = cm.ScalarMappable(norm=norm, cmap=new_cmap)
    return mapper, np.array([(mapper.to_rgba(v)) for v in R])

def truncate_colormap(cmap, minval=0.0, maxval=16, n=100):
    new_cmap = LinearSegmentedColormap.from_list(
        'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval,
                                            b=maxval),
        cmap(np.linspace(minval, maxval, n)))
    return new_cmap

def lick_tests(wdir, velscale=None):
    """ Calculates Lick indices on samples from full spectrum fitting. """
    velscale = context.velscale if velscale is None else velscale
    bandsfile = os.path.join(os.path.split(os.path.abspath(__file__))[0],
                                "tables/spindex_CS.dat")
    bandsz0 = np.loadtxt(bandsfile, usecols=(3,4,1,2,5,6)) * u.AA
    units_bin = np.loadtxt(bandsfile, usecols=(7,))
    units = np.where(units_bin, u.Unit("mag"), u.Unit("angstrom"))
    names = np.loadtxt(bandsfile, usecols=(8,), dtype=str)
    names = [_.replace("_", "\_") for _ in names]
    # Setting input data
    emcee_dir = os.path.join(wdir, "EMCEE")
    tables = sorted([_ for _ in os.listdir(emcee_dir) if _.endswith(
        "results.fits")])
    lick = Lick(bandsz0, units=units)
    norm = Normalize(vmin=0, vmax=16, clip=True)
    cmap = plt.get_cmap('Blues_r')
    new_cmap = truncate_colormap(cmap, 0.0, 0.5)
    # sed = build_sed_model(sample="test")[0]
    mapper = cm.ScalarMappable(norm=norm, cmap=new_cmap)
    table = Table.read(os.path.join(wdir, "results.fits"))
    for i, gal in enumerate(table):
        vel = gal["V"] * u.km / u.s
        sci_table = os.path.join(wdir, "sci",
                        "{0[0]}_sn250_{0[1]}.fits".format(
                         gal["BIN"].split("_")))
        sci = Table.read(sci_table)
        if i == 0:
            wave = sci["wave"].data
            sed = build_sed_model(wave)[0]
            # Masking wavelengths with sky lines/ residuals
            skylines = np.array(
                [4792, 4860, 4923, 5071, 5239, 5268, 5577, 5889.99,
                 5895, 5888, 5990, 5895, 6300, 6363, 6386, 6562,
                 6583, 6717, 6730, 7246, 8286, 8344, 8430, 8737,
                 8747, 8757, 8767,
                 8777, 8787, 8797, 8827, 8836, 8919, 9310])
            goodpixels = np.arange(len(wave))
            for line in skylines:
                sky = np.argwhere(
                    (wave <= line - 2) | (wave >= line + 2)).ravel()
                goodpixels = np.intersect1d(goodpixels, sky)
        # Performing Lick measurements
        # Original data
        flam = sci["flam"].data
        ew_obs= lick(wave, flam, vel=vel)[1].ravel()
        # Interpolated, observed data
        interpolator = interp1d(wave[goodpixels], flam[goodpixels])
        ew_obs_interp = lick(wave, interpolator(wave), vel=vel)[1].ravel()
        # SSP model
        params = np.array([gal[p] for p in sed.parnames])
        ew_model= lick(wave, sed(params), vel=vel)[1].ravel()
        # SSP interpolated
        ew_model_interp = lick(wave[goodpixels], sed(params)[goodpixels],
                               vel=vel)[1].ravel()
        for j, ew in enumerate([ew_obs, ew_obs_interp, ew_model,
                               ew_model_interp]):
            plt.plot(names, ew, "o-", label="{}".format(j))
        plt.legend()
        plt.show()
    plt.legend()
    plt.show()

if __name__ == "__main__":
    wdir = os.path.join(context.data_dir, "MUSE/voronoi/sn250")
    lick_tests(wdir)

import numpy,from scipy,import astropy,from astropy


"""
    mesh
    ~~~~

    This module contains the Mesh, MeshData, and Material classes.
    This documentation was auto-generated from the mesh.py file.
"""

import numpy as np
from .utils import vertices as vertutils
from .utils import NameLabelMixin
from . import physical, shader
from .texture import Texture
from .vertex import VAO, VBO
import pyglet.gl as gl
from copy import deepcopy

def gen_fullscreen_quad(name='FullScreenQuad'):
    verts = np.array([[-1, -1, -.5], [-1, 1, -.5], [1, 1, -.5], [-1, -1, -.5], [1, 1, -.5], [1, -1, -.5]], dtype=np.float32)
    normals=np.array([[0, 0, 1]] * 6, dtype=np.float32)
    texcoords=np.array([[0, 0], [0, 1], [1, 1], [0, 0], [1, 1], [1, 0]], dtype=np.float32)
    return Mesh(name=name, arrays=(verts, normals, texcoords), mean_center=False)

class EmptyEntity(shader.HasUniformsUpdater, physical.PhysicalGraph, NameLabelMixin):
    """An object that occupies physical space and uniforms, but doesn't actually draw anything when draw() is called."""

    def draw(self, *args, **kwargs):
        pass

    def reset_uniforms(self):
        pass

class Mesh(shader.HasUniformsUpdater, physical.PhysicalGraph, NameLabelMixin):

    triangles = gl.GL_TRIANGLES
    points = gl.GL_POINTS

    def __init__(self, arrays, textures=(), mean_center=True,
                 gl_states=(), drawmode=gl.GL_TRIANGLES, point_size=15, dynamic=False, visible=True, **kwargs):
        """
        Returns a Mesh object, containing the position, rotation, and color info of an OpenGL Mesh.

        Meshes have two coordinate system, the "local" and "world" systems, on which the transforms are performed
        sequentially.  This allows them to be placed in the scene while maintaining a relative position to one another.

        .. note:: Meshes are not usually instantiated directly, but from a 3D file, like the WavefrontReader .obj and .mtl files.

        Args:
            name (str): the mesh's name.
            arrays (tuple): a list of 2D arrays to be rendered.  All arrays should have same number of rows. Arrays will be accessible in shader in same attrib location order.
            texture (Texture): a Texture instance, which is linked when the Mesh is rendered.
            visible (bool): whether the Mesh is available to be rendered.  To make hidden (invisible), set to False.

        Returns:
            Mesh instance
        """

        super(Mesh, self).__init__(**kwargs)
        self.reset_uniforms()

        arrays = tuple(np.array(array, dtype=np.float32) for array in arrays)
        self.arrays, self.array_indices = vertutils.reindex_vertices(arrays)

        # Mean-center vertices and move position to vertex mean.
        vertex_mean = self.arrays[0][self.array_indices, :].mean(axis=0)

        if mean_center:
            self.arrays[0][:] -= vertex_mean
        if 'position' in kwargs:
            self.position.xyz = kwargs['position']
        elif mean_center:
            self.position.xyz = vertex_mean
        self._mean_center = mean_center
        # self.position.xyz = vertex_mean if not 'position' in kwargs else kwargs['position']

        # Change vertices from an Nx3 to an Nx4 array by appending ones.  This makes some calculations more efficient.
        arrays = list(self.arrays)
        arrays[0] = np.append(self.arrays[0], np.ones((self.arrays[0].shape[0], 1), dtype=np.float32), axis=1)
        self.arrays = tuple(arrays)

        self.textures = list(textures)
        self.vao = None  # Will be created upon first draw, when OpenGL context is available.
        self.gl_states = gl_states
        self.drawmode = drawmode
        self.point_size = point_size
        self.dynamic = dynamic
        self.visible = visible
        self.vbos = []

    def __repr__(self):
        return "<Mesh(name='{self.name}', position_rel={self.position}, position_glob={self.position_global}, rotation={self.rotation})".format(self=self)

    def copy(self):
        return Mesh(arrays=deepcopy([arr.copy() for arr in [self.vertices, self.normals, self.texcoords]]), texture=self.textures, mean_center=deepcopy(self._mean_center),
                    position=self.position.xyz, rotation=self.rotation.__class__(*self.rotation[:]), scale=self.scale.xyz,
                    drawmode=self.drawmode, point_size=self.point_size, dynamic=self.dynamic, visible=self.visible,
                    gl_states=deepcopy(self.gl_states))

    def reset_uniforms(self):
        self.uniforms['model_matrix'] = self.model_matrix_global.view()
        self.uniforms['normal_matrix'] = self.normal_matrix_global.view()

    @property
    def dynamic(self):
        return self._dynamic

    @dynamic.setter
    def dynamic(self, value):
        for array in self.arrays:
            array.setflags(write=True if value else False)
        self._dynamic = value

    @property
    def vertices(self):
        """Mesh vertices, centered around 0,0,0"""
        return self.arrays[0][:, :3].view()

    @vertices.setter
    def vertices(self, value):
        self.arrays[0][:, :3] = value

    @property
    def normals(self):
        """Mesh normals array."""
        return self.arrays[1][:, :3].view()

    @normals.setter
    def normals(self, value):
        self.arrays[1][:, :3] = value

    @property
    def texcoords(self):
        return self.arrays[2][:, :2].view()

    @texcoords.setter
    def texcoords(self, value):
        self.arrays[2][:, :2] = value

    @property
    def vertices_local(self):
        """Vertex position, in local coordinate space (modified by model_matrix)"""
        return np.dot(self.model_matrix, self.vertices)

    @property
    def vertices_global(self):
        """Vertex position, in world coordinate space (modified by model_matrix)"""
        return np.dot(self.model_matrix_global, self.vertices)

    @property
    def texture(self):
        raise DeprecationWarning("Mesh.texture no longer exists.  Instead, please append textures to the Mesh.textures list.")

    @texture.setter
    def texture(self, value):
        raise DeprecationWarning("Mesh.texture no longer exists.  Instead, please append textures to the Mesh.textures list.")

    @classmethod
    def from_incomplete_data(cls, vertices, normals=(), texcoords=(), **kwargs):
        """Return a Mesh with (vertices, normals, texcoords) as arrays, in that order.
           Useful for when you want a standardized array location format across different amounts of info in each mesh."""
        normals = normals if hasattr(texcoords, '__iter__') and len(normals) else vertutils.calculate_normals(vertices)
        texcoords = texcoords if hasattr(texcoords, '__iter__') and len(texcoords) else np.zeros((vertices.shape[0], 2), dtype=np.float32)
        return cls(arrays=(vertices, normals, texcoords), **kwargs)

    def _fill_vao(self):
        """Put array location in VAO for shader in same order as arrays given to Mesh."""
        with self.vao:
            self.vbos = []
            for loc, verts in enumerate(self.arrays):
                vbo = VBO(verts)
                self.vbos.append(vbo)
                self.vao.assign_vertex_attrib_location(vbo, loc)

    def draw(self):
        if not self.vao:
            self.vao = VAO(indices=self.array_indices)
            self._fill_vao()

        if self.visible:
            if self.dynamic:
                for vbo in self.vbos:
                    vbo._buffer_subdata()

            if self.drawmode == gl.GL_POINTS:
                gl.glPointSize(self.point_size)

            for texture in self.textures:
                texture.bind()

            with self.vao as vao:
                self.uniforms.send()
                vao.draw(mode=self.drawmode)

            for texture in self.textures:
                texture.unbind()


import time
import numpy as np
import tensorflow as tf
from offrl_evaluate import evaluate_policy
from .mbrl_base import BNN, FakeEnv # pylint: disable=relative-beyond-top-level
from .base import OnlineReplayBuffer, SquahedGaussianActor, QNetwork # pylint: disable=relative-beyond-top-level

class REPBRL(tf.keras.layers.Layer):
    def __init__(self, state_dim, action_dim, termination_fn, rollout_steps=5,
                 max_model_t=None, ipm_coef=0, mopo_penalty_coeff=0.0, my_penalty_coeff=0.0,
                 use_pre_mean=True, **_):
        self.state_dim, self.action_dim, self.max_model_t \
            = state_dim, action_dim, max_model_t
        self.obs_ph, self.action_ph, self.reward_ph, self.terminal_ph, self.next_obs_ph \
            = [tf.keras.layers.Input(d) for d in [state_dim, action_dim, 1, 1, state_dim]]
        self.data_obs_ph, self.data_action_ph, self.data_next_obs_ph, self.data_init_obs_ph,\
            self.data_terminal_ph = [
                tf.keras.layers.Input(d) for d in [state_dim, action_dim, state_dim, state_dim, 1]]

        # Model params
        self.model_hidden_dim = 200
        self.model = BNN(input_dim=state_dim + action_dim, output_dim=state_dim + 1,
                         hidden_dim=self.model_hidden_dim, use_pre_mean=use_pre_mean)
        self.fake_env = FakeEnv(self.model, termination_fn, mopo_penalty_coeff=mopo_penalty_coeff,
                                my_penalty_coeff=my_penalty_coeff)

        # policy optimization params
        self.rollout_steps = rollout_steps # 1 or 5 in MOPO
        self.sac_hidden_dim, self.batch_size = 256, 256
        self.rollout_batch_size = int(1e4) if rollout_steps == 1000 else int(1e6 / rollout_steps)
        self.model_retain_epochs = 5
        self.model_train_freq, self.evaluate_freq = 10000, 10000
        self.real_ratio, self.gamma, self.tau, self.learning_rate = 0.05, 0.99, 5e-3, 3e-4
        self.env_batch_size = int(self.batch_size * self.real_ratio)
        self.model_batch_size = self.batch_size - self.env_batch_size
        self.model_buffer_size = int(self.model_retain_epochs * rollout_steps \
            * self.rollout_batch_size * self.evaluate_freq / self.model_train_freq)
        self.target_entropy = -self.action_dim

        # IPM params
        self.ipm_coef = ipm_coef
        self.sess = tf.keras.backend.get_session()

        # build SAC
        self.actor = actor = SquahedGaussianActor(action_dim, hidden_dim=self.sac_hidden_dim)
        critic_online = QNetwork(hidden_dim=self.sac_hidden_dim)
        critic_target = QNetwork(hidden_dim=self.sac_hidden_dim)
        log_ent_coef = tf.keras.backend.variable(1.0, dtype=tf.float32, name='log_ent_coef')
        ent_coef = tf.exp(log_ent_coef)

        # Actor training (with kl regularization)
        self.det_actions, self.cur_actions, cur_logp, dist = actor([self.obs_ph])
        _, next_actions, next_logp, _ = actor([self.next_obs_ph])
        ipm = self._build_ipm(actor)

        bc_loss = actor.nlogp(dist, self.action_ph)
        cur_qpis = critic_online([self.obs_ph, self.cur_actions])
        actor_loss = tf.reduce_mean(ent_coef * cur_logp - tf.reduce_min(cur_qpis, axis=0))
        if self.ipm_coef > 0:
            actor_loss += tf.reduce_mean(ipm_coef * ipm)
        ent_coef_loss = -tf.reduce_mean(
            log_ent_coef * tf.stop_gradient(cur_logp + self.target_entropy))

        cur_qs = critic_online([self.obs_ph, self.action_ph])
        next_qs = critic_target([self.next_obs_ph, next_actions])
        next_q = tf.reduce_min(next_qs, axis=0) - ent_coef * next_logp
        q_backup = tf.stop_gradient(self.reward_ph + (1 - self.terminal_ph) * self.gamma * next_q)
        q_loss = tf.reduce_sum([tf.losses.mean_squared_error(q_backup, cur_q) for cur_q in cur_qs])

        bc_optimizer = tf.train.AdamOptimizer(self.learning_rate)
        self.behavior_clone_op = bc_optimizer.minimize(bc_loss, var_list=actor.trainable_variables)
        actor_opt = tf.train.AdamOptimizer(self.learning_rate)
        actor_train_op = actor_opt.minimize(actor_loss, var_list=actor.trainable_variables)
        critic_opt = tf.train.AdamOptimizer(self.learning_rate)
        critic_train_op = critic_opt.minimize(q_loss, var_list=critic_online.trainable_variables)
        entropy_opt = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
        entropy_train_op = entropy_opt.minimize(ent_coef_loss, var_list=log_ent_coef)

        with tf.control_dependencies([critic_train_op]):
            # Update target network
            source_params = critic_online.trainable_variables
            target_params = critic_target.trainable_variables
            target_update_op = [
                tf.assign(target, (1 - self.tau) * target + self.tau * source)
                for target, source in zip(target_params, source_params)]

        self.sess.run(tf.variables_initializer(
            actor_opt.variables() + critic_opt.variables() + entropy_opt.variables()
            + bc_optimizer.variables()))
        critic_target.set_weights(critic_online.get_weights())
        self.step_ops = [actor_train_op, critic_train_op, entropy_train_op, target_update_op]
        self.step_ops += [actor_loss, q_loss, tf.reduce_mean(cur_qs), ent_coef,
                          tf.reduce_mean(dist.entropy()), tf.reduce_mean(cur_logp), ipm]
        self.info_labels = ['actor_loss', 'q_loss', 'mean(cur_qs)', 'ent_coef', 'entropy',
                            'cur_logp', 'ipm']

        self.standardizer = None
        self.unstandardizer = None

    def _build_ipm(self, actor):
        _, initial_actions, _, _ = actor([self.data_init_obs_ph])
        _, next_actions, _, _ = actor([self.data_next_obs_ph])
        rep = self.model.representation(tf.concat([self.data_obs_ph, self.data_action_ph], axis=-1))
        nrep = self.model.representation(tf.concat([self.data_next_obs_ph, next_actions], axis=-1))
        irep = self.model.representation(
            tf.concat([self.data_init_obs_ph, initial_actions], axis=-1))

        def dot_k(rep1, rep2):
            # use dot kernel
            return tf.reduce_mean(tf.matmul(rep1, rep2, transpose_b=True), axis=0)
        mask = 1 - self.data_terminal_ph
        ipm = dot_k(rep, rep) + (1 - self.gamma) ** 2 * dot_k(irep, irep) \
            + self.gamma ** 2 * dot_k(nrep, nrep) * mask * tf.transpose(mask)\
            - 2 * (1 - self.gamma) * dot_k(irep, rep) \
            - 2 * self.gamma * dot_k(nrep, rep) * mask \
            + 2 * (1 - self.gamma) * self.gamma * dot_k(nrep, irep) * mask
        return tf.reduce_mean(ipm)

    def _pretrain_actors(self, replay_buffer, iterations):
        for _ in range(iterations):
            obs, act, *_ = replay_buffer.sample(self.batch_size)
            self.sess.run(self.behavior_clone_op, feed_dict={self.obs_ph: obs, self.action_ph: act})

    def train_step(self, replay_buffer, model_buffer):
        if self.model_batch_size == 0:
            env_batch = replay_buffer.sample(self.env_batch_size)
            obs, action, reward, next_obs, terminal = env_batch
        elif self.env_batch_size == 0:
            model_batch = model_buffer.sample(self.model_batch_size)
            obs, action, reward, next_obs, terminal = model_batch
        else:
            env_batch = replay_buffer.sample(self.env_batch_size)
            model_batch = model_buffer.sample(self.model_batch_size)
            obs, action, reward, next_obs, terminal = [np.concatenate([each1, each2], axis=0) \
                for each1, each2 in zip(env_batch, model_batch)]
        feed_dict = {self.obs_ph: obs, self.action_ph: action, self.reward_ph: reward,
                     self.next_obs_ph: next_obs, self.terminal_ph: terminal}
        # sample from obs instead of initial obs
        #feed_dict[self.initial_obs_ph] = replay_buffer.sample_obs(self.batch_size)
        obs, action, _, next_obs, terminal = replay_buffer.sample(self.batch_size)
        init_obs = replay_buffer.sample_obs(self.batch_size)
        feed_dict.update({self.data_obs_ph:obs, self.data_action_ph:action, \
            self.data_next_obs_ph: next_obs, self.data_terminal_ph: terminal, \
                self.data_init_obs_ph: init_obs})

        step_result = self.sess.run(self.step_ops, feed_dict=feed_dict)
        return step_result[-len(self.info_labels):]

    def batch_learn(self, buffer, vec_env, total_timesteps, result_filepath=None, model_path=None,
                    policy_path=None, **_):
        model_buffer = OnlineReplayBuffer(self.state_dim, self.action_dim, self.model_buffer_size)
        replay_buffer = buffer
        self.standardizer = replay_buffer.standardizer

        start_time = time.time()
        result = {'timestep': [], 'evals': [], 'infos': []}
        v_max, v_min = 10 * np.max(buffer.reward) / (1 - self.gamma), \
            10 * np.min(buffer.reward) / (1 - self.gamma)
        v_max, v_min = max(v_max, -v_min), min(-v_max, v_min)
        #self._pretrain_actors(replay_buffer, 30000)

        if self.real_ratio < 1.0: # train model
            if model_path is None:
                print('[ REPBRL ] Train Model')
                train_inputs, train_targets, terminals = replay_buffer.format_for_model_training()
                self.model.train(
                    train_inputs, train_targets, terminals, batch_size=self.batch_size,
                    max_epochs=None, max_t=self.max_model_t) #train models on standardized input outputs
            else:
                train_inputs, train_targets, terminals = replay_buffer.format_for_model_training()
                self.model.set_weights(np.load(model_path, allow_pickle=True)[()])
                self.model.set_model_inds(train_inputs, train_targets)
        infos = []
        for timestep in range(total_timesteps):
            if timestep % self.model_train_freq == 0 and self.real_ratio < 1.0:
                self._rollout_model(replay_buffer, model_buffer)

            infos.append(self.train_step(replay_buffer, model_buffer))

            if timestep % self.evaluate_freq == 0:
                current_info = np.mean(infos, axis=0)
                infos = []
                evaluation = evaluate_policy(vec_env, self)
                result['timestep'].append(timestep)
                result['evals'].append(evaluation)
                result['infos'].append(current_info)
                print(f't={timestep}: {evaluation} (elapsed_time={time.time() - start_time})')
                print('\n============================')
                for label, value in zip(self.info_labels, current_info):
                    print(f'{label:>12}: {value:>10.5f}')
                print('============================\n', flush=True)
                np.save(policy_path, self.actor.get_weights())
                #if current_info[2] > v_max or current_info[2] < v_min:
                    #print('Q diverged! terminating training')
                    #break

                if result_filepath:
                    np.save(result_filepath + '.tmp.npy', result)
        return result

    def _rollout_model(self, replay_buffer, model_buffer):
        print(f'[ Model Rollout ] Starting | '
              f'Rollout steps: {self.rollout_steps} | Batch size: {self.rollout_batch_size}')
        # get standardized obs from real samples
        obs = replay_buffer.sample_obs(self.rollout_batch_size)
        steps_added = []
        for i in range(self.rollout_steps):
            act = self.sess.run(self.cur_actions, feed_dict={self.obs_ph:obs})
            next_obs, rew, term, _ = self.fake_env.step(obs, act, replay_buffer)
            # fake env gives standardized obs
            steps_added.append(len(obs))
            model_buffer.add_samples(obs, act, rew, next_obs, term)
            # model_buffer contains standardized obs

            nonterm_mask = ~term.squeeze(-1)
            if nonterm_mask.sum() == 0:
                print(f'[ Model Rollout ] Breaking early: {i} | '
                      f'{nonterm_mask.sum()} / {nonterm_mask.shape}')
                break

            obs = next_obs[nonterm_mask]

        mean_rollout_length = sum(steps_added) / self.rollout_batch_size
        print(f'[ Model Rollout ] Added: {sum(steps_added):.1e} | Model pool: '
              f'{model_buffer.size:.1e} (max {model_buffer.buffer_size:.1e}) | '
              f'Length: {mean_rollout_length}')

    def predict(self, obs, deterministic=False):
        assert len(obs.shape) == 2
        if self.standardizer is not None:
            obs = self.standardizer(obs)
        if deterministic:
            return self.sess.run(self.det_actions, feed_dict={self.obs_ph:obs})
        else:
            return self.sess.run(self.cur_actions, feed_dict={self.obs_ph:obs})

    def _initial_exploration(self, train_env, replay_buffer):
        # initial exploration, 500 steps in case of pendulum
        obs = train_env.reset()
        for _ in range(5000):
            action = np.random.uniform(-1.0, 1.0, size=self.action_dim)
            next_obs, reward, done, _ = train_env.step(action)
            replay_buffer.add_sample(obs, action, [reward], next_obs, [float(done)])
            obs = next_obs
            if done:
                obs = train_env.reset()

    def learn(self, train_env, vec_env, total_epochs, result_filepath=None, epoch_length=1000,
              update_per_step=20):
         # epoch length: 250 in case of pendulum
         # update_per_step: depends on domain
        replay_buffer = OnlineReplayBuffer(self.state_dim, self.action_dim, int(1e6))
        model_buffer = OnlineReplayBuffer(self.state_dim, self.action_dim, self.model_buffer_size)
        self._initial_exploration(train_env, replay_buffer)
        episode_rewards = [0.0]

        start_time = time.time()
        result = {'timestep': [], 'evals': [], 'infos': []}
        obs = train_env.reset()
        for epoch in range(total_epochs):
            infos_values = []
            for timestep in range(epoch_length):
                if timestep % self.model_train_freq == 0 and self.real_ratio < 1.0:
                    print(f'[ REPBRL ] Freq {self.model_train_freq} | '
                          f'timestep {timestep}/{epoch_length}')
                    train_inputs, train_targets = replay_buffer.format_for_model_training()
                    self.model.train(
                        train_inputs, train_targets, batch_size=self.batch_size, holdout_ratio=0.2,
                        max_t=self.max_model_t)
                    self._rollout_model(replay_buffer, model_buffer)
                # Take an action and store transition in the replay buffer.
                action = self.predict(np.array([obs]), deterministic=False)[0].flatten()
                next_obs, reward, done, _ = train_env.step(action)
                replay_buffer.add_sample(obs, action, [reward], next_obs, [float(done)])
                obs = next_obs
                episode_rewards[-1] += reward
                if done:
                    obs = train_env.reset()
                    episode_rewards.append(0.0)
                if replay_buffer.can_sample(self.env_batch_size):
                    infos_values.append(self._sac_train(
                        replay_buffer, model_buffer, update_per_step, verbose=0))
            evaluation = evaluate_policy(vec_env, self, deterministic=True)
            infos_value = np.mean(infos_values, axis=0)
            result['timestep'].append(epoch * epoch_length)
            result['evals'].append(evaluation)
            result['infos'].append(infos_value)
            print(f't={epoch * epoch_length}: {evaluation} '
                  f'(elapsed_time={time.time() - start_time})')
            print('\n============================')
            for label, value in zip(self.info_labels, infos_value):
                print(f'{label:>12}: {value:>10.3f}')
            print('============================\n')

            if result_filepath:
                np.save(result_filepath + '.tmp.npy', result)

        return episode_rewards


import serial
import enum
from crccheck.crc import Crc16
import os
import pickle
import sklearn
import numpy as np
import zlib
from subprocess import check_output, CalledProcessError
import pandas as pd
import ast
import random
import surprise

STARTING_BYTE = 0x01

# rules_df = pd.read_csv(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "projects", "SmartGroceries", "data", 'rules_df.csv'))
# rules_df.Antecedent = rules_df.Antecedent.apply(lambda x: ast.literal_eval(x))
# rules_df.Consequents = rules_df.Consequents.apply(lambda x: ast.literal_eval(x))
pickls_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "projects", "SmartGroceries", "pickles")
rules_df_categorized = pickle.load(open(os.path.join(pickls_path, 'rules_df_categorized.pkl'), 'rb'))
print("rules_df_categorized loaded")
product_to_category = pickle.load(open(os.path.join(pickls_path, 'product_to_category.pkl'), 'rb'))
# collab_filtering_predictions = pickle.load(open(os.path.join(pickls_path, 'collab_filtering_predictions.pkl'), 'rb'))
print("product_to_category loaded")
users_categorized_top_items = pickle.load(open(os.path.join(pickls_path, 'reduced_users_categorized_top_items.pkl'), 'rb'))
print("reduced_users_categorized_top_items loaded")
top_100_items = [24852, 13176, 21137, 21903, 47209, 47766, 47626, 16797, 26209, 27845, 27966, 22935, 24964, 45007, 39275, 49683, 28204, 5876, 8277, 40706, 4920, 30391, 45066, 42265, 49235, 44632, 19057, 4605, 37646, 21616, 17794, 27104, 30489, 31717, 27086, 44359, 28985, 46979, 8518, 41950, 26604, 5077, 34126, 22035, 39877, 35951, 43352, 10749, 19660, 9076, 21938, 43961, 24184, 34969, 46667, 48679, 25890, 31506, 12341, 39928, 24838, 5450, 22825, 5785, 35221, 28842, 33731, 27521, 44142, 33198, 8174, 20114, 8424, 27344, 11520, 29487, 18465, 28199, 15290, 46906, 9839, 27156, 3957, 43122, 23909, 34358, 4799, 9387, 16759, 196, 42736, 38689, 4210, 41787, 41220, 47144, 7781, 33000, 20995, 21709]

class SerialState(enum.Enum):
    WaitingToStart = 0
    WaitingForString = 1
    WaitingForChecksum1 = 2
    WaitingForChecksum2 = 3
    CommandDone = 4 #Command Done -> Take action, send ACK and Go to 0

class ExecState(enum.Enum):
    NotConnected = 0
    Connected = 1 #Waiting to Select The Lab
    LabSelected = 2 #Ready to load model
    ModelLoaded = 3 #Ready to start processing
    Processing = 4
    Done = 5

class LabCode(enum.Enum):
    LabTest = 0
    Lab1 = 1
    Lab2 = 2
    Lab3 = 3
    Lab4 = 4

class PiExecuter():
    def __init__(self, serialPort):
        self.port = serialPort
        self.execState = ExecState.Connected
        self.serialState = SerialState.WaitingToStart

        self._currentLab = ""
        self._currentModelPath = ""

        self._currentSerialString = ""
        self._currentCheckSum = bytearray(2)

    def readSerial(self):
        if self.serialState == SerialState.WaitingToStart:
            newByte = self.port.read()
            if len(newByte) == 1:
                if newByte[0] == STARTING_BYTE:
                    self.serialState = SerialState.WaitingForString

        if self.serialState == SerialState.WaitingForString:
            newBytes = self.port.read_until(b'\0')
            if len(newBytes) >= 1:
                for i in range (len(newBytes)):
                    if newBytes[i] == STARTING_BYTE:
                        pass
                    else:
                        self._currentSerialString = self._currentSerialString + newBytes[i:].decode("utf-8")
                        if newBytes[-1] == 0x00:
                            self.serialState = SerialState.WaitingForChecksum1
                        break
            print(len(self._currentSerialString))

        if self.serialState == SerialState.WaitingForChecksum1:
            newByte = self.port.read()
            if len(newByte) == 1:
                self._currentCheckSum[0] = newByte[0]
                self.serialState = SerialState.WaitingForChecksum2

        if self.serialState == SerialState.WaitingForChecksum2:
            newByte = self.port.read()
            if len(newByte) == 1:
                self._currentCheckSum[1] = newByte[0]
                self.serialState = SerialState.CommandDone

        if self.serialState == SerialState.CommandDone:
            # check the command integrity
            receivedCommandCrc = Crc16()
            receivedCommandCrc.process(self._currentSerialString.encode('utf-8'))
            receivedCommandCrcBytes = receivedCommandCrc.finalbytes()
            print("Checksum Calc based on {}".format(self._currentSerialString.encode('utf-8')))
            print("Checksum Received: {}, Calculated: {}".format(self._currentCheckSum, receivedCommandCrcBytes))
            if receivedCommandCrcBytes == self._currentCheckSum:
                self.processSerialCommand(self._currentSerialString)

            self._currentSerialString = ""
            self._currentCheckSum[0] = 0x00
            self._currentCheckSum[1] = 0x00

            self.serialState = SerialState.WaitingToStart

    def processSerialCommand(self, commandStr):
        ackPayload = "None"

        (command, payload) = commandStr.rstrip('\t\r\n\0').split(':')
        print("Received command: {}, with payload: {} in State {}".format(command, payload, self.execState))
        if command == "RESET":
            self.execState = ExecState.Connected
            self._currentSerialString = ""
            self._currentCheckSum[0] = 0x00
            self._currentCheckSum[1] = 0x00
            self.serialState = SerialState.WaitingToStart

        elif command == "GET_IP":
            ackPayload = check_output(['hostname','-I']).decode('utf-8').split(" ")[0]

        elif command == "UPDATE_SCRIPT":
            try:
                ackPayload = check_output(['git','pull'], cwd='/home/pi/SFU_ML').decode('utf-8')
            except CalledProcessError as e:
                ackPayload = "FAILED: " + e.output()
        elif command == "PROCESS_PROJECT_GROUP_2":
            ackPayload = self.processProjectGroup2(payload)

        elif self.execState == ExecState.NotConnected:
            raise Exception("Wrong Exec State reached somehow: {}".format(self.execState))

        elif self.execState == ExecState.Connected:
            if command != "SELECT_LAB":
                raise Exception("You need to select the lab in the current state of: {}".format(self.execState))
            else:
                self._currentLab = LabCode[payload]
                self.execState = ExecState.LabSelected

        elif self.execState == ExecState.LabSelected:
            if not command in ["SAVE_MODEL", "LOAD_MODEL"]:
                raise Exception("You need to send or load a default model in the current state of: {}".format(self.execState))
            else:
                if self._currentLab != LabCode.LabTest:
                    self._currentModelPath = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'saved_models')
                    if command == "SAVE_MODEL":
                        savedModelStrings = payload.split(" ")
                        savedModel = bytearray([int(s) for s in savedModelStrings])
                        savedModel = zlib.decompress(savedModel)
                        # print(savedModel)
                        self._currentModelPath = os.path.join(self._currentModelPath, "lastModelSaved")
                        with open(self._currentModelPath, "wb") as modelFile:
                            modelFile.write(bytearray(savedModel))
                    elif command == "LOAD_MODEL":
                        self._currentModelPath = os.path.join(self._currentModelPath, payload)

                    try:
                        self._loadedModel = pickle.load(open(self._currentModelPath, 'rb'))
                        print("Loaded the Model File successfully")
                    except Exception as e:
                        raise Exception("Problem opening the current model with path {} because of exception: {}".format(str(self._currentModelPath), str(e)))

            self.execState = ExecState.ModelLoaded
            self.execState = ExecState.Processing

        elif self.execState == ExecState.ModelLoaded:
            #for any preprocessing required
            pass #NotImplemented

        elif self.execState == ExecState.Processing:
            if not command in ["PROCESS", "PROCESSING_DONE"]:
                raise Exception("You need to send a PROCESS or PROCESSING_DONE command in the current state of: {}".format(self.execState))
            else:
                if command == "PROCESS":
                    if self._currentLab == LabCode.LabTest:
                        ackPayload = self.processLabTest(payload)
                    elif self._currentLab == LabCode.Lab1:
                        ackPayload = self.processLab1andLab2(payload)
                    elif self._currentLab == LabCode.Lab2:
                        ackPayload = self.processLab1andLab2(payload)
                    else:
                        raise Exception("The lab selected is not a valid lab. This is likely a software implementation issue")
                elif command == "PROCESSING_DONE":
                    self.execState = ExecState.Connected

        self.sendSerialAck(ackPayload)

    def processLabTest(self, payload):
        inputs = [float(i) for i in payload.split(',')]
        outputs = [0, 0]
        outputs[0] = sum(inputs)
        outputs[1] = 1
        for input in inputs:
            outputs[1] *= input
        strOutputs = [str(o) for o in outputs]
        outputPayload = ', '.join(strOutputs)
        print("The Acknowledgment Payload is:"+outputPayload)
        return outputPayload

    def processLab1andLab2(self, payload):
        # print("Went into process lab 1 or Lab 2")
        input_list = [float(i) for i in payload.split(',')]
        y_pred = self._loadedModel.predict(np.array(input_list).reshape(1, -1)).astype('float32')
        outputPayload = ', '.join([str(i) for i in list(y_pred.flatten())])
        print("The Acknowledgment Payload is:"+outputPayload)
        if outputPayload == "":
            raise Exception("Problem with getting a prediction for the input data, \
please check the provided model and restart both the application and Raspberry Pi")
        return outputPayload

    def processProjectGroup2(self, payload):
        user_id = int(payload.split(";")[0])
        input_items = [ int(id.strip()) for id in payload.split(";")[1].split(",")]
        try:
            try_to_execlude = [ int(id.strip()) for id in payload.split(";")[2].split(",")]
        except:
            try_to_execlude = []

        def association_rules_algo(current_items, k=5, rules_df=None, no_popular_items=False):
            # print(current_items)
            if rules_df is None:
                raise RuntimeError("Rules Data Frame needs to be provided (acts as a trained model)")
            possible_predictions_df = rules_df[rules_df['Antecedent'].map(lambda x: all(id in current_items for id in x))]
            possible_predictions_df = possible_predictions_df.sort_values('Confidence', ascending=False)
            possible_predictions = []
            for row in possible_predictions_df['Consequents']:
                [possible_predictions.append(x) for x in row if not (x in possible_predictions or x in current_items)]
            print(possible_predictions)
            if no_popular_items:
                possible_predictions = [x for x in possible_predictions if (x not in top_100_items)]
            print(possible_predictions)
            filtered_predictions = possible_predictions.copy()
            possible_predictions.reverse()
            for item in possible_predictions:
                if (item in try_to_execlude) and len(filtered_predictions) > k:
                    filtered_predictions.remove(item)
            print(filtered_predictions)
            return filtered_predictions[:min(k, len(filtered_predictions))]

        def association_rules_categories_algo(current_items, k=5, rules_df=None, user_id = None, no_popular_items=False):
            if rules_df is None:
                raise RuntimeError("Rules Data Frame needs to be provided (acts as a trained model)")
            if user_id is None:
                raise RuntimeError("User ID needs to be provided for this algorithm to work")
            current_categories = set([product_to_category[id] for id in current_items if id in product_to_category.keys()])
            possible_predictions_df = rules_df[rules_df['Antecedent'].map(lambda x: all(category in current_categories for category in x))]
            possible_predictions_df = possible_predictions_df.sort_values('Confidence', ascending=False)
            possible_predictions = []
            for row in possible_predictions_df['Consequents']:
                [possible_predictions.append(x) for x in row if not (x in possible_predictions or x in current_categories)]
            final_predictions = []
            for category in possible_predictions:
                if category in users_categorized_top_items[user_id].keys():
                    for item, count in users_categorized_top_items[user_id][category]:
                        if not no_popular_items or item not in top_100_items:
                            # if item not in try_to_execlude :
                            final_predictions.append(item)
                            break
                    # if len(final_predictions) == k:
                    #     break
            filtered_predictions = final_predictions.copy()
            final_predictions.reverse()
            for item in final_predictions:
                if (item in try_to_execlude) and len(filtered_predictions) > k:
                    filtered_predictions.remove(item)
            return filtered_predictions[:min(k, len(filtered_predictions))]

        def get_predictions_for_user(predictions, uid):
            top_items = []
            for uid, iid, _, est, _ in predictions:
                if uid == user_id:
                    top_items.append((iid, est))
            top_items.sort(key=lambda x: x[1], reverse=True)
            return top_items

        def collab_filtering_algo(cart_current_items, cf_predictions_tuple=None, k=5, user_id=None):
            if user_id is None:
                raise RuntimeError("User ID needs to be provided for this algorithm to work")
            if cf_predictions_tuple is None:
                raise RuntimeError("Provide the saved predictions in the Surprise Predictions format")
            possible_recommendations = [iid for (iid, _) in get_predictions_for_user(cf_predictions_tuple, user_id) if iid not in (cart_current_items)] #try_to_execlude
            return possible_recommendations[:min(k, len(possible_recommendations))]

        def hybrid_algo(cart_current_items, k=5, cf_predictions_tuple=None, user_id=None, rules_df=None, no_popular_items=True):
            if user_id is None or cf_predictions_tuple is None or rules_df is None:
                raise RuntimeError("Please provide all needed keyword arguments user_id, rules_df, and cf_predictions_tuple")
            s = random.randrange(1, k)
            assoc_rules_predictions = association_rules_categories_algo(cart_current_items, k=k-s, rules_df=rules_df, user_id=user_id, no_popular_items=True)
            m = k-len(assoc_rules_predictions)
            collab_filtering_predictions = collab_filtering_algo(cart_current_items, cf_predictions_tuple=cf_predictions_tuple, k=m,user_id=user_id)
            print("assoc", len(assoc_rules_predictions), "collab", len(collab_filtering_predictions))
            return assoc_rules_predictions+collab_filtering_predictions

        # recommendations = hybrid_algo(input_items, user_id=user_id, k=5, rules_df=rules_df_categorized, cf_predictions_tuple=collab_filtering_predictions)
        recommendations = association_rules_categories_algo(input_items, user_id=user_id, k=5, rules_df=rules_df_categorized, no_popular_items=True)
        print(recommendations)

        return ",".join([str(x) for x in recommendations])

    def sendSerialAck(self, result=None):
        outBuffer = bytearray()
        outBuffer.append(STARTING_BYTE)
        if result == None:
            result = "None"
        ackBytes = ("ACK:"+ result).encode("utf-8")
        outBuffer.extend(ackBytes)
        outBuffer.append(0x00)
        newChecksum = Crc16()
        newChecksum.process(outBuffer[1:])
        checksumBytes = newChecksum.finalbytes()
        outBuffer.extend(checksumBytes)
        self.port.write(outBuffer)
        # print("sent Ack: {}".format(outBuffer))

    @property
    def execState(self):
        return self._execState

    @execState.setter
    def execState(self, newVal):
        print("Switched to Exec State: {}".format(newVal))
        self._execState = newVal

    @property
    def serialState(self):
        return self._serialState

    @serialState.setter
    def serialState(self, newVal):
        # print("Switched to Serial State: {}".format(newVal))
        self._serialState = newVal


import numpy as np
import cv2 as cv
import time
import RPi.GPIO as GPIO
import Adafruit_PCA9685
import threading
import importlib,sys
#导入必要的库文件
importlib.reload(sys)
servo_pwm = Adafruit_PCA9685.PCA9685()
#设置二维舵机对象
servo_pwm.set_pwm_freq(50)  # 设置频率
servo_pwm.set_pwm(4,0,325)  # 底座舵机
servo_pwm.set_pwm(9,0,325)  # 倾斜舵机
face_cascade = cv.CascadeClassifier('./haar/haarcascade_frontalface_default.xml')
#设置调用级联分级器对象并调用opencv自带的Haar文件
recognizer = cv.face.LBPHFaceRecognizer_create()
recognizer.read('./Train.yml')
#设置比对方法和训练文件
cap = cv.VideoCapture(0)
#初始化摄像头
usb_cap.set(3, 640)
usb_cap.set(4, 480)
#设置视频分辨率
pid_x=0
pid_y=0
pid_w=0
pid_h=0
#舵机云台的每个自由度需要4个变量
pid_thisError_x=0   #当前误差值
pid_lastError_x=0   #上次误差值
pid_thisError_y=0
pid_lastError_y=0
#舵机初始角度
pid_X_P = 330
pid_Y_P = 330
#下达舵机转动角度
pid_flag=0
size=0
#图片尺寸，用于判断大小
makerobo_facebool = False
#判断要不要让舵机执行转向
GPIO.setwarnings(False)
# 机器人舵机旋转
def Robot_servo():
    while True:
        servo_pwm.set_pwm(4,0,595-pid_X_P)
        servo_pwm.set_pwm(9,0,80+pid_Y_P)
#定义舵机初始位置
servo_tid=threading.Thread(target=Robot_servo)  # 多线程
servo_tid.setDaemon(True)
servo_tid.start()                               # 开启线程
while (usb_cap.isOpened()):                     #先看看摄像头开了没
    size=0
    ret,img = usb_cap.read()                    #获取当前帧
    face = np.zeros((640,640,3),np.uint8)       #统一用三通道方图处理图片
    face[80:560,0:640] = img[0:480,0:640]       #图片无损接入
    gray = cv.cvtColor(face, cv.COLOR_BGR2GRAY)#灰度处理
    faces = face_cascade.detectMultiScale(gray,1.3,3,0,(50,50),(400,400))#设置好合适的人脸参数
    if len(faces)==0:#faces容器内没有任何元素，没有检测到人任何一张脸
        face=np.rot90(face)                     #由于手机支架只有横竖两种工作模式，尝试将图片矩阵反转90°
        gray = cv.cvtColor(face, cv.COLOR_BGR2GRAY)  # 灰度处理
        faces = face_cascade.detectMultiScale(gray,1.3,5,0,(50,50),(400,400))
        if len(faces)==0:
            pid_x,pid_y,pid_w,pid_h = 320,240,0,0
            #如果横竖都没有检测到人脸就定位在中央，舵机不再转动
    if len(faces)!=0:
        idnum, confidence = recognizer.predict(gray[y:y + h, x:x + w])
        if confidence > 30:
            pid_x, pid_y, pid_w, pid_h = 320, 240, 0, 0
            #当前图片即使有人脸，但也不是我们需要追踪的对象
        else
            for x,y,w,h in faces:
                if w*h > size:
                    pid_x,pid_y,pid_w,pid_h = x,y,w,h
                    size = w*h
    #找到最大人脸坐标
    face=np.array(face)
    cv.rectangle(face,(pid_x,pid_y),(pid_x+pid_w,pid_y+pid_h),(0,255,0),1)#用矩形方框标记出人脸
    result=(pid_x,pid_y,pid_w,pid_h)
    pid_x=result[0]+pid_w/2
    pid_y=result[1]+pid_h/2
    makerobo_facebool = True

    # 误差值处理，根据舵机实际情况调整
    pid_thisError_x=pid_x-160
    pid_thisError_y=pid_y-120

    #自行对P和D两个值进行调整，以达到最稳定的追踪状态
    pwm_x = pid_thisError_x*5+1*(pid_thisError_x-pid_lastError_x)
    pwm_y = pid_thisError_y*5+1*(pid_thisError_y-pid_lastError_y)

    #迭代误差值操作
    pid_lastError_x = pid_thisError_x
    pid_lastError_y = pid_thisError_y

    #舵机拟需转动角度
    pid_XP=pwm_x/100
    pid_YP=pwm_y/100

    # pid_X_P pid_Y_P 为最终PID值
    pid_X_P=pid_X_P+int(pid_XP)
    pid_Y_P=pid_Y_P+int(pid_YP)

    #限值舵机在一定的范围之内
    if pid_X_P>650:
        pid_X_P=650
    if pid_X_P<0:
        pid_X_P=0
    if pid_Y_P>650:
        pid_Y_P=650
    if pid_Y_P<0:
        pid_Y_P=0
    cv.imshow("摄像头实时页面", face)
    if cv.waitKey(1)==119:
        break
print("\n [INFO] Exiting Program and cleanup stuff \n")
#释放资源
usb_cap.release()
cv.destroyAllWindows()


import sys
import torch
import numpy as np

from denoiser import Denoiser
import glow

_waveglow_path = sys.path[0]

class WaveglowWrapper:
    def __init__(self, model_path, device, sigma=0.666, strength=0.1):
        self.device = torch.device("cpu" if not torch.cuda.is_available() else device)
        self.dtype = torch.float if self.device.type == "cpu" else torch.half

        self.model = torch.load(model_path, map_location=self.device)["model"]
        self.model.device = self.device

        for m in self.model.modules():
            if "Conv" in str(type(m)):
                setattr(m, "padding_mode", "zeros")

        self.model.eval().to(device=self.device, dtype=self.dtype)

        for k in self.model.convinv:
            k.float()

        self.denoiser = Denoiser(self.model, device=self.device)

        self.sigma = sigma
        self.strength = strength

    def __call__(self, spectrogram):
        with torch.no_grad():
            audio = self.model.infer(spectrogram, self.sigma)

        return audio

    def denoise(self, audio):
        if type(audio) == np.ndarray:
            audio = torch.tensor(audio).to(self.device, self.dtype)

        if audio.ndim == 1:
            audio = audio.view(1, -1)
        audio = self.denoiser(audio, self.strength)[:, 0]

        return audio.data.cpu().numpy()

    @staticmethod
    def clear_cache():
        torch.cuda.empty_cache()


import os
import cv2
import pickle
import imutils
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.models import load_model

def load_history_file(history_filename):
    pickle_in = open(history_filename, 'rb')
    saved_history = pickle.load(pickle_in)
    return saved_history

def get_faces_coords(img, classifier):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    faces = classifier.detectMultiScale(gray, 1.05, 10)
    return faces

def get_prediction(face, model):
    img = cv2.resize(face, (128, 128), interpolation=cv2.INTER_CUBIC)
    return model.predict(np.expand_dims(img, axis=0))

def main():
    classes = ['mask_weared_incorrect', 'with_mask', 'without_mask']
    model = load_model('./mask_status.h5')
    cascade_classifier = cv2.CascadeClassifier('./cascade_classifier/haarcascade_frontalface_alt2.xml')

    for img in os.listdir('./Test/'):
        img_path = './Test/' + img
        img = cv2.imread(img_path)
        face_coords = get_faces_coords(img=img, classifier=cascade_classifier)
        for face in face_coords:
            x, y, w, h = face
            roi = img[y:y+h, x:x+w] / 255
            prediction = get_prediction(face=roi, model=model)
            prediction = classes[np.argmax(prediction)]
            if prediction == 'with_mask':
                color = (0, 255, 0)
            else:
                color = (0, 0, 255)
            cv2.rectangle(img, pt1=(x, y), pt2=(x+w, y+h), color=color, thickness=2)
            cv2.putText(img=img, text=prediction, org=(x, y), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=color, thickness=2)

        plt.imshow(imutils.opencv2matplotlib(img))
        plt.show()

if __name__ == '__main__':
    main()


#!/usr/bin/env python

import time
import numpy as np
import os
from softlearning.policies.safe_utils.logx import load_policy
from softlearning.policies.safe_utils.logx import EpochLogger
from mbpo.models import fake_env, bnn, constructor
import tensorflow as tf
import matplotlib.pyplot as plt

import time
import random

def run_policy(env, get_action, fpath, max_ep_len=None, num_episodes=10, render=True):

    assert env is not None, \
        "Environment not found!\n\n It looks like the environment wasn't saved, " + \
        "and we can't run the agent in it. :("

    params1= {
            'name':'BNN_10262',
            'load_model':True,
            'model_dir':'/home/mo/ray_mbpo/AntSafe/defaults/seed:2608_2020-10-30_02-55-39h74qa9rw/checkpoint_340/',
            'num_networks':7,
            'num_elites':5,
            'loss':'MSE',
            'use_scaler_in':True,
            'use_scaler_out':True,
            }
    model1 = bnn.BNN(params1)

    params2= {
            'name':'BNN_923',
            'load_model':True,
            'model_dir':'/home/mo/ray_mbpo/AntSafe/defaults/seed:2608_2020-10-30_02-55-39h74qa9rw/checkpoint_40/',
            'num_networks':7,
            'num_elites':5,
            'loss':'MSE',
            'use_scaler_in':True,
            'use_scaler_out':True,
            }
    model2 = bnn.BNN(params2)

    model2.finalize(tf.train.AdamOptimizer)
    model1.finalize(tf.train.AdamOptimizer)

    logger = EpochLogger()
    o, r, d, ep_ret, ep_cost, ep_len, n = env.reset(), 0, False, 0, 0, 0, 0
    xdata = []
    ydata1 = []
    ydata2 = []
    plt.show()

    axes = plt.gca()
    axes.set_xlim(0, 700)
    axes.set_ylim(0, 0.1)
    axes.axvline(x=150, ymin=0, ymax=10, color='black', linewidth=2.5)
    axes.axvline(x=350, ymin=0, ymax=10, color='black', linewidth=2.5)
    axes.set_title('Ensemble Disagreement')
    # axes.set_yscale('log')
    line1, = axes.plot(xdata, ydata1, color='green', label='Model trained under current policy') #'r-')
    line2, = axes.plot(xdata, ydata2, color='red', label='Model trained under older policy')#'r-')
    axes.legend(fontsize='x-large', loc='upper left')

    while n < num_episodes:
        if render:
            env.render(mode='human')
            # time.sleep(1e-3)
        # @TODO anyboby
        ### bad hardcoding
        o_main = o[:29][None]

        a_main = get_action(o_main)

        if ep_len==150 or ep_len==350:
            time.sleep(3)
            m_steps = 0
            m_obs1 = o_main
            m_obs2 = o_main
            while True:
                ep_len += 1
                m_steps += 1

                a_m1 = get_action(m_obs1)

                ######### model 1
                inputs1 = np.concatenate([m_obs1,a_m1], axis=-1)
                n_obs1 = model1.predict(inputs1, factored=True, inc_var=False)

                unc1 = np.mean(np.var(n_obs1, axis=0))

                n_obs1 = n_obs1[np.random.randint(7)]
                n_obs1[...,:-1] += m_obs1
                n_obs_qpos1 = n_obs1[...,:13]
                n_obs_qvel1 = n_obs1[...,13:27]
                n_obs_x1 = n_obs1[...,27:28]*5
                n_obs_y_off1 = n_obs1[...,28:29]

                if n_obs_x1<20:
                    n_obs_y1 = n_obs_y_off1 + n_obs_x1*np.tan(30/360*2*np.pi)
                elif n_obs_x1>20 and n_obs_x1<60:
                    n_obs_y1 = n_obs_y_off1 - (n_obs_x1-40)*np.tan(30/360*2*np.pi)
                elif n_obs_x1>60 and n_obs_x1<100:
                    n_obs_y1 = n_obs_y_off1 + (n_obs_x1-80)*np.tan(30/360*2*np.pi)
                else:
                    n_obs_y1 = n_obs_y_off1 + 20*np.tan(30/360*2*np.pi)

                ########## model 2
                a_m2 = get_action(m_obs2)
                inputs2 = np.concatenate([m_obs2,a_m2], axis=-1)

                n_obs2 = model2.predict(inputs2, factored=True, inc_var=False)

                unc2 = np.mean(np.var(n_obs2, axis=0))

                n_obs2 = n_obs2[np.random.randint(7)]
                n_obs2[...,:-1] += m_obs2
                n_obs_qpos2 = n_obs2[...,:13]
                n_obs_qvel2 = n_obs2[...,13:27]
                n_obs_x2 = n_obs2[...,27:28]*5
                n_obs_y_off2 = n_obs2[...,28:29]

                if n_obs_x2<20:
                    n_obs_y2 = n_obs_y_off2 + n_obs_x2*np.tan(30/360*2*np.pi)
                elif n_obs_x2>20 and n_obs_x2<60:
                    n_obs_y2 = n_obs_y_off2 - (n_obs_x2-40)*np.tan(30/360*2*np.pi)
                elif n_obs_x2>60 and n_obs_x2<100:
                    n_obs_y2 = n_obs_y_off2 + (n_obs_x2-80)*np.tan(30/360*2*np.pi)
                else:
                    n_obs_y2 = n_obs_y_off2 + 20*np.tan(30/360*2*np.pi)
                ###########

                xdata.append(ep_len)
                ydata1.append(unc1)
                ydata2.append(unc2)
                line1.set_xdata(xdata)
                line1.set_ydata(ydata1)
                line2.set_xdata(xdata)
                line2.set_ydata(ydata2)

                plt.draw()
                plt.pause(1e-17)

                a = get_action(o_main)

                o, _, _, _ = env.step(np.concatenate([a, a_m1, a_m2], axis=-1))
                o_main=o[:29][None]
                real_qpos = o_main[...,:13]
                real_qvel = o_main[...,13:27]
                real_x = o_main[...,27:28]*5
                real_y_off = o_main[...,28:29]
                if real_x<20:
                    real_y = real_y_off + real_x*np.tan(30/360*2*np.pi)
                elif real_x>20 and real_x<60:
                    real_y = real_y_off - (real_x-40)*np.tan(30/360*2*np.pi)
                elif real_x>60 and real_x<100:
                    real_y = real_y_off + (real_x-80)*np.tan(30/360*2*np.pi)
                else:
                    real_y = real_y_off + 20*np.tan(30/360*2*np.pi)

                qpos = np.concatenate([real_x, real_y, real_qpos, n_obs_x2, n_obs_y2, n_obs_qpos2, n_obs_x1, n_obs_y1, n_obs_qpos1], axis=-1)
                qvel = np.concatenate([real_qvel, n_obs_qvel2, n_obs_qvel1], axis=-1)

                init_qpos = env._env.env.init_qpos.copy()
                init_qpos[:-42] = qpos
                init_qvel = env._env.env.init_qvel.copy()
                init_qvel[:-36] = qvel

                env._env.env.set_state(qpos=init_qpos, qvel=init_qvel)
                env.render(mode='human')

                m_obs1 = n_obs1[...,:-1]
                m_obs2 = n_obs2[...,:-1]

                if m_steps > 198:
                    # time.sleep(3)
                    break

        pred1 = model1.predict(np.concatenate((o_main,a_main), axis=-1), factored=True, inc_var=False)
        unc1 = np.mean(np.var(pred1, axis=0))

        pred2 = model2.predict(np.concatenate((o_main,a_main), axis=-1), factored=True, inc_var=False)
        unc2 = np.mean(np.var(pred2, axis=0))

        o, r, d, info = env.step(np.concatenate([a_main, a_main, a_main], axis=-1))#a)
        ep_ret += r #r[0]
        ep_cost +=  info.get('cost',0) #info[0].get('cost', 0)
        ep_len += 1

        #### restore positions
        o_main=o[:29][None]
        real_qpos = o_main[...,:13]
        real_qvel = o_main[...,13:27]
        real_x = o_main[...,27:28]*5
        real_y_off = o_main[...,28:29]
        if real_x<20:
            real_y = real_y_off + real_x*np.tan(30/360*2*np.pi)
        elif real_x>20 and real_x<60:
            real_y = real_y_off - (real_x-40)*np.tan(30/360*2*np.pi)
        elif real_x>60 and real_x<100:
            real_y = real_y_off + (real_x-80)*np.tan(30/360*2*np.pi)
        else:
            real_y = real_y_off + 20*np.tan(30/360*2*np.pi)

        qpos = np.concatenate([real_x, real_y, real_qpos, real_x, real_y, real_qpos, real_x, real_y, real_qpos], axis=-1)
        qvel = np.concatenate([real_qvel, real_qvel, real_qvel], axis=-1)

        init_qpos = env._env.env.init_qpos.copy()
        init_qpos[:-42] = qpos
        init_qvel = env._env.env.init_qvel.copy()
        init_qvel[:-36] = qvel

        env._env.env.set_state(qpos=init_qpos, qvel=init_qvel)

        #### plot
        xdata.append(ep_len)
        ydata1.append(unc1)
        ydata2.append(unc2)
        line1.set_xdata(xdata)
        line1.set_ydata(ydata1)
        line2.set_xdata(xdata)
        line2.set_ydata(ydata2)
        plt.draw()
        plt.pause(1e-17)
        # time.sleep(0.1)

        if d or (ep_len >= 500): #(ep_len == max_ep_len):
            #logger.store(EpRet=ep_ret, EpCost=ep_cost, EpLen=ep_len)
            print('Episode %d \t EpRet %.3f \t EpCost %.3f \t EpLen %d'%(n, ep_ret, ep_cost, ep_len))
            o, r, d, ep_ret, ep_cost, ep_len = env.reset(), 0, False, 0, 0, 0
            n += 1
            xdata = []
            ydata1 = []
            ydata2 = []

    # logger.log_tabular('EpRet', with_min_and_max=True)
    # logger.log_tabular('EpCost', with_min_and_max=True)
    # logger.log_tabular('EpLen', average_only=True)
    # logger.dump_tabular(output_dir = os.path.join(fpath, '/Inference'))

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('fpath', type=str)
    parser.add_argument('--len', '-l', type=int, default=1000)
    parser.add_argument('--episodes', '-n', type=int, default=100)
    parser.add_argument('--norender', '-nr', action='store_true')
    parser.add_argument('--itr', '-i', type=int, default=-1)
    parser.add_argument('--deterministic', '-d', action='store_true')
    args = parser.parse_args()
    env, get_action, sess = load_policy(args.fpath,
                                        args.itr if args.itr >=0 else 'last',
                                        args.deterministic)
    run_policy(env, get_action, args.fpath, args.len, args.episodes, not(args.norender))


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

import re
import numpy as np

from nomad.parsing.file_parser import BasicParser
from nomad.units import ureg

class QboxParser(BasicParser):
    def __init__(self):
        re_f = r'\-*\d+\.\d+E*e*\-*\+*\d*'

        def get_positions(val):
            labels = re.findall(r'atom name\=\"([A-Z][a-z]*)\w*\"', val)
            positions = re.findall(rf'\<position\> *({re_f} +{re_f} +{re_f})', val)
            forces = re.findall(rf'\<force\> *({re_f} +{re_f} +{re_f})', val)
            try:
                positions = np.array([p.split() for p in positions], dtype=np.dtype(np.float64))
                forces = np.array([f.split() for f in forces], dtype=np.dtype(np.float64))
            except Exception:
                pass
            return dict(atom_labels=labels, atom_positions=positions, atom_forces=forces)

        super().__init__(
            specifications=dict(
                name='parsers/qbox', code_name='qbox', code_homepage='http://qboxcode.org/',
                domain='dft', mainfile_mime_re=r'(application/xml)|(text/.*)',
                mainfile_contents_re=(r'http://qboxcode.org')),
            units_mapping=dict(length=ureg.bohr, energy=ureg.hartree),
            program_version=r'I qbox ([\d\.]+)',
            lattice_vectors=rf'\<unit\_cell\s+(a\=[\s\S]+?)\/\>',
            atom_labels_atom_positions_atom_forces=(
                r'(\<atom name\=[\s\S]+?)\<\/atomset\>', get_positions),
            electronic_kinetic_energy=rf'\<ekin\> *({re_f}) *\<\/ekin\>',
            energy_total=rf'\<etotal\> *({re_f}) *\<\/etotal\>',
            energy_XC=rf'\<exc\> *({re_f}) *\<\/exc\>')


import cv2
import numpy as np
import winsound
import os

# Capturing frames
cap = cv2.VideoCapture(0)

while True:

    # Reading the captured frames-> frame1, frame2
    ret,frame1 = cap.read()
    ret,frame2 = cap.read()

    # Checking Difference in motion in 2 frames
    diff_frame = cv2.absdiff(frame1,frame2)

    # Converting BGR-> Gray and making frame blur
    gray = cv2.cvtColor(diff_frame,cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray,(5,5),0)
    _,thresh = cv2.threshold(blur,20,255,cv2.THRESH_BINARY)

    # Dilating the frame to make get uniform contour
    dilated = cv2.dilate(thresh,None,iterations=10)

    #  Finding the contours in the difference to get coordinates of motion
    contour,x = cv2.findContours(dilated,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)

    for c in contour:

        # If the contourArea < 1000 ignore the frame
        if cv2.contourArea(c) < 1000:
            continue

        # Drawing rectangle around the moving object
        x,y,w,h = cv2.boundingRect(c)
        cv2.rectangle(frame1,(x,y),(x+w,y+h),(0,255,0),2)
        cv2.imshow("The Frame",frame1)

    # Press 'q' to exit the frame
    if cv2.waitKey(1) & 0xff == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()


import multiprocessing  # 多线程模块
import threading  # 线程模块
import queue
import tensorflow as tf
import numpy as np
import gym
import os
import shutil  # 拷贝文件用
import matplotlib.pyplot as plt
from FeudalBatchProcessor import FeudalBatchProcessor
import policy_utils
from LSTMmodel import SingleStepLSTM

Game = 'CartPole-v0'
N_workers = multiprocessing.cpu_count()  # 独立玩家个体数为cpu数
# MAX_GLOBAL_EP = 2000  # 中央大脑最大回合数
MAX_GLOBALE_STEP = 100000  # 中央大脑最大步数
GLOBAL_NET_SCOPE = 'Global_Net'  # 中央大脑的名字
UPDATE_GLOBALE_ITER = 10  # 中央大脑每N次提升一次
GAMMA = 0.9  # 衰减度
LR_A = 0.0001  # Actor网络学习率
LR_C = 0.001  # Critic 网络学习率
beta_start = 0.01
beta_end = 0.001
decay_steps = 50000

GLOBALE_RUNNING_R = []  # 存储总的reward
# GLOBALE_EP = 0  # 中央大脑步数
GLOBALE_STEP = 0  # 中央大脑步数

env = gym.make(Game)  # 定义游戏环境

N_S = env.observation_space.shape[0]  # 观测值个数
N_A = env.action_space.n  # 行为值个数

class ACnet(object):  # 这个class即可用于生产global net，也可生成 worker net，因为结构相同
    def __init__(self, scope, globalAC=None, global_step=None):  # scope 用于确定生成什么网络
        # global GLOBALE_STEP
        # self.global_step = GLOBALE_STEP
        if scope == GLOBAL_NET_SCOPE:  # 创建中央大脑
            with tf.variable_scope(scope):
                self.global_step = tf.get_variable("global_step", [], tf.int32,
                                                   initializer=tf.constant_initializer(0, dtype=tf.int32),
                                                   trainable=False)
                self.obs_space = N_S
                self.act_space = N_A
                self.k = 16
                self.g_dim = 256
                self.c = 10
                self.vf_hidden_size = 128  # for value function network
                self.alpha = 0.5  # for build loss
                self.batch_processor = FeudalBatchProcessor(self.c)
                self.build_model()  # build feudal policy model

        else:  # 创建worker两个网络的具体步骤
            with tf.variable_scope(scope):  # 这里的scope传入的是worker的名字
                self.global_step = globalAC.global_step
                self.obs_space = N_S
                self.act_space = N_A
                self.k = 16
                self.g_dim = 256
                self.c = 10
                self.vf_hidden_size = 128  # for value function network
                self.alpha = 0.5  # for build loss
                self.batch_processor = FeudalBatchProcessor(self.c)
                self.build_model()  # build feudal policy model

            with tf.name_scope('local_grad'):
                grads = tf.gradients(self.loss, self.var_list)
                grads, _ = tf.clip_by_global_norm(grads, 40)

            with tf.name_scope('sync'):  # worker和global的同步过程
                with tf.name_scope('pull'):  # 获取global参数,复制到local—net
                    self.pull_params_op = tf.group(*[v1.assign(v2)
                                                     for v1, v2 in zip(self.var_list, globalAC.var_list)])
                with tf.name_scope('push'):  # 将参数传送到gloabl中去
                    self.update_params_op = OPT.apply_gradients(zip(grads, globalAC.var_list))
                    # 其中传送的是local—net的actor和critic的参数梯度grads,具体计算在上面定义
                    # apply_gradients是tf.train.Optimizer中自带的功能函数，将求得的梯度参数更新到global中
            self.inc_step = self.global_step.assign_add(tf.shape(self.obs)[0])
            self.train_op = tf.group(self.update_params_op, self.inc_step)
            # GLOBALE_STEP += tf.shape(self.obs)[0]

    def build_model(self):
        """
        Builds the manager and worker models.
        """
        with tf.variable_scope('FeUdal'):
            self.build_placeholders()
            self.build_perception()
            self.build_manager()
            self.build_worker()
            self.build_loss()
            self.var_list = tf.get_collection(
                tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name)
        # for v in self.var_list:
        #     print v.name

        self.state_in = [self.worker_lstm.state_in[0],
                         self.worker_lstm.state_in[1],
                         self.manager_lstm.state_in[0],
                         self.manager_lstm.state_in[1]
                         ]
        self.state_out = [self.worker_lstm.state_out[0],
                          self.worker_lstm.state_out[1],
                          self.manager_lstm.state_out[0],
                          self.manager_lstm.state_out[1]
                          ]
        # for v in self.var_list:
        #     print v

    def build_placeholders(self):
        # standard for all policies
        self.obs = tf.placeholder(tf.float32, [None,
                                               self.obs_space])  # ! self.obs = tf.placeholder(tf.float32, [None] + list(self.obs_space))
        # ! self.obs_space = env.observation_space.shape
        self.r = tf.placeholder(tf.float32, (None,1))
        self.ac = tf.placeholder(tf.float32, (None, self.act_space))
        self.adv = tf.placeholder(tf.float32, [None])  # unused

        # specific to FeUdal
        self.prev_g = tf.placeholder(tf.float32, (None, None, self.g_dim))
        self.ri = tf.placeholder(tf.float32, (None,))
        self.s_diff = tf.placeholder(tf.float32, (None, self.g_dim))

    def build_perception(self):
        self._obs = tf.expand_dims(self.obs, -1)  # !
        self._obs = tf.expand_dims(self._obs, -1)  # !
        conv1 = tf.layers.conv2d(inputs=self._obs,
                                 filters=16,
                                 kernel_size=[2, 1],  # ! kernel_size = [8,8]
                                 activation=tf.nn.elu,
                                 strides=1)  # ! strides = 4
        conv2 = tf.layers.conv2d(inputs=conv1,
                                 filters=32,
                                 kernel_size=[2, 1],  # ! kernel_size = [4,4]
                                 activation=tf.nn.elu,
                                 strides=1)  # ! strides = 2

        flattened_filters = policy_utils.flatten(conv2)
        self.z = tf.layers.dense(inputs=flattened_filters,
                                 units=256,
                                 activation=tf.nn.elu)

    def build_manager(self):
        with tf.variable_scope('manager'):
            # Calculate manager internal state
            self.s = tf.layers.dense(inputs=self.z,
                                     units=self.g_dim,
                                     activation=tf.nn.elu)

            # Calculate manager output g
            x = tf.expand_dims(self.s, [0])
            self.manager_lstm = SingleStepLSTM(x,
                                               self.g_dim,
                                               step_size=tf.shape(self.obs)[:1])
            g_hat = self.manager_lstm.output
            self.g = tf.nn.l2_normalize(g_hat, dim=1)

            self.manager_vf = self.build_value(g_hat)

    def build_worker(self):
        with tf.variable_scope('worker'):
            num_acts = self.act_space

            # Calculate U
            self.worker_lstm = SingleStepLSTM(tf.expand_dims(self.z, [0]),
                                              size=num_acts * self.k,
                                              step_size=tf.shape(self.obs)[:1])
            flat_logits = self.worker_lstm.output

            self.worker_vf = self.build_value(flat_logits)

            U = tf.reshape(flat_logits, [-1, num_acts, self.k])

            # Calculate w
            cut_g = tf.stop_gradient(self.g)
            cut_g = tf.expand_dims(cut_g, [1])
            gstack = tf.concat([self.prev_g, cut_g], axis=1)

            self.last_c_g = gstack[:, 1:]
            # print self.last_c_g
            gsum = tf.reduce_sum(gstack, axis=1)
            phi = tf.get_variable("phi", (self.g_dim, self.k))
            w = tf.matmul(gsum, phi)
            w = tf.expand_dims(w, [2])
            # Calculate policy and sample
            logits = tf.reshape(tf.matmul(U, w), [-1, num_acts])
            self.pi = tf.nn.softmax(logits)
            self.log_pi = tf.nn.log_softmax(logits)
            self.sample = policy_utils.categorical_sample(
                tf.reshape(logits, [-1, num_acts]), num_acts)[0, :]

    def build_value(self, _input):
        with tf.variable_scope('VF'):
            hidden = tf.layers.dense(inputs=_input,
                                     units=self.vf_hidden_size,
                                     activation=tf.nn.elu)

            w = tf.get_variable("weights", (self.vf_hidden_size, 1))
            return tf.matmul(hidden, w)

    def build_loss(self):
        cutoff_vf_manager = tf.reshape(tf.stop_gradient(self.manager_vf), [-1])
        dot = tf.reduce_sum(tf.multiply(self.s_diff, self.g), axis=1)
        gcut = tf.stop_gradient(self.g)
        mag = tf.norm(self.s_diff, axis=1) * tf.norm(gcut, axis=1) + .0001
        dcos = dot / mag
        manager_loss = -tf.reduce_sum((self.r - cutoff_vf_manager) * dcos)

        cutoff_vf_worker = tf.reshape(tf.stop_gradient(self.worker_vf), [-1])
        log_p = tf.reduce_sum(self.log_pi * self.ac, [1])
        worker_loss = (self.r + self.alpha * self.ri - cutoff_vf_worker) * log_p
        worker_loss = -tf.reduce_sum(worker_loss, axis=0)

        Am = self.r - self.manager_vf
        manager_vf_loss = .5 * tf.reduce_sum(tf.square(Am))

        Aw = (self.r + self.alpha * self.ri) - self.worker_vf
        worker_vf_loss = .5 * tf.reduce_sum(tf.square(Aw))

        entropy = -tf.reduce_sum(self.pi * self.log_pi)

        beta = tf.train.polynomial_decay(beta_start, self.global_step,
                                         end_learning_rate=beta_end,
                                         decay_steps=decay_steps,
                                         power=1)

        # worker_loss = tf.Print(worker_loss,[manager_loss,worker_loss,manager_vf_loss,worker_vf_loss,entropy])
        self.loss = worker_loss + manager_loss + \
                    worker_vf_loss + manager_vf_loss - \
                    entropy * beta

    def update_global(self, feed_dict):  # 定义更新global参数函数
        SESS.run([self.update_params_op], feed_dict)  # 分别更新actor和critic网络

    def pull_global(self):  # 定义更新local参数函数
        SESS.run([self.pull_params_op])

    def action(self, ob, g, cw, hw, cm, hm):  # 定义选择动作函数
        # ob = ob[np.newaxis, :]
        ob = ob.reshape([-1, self.obs_space])
        return SESS.run([self.sample, self.manager_vf, self.g, self.s, self.last_c_g] + self.state_out,
                        feed_dict={self.obs: ob, self.state_in[0]: cw, self.state_in[1]: hw, self.state_in[2]: cm,
                                   self.state_in[3]: hm, self.prev_g: g})
        # return np.random.choice(range(probs.shape[1]), p=probs.ravel())  # 从probs中按概率选取出某一个动作

    def value(self, ob, g, cw, hw, cm, hm):
        sess = tf.get_default_session()
        return sess.run(self.manager_vf,
                        {self.obs: [ob], self.state_in[0]: cw, self.state_in[1]: hw,
                         self.state_in[2]: cm, self.state_in[3]: hm,
                         self.prev_g: g})[0]

    def get_initial_features(self):
        return np.zeros((1, 1, self.g_dim), np.float32), self.worker_lstm.state_init + self.manager_lstm.state_init

    def update_batch(self, batch):
        return self.batch_processor.process_batch(batch)

class Worker(object):
    def __init__(self, name, globalAC):  # 传入的name是worker的名字，globalAC是已经建立好的中央大脑GLOBALE—AC
        self.env = gym.make(Game).unwrapped
        self.name = name  # worker的名字
        self.global_AC = globalAC
        self.local_AC = ACnet(scope=name, globalAC=globalAC)  # 第二个参数当传入的是已经建立好的GLOBALE—AC时创建的是local net
        # 建立worker的AC网络
        self.runner = policy_utils.RunnerThread(self.env, self.local_AC, 20, visualise=0)

    def pull_batch_from_queue(self):
        """
        self explanatory:  take a rollout from the queue of the thread runner.
        """
        rollout = self.runner.queue.get(timeout=600.0)
        while not rollout.terminal:
            try:
                rollout.extend(self.runner.queue.get_nowait())
            except queue.Empty:
                break
        return rollout

    def start(self, sess, summary_writer):
        self.runner.start_runner(sess, summary_writer)

    def work(self):  # 定义worker运行的的具体过程
        global GLOBALE_STEP, MAX_GLOBALE_STEP
        # global GLOBALE_RUNNING_R, GLOBALE_EP  # 两个全局变量，R是所有worker的总reward，ep是所有worker的总episode
        # total_step = 1  # 本worker的总步数
        # buffer_s, buffer_a, buffer_r = [], [], []  # state,action,reward的缓存
        SESS.run(self.local_AC.pull_params_op)
        self.start(SESS, summary_writer=0)
        global_step = SESS.run(self.global_AC.global_step)
        # print(type(GLOBALE_STEP < MAX_GLOBALE_STEP))
        while not COORD.should_stop() and global_step < MAX_GLOBALE_STEP:  # 停止本worker运行的条件
            # 本循环一次是一个回合

            # s = self.env.reset()  # 初始化环境
            if self.name == 'W_0':  # 只有worker0才将动画图像显示
                self.env.render()
            ep_r = 0  # 本回合总的reward
            SESS.run(self.local_AC.pull_params_op)
            rollout = self.pull_batch_from_queue()
            batch = policy_utils.process_rollout(rollout, gamma=.99)
            batch = self.local_AC.update_batch(batch)
            # batch.ri = [item for sublist in batch.ri for item in sublist]
            # returns = [item for sublist in batch.returns for item in sublist]
            # batch._replace(returns=returns)
            # print("batch.returns.shape:",batch.returns.shape)
            # print("batch.ri.shape:",batch.ri.le)
            fetches = [self.local_AC.train_op]
            feed_dict = {
                self.local_AC.obs: batch.obs,
                self.global_AC.obs: batch.obs,

                self.local_AC.ac: batch.a,
                self.global_AC.ac: batch.a,

                self.local_AC.r: batch.returns,
                self.global_AC.r: batch.returns,

                self.local_AC.s_diff: batch.s_diff,
                self.global_AC.s_diff: batch.s_diff,

                self.local_AC.prev_g: batch.gsum,
                self.global_AC.prev_g: batch.gsum,

                self.local_AC.ri: batch.ri,
                self.global_AC.ri: batch.ri
            }

            for i in range(len(self.local_AC.state_in)):
                feed_dict[self.local_AC.state_in[i]] = batch.features[i]
                feed_dict[self.global_AC.state_in[i]] = batch.features[i]

            fetched = SESS.run(fetches, feed_dict=feed_dict)
            # while True:  # 本循环一次是一步
            #     if self.name == 'W_0':  # 只有worker0才将动画图像显示
            #         self.env.render()
            #
            #     fetched = self.AC.action(last_state, *last_features)  # 将当前状态state传入AC网络选择动作action
            #     action, value_, g, s, last_c_g, features = fetched[0], fetched[1], \
            #                                                fetched[2], fetched[3], \
            #                                                fetched[4], fetched[5:]
            #     a = action.argmax()
            #     state, reward, done, info = self.env.step(a)  # 行动并获得新的状态和回报等信息
            #     rollout.add(last_state,action,reward,value_,g,s,done,last_features)
            #
            #     if done: reward = -5  # 如果结束了，reward给一个惩罚数
            #
            #     ep_r += reward  # 记录本回合总体reward
            #     # buffer_s.append(s)  # 将当前状态，行动和回报加入缓存
            #     # buffer_a.append(a)
            #     # buffer_r.append(r)
            #     last_state = state
            #     last_features = features
            #     if total_step % UPDATE_GLOBALE_ITER == 0 or done:  # 每iter步完了或者或者到达终点了，进行同步sync操作
            #         if done:
            #             v_s_ = 0  # 如果结束了，设定对未来的评价值为0
            #         else:
            #             v_s_ = SESS.run(self.AC.v, feed_dict={self.AC.s: s_[np.newaxis, :]})[
            #                 0, 0]  # 如果是中间步骤，则用AC网络分析下一个state的v评价
            #
            #         buffer_v_target = []
            #         for r in buffer_r[::-1]:  # 将下一个state的v评价进行一个反向衰减传递得到每一步的v现实
            #             v_s_ = r + GAMMA * v_s_
            #             buffer_v_target.append(v_s_)  # 将每一步的v现实都加入缓存中
            #         buffer_v_target.reverse()  # 反向后，得到本系列操作每一步的v现实(v-target)
            #
            #         buffer_s, buffer_a, buffer_v_target = np.vstack(buffer_s), np.vstack(buffer_a), np.vstack(
            #             buffer_v_target)
            #
            #         feed_dict = {
            #             self.AC.obs: buffer_s,  # 本次走过的所有状态，用于计算v估计
            #             self.AC.ac: buffer_a,  # 本次进行过的所有操作，用于计算a—loss
            #             self.AC.v: buffer_v_target  # 走过的每一个state的v现实值，用于计算td
            #         }
            #
            #         self.AC.update_global(feed_dict)  # update—global的具体过程在AC类中定义，feed-dict如上
            #
            #         buffer_s, buffer_a, buffer_r = [], [], []  # 清空缓存
            #
            #         self.AC.pull_global()  # 从global—net提取出参数赋值给local—net
            #
            #     s = s_  # 跳转到下一个状态
            #     total_step += 1  # 本回合总步数加1
            #
            #     if done:  # 如果本回合结束了
            #         if len(GLOBALE_RUNNING_R) == 0:  # 如果尚未记录总体running
            #             GLOBALE_RUNNING_R.append(ep_r)
            #         else:
            #             GLOBALE_RUNNING_R.append(0.9 * GLOBALE_RUNNING_R[-1] + 0.1 * ep_r)
            #
            #         print(self.name, 'EP:', GLOBALE_EP)
            #         GLOBALE_EP += 1  # 加一回合
            #         break  # 结束本回合

            # global_step = SESS.run(self.global_AC.global_step)

if __name__ == '__main__':
    SESS = tf.Session()

    with tf.device('/cpu:0'):
        OPT = tf.train.AdamOptimizer(1e-4)  # 后续主要是使用该optimizer中的apply—gradients操作
        # OPT_C = tf.train.RMSPropOptimizer(LR_C, name='RMSPropC')  # 定义critic训练过程
        GLOBAL_AC = ACnet(scope=GLOBAL_NET_SCOPE)  # 创建中央大脑GLOBALE_AC，只创建结构(A和C的参数)
        workers = []
        for i in range(N_workers):  # N—workers等于cpu数量
            i_name = 'W_%i' % i  # worker name
            workers.append(Worker(name=i_name, globalAC=GLOBAL_AC))  # 创建独立的worker

        COORD = tf.train.Coordinator()  # 多线程
        SESS.run(tf.global_variables_initializer())  # 初始化所有参数

        worker_threads = []
        for worker in workers:  # 并行过程
            job = lambda: worker.work()  # worker的工作目标,此处调用Worker类中的work
            t = threading.Thread(target=job)  # 每一个线程完成一个worker的工作目标
            t.start()  # 启动每一个worker
            worker_threads.append(t)  # 每一个worker的工作都加入thread中
        COORD.join(worker_threads)  # 合并几个worker,当每一个worker都运行完再继续后面步骤

        plt.plot(np.arange(len(GLOBALE_RUNNING_R)), GLOBALE_RUNNING_R)  # 绘制reward图像
        plt.xlabel('step')
        plt.ylabel('Total moving reward')
        plt.show()