import warnings
from itertools import zip_longest

from abc import abstractmethod
import numpy as np
import tensorflow as tf
from stable_baselines.a2c.utils import conv, linear, conv_to_fc, batch_to_seq, seq_to_batch, lstm
from stable_baselines.common.policies import BasePolicy
from stable_baselines.common.input import observation_input
from stable_baselines.poar.utils import nature_autoencoder, bn_autoencoder, inverse_net, \
    forward_net, autoencoderMP, naive_autoencoder, transition_net, encoder, reward_net, natural_autoencoder
from gym.spaces.discrete import Discrete

from stable_baselines.common.distributions import make_proba_dist_type, CategoricalProbabilityDistribution, \
    MultiCategoricalProbabilityDistribution, DiagGaussianProbabilityDistribution, BernoulliProbabilityDistribution

from ipdb import set_trace as tt

_BATCH_NORM_DECAY = 0.997
_BATCH_NORM_EPSILON = 1e-5

def nature_cnn(scaled_images, **kwargs):
    """
    CNN from Nature paper.

    :param scaled_images: (TensorFlow Tensor) Image input placeholder
    :param kwargs: (dict) Extra keywords parameters for the convolutional layers of the CNN
    :return: (TensorFlow Tensor) The CNN output layer
    """
    activ = tf.nn.relu
    layer_1 = activ(conv(scaled_images, 'c1', n_filters=32, filter_size=8, stride=4, init_scale=np.sqrt(2), **kwargs))
    layer_2 = activ(conv(layer_1, 'c2', n_filters=64, filter_size=4, stride=2, init_scale=np.sqrt(2), **kwargs))
    layer_3 = activ(conv(layer_2, 'c3', n_filters=64, filter_size=3, stride=1, init_scale=np.sqrt(2), **kwargs))
    layer_3 = conv_to_fc(layer_3)
    return activ(linear(layer_3, 'fc1', n_hidden=512, init_scale=np.sqrt(2)))

def mlp_extractor(flat_observations, net_arch, act_fun):
    """
    Constructs an MLP that receives observations as an input and outputs a latent representation for the policy and
    a value network. The ``net_arch`` parameter allows to specify the amount and size of the hidden layers and how many
    of them are shared between the policy network and the value network. It is assumed to be a list with the following
    structure:

    1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer.
       If the number of ints is zero, there will be no shared layers.
    2. An optional dict, to specify the following non-shared layers for the value network and the policy network.
       It is formatted like ``dict(vf=[<value layer sizes>], pi=[<policy layer sizes>])``.
       If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed.

    For example to construct a network with one shared layer of size 55 followed by two non-shared layers for the value
    network of size 255 and a single non-shared layer of size 128 for the policy network, the following layers_spec
    would be used: ``[55, dict(vf=[255, 255], pi=[128])]``. A simple shared network topology with two layers of size 128
    would be specified as [128, 128].

    :param flat_observations: (tf.Tensor) The observations to base policy and value function on.
    :param net_arch: ([int or dict]) The specification of the policy and value networks.
        See above for details on its formatting.
    :param act_fun: (tf function) The activation function to use for the networks.
    :return: (tf.Tensor, tf.Tensor) latent_policy, latent_value of the specified network.
        If all layers are shared, then ``latent_policy == latent_value``
    """
    latent = flat_observations
    policy_only_layers = []  # Layer sizes of the network that only belongs to the policy network
    value_only_layers = []  # Layer sizes of the network that only belongs to the value network

    # Iterate through the shared layers and build the shared parts of the network
    for idx, layer in enumerate(net_arch):
        if isinstance(layer, int):  # Check that this is a shared layer
            layer_size = layer
            latent = act_fun(linear(latent, "shared_fc{}".format(idx), layer_size, init_scale=np.sqrt(2)))
        else:
            assert isinstance(layer, dict), "Error: the net_arch list can only contain ints and dicts"
            if 'pi' in layer:
                assert isinstance(layer['pi'], list), "Error: net_arch[-1]['pi'] must contain a list of integers."
                policy_only_layers = layer['pi']

            if 'vf' in layer:
                assert isinstance(layer['vf'], list), "Error: net_arch[-1]['vf'] must contain a list of integers."
                value_only_layers = layer['vf']
            break  # From here on the network splits up in policy and value network

    # Build the non-shared part of the network
    latent_policy = latent
    latent_value = latent
    for idx, (pi_layer_size, vf_layer_size) in enumerate(zip_longest(policy_only_layers, value_only_layers)):
        if pi_layer_size is not None:
            assert isinstance(pi_layer_size, int), "Error: net_arch[-1]['pi'] must only contain integers."
            latent_policy = act_fun(linear(latent_policy, "pi_fc{}".format(idx), pi_layer_size, init_scale=np.sqrt(2)))

        if vf_layer_size is not None:
            assert isinstance(vf_layer_size, int), "Error: net_arch[-1]['vf'] must only contain integers."
            latent_value = act_fun(linear(latent_value, "vf_fc{}".format(idx), vf_layer_size, init_scale=np.sqrt(2)))

    return latent_policy, latent_value

class SRLActorCriticPolicy(BasePolicy):
    """
    Policy object that implements actor critic

    :param sess: (TensorFlow session) The current TensorFlow session
    :param ob_space: (Gym Space) The observation space of the environment
    :param ac_space: (Gym Space) The action space of the environment
    :param n_env: (int) The number of environments to run
    :param n_steps: (int) The number of steps to run for each environment
    :param n_batch: (int) The number of batch to run (n_envs * n_steps)
    :param reuse: (bool) If the policy is reusable or not
    :param scale: (bool) whether or not to scale the input
    """

    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, split_dim=(), reuse=False, scale=False):
        super(SRLActorCriticPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=reuse,
                                                   scale=scale)
        self._pdtype = make_proba_dist_type(ac_space)
        self._policy = None
        self._proba_distribution = None
        self._value_fn = None
        self._action = None
        self._deterministic_action = None
        self.reconstruct_obs = None
        self.next_reconstruct_obs = None
        self.next_processed_obs = None
        self.srl_action = None
        self.srl_state = None
        self.next_latent_obs = None
        self.latent_obs = None
        self.srl_reward = None
        self._split_dim = split_dim

    def _setup_init(self):
        """
        sets up the distibutions, actions, and value
        """
        with tf.variable_scope("output", reuse=True):
            assert self.policy is not None and self.proba_distribution is not None and self.value_fn is not None
            self._action = self.proba_distribution.sample()
            self._deterministic_action = self.proba_distribution.mode()
            self._neglogp = self.proba_distribution.neglogp(self.action)
            if isinstance(self.proba_distribution, CategoricalProbabilityDistribution):
                self._policy_proba = tf.nn.softmax(self.policy)
            elif isinstance(self.proba_distribution, DiagGaussianProbabilityDistribution):
                self._policy_proba = [self.proba_distribution.mean, self.proba_distribution.std]
            elif isinstance(self.proba_distribution, BernoulliProbabilityDistribution):
                self._policy_proba = tf.nn.sigmoid(self.policy)
            elif isinstance(self.proba_distribution, MultiCategoricalProbabilityDistribution):
                self._policy_proba = [tf.nn.softmax(categorical.flatparam())
                                      for categorical in self.proba_distribution.categoricals]
            else:
                self._policy_proba = []  # it will return nothing, as it is not implemented
            self._value_flat = self.value_fn[:, 0]

    @property
    def split_dim(self):
        return self._split_dim

    @property
    def pdtype(self):
        """ProbabilityDistributionType: type of the distribution for stochastic actions."""
        return self._pdtype

    @property
    def policy(self):
        """tf.Tensor: policy output, e.g. logits."""
        return self._policy

    @property
    def proba_distribution(self):
        """ProbabilityDistribution: distribution of stochastic actions."""
        return self._proba_distribution

    @property
    def value_fn(self):
        """tf.Tensor: value estimate, of shape (self.n_batch, 1)"""
        return self._value_fn

    @property
    def value_flat(self):
        """tf.Tensor: value estimate, of shape (self.n_batch, )"""
        return self._value_flat

    @property
    def action(self):
        """tf.Tensor: stochastic action, of shape (self.n_batch, ) + self.ac_space.shape."""
        return self._action

    @property
    def deterministic_action(self):
        """tf.Tensor: deterministic action, of shape (self.n_batch, ) + self.ac_space.shape."""
        return self._deterministic_action

    @property
    def neglogp(self):
        """tf.Tensor: negative log likelihood of the action sampled by self.action."""
        return self._neglogp

    @property
    def policy_proba(self):
        """tf.Tensor: parameters of the probability distribution. Depends on pdtype."""
        return self._policy_proba

    @abstractmethod
    def step(self, obs, state=None, mask=None, deterministic=False):
        """
        Returns the policy for a single step

        :param obs: ([float] or [int]) The current observation of the environment
        :param state: ([float]) The last states (used in recurrent policies)
        :param mask: ([float]) The last masks (used in recurrent policies)
        :param deterministic: (bool) Whether or not to return deterministic actions.
        :return: ([float], [float], [float], [float]) actions, values, states, neglogp
        """
        raise NotImplementedError

    @abstractmethod
    def value(self, obs, state=None, mask=None):
        """
        Returns the value for a single step

        :param obs: ([float] or [int]) The current observation of the environment
        :param state: ([float]) The last states (used in recurrent policies)
        :param mask: ([float]) The last masks (used in recurrent policies)
        :return: ([float]) The associated value of the action
        """
        raise NotImplementedError

class FeedForwardPolicy(SRLActorCriticPolicy):
    """
    Policy object that implements actor critic, using a feed forward neural network.

    :param sess: (TensorFlow session) The current TensorFlow session
    :param ob_space: (Gym Space) The observation space of the environment
    :param ac_space: (Gym Space) The action space of the environment
    :param n_env: (int) The number of environments to run
    :param n_steps: (int) The number of steps to run for each environment
    :param n_batch: (int) The number of batch to run (n_envs * n_steps)
    :param reuse: (bool) If the policy is reusable or not
    :param layers: ([int]) (deprecated, use net_arch instead) The size of the Neural network for the policy
        (if None, default to [64, 64])
    :param net_arch: (list) Specification of the actor-critic policy network architecture (see mlp_extractor
        documentation for details).
    :param act_fun: (tf.func) the activation function to use in the neural network.
    :param cnn_extractor: (function (TensorFlow Tensor, ``**kwargs``): (TensorFlow Tensor)) the CNN feature extraction
    :param feature_extraction: (str) The feature extraction type ("cnn" or "mlp")
    :param kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
    """

    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, layers=None, net_arch=None,
                 act_fun=tf.tanh, feature_extraction="cnn", structure=None, **kwargs):
        super(FeedForwardPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=reuse,
                                                scale=(feature_extraction == "cnn"))
        self._kwargs_check(feature_extraction, kwargs)

        if layers is not None:
            warnings.warn("Usage of the `layers` parameter is deprecated! Use net_arch instead "
                          "(it has a different semantics though).", DeprecationWarning)
            if net_arch is not None:
                warnings.warn("The new `net_arch` parameter overrides the deprecated `layers` parameter!",
                              DeprecationWarning)
        if net_arch is None:
            if layers is None:
                layers = [64, 64]
            net_arch = [dict(vf=layers, pi=layers)]
        with tf.variable_scope("model", reuse=reuse):
            # By default, we consider the inputs are raw_pixels
            if structure == 'autoencoder':
                self.reconstruct_obs, latent_obs = nature_autoencoder(self.processed_obs, state_dim=512)
                pi_latent = vf_latent = latent_obs
            elif structure == 'autoencoder_bn':
                self.reconstruct_obs, latent_obs = bn_autoencoder(self.processed_obs, state_dim=512)
                pi_latent = vf_latent = latent_obs
            elif structure == 'autoencoder_mlp':
                self.reconstruct_obs, latent_obs = nature_autoencoder(self.processed_obs, state_dim=512)
                pi_latent, vf_latent = mlp_extractor(latent_obs, net_arch, act_fun)
            else:
                self.reconstruct_obs = self.processed_obs
                pi_latent = vf_latent = nature_cnn(self.processed_obs, **kwargs)

            self._value_fn = linear(vf_latent, 'vf', 1)

            self._proba_distribution, self._policy, self.q_value = \
                self.pdtype.proba_distribution_from_latent(pi_latent, vf_latent, init_scale=0.01)

        self._setup_init()

    def step(self, obs, state=None, mask=None, deterministic=False):
        if deterministic:
            action, value, neglogp = self.sess.run([self.deterministic_action, self.value_flat, self.neglogp],
                                                   {self.obs_ph: obs})
        else:
            action, value, neglogp = self.sess.run([self.action, self.value_flat, self.neglogp],
                                                   {self.obs_ph: obs})
        return action, value, self.initial_state, neglogp

    def proba_step(self, obs, state=None, mask=None):
        return self.sess.run(self.policy_proba, {self.obs_ph: obs})

    def value(self, obs, state=None, mask=None):
        return self.sess.run(self.value_flat, {self.obs_ph: obs})

class NatureCnnPolicy(FeedForwardPolicy):
    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, **_kwargs):
        super(NatureCnnPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                              feature_extraction="cnn", **_kwargs)

class AEPolicy(FeedForwardPolicy):
    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, **_kwargs):
        super(AEPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                       feature_extraction='cnn', structure='autoencoder', **_kwargs)

class AEBNPolicy(FeedForwardPolicy):
    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, **_kwargs):
        super(AEBNPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                         feature_extraction='cnn', structure='autoencoder_bn', **_kwargs)

class AEMlpPolicy(FeedForwardPolicy):
    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, **_kwargs):
        super(AEMlpPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                          feature_extraction='cnn', structure='autoencoder_mlp', **_kwargs)

class MlpPolicy(FeedForwardPolicy):
    """
    Policy object that implements actor critic, using a MLP (2 layers of 64)

    :param sess: (TensorFlow session) The current TensorFlow session
    :param ob_space: (Gym Space) The observation space of the environment
    :param ac_space: (Gym Space) The action space of the environment
    :param n_env: (int) The number of environments to run
    :param n_steps: (int) The number of steps to run for each environment
    :param n_batch: (int) The number of batch to run (n_envs * n_steps)
    :param reuse: (bool) If the policy is reusable or not
    :param _kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
    """

    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, **_kwargs):
        super(MlpPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                        feature_extraction="mlp", **_kwargs)

class SRLPolicy(SRLActorCriticPolicy):
    """
    SRL policy that can combine or split dimension for SRL model. An implementation of SRL in tensorflow
    """
    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, layers=None, net_arch=None,
                 act_fun=tf.tanh, split_dim=200, feature_extraction="cnn", structure='autoencoder', **kwargs):
        super(SRLPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=reuse,
                                        scale=(feature_extraction == "cnn"))
        self._kwargs_check(feature_extraction, kwargs)

        # placeholder for the next observation
        with tf.variable_scope('Next_observation'):
            self.next_obs_ph, self.next_processed_obs = observation_input(ob_space, n_batch, name='NextObs',
                                                                          scale=(feature_extraction == "cnn"))
            self._action_ph = tf.placeholder(dtype=ac_space.dtype, shape=(n_batch,) + ac_space.shape, name="action_ph")

        if layers is not None:
            warnings.warn("Usage of the `layers` parameter is deprecated! Use net_arch instead "
                          "(it has a different semantics though).", DeprecationWarning)
            if net_arch is not None:
                warnings.warn("The new `net_arch` parameter overrides the deprecated `layers` parameter!",
                              DeprecationWarning)
        if net_arch is None:
            if layers is None:
                layers = [64, 64]
            # preserved for further use, for example, pi_latent could be passed in to a MLP strutre
            net_arch = [dict(vf=layers, pi=layers)]
        with tf.variable_scope("model", reuse=reuse):
            # By default, we consider the inputs are raw_pixels
            pi_latent = vf_latent = self.srl_scope(split_dim, ac_space, structure)
            # latent_obs = self.srl_scope(split_dim, ac_space, structure)
            # pi_latent, vf_latent = mlp_extractor(latent_obs, net_arch, act_fun)
            self._value_fn = linear(vf_latent, 'vf', 1)
            self._proba_distribution, self._policy, self.q_value = \
                self.pdtype.proba_distribution_from_latent(pi_latent, vf_latent, init_scale=0.01)

        self._setup_init()

    def srl_scope(self, split_dim, ac_space, structure='autoencoder'):
        """
        The manager method that include all SRL models
        :param split_dim: dict with indicated dimension for each model
        :param ac_space:
        :param structure: the autoencoder structure
        :return:
        """
        if structure == 'autoencoder':
            encoder_fn = nature_autoencoder
        elif structure == 'natural':
            encoder_fn = natural_autoencoder
        elif structure == 'naive_autoencoder':
            encoder_fn = naive_autoencoder
        else:
            encoder_fn = encoder
        state_dim = 0
        previous_dim = 0
        dim_attr_dict = {}
        for key in split_dim:
            if split_dim[key] < 0:
                assert state_dim != previous_dim, "Error: the first dimension should not be 0"
                dim_attr_dict[key] = (previous_dim, state_dim)
            else:
                previous_dim = state_dim
                state_dim += split_dim[key]
                dim_attr_dict[key] = (previous_dim, state_dim)
        if 'domain' in split_dim:
            self.domain_dim = dim_attr_dict['domain']
        self.next_reconstruct_obs, self.next_latent_obs = encoder_fn(self.next_processed_obs, state_dim=state_dim)
        self.reconstruct_obs, self.latent_obs = encoder_fn(self.processed_obs, state_dim=state_dim)
        # We make next_latent_obs to be observable from outside to compute the loss with srl_state

        with tf.variable_scope('SRL'):
            if "forward" in split_dim:
                self.srl_state = forward_net(
                    self.latent_obs[..., dim_attr_dict["forward"][0]:dim_attr_dict["forward"][1]],
                    self._action_ph, ac_space, state_dim=state_dim)
            if "inverse" in split_dim:
                self.srl_action = inverse_net(
                    self.latent_obs[..., dim_attr_dict["inverse"][0]:dim_attr_dict["inverse"][1]],
                    self.next_latent_obs[..., dim_attr_dict["inverse"][0]:dim_attr_dict["inverse"][1]], ac_space)
            if "reward" in split_dim:
                # if dim_attr_dict["reward"][1] - dim_attr_dict["reward"][0] == dim_attr_dict["inverse"][1]-dim_attr_dict["inverse"][0]:
                #     self.srl_reward = reward_net(
                #         self.latent_obs[..., dim_attr_dict["reward"][0]:dim_attr_dict["reward"][1]] -
                #         self.latent_obs[..., dim_attr_dict["inverse"][0]:dim_attr_dict["inverse"][1]],
                #         self.next_latent_obs[..., dim_attr_dict["reward"][0]:dim_attr_dict["reward"][1]] -
                #         self.next_latent_obs[..., dim_attr_dict["inverse"][0]:dim_attr_dict["inverse"][1]],
                #         reward_dim=2
                #     )
                # else:
                self.srl_reward = reward_net(
                    self.latent_obs[..., dim_attr_dict["reward"][0]:dim_attr_dict["reward"][1]],
                    self.next_latent_obs[..., dim_attr_dict["reward"][0]:dim_attr_dict["reward"][1]], reward_dim=10)
        return self.latent_obs

    def step(self, obs, next_obs=None, state=None, mask=None, deterministic=False):
        """
        for the model to take action, predict value function
        :param obs:
        :param next_obs:
        :param state:
        :param mask:
        :param deterministic:
        :return:
        """
        if deterministic:
            action, value, neglogp = self.sess.run([self.deterministic_action, self.value_flat, self.neglogp],
                                                   {self.obs_ph: obs})
        else:
            action, value, neglogp = self.sess.run([self.action, self.value_flat, self.neglogp],
                                                   {self.obs_ph: obs})

        p_obs, ae_obs, latent = self.sess.run([self.processed_obs, self.reconstruct_obs, self.latent_obs],
                                              {self.obs_ph: obs})
        return action, value, self.initial_state, neglogp, ae_obs, latent

    def proba_step(self, obs, state=None, mask=None):
        return self.sess.run(self.policy_proba, {self.obs_ph: obs})

    def value(self, obs, state=None, mask=None):
        return self.sess.run(self.value_flat, {self.obs_ph: obs})

class NAEPolicy(SRLPolicy):
    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, **_kwargs):
        super(NAEPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                        feature_extraction='cnn', structure='naive_autoencoder', **_kwargs)

class AESRLPolicy(SRLPolicy):
    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, **_kwargs):
        super(AESRLPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                          feature_extraction='cnn', structure='autoencoder', **_kwargs)

class NAESRLPolicy(SRLPolicy):
    def __init__(self, sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=False, **_kwargs):
        super(NAESRLPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                          feature_extraction='cnn', structure='natural', **_kwargs)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Jan 10 13:41:44 2021

@author: winter_camp
"""
import pyaudio
import numpy as np
#import matplotlib.pyplot as plt
import time

from scipy.fftpack import fft

class Receiver(object):

    CHUNK               = 1024 * 2
    FORMAT              = pyaudio.paInt16
    CHANNELS            = 1
    RATE                = 44100

    FRAMES_PER_FFT      = 16 # FFT takes average across how many frames
    SAMPLES_PER_FFT     = CHUNK * FRAMES_PER_FFT
    FREQ_STEP           = float(RATE)/SAMPLES_PER_FFT

    CHANNEL_NUMBER      = 8
    SHARED_CHANNEL      = 8
    TRACK_NUM           = 1

    FRAME_TIME          = 0.2

    active_freq_bin     = [55,193,323]
    ending_freq_bin     = [56,196,323]
    ending_channel_num  = [0,4,7]
    '''
    d_channel_1 = [[53,57,58],[59,60],[61,62],[63,64],[65,66],[67,68],[69,70],[71,72],[73,74],[75,76],[77,78],[79,80],[81,82],[83,84],[85,86],[87,88]]
    d_channel_2 = [[89,90],[91,92],[93,94],[95,96],[97,98],[99,100],[101,102],[103,104],[105,106],[107,108],[109,110],[111,112],[113,114],[115,116],[117,118],[119,120]]
    d_channel_3 = [[121,122],[123,124],[125,126],[127,128],[129,130],[131,132],[133,134],[135,136],[137,138],[139,140],[141,142],[143,144],[145,146],[147,148],[149,150],[151,152]]
    d_channel_4 = [[153,154],[155,156],[157,158],[159,160],[161,162],[163,164],[165,166],[167,168],[169,170],[171,172],[173,174],[175,176],[177,178],[179,180],[181,182],[183,184]]
    d_channel_5 = [[185,186],[187,188],[189,190],[191,192],[193,194],[195,196],[197,198],[199,200],[201,202],[203,204],[205,206],[207,208],[209,210],[211,212],[213,214],[215,216]]
    d_channel_6 = [[217,218],[219,220],[221,222],[223,224],[225,226],[227,228],[229,230],[231,232],[233,234],[235,236],[237,238],[239,240],[241,242],[243,244],[245,246],[247,248]]
    d_channel_7 = [[249,250],[251,252],[253,254],[255,256],[257,258],[259,260],[261,262],[263,264],[265,266],[267,268],[269,270],[271,272],[273,274],[275,276],[277,278],[279,280]]
    d_channel_8 = [[281,282],[283,284],[285,286],[287,288],[289,290],[291,292],[293,294],[295,296],[297,298],[299,300],[301,302],[303,304],[305,306],[307,308],[309,310],[311,312,313]]
    d_channel_9 = [[314]]
    '''
    d_channel_1 = [[53,57,58],[59,60],[61,62],[63,64],[65,66],[67,68],[69,70],[71,72],[73,74],[75,76],[77,78],[79,80],[81,82],[83,84],[85,86],[87,88]]
    d_channel_2 = [[91,92],[93,94],[95,96],[97,98],[99,100],[101,102],[103,104],[105,106],[107,108],[109,110],[111,112],[113,114],[115,116],[117,118],[119,120],[121,122]]
    d_channel_3 = [[125,126],[127,128],[129,130],[131,132],[133,134],[135,136],[137,138],[139,140],[141,142],[143,144],[145,146],[147,148],[149,150],[151,152],[153,154],[155,156]]
    d_channel_4 = [[159,160],[161,162],[163,164],[165,166],[167,168],[169,170],[171,172],[173,174],[175,176],[177,178],[179,180],[181,182],[183,184],[185,186],[187,188],[189,190]]
    d_channel_5 = [[191,192],[193,194],[195,196],[197,198],[199,200],[201,202],[203,204],[205,206],[207,208],[209,210],[211,212],[213,214],[215,216],[217,218],[219,220],[221,222]]
    d_channel_6 = [[223,224],[225,226],[227,228],[229,230],[231,232],[233,234],[235,236],[237,238],[239,240],[241,242],[243,244],[245,246],[247,248],[249,250],[251,252],[253,254]]
    d_channel_7 = [[257,258],[259,260],[261,262],[263,264],[265,266],[267,268],[269,270],[271,272],[273,274],[275,276],[277,278],[279,280],[281,282],[283,284],[285,286],[287,288]]
    d_channel_8 = [[291,292],[293,294],[295,296],[297,298],[299,300],[301,302],[303,304],[305,306],[307,308],[309,310],[311,312],[313,314],[315,316],[317,318],[319,320],[321,322,323]]
    d_channel_9 = [[324]]

    d_channel = []
    chunk_list = ['0000', '0001', '0010', '0011', '0100', '0101', '0110', '0111', '1000', '1001', '1010', '1011', '1100', '1101', '1110', '1111']

    activation_info = [[],[],[],[],[]]
    received_info   = [] #i0: Estimated length
    ending_info     = [[],[],[],[],[],[],[],[]]
    ending_mark     = [0,0] #i0: ending pointer, i1: Estimated length

    status = 0
    #speed_info = [[0,0],[0,2],[1,3]]
    speed_info = [[0,0],[4,0],[7,0]]

    def __init__(self, actived_channel = 8, sensitivity = 'medium'):
        self.d_channel.append(self.d_channel_1)
        self.d_channel.append(self.d_channel_2)
        self.d_channel.append(self.d_channel_3)
        self.d_channel.append(self.d_channel_4)
        self.d_channel.append(self.d_channel_5)
        self.d_channel.append(self.d_channel_6)
        self.d_channel.append(self.d_channel_7)
        self.d_channel.append(self.d_channel_8)
        self.d_channel.append(self.d_channel_9)

        self.p = pyaudio.PyAudio()

        self.stream = self.p.open(
            format = self.FORMAT,
            channels = self.CHANNELS,
            rate = self.RATE,
            input = True,
            output = True,
            frames_per_buffer = self.CHUNK
            )

        self.current_bins = []
        self.pointers = []
        self.recieved_bins = []
        self.fft = []
        self.scheduled_pointer = 0
        self.retrieved_data = []
        '''
        if speed == 'fast':
            self.SHARED_CHANNEL = 8
            self.TRACK_NUM = 1
            self.speed_info = [[0,0],[4,0],[7,0]]
        '''
        if sensitivity == 'medium':
            self.sensitivity = 2
        elif sensitivity == 'low':
            self.sensitivity = 3
        elif sensitivity == 'high':
            self.sensitivity = 1
        else:
            raise ParameterError(message = 'sensitivity could only be low, medium, or high')

        for i in range(self.SHARED_CHANNEL):
            self.pointers.append(0)
            self.current_bins.append([0,0,0,0,0,0,0])
            self.recieved_bins.append([])

        #print(self.stream)
    def refresh_audio(self):
        self.p = pyaudio.PyAudio()

        self.stream = self.p.open(
            format = self.FORMAT,
            channels = self.CHANNELS,
            rate = self.RATE,
            input = True,
            output = True,
            frames_per_buffer = self.CHUNK
            )

    def update_speed_info(self):
        if (self.SHARED_CHANNEL == 8):
            self.speed_info = [[0,0],[4,0],[7,0]]
        elif (self.SHARED_CHANNEL == 2):
            self.speed_info = [[0,0],[0,2],[1,3]]
        else:
            self.speed_info = [[0,0],[0,1],[3,1]]

    def callback(self, input_data, frame_count, time_info, flags):

        return (input_data, pyaudio.paContinue)

    def read_block(self, standby_time):
        '''
        Read the input audio (Blocking Mode):
            Standby and listen until all the equested frames have been recorded.

        Parameters
        ----------
        standby_time : int
            The requested time (in second).

        Returns
        -------
        retrieved_data: list
            Every retrieved data during the standby time

        '''
        frame_count = 0
        start_time = time.time()

        #self.current_bins = []
        #self.pointers = []
        #self.recieved_bins = []
        self.scheduled_pointer = 0
        '''
        for i in range(self.SHARED_CHANNEL):
            self.pointers.append(0)
            self.current_bins.append([0,0,0,0,0,0,0])
            self.recieved_bins.append([])
            #self.scheduled_bins.append([])
        '''
        while (time.time() - start_time < standby_time):
            bit = self.read_frame()
            print(bit, end=(''))
            frame_count += 1

        return self.retrieved_data

    def read_frame_audio(self, audio, log = False):
        '''
        Takes 2048 frames of audio data and perform analysis.

        Parameters
        ----------
        audio : np array with a size of (2048,)
            The array version of the audio buffer.
        log : bool, optional
            If true, it will add a status flag to the returns, accompany with the retrived result.
            The default is False.

        Returns
        -------
        retrieved_data: String
            The retrieved and demodulated data, empty if the receiver have not detected anything yet.

        status: int
            0: Unactivated
            1: Activating
            2: Activated, preparing
            3: Activation Failed, rollback to unactivated
            4: Listening
            5: Cycle Terminated, received auccessfully and back to Unactivated (status 0)
            6: Cycle Terminated, received failed and back to Unactivated (status 0)

        '''
        try:
            if (audio.shape[0]==2048):
                self.fft = fft(audio)
                #######################

                freq_bins = []
                for i in range(self.SHARED_CHANNEL):
                    candidate_freq = []
                    for j in range(int(self.CHANNEL_NUMBER/self.SHARED_CHANNEL)):
                        freq_bin = np.abs(self.fft[ self.d_channel[j*self.SHARED_CHANNEL+i][0][0]  : self.d_channel[j*self.SHARED_CHANNEL+i][15][-1]+1 ]).argmax() + self.d_channel[j*self.SHARED_CHANNEL+i][0][0]
                        candidate_freq.append(freq_bin)
                    freq_bins.append(candidate_freq)
                #print(freq_bins)
                try:
                    self.status = self.update_statue(freq_bins,self.status)
                except ActivationError as ae:
                    print(ae)
                bit = ''
                if self.status==4 and len(self.recieved_bins[0]) != self.scheduled_pointer:
                    length = True
                    for i in range(self.SHARED_CHANNEL-1):
                        if len(self.recieved_bins[i])==len(self.recieved_bins[i+1]):
                            continue
                        else:
                            length = False
                            break
                    if length:
                        bit = self.convert_result([[item[-1]] for item in self.recieved_bins])
                        self.scheduled_pointer = len(self.recieved_bins[0])
                if log:
                    if self.status == 3:
                        self.status = 0
                        return bit, 3
                    elif self.status == 5:
                        self.status = 0
                        return self.retrieved_data[-1], 5
                    elif self.status == 6:
                        self.status = 0
                        return bit, 6
                    else:
                        return bit, self.status
                else:
                    if self.status == 3:
                        self.status = 0
                        return ''
                    elif self.status == 5:
                        self.status = 0
                        return self.retrieved_data[-1]
                    elif self.status == 6:
                        self.status = 0
                        return ''
                    else:
                        return bit
            else:
                raise AttributeError

                #######################
        except AttributeError:
                print('The number of frames need to be 2048')

    def read_frame(self, log = False):
        '''
        Read the input audio (Callback Mode):
            Called each frame to listen to the audio.

        Parameters
        ----------
        log : bool, optional
            If true, it will add a status flag to the returns, accompany with the retrived result.
            The default is False.

        Returns
        -------
        retrieved_data: String
            The retrieved and demodulated data, empty if the receiver have not detected anything yet.

        status: int
            0: Unactivated
            1: Activating
            2: Activated, preparing
            3: Activation Failed, rollback to unactivated
            4: Listening
            5: Terminated, received auccessfully
            6: Terminated, received failed

        '''
        data = self.stream.read(self.CHUNK, exception_on_overflow = False)
        data_int = np.frombuffer(data, dtype = np.int16)
        self.fft = fft(data_int)

        #######################

        freq_bins = []
        for i in range(self.SHARED_CHANNEL):
            candidate_freq = []
            for j in range(int(self.CHANNEL_NUMBER/self.SHARED_CHANNEL)):
                freq_bin = np.abs(self.fft[ self.d_channel[j*self.SHARED_CHANNEL+i][0][0]  : self.d_channel[j*self.SHARED_CHANNEL+i][15][-1]+1 ]).argmax() + self.d_channel[j*self.SHARED_CHANNEL+i][0][0]
                candidate_freq.append(freq_bin)
            freq_bins.append(candidate_freq)
        #print(freq_bins)
        try:
            self.status = self.update_statue(freq_bins,self.status)
        except ActivationError as ae:
            print(ae)
        bit = ''
        if self.status==4 and len(self.recieved_bins[0]) != self.scheduled_pointer:
            length = True
            for i in range(self.SHARED_CHANNEL-1):
                if len(self.recieved_bins[i])==len(self.recieved_bins[i+1]):
                    continue
                else:
                    length = False
                    break
            if length:
                try:
                    bit = self.convert_result([[item[-1]] for item in self.recieved_bins])
                    self.scheduled_pointer = len(self.recieved_bins[0])
                except:
                    print('transmission interrupted')
                    self.status = 6
                    self.clear_session()

        if log:
            if self.status == 3:
                self.status = 0
                return bit, 3
            elif self.status == 5:
                self.status = 0
                return self.retrieved_data[-1], 5
            elif self.status == 6:
                self.status = 0
                return bit, 6
            else:
                return bit, self.status
        else:
            if self.status == 3:
                self.status = 0
                return ''
            elif self.status == 5:
                self.status = 0
                return self.retrieved_data[-1]
            elif self.status == 6:
                self.status = 0
                return ''
            else:
                return bit

        #######################

    #def refresh_result(self):

    def most_frequent(self, List):
        counter = 0
        num = List[0]
        for i in List:
            curr_frequency = List.count(i)
            if(curr_frequency> counter):
                counter = curr_frequency
                num = i
        return num
    def get_bin_num(self, freq_bin,n):
        for i in range(16):
            if freq_bin in self.d_channel[n][i]:
                return i
        print('request {} in number {} channel'.format(freq_bin,n))
        return 99

    '''
    Status:
        0: Unactivated
        1: Activating
        2: Activated, Preparing
        3: Activation Failed, rollback to unactivated
        4: Listening
            4.5 (Hide): Terminating
        5: Terminated, Received Successfully
        6: Terminated, Received Failed
    '''
    def update_statue(self, freq_bins,status):
         # if the activation frequency is been detected three times,
        if (status == 0):
            if (freq_bins[self.speed_info[0][0]][self.speed_info[0][1]] == self.active_freq_bin[0] and freq_bins[self.speed_info[1][0]][self.speed_info[1][1]] == self.active_freq_bin[1] and freq_bins[self.speed_info[2][0]][self.speed_info[2][1]] == self.active_freq_bin[2]):
                self.activation_info = [[],[],[],[],[]]
                self.received_info = []
                self.pointers[0] = 1
                status = 1
                print('activating...')
                ########## TODO ##########
                self.activation_info[0].append(freq_bins[1][0])
                self.activation_info[1].append(freq_bins[2][0])
                self.activation_info[2].append(freq_bins[3][0])
                self.activation_info[3].append(self.get_bin_num(freq_bins[5][0],5))
                self.activation_info[4].append(self.get_bin_num(freq_bins[6][0],6))

        elif (status == 1):
            if (freq_bins[self.speed_info[0][0]][self.speed_info[0][1]] == self.active_freq_bin[0] and freq_bins[self.speed_info[1][0]][self.speed_info[1][1]] == self.active_freq_bin[1] and freq_bins[self.speed_info[2][0]][self.speed_info[2][1]] == self.active_freq_bin[2]):
                self.pointers[0] += 1
                ########## TODO ##########
                self.activation_info[0].append(freq_bins[1][0])
                self.activation_info[1].append(freq_bins[2][0])
                self.activation_info[2].append(freq_bins[3][0])
                self.activation_info[3].append(self.get_bin_num(freq_bins[5][0],5))
                self.activation_info[4].append(self.get_bin_num(freq_bins[6][0],6))

                if (self.pointers[0] == self.sensitivity):
                    self.SHARED_CHANNEL = self.most_frequent(self.activation_info[3]+self.activation_info[4])
                    if self.SHARED_CHANNEL!=4 and self.SHARED_CHANNEL!=8 and self.SHARED_CHANNEL!=2:
                        self.clear_session()
                        status = 3
                        raise ActivationError

                    self.TRACK_NUM = int(self.CHANNEL_NUMBER / self.SHARED_CHANNEL)
                    self.update_speed_info()
                    self.current_bins = []
                    self.recieved_bins = []
                    self.pointers = []
                    for i in range(self.SHARED_CHANNEL):
                        self.current_bins.append([0,0,0,0,0,0,0])
                        self.recieved_bins.append([])
                        self.pointers.append(0)
                    status = 2
                    #self.pointers[0] = 0
                    print("Activated, on preparing")

            else:
                self.clear_session()
                status = 3

                print('activation failed')
        elif (status == 2):
            if (freq_bins[self.speed_info[0][0]][self.speed_info[0][1]] != self.active_freq_bin[0] and freq_bins[self.speed_info[2][0]][self.speed_info[2][1]] != self.active_freq_bin[2]):
                #print(self.activation_info)
                self.received_info.append(100*self.get_bin_num(self.most_frequent(self.activation_info[0]),1) + 10*self.get_bin_num(self.most_frequent(self.activation_info[1]),2) + self.get_bin_num(self.most_frequent(self.activation_info[2]),3))

                print('Estimated length: {}'.format(self.received_info[0]))
                print('On recieving...')
                self.d_channel[0][0] = [57,58]
                #self.d_channel[7][15] = [311,312]
                status = self.check_channels(freq_bins)
            else:
                ########## TODO ##########
                if self.SHARED_CHANNEL == 8:
                    self.activation_info[0].append(freq_bins[1][0])
                    self.activation_info[1].append(freq_bins[2][0])
                    self.activation_info[2].append(freq_bins[3][0])
                elif self.SHARED_CHANNEL == 2:
                    self.activation_info[0].append(freq_bins[1][0])
                    self.activation_info[1].append(freq_bins[0][1])
                    self.activation_info[2].append(freq_bins[1][1])
                else:
                    self.activation_info[0].append(freq_bins[1][0])
                    self.activation_info[1].append(freq_bins[2][0])
                    self.activation_info[2].append(freq_bins[3][0])

        elif (status == 4):
            status = self.check_channels(freq_bins)

            if (self.ending_mark[0]>=1):
                self.ending_mark[0] += 1
                ########## TODO ##########
                if self.SHARED_CHANNEL == 8:
                    self.ending_info[0].append(freq_bins[0][0])
                    self.ending_info[1].append(freq_bins[1][0])
                    self.ending_info[2].append(freq_bins[2][0])
                    self.ending_info[3].append(freq_bins[3][0])
                    self.ending_info[4].append(freq_bins[4][0])
                    self.ending_info[7].append(freq_bins[7][0])
                elif self.SHARED_CHANNEL == 2:
                    self.ending_info[0].append(freq_bins[0][0])
                    self.ending_info[1].append(freq_bins[1][0])
                    self.ending_info[2].append(freq_bins[0][1])
                    self.ending_info[3].append(freq_bins[1][1])
                    self.ending_info[4].append(freq_bins[0][2])
                    self.ending_info[7].append(freq_bins[1][3])
                else:
                    self.ending_info[0].append(freq_bins[0][0])
                    self.ending_info[1].append(freq_bins[1][0])
                    self.ending_info[2].append(freq_bins[2][0])
                    self.ending_info[3].append(freq_bins[3][0])
                    self.ending_info[4].append(freq_bins[0][1])
                    self.ending_info[7].append(freq_bins[3][1])

                if self.ending_mark[0] >= 5:

                    validated = 0
                    for i in range(3):
                        count = 0
                        for info in self.ending_info[self.ending_channel_num[i]]:
                            if info == self.ending_freq_bin[i]:
                                count += 1
                                if count == 2:
                                    validated += 1
                                    break
                    if validated == 3:
                        #if validated ended
                        print('')
                        #print('Recieved: {}, length: {}, {}'.format(self.copy_recieved_bins,len(self.copy_recieved_bins[0]),len(self.copy_recieved_bins[1])))
                        self.d_channel[0][0] = [53,57,58]
                        self.ending_mark[0] = 0

                        ##################CHECK
                        self.ending_mark[1]=100*self.get_bin_num(self.most_frequent(self.ending_info[1]),1) + 10*self.get_bin_num(self.most_frequent(self.ending_info[2]),2) + self.get_bin_num(self.most_frequent(self.ending_info[3]),3)
                        print('Ending marks estimated length: {}'.format(self.ending_mark[1]))

                        ##################
                        if(self.check_result(self.received_info[0], self.ending_mark[1], self.copy_recieved_bins) == 1):
                            try:
                                result = self.convert_result(self.copy_recieved_bins, self.trim)
                            except UnicodeDecodeError:
                                self.clear_session()
                                return 6
                            self.retrieved_data.append(result)
                            #print(result)

                            self.clear_session()

                            return 5
                        else:
                            self.clear_session()
                            return 6
                    else:
                        #if not ended
                        self.d_channel[0][0] = [57,58]
                        self.ending_mark[0] = 0

            elif (freq_bins[self.speed_info[2][0]][self.speed_info[2][1]] == self.ending_freq_bin[2]):

                #print('detecting ending...')
                #if one of the ending bit appears, expande the fft scope, start monitoring for 5 frames
                #create copy of recieved_bin
                #undo the pointer
                self.ending_mark[0] = 1
                self.d_channel[0][0] = [56,57,58]
                self.ending_info = [[],[],[],[],[],[],[],[]]

                self.copy_recieved_bins = []
                for bins in self.recieved_bins:
                    copy_bins = bins.copy()
                    self.copy_recieved_bins.append(copy_bins)
                for pointer, current_bin, recieved_bin in zip(self.pointers, self.current_bins, self.copy_recieved_bins):
                    if pointer >= 3:
                        recieved_bin.append(current_bin[pointer-1])
        return status

    def clear_session(self):
        self.activation_info = [[],[],[],[],[]]
        self.received_info = []
        self.SHARED_CHANNEL = 8
        self.scheduled_pointer = 0
        self.trim = 0
        self.update_speed_info()
        self.d_channel[0][0] = [53,57,58]
        self.ending_mark[0] = 0
        self.refresh_audio()
        #self.TRACK_NUM = int(self.CHANNEL_NUMBER / self.SHARED_CHANNEL)

    def check_channels(self, freq_bins):
        frame_result = []
        for freq_bin, current_bin, recieved_bin,i in zip(freq_bins,self.current_bins,self.recieved_bins,range(self.SHARED_CHANNEL)):
            freq_bin_nums = []
            for j in range(self.TRACK_NUM):
                freq_bin_nums.append(self.get_bin_num(freq_bin[j],j*self.SHARED_CHANNEL+i))
            frame_result.append(freq_bin_nums)
            freq_bin_num = self.most_frequent(freq_bin_nums)
            if (self.pointers[i]==0):
                current_bin[self.pointers[i]] = freq_bin_num
                self.pointers[i] += 1
            elif ( freq_bin_num == current_bin[self.pointers[i]-1]) or (self.pointers[i] < 3):
                #if this bit is the same as the last bit
                current_bin[self.pointers[i]] = freq_bin_num
                self.pointers[i] += 1
                if (self.pointers[i] == 7):
                    recieved_bin.append(current_bin[self.pointers[i]-1])
                    current_bin[0] = freq_bin_num
                    self.pointers[i] = 3
            else:
                #if new bit appears,
                number = self.most_frequent(current_bin[0:self.pointers[i]-1])
                recieved_bin.append(number)
                current_bin[0] = freq_bin_num
                self.pointers[i] = 1
        #print(frame_result)

        return 4

    def bin_to_ascii(self,bin_data):
        st = ''
        for i in range(len(bin_data)):
            st += str(int(bin_data[i]))
        st = st[:1]+'b'+st[1:]
        n = int(st, 2)
        return n.to_bytes((n.bit_length() + 7) // 8, 'big').decode()

    def check_result(self, a, e, received):

        ###########CHECK
        minimum_block_length = self.SHARED_CHANNEL * 4
        full_length_a = (int((a-1) / minimum_block_length) + 1)
        full_length_e = (int((e-1) / minimum_block_length) + 1)

        for i in range(self.SHARED_CHANNEL):
            if len(received[i]) == full_length_a:
                self.trim = a
                continue
            elif len(received[i]) == full_length_e:
                self.trim = e
                continue
            else:
                print('recieve failed')
                return 0

        return 1
        ###########

    '''
        if (len(self.copy_recieved_bins[0]) != len(self.copy_recieved_bins[1])):
            print('recieve failed')
            return 0
        else:
            return 1
        '''

    def convert_result(self, received, trim = -1):
        '''
        Convert received binaries from all channels to demodulated message

        Parameters
        ----------
        received : ndarray (binaries)
            received binaries from all channels.

        Returns
        -------
        result : string
            converted message.

        '''
        binary = ''

        for i in range(len(received[0])):
            for j in range(self.SHARED_CHANNEL):
                binary += self.chunk_list[received[j][i]]
        if trim == -1:
            trim = len(binary)
        #print(binary)
        return self.bin_to_ascii(binary[0:trim])

    def fft(self):
        return self.fft

    def clear(self):
        '''
        Clear the retrieved data

        Returns
        -------
        None.

        '''
        self.status = 0
        self.current_bins = []
        self.pointers = []
        self.recieved_bins = []
        self.fft = []
        self.retrieved_data = []
        for i in range(self.SHARED_CHANNEL):
            self.pointers.append(0)
            self.current_bins.append([0,0,0,0,0,0,0])
            self.recieved_bins.append([])

    def received_data(self):
        '''
        Return all received data

        Returns
        -------
        retrieved_data: list
            All retrieved data during the standby time

        '''
        return self.retrieved_data

#%%

class Error(Exception):
    """Base class for exceptions in this module."""
pass

class ParameterError(Error):

    def __init__(self, message = 'speed could only be slow, medium or fast'):
        self.message = message
        super().__init__(self.message)

class ActivationError(Error):

    def __init__(self, message = 'Activation failed, increase volume'):
        self.message = message
        super().__init__(self.message)

class ASCIIError(Error):
    def __init___(self, message = 'Invalid data received'):
        super().__init__(self.message)


import unittest
from ai.go.go_game import GoGame
from ai.go.go_engine import GoEngine, GoState, GoColor
from ai.rl.mcts import MCPlayer, MCSearch
from ai.rl.models import CNNModel
import numpy as np
import tensorflow as tf

class TestGoBoard(unittest.TestCase):
    def setUp(self):
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True

        board_size = 19
        komi = 7.5
        self.go_game = GoGame(board_size=board_size)
        model = CNNModel(board_size=board_size)
        history = np.zeros((2, board_size, board_size), dtype=int)
        root_state = GoState(komi, GoColor.black, history[-1], history)
        values, win_rate = model.predict([root_state])
        search = MCSearch(GoEngine(), model, root_state, values, win_rate)
        self.white_player = MCPlayer(search)
        self.black_player = MCPlayer(search)

    def test_play_black_stone(self):
        score = self.go_game.play_game(self.black_player, self.white_player)
        print(score)
        print(self.go_game.history[len(self.go_game.moves)])
        self.go_game.to_sgf('test.sgf')


#!/usr/bin/python
import sys
import os
import subprocess
from os.path import join, isdir
import numpy as np
import fileinput
import json
import random
from itertools import chain
from numpy.random import permutation
##------------------------------------------------------------------
import torch
from torch.autograd import Variable
#----------------------------------------
import torch.nn as nn
from torch import autograd, nn, optim
os.environ['PYTHONUNBUFFERED'] = '0'
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim

from random import shuffle
from statistics import mean
import matplotlib
import matplotlib.pyplot as plt
plt.switch_backend('agg')
matplotlib.pyplot.viridis()
import glob

#*************************************************************************************************************************
####### Loading the Parser and default arguments
sys.path.insert(0,'/mnt/matylda3/vydana/HOW2_EXP/MT_Transformer/MT_TransV1/')
import MT_Transformer_arg
from MT_Transformer_arg import parser
args = parser.parse_args()

###save architecture for decoding
model_path_name=join(args.model_dir,'model_architecture_')
with open(model_path_name, 'w') as f:
    json.dump(args.__dict__, f, indent=2)
print(args)
# #####setting the gpus in the gpu cluster
# #**********************************
from Set_gpus import Set_gpu
if args.gpu:
    Set_gpu()
###----------------------------------------
#==============================================================
from Dataloader_for_MT_v2_check import DataLoader
from TRANSFORMER_MT_V1 import Transformer
from Initializing_Transformer_MT import Initialize_Att_model
from Training_loop_MT import train_val_model
from Load_sp_model import Load_sp_models
##==================================
#==============================================================
############################################
#=============================================================
def main():
        ##Load setpiece models for Dataloaders
        Src_model=Load_sp_models(args.Src_model_path)
        Tgt_model=Load_sp_models(args.Tgt_model_path)
        ###initilize the model
        #============================================================
        #------------------------------------------------------------
        #/mnt/matylda3/vydana/HOW2_EXP/IWSLT2021_MT_systems/Data_files_Mustc/train_splits
        #dev_splits
        #train_splits_V2
        #train_splits_V2/*
        train_gen = DataLoader({'files': glob.glob(args.data_dir + "dev_scp_sorted"),
                                'max_batch_label_len': args.max_batch_label_len,
                                'max_batch_len': args.max_batch_len,
                                'max_feat_len': args.max_feat_len,
                                'max_label_len': args.max_label_len,
                                'Src_model': Src_model,
                                'Tgt_model': Tgt_model,
                                'queue_size': 100,
                                'apply_cmvn': 1,
                                'min_words': 0,
                                'max_words': args.max_words,
                                'min_len_ratio': args.min_len_ratio})

        sentences=0
        for i in range(166563):
                B1 = train_gen.next()
                assert B1 is not None, "None should never come out of the DataLoader"
                smp_Src_labels = B1.get('smp_Src_labels')
                smp_Tgt_labels = B1.get('smp_Tgt_labels')
                sentences+=smp_Src_labels.shape[0]
                print(':===>',i, smp_Src_labels.shape, smp_Tgt_labels.shape,"sentences:-->",sentences)

#=======================================================
#=============================================================================================
if __name__ == '__main__':
    main()


import numpy as np
import random
import torch

import time

from torch.distributions import Normal
from collections import deque
import torch.nn.functional as F
from tensorboardX import SummaryWriter

class Memory(object):
    def __init__(self, memory_size: int) -> None:
        self.memory_size = memory_size
        self.buffer = deque(maxlen=self.memory_size)

    def add(self, experience) -> None:
        self.buffer.append(experience)

    def sample(self, batch_size: int, continuous: bool = True):
        if batch_size > len(self.buffer):
            batch_size = len(self.buffer)
        if continuous:
            rand = random.randint(0, len(self.buffer) - batch_size)
            return [self.buffer[i] for i in range(rand, rand + batch_size)]
        else:
            indexes = np.random.choice(np.arange(len(self.buffer)), size=batch_size, replace=False)
            return [self.buffer[i] for i in indexes]

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

class GraspTrain(object):
    def __init__(self, mlp_policy_model, mlp_value_model, mlp_orientation, mlp_orientation_value,
                 cnn_policy_model, cnn_orientation, action_dim, state_dim,
                 lr=0.0001, update_steps=4, buff_capability=32, batch_size=32, gamma=0.9, ppo_epoch=10):
        self.mlp_policy = mlp_policy_model(state_dim-1, action_dim-1).cuda()
        self.mlp_old_policy = mlp_policy_model(state_dim-1, action_dim-1).cuda()
        self.mlp_value = mlp_value_model(state_dim-1).cuda()

        self.mlp_old_orientation = mlp_orientation(1, 1).cuda()
        self.mlp_orientation = mlp_orientation(1, 1).cuda()
        self.mlp_orientation_value = mlp_orientation_value(1).cuda()

        self.cnn_policy = cnn_policy_model(action_dim-1).cuda()
        self.cnn_orientation = cnn_orientation(1).cuda()

        self.steps = 0
        self.buffer = Memory(buff_capability)
        self.buffer_capability = buff_capability
        self.buffer_count = 0
        self.update_steps = update_steps
        self.gamma = gamma
        self.policy_optim = torch.optim.Adam(self.mlp_policy.parameters(), lr=1e-4)
        self.valu_optim = torch.optim.Adam(self.mlp_value.parameters(), lr=2e-4)

        self.orientation_optim = torch.optim.Adam(self.mlp_orientation.parameters(), lr=1e-4)
        self.orientation_value_optim = torch.optim.Adam(self.mlp_orientation_value.parameters(), lr=2e-4)

        self.cnn_policy_optim = torch.optim.Adam(self.cnn_policy.parameters(), lr=1e-4)
        self.cnn_orientation_optim = torch.optim.Adam(self.cnn_orientation.parameters(), lr=1e-3)

        # self.optimizer2 = torch.optim.Adam([{'params': self.mlp_policy.parameters(), 'lr': 1e-4},
        #                                     {'params': self.mlp_value.parameters(), 'lr': 2e-4}])

        self.loss_steps = 0
        self.loss_mlp_steps = 0
        self.orientation_mlp_steps = 0
        self.orientation_steps = 0
        self.ppo_epoch = ppo_epoch
        self.cnn_epoch = 30
        self.batch_size = batch_size

        self.writer = SummaryWriter('./logs_grasp_divide')
        # self.logger = Logger('./logs_grasp')

        # self.mlplogger = Logger('./logs_grasp')

    def learn(self):
        self.steps += 1
        experiences = self.buffer.sample(self.batch_size)
        self.mlp_old_policy.load_state_dict(self.mlp_policy.state_dict())
        self.mlp_old_orientation.load_state_dict(self.mlp_orientation.state_dict())

        # state , action, orientation_action, next_state, reward, done, observation, next_observation, current_position, next_position
        state, action, orientation_action, next_state, reward, done, ob, next_ob, pos, next_pos = zip(*experiences)
        state = np.asarray(state)
        next_state = np.asarray(next_state)
        reward = np.asarray(reward)
        pos = np.asarray(pos)
        orientation_action = np.asarray(orientation_action)

        # print(state.shape, orientation_action.shape)
        orientation = state[:, -1]
        state = state[:, 0:2]
        next_orientation = next_state[:, -1]
        next_state = next_state[:, 0:2]
        orientation_reward = reward[:, -1]
        reward = reward[:, 0]
        orientation_pos = pos[:, -1]
        pos = pos[:, 0:2]

        orientation = torch.FloatTensor(orientation).cuda().unsqueeze(1)
        state = torch.FloatTensor(state).cuda()
        action = torch.FloatTensor(action).cuda()
        orientation_action = torch.FloatTensor(orientation_action).cuda().unsqueeze(1)
        next_state = torch.FloatTensor(next_state).cuda()
        next_orientation = torch.FloatTensor(next_orientation).cuda().unsqueeze(1)
        orientation_reward = torch.FloatTensor(orientation_reward).cuda().unsqueeze(1)
        reward = torch.FloatTensor(reward).cuda().unsqueeze(1)
        done = torch.FloatTensor(done).cuda().unsqueeze(1)

        ob = torch.FloatTensor(ob).cuda().unsqueeze(1)
        pos = torch.FloatTensor(pos).cuda()
        orientation_pos = torch.FloatTensor(orientation_pos).cuda().unsqueeze(1)

        # print(state.shape, action.shape, next_state.shape,
        #       reward.shape, ob.shape,
        #       pos.shape, orientation_pos.shape,
        #       orientation_reward.shape, orientation_action.shape, orientation.shape)
        # print(1-done)

        with torch.no_grad():
            old_mean, old_std = self.mlp_old_policy(state)
            old_normal = Normal(old_mean, old_std)

            target_v = reward + self.gamma * (1-done) * self.mlp_value(next_state)
            advantage = (target_v - self.mlp_value(state))
        # print('target v : ', target_v.shape)
        # print('advantage : ', advantage)

        for _ in range(self.ppo_epoch):
            self.loss_mlp_steps += 1
            mean, std = self.mlp_policy(state)
            n = Normal(mean, std)
            action_log_prob = n.log_prob(action)
            old_action_prob = old_normal.log_prob(action)
            # print('before : ', action_log_prob.shape, old_action_prob.shape)
            action_log_prob = torch.sum(action_log_prob, dim=1, keepdim=True)
            old_action_prob = torch.sum(old_action_prob, dim=1, keepdim=True)
            # print('position : ', action_log_prob.shape, old_action_prob.shape)
            ratio = torch.exp(action_log_prob - old_action_prob)
            # print(ratio)
            # ratio = torch.mean(ratio, 1, keepdim=True)
            # print(ratio.shape)
            # print('advantage [index] : ', advantage[index].shape)
            L1 = ratio * advantage
            # print(ratio)
            L2 = torch.clamp(ratio, 0.8, 1.2) * advantage
            action_loss = torch.min(L1, L2)  # + 1e-3 * n.entropy()
            action_loss = - action_loss.mean()

            value_loss = F.mse_loss(self.mlp_value(state), target_v)

            self.policy_optim.zero_grad()
            action_loss.backward()
            self.policy_optim.step()
            self.valu_optim.zero_grad()
            value_loss.backward()
            self.valu_optim.step()

            self.writer.add_scalar('action loss mlp', action_loss.item(), self.loss_mlp_steps)
            self.writer.add_scalar('value loss mlp', value_loss.item(), self.loss_mlp_steps)
            #
            # self.optimizer2.zero_grad()
            # action_loss.backward()
            # value_loss.backward()
            # self.optimizer2.step()

        with torch.no_grad():
            old_mean, old_std = self.mlp_old_orientation(orientation)
            old_normal = Normal(old_mean, old_std)

            target_v = orientation_reward + self.gamma * (1-done) * self.mlp_orientation_value(next_orientation)
            advantage = (target_v - self.mlp_orientation_value(orientation))
        # print('target v : ', target_v.shape)
        # print('advantage : ', advantage)

        for _ in range(self.ppo_epoch):
            self.orientation_mlp_steps += 1
            mean, std = self.mlp_orientation(orientation)
            n = Normal(mean, std)
            action_log_prob = n.log_prob(orientation_action)
            old_action_prob = old_normal.log_prob(orientation_action)
            # print('before : ', action_log_prob.shape, old_action_prob.shape)
            action_log_prob = torch.sum(action_log_prob, dim=1, keepdim=True)
            old_action_prob = torch.sum(old_action_prob, dim=1, keepdim=True)
            # print(action_log_prob.shape, old_action_prob.shape)
            ratio = torch.exp(action_log_prob - old_action_prob)
            # print(ratio)
            # ratio = torch.mean(ratio, 1, keepdim=True)
            # print(ratio.shape)
            # print('advantage [index] : ', advantage[index].shape)
            L1 = ratio * advantage
            # print(ratio)
            L2 = torch.clamp(ratio, 0.8, 1.2) * advantage
            action_loss = torch.min(L1, L2)  # + 1e-3 * n.entropy()
            action_loss = - action_loss.mean()

            value_loss = F.mse_loss(self.mlp_orientation_value(orientation), target_v)

            self.orientation_optim.zero_grad()
            action_loss.backward()
            self.orientation_optim.step()
            self.orientation_value_optim.zero_grad()
            value_loss.backward()
            self.orientation_value_optim.step()

            self.writer.add_scalar('orientation loss mlp', action_loss.item(), self.orientation_mlp_steps)
            self.writer.add_scalar('orientation value loss mlp', value_loss.item(), self.orientation_mlp_steps)
            #
            # self.optimizer2.zero_grad()
            # action_loss.backward()
            # value_loss.backward()
            # self.optimizer2.step()

        with torch.no_grad():
            old_mean, old_std = self.mlp_policy(state)
            old_normal = Normal(old_mean, old_std)
            # print('position \n old mean: ', old_mean, '\n old std: ', old_std)
        for _ in range(self.cnn_epoch):
            self.loss_steps += 1
            mean, std = self.cnn_policy(ob, pos)
            n = Normal(mean, std)
            # print('position \n new mean: ', mean, '\n new std: ', std)
            kl_loss = torch.distributions.kl.kl_divergence(n, old_normal)
            # print('position : ', kl_loss)
            kl_loss = kl_loss.sum(dim=1, keepdim=True)

            # print(kl_loss.shape)
            kl_loss = kl_loss.mean()
            # print(kl_loss)
            # kl_loss.register_hook(lambda g: self.writer.add_scalar('position gradient ', g.item(), self.orientation_steps))

            self.cnn_policy_optim.zero_grad()
            kl_loss.backward()
            self.cnn_policy_optim.step()
            self.writer.add_scalar('action loss', kl_loss.item(), self.loss_steps)

            # info = {'action loss': action_loss.item(), 'value loss': value_loss.item(),
            #         'distribution ratio': ratio.mean().item()}
            # for tag, value in info.items():
            #     self.logger.scalar_summary(tag, value, step=self.loss_steps)
        with torch.no_grad():
            old_mean, old_std = self.mlp_orientation(orientation)
            old_normal = Normal(old_mean, old_std)
            # print('orientation \n old mean: ', old_mean, '\n old std: ', old_std)
        for _ in range(self.cnn_epoch):
            self.orientation_steps += 1
            mean, std = self.cnn_orientation(ob, orientation_pos)
            n = Normal(mean, std)
            # print('orientation \n new mean: ', mean, '\n new std: ', std)
            kl_loss = torch.distributions.kl.kl_divergence(n, old_normal)
            # print('orientation : ', kl_loss)
            kl_loss = kl_loss.sum(dim=1, keepdim=True)
            # print(kl_loss.shape)
            kl_loss = kl_loss.mean()
            # print(kl_loss)
            # kl_loss.register_hook(lambda g: self.writer.add_scalar('orientation gradient ', g.item(), self.orientation_steps))
            self.cnn_orientation_optim.zero_grad()
            kl_loss.backward()
            self.cnn_orientation_optim.step()
            self.writer.add_scalar('orientation action loss', kl_loss.item(), self.orientation_steps)

            # info = {'action loss': action_loss.item(), 'value loss': value_loss.item(),
            #         'distribution ratio': ratio.mean().item()}
            # for tag, value in info.items():
            #     self.logger.scalar_summary(tag, value, step=self.loss_steps)

    def learn_orientation(self):
        self.steps += 1
        experiences = self.buffer.sample(self.batch_size)
        self.mlp_old_orientation.load_state_dict(self.mlp_orientation.state_dict())

        # state , action, orientation_action, next_state, reward, done, observation, next_observation, current_position, next_position
        state, action, orientation_action, next_state, reward, done, ob, next_ob, pos, next_pos = zip(*experiences)
        state = np.asarray(state)
        next_state = np.asarray(next_state)
        reward = np.asarray(reward)
        pos = np.asarray(pos)
        orientation_action = np.asarray(orientation_action)

        # print(state.shape, orientation_action.shape)
        orientation = state[:, -1]
        state = state[:, 0:2]
        next_orientation = next_state[:, -1]
        next_state = next_state[:, 0:2]
        orientation_reward = reward[:, -1]
        reward = reward[:, 0]
        orientation_pos = pos[:, -1]
        pos = pos[:, 0:2]

        orientation = torch.FloatTensor(orientation).cuda().unsqueeze(1)
        state = torch.FloatTensor(state).cuda()
        action = torch.FloatTensor(action).cuda()
        orientation_action = torch.FloatTensor(orientation_action).cuda().unsqueeze(1)
        next_state = torch.FloatTensor(next_state).cuda()
        next_orientation = torch.FloatTensor(next_orientation).cuda().unsqueeze(1)
        orientation_reward = torch.FloatTensor(orientation_reward).cuda().unsqueeze(1)
        reward = torch.FloatTensor(reward).cuda().unsqueeze(1)
        done = torch.FloatTensor(done).cuda().unsqueeze(1)

        ob = torch.FloatTensor(ob).cuda().unsqueeze(1)
        pos = torch.FloatTensor(pos).cuda()
        orientation_pos = torch.FloatTensor(orientation_pos).cuda().unsqueeze(1)

        with torch.no_grad():
            old_mean, old_std = self.mlp_old_orientation(orientation)
            old_normal = Normal(old_mean, old_std)

            target_v = orientation_reward + self.gamma * (1-done) * self.mlp_orientation_value(next_orientation)
            advantage = (target_v - self.mlp_orientation_value(orientation))
        # print('target v : ', target_v.shape)
        # print('advantage : ', advantage)

        for _ in range(self.ppo_epoch):
            self.orientation_mlp_steps += 1
            mean, std = self.mlp_orientation(orientation)
            n = Normal(mean, std)
            action_log_prob = n.log_prob(orientation_action)
            old_action_prob = old_normal.log_prob(orientation_action)
            # print('before : ', action_log_prob.shape, old_action_prob.shape)
            action_log_prob = torch.sum(action_log_prob, dim=1, keepdim=True)
            old_action_prob = torch.sum(old_action_prob, dim=1, keepdim=True)
            # print(action_log_prob.shape, old_action_prob.shape)
            ratio = torch.exp(action_log_prob - old_action_prob)
            # print(ratio)
            # ratio = torch.mean(ratio, 1, keepdim=True)
            # print(ratio.shape)
            # print('advantage [index] : ', advantage[index].shape)
            L1 = ratio * advantage
            # print(ratio)
            L2 = torch.clamp(ratio, 0.8, 1.2) * advantage
            action_loss = torch.min(L1, L2)  # + 1e-3 * n.entropy()
            action_loss = - action_loss.mean()

            value_loss = F.mse_loss(self.mlp_orientation_value(orientation), target_v)

            self.orientation_optim.zero_grad()
            action_loss.backward()
            self.orientation_optim.step()
            self.orientation_value_optim.zero_grad()
            value_loss.backward()
            self.orientation_value_optim.step()

            self.writer.add_scalar('orientation loss mlp', action_loss.item(), self.orientation_mlp_steps)
            self.writer.add_scalar('orientation value loss mlp', value_loss.item(), self.orientation_mlp_steps)
            #
            # self.optimizer2.zero_grad()
            # action_loss.backward()
            # value_loss.backward()
            # self.optimizer2.step()

        with torch.no_grad():
            old_mean, old_std = self.mlp_orientation(orientation)
            old_normal = Normal(old_mean, old_std)

        for _ in range(self.cnn_epoch):
            self.orientation_steps += 1
            mean, std = self.cnn_orientation(ob, orientation_pos)
            n = Normal(mean, std)

            kl_loss = torch.distributions.kl.kl_divergence(n, old_normal)
            # kl_loss.register_hook(lambda g: print('origin gradient: ', g))
            # print(kl_loss.shape)
            # kl_loss = kl_loss.sum(dim=1, keepdim=True)
            # print(kl_loss.shape)
            kl_loss = kl_loss.mean()
            # kl_loss.register_hook(lambda g: print('mean gradient: ', g))
            # print(kl_loss)
            # time.sleep(1)
            self.cnn_orientation_optim.zero_grad()
            kl_loss.backward()
            self.cnn_orientation_optim.step()
            self.writer.add_scalar('orientation action loss', kl_loss.item(), self.orientation_steps)

            # info = {'action loss': action_loss.item(), 'value loss': value_loss.item(),
            #         'distribution ratio': ratio.mean().item()}
            # for tag, value in info.items():
            #     self.logger.scalar_summary(tag, value, step=self.loss_steps)

    def select_action(self, state):
        state = torch.FloatTensor(state).cuda().unsqueeze(0)
        # print('state : ', state)
        with torch.no_grad():
            mean, std = self.mlp_policy(state)
        # print('mean, std', mean.shape, std.shape, mean, std)
        dist = Normal(mean, std)
        action = dist.sample()
        # action_log_prob = dist.log_prob(action)
        action = action.clamp(-1, 1)
        # print(action_log_prob.shape, action.shape)
        # print(action)
        return action.cpu().squeeze().numpy()
        # return action.cpu().numpy()

    def select_orientation(self, state):
        state = torch.FloatTensor(state).cuda().unsqueeze(0)
        # print('state : ', state)
        with torch.no_grad():
            mean, std = self.mlp_orientation(state)
        # print('mean, std', mean.shape, std.shape, mean, std)
        dist = Normal(mean, std)
        action = dist.sample()
        # action_log_prob = dist.log_prob(action)
        action = action.clamp(-1, 1)
        # print(action_log_prob.shape, action.shape)
        # print(action)
        return action.cpu().squeeze().numpy()
        # return action.cpu().numpy()

    def select_action_cnn(self, frame, pos):
        frame = torch.FloatTensor(frame).cuda().unsqueeze(0).unsqueeze(1)
        pos = torch.FloatTensor(pos).cuda().unsqueeze(0)
        # print('pos : ', pos, pos.shape)
        # print('state : ', state)
        with torch.no_grad():
            mean, std = self.cnn_policy(frame, pos)
        # print('mean, std', mean.shape, std.shape, mean, std)
        dist = Normal(mean, std)
        action = dist.sample()
        # action_log_prob = dist.log_prob(action)
        action = action.clamp(-1, 1)
        # print(action_log_prob.shape, action.shape)
        # print(action)
        return action.cpu().squeeze().numpy()
        # return action.cpu().numpy()

    def select_orientation_cnn(self, frame, pos):
        frame = torch.FloatTensor(frame).cuda().unsqueeze(0).unsqueeze(1)
        pos = torch.FloatTensor(pos).cuda().unsqueeze(0)
        # print('pos : ', pos, pos.shape)
        # print('state : ', state)
        with torch.no_grad():
            mean, std = self.cnn_orientation(frame, pos)
        # print('mean, std', mean.shape, std.shape, mean, std)
        dist = Normal(mean, std)
        action = dist.sample()
        # action_log_prob = dist.log_prob(action)
        action = action.clamp(-1, 1)
        # print(action_log_prob.shape, action.shape)
        # print(action)
        return action.cpu().squeeze().numpy()

    def save_model(self, policy_path, value_path, orientation, orientation_value, cnn_policy_path, cnn_orientation):
        save_path = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + policy_path
        save_value = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + value_path
        save_path_orientation = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + orientation
        save_value_orientation = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + orientation_value
        save_path_cnn = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + cnn_policy_path
        save_orientation_cnn = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + cnn_orientation
        torch.save(self.mlp_policy.state_dict(), save_path)
        torch.save(self.mlp_value.state_dict(), save_value)
        torch.save(self.mlp_orientation.state_dict(), save_path_orientation)
        torch.save(self.mlp_orientation_value.state_dict(), save_value_orientation)
        torch.save(self.cnn_policy.state_dict(), save_path_cnn)
        torch.save(self.cnn_orientation.state_dict(), save_orientation_cnn)
        print('model_divide saved in ' + str(time.asctime(time.localtime(time.time()))))

    def save_orientation(self, cnn_orientation):
        save_orientation_cnn = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + cnn_orientation
        torch.save(self.cnn_orientation.state_dict(), save_orientation_cnn)
        print('model_divide saved in ' + str(time.asctime(time.localtime(time.time()))))

    def load_model(self, policy_path, value_path, orientation, orientation_value, cnn_policy_path, cnn_orientation):
        save_path = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + policy_path
        save_value = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + value_path
        save_path_orientation = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + orientation
        save_value_orientation = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + orientation_value
        save_path_cnn = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + cnn_policy_path
        save_orientation_cnn = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + cnn_orientation
        self.mlp_policy.load_state_dict(torch.load(save_path))
        self.mlp_value.load_state_dict(torch.load(save_value))
        self.mlp_orientation.load_state_dict(torch.load(save_path_orientation))
        self.mlp_orientation_value.load_state_dict(torch.load(save_value_orientation))

        self.cnn_policy.load_state_dict(torch.load(save_path_cnn))
        self.cnn_orientation.load_state_dict(torch.load(save_orientation_cnn))
        print(policy_path + ' model_divide is loaded!')

    def load_cnn_model(self, cnn_policy_path, cnn_orientation):
        save_path_cnn = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + cnn_policy_path
        save_orientation_cnn = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model_divide/' + cnn_orientation
        self.cnn_policy.load_state_dict(torch.load(save_path_cnn))
        self.cnn_orientation.load_state_dict(torch.load(save_orientation_cnn))
        print('cnn model loaded!')

    def load_part_model(self, cnn_pos_policy):
        save_path_cnn = '/home/chen/PycharmProjects/Reinforcement/VisualGrasp/model/' + cnn_pos_policy
        part_model_dict = torch.load(save_path_cnn)
        orientation_dict = self.cnn_orientation.state_dict()
        part_model_dict = {k: v for k, v in part_model_dict.items() if k in orientation_dict}
        orientation_dict.update(part_model_dict)
        self.cnn_orientation.load_state_dict(orientation_dict)
        print('part model loaded!')


import torch
import numpy as np
import gym
from algorithms.q_learning_cnn import Q_learning
from custom_envs.custom_acrobat import AcrobotEnv
from custom_envs.custom_mountain_car_discrete import MountainCarEnv
import gym_sokoban
from atari_wrapper import make_atari, wrap_deepmind
import gym_miniworld

from model import NN_Paramters
from parameters import Algo_Param, Save_Paths, Load_Paths
# Use the Baseline Atari environment because of Deepmind helper functions
"""
env = gym.make('MiniWorld-CollectHealth-v0')
env2 = gym.make('MiniWorld-CollectHealth-v0')
state = env.reset()
print(np.shape(state))
"""
# Warp the frames, grey scale, stake four frame and scale to smaller ratio
#env = wrap_deepmind(env, frame_stack=True, scale=True)

#env = MountainCarEnv()
#env2 = MountainCarEnv()
env = gym.make('MiniGrid-Empty-8x8-v0')
env2 = gym.make('MiniGrid-Empty-8x8-v0')

print(env.action_space.shape)
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n

q_nn_param = NN_Paramters(state_dim, action_dim, hidden_layer_dim=[512],
                          non_linearity=torch.relu, device=torch.device("cuda"), l_r=0.001)
algo_param = Algo_Param()
algo_param.gamma = 0.99

max_episodes = 200

Q = Q_learning(env, q_nn_param=q_nn_param, algo_param=algo_param)
print(Q.Q)
#env2 = env
#Q.load("q", "target_q")
update_interval = 10
save_interval = 1000
eval_interval = 10000
state = Q.initalize()

for i in range(100000):
    print(i)
    if i%100 == 0:
        print(i)
    Q.update()
    state = Q.step(state)
    if i%update_interval == 0:
        Q.hard_update()
    if i%save_interval == 0:
        print("saving")
        Q.save("q", "target_q")
    if i%eval_interval == 0:
        s = env2.reset()

        i_s = s
        rew = 0
        for j in range(max_episodes):
            a = Q.get_action(s)
            s, r, d, _ = env2.step(a)
            rew += r
            env2.render()
            if j == max_episodes-1:
                d = True
            if d == True:
                break
        print("reward at itr " + str(i) + " = " + str(rew) )
#Q.memory = torch.load("mem")
torch.save(Q.memory, "mem")


# compile cython modules
import os
os.system('python experience_replay_setup.py build_ext --inplace')

# load dependencies
import tensorflow as tf
physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
tf.config.experimental.set_memory_growth(physical_devices[0], True)

from tensorflow.keras.models import load_model
from tensorflow.keras.optimizers import RMSprop, Adam
from tensorflow.keras import initializers

import gym

import numpy as np

from deep_q_agents import EpsAnnDQNAgent
from deep_q_networks import DeepQNetwork
from experience_replay import PrioritizedExperienceReplay
from atari_preprocessing import atari_montezuma_processor, ProcessedAtariEnv
from openai_baseline_wrappers import make_atari, wrap_deepmind
from load_data import LoadAtariHeadData

#create environment
frame_processor = atari_montezuma_processor
game_id = 'MontezumaRevengeNoFrameskip-v4'
game_name = 'montezuma_revenge'
env = make_atari(game_id)
env = wrap_deepmind(env)
env = ProcessedAtariEnv(env, frame_processor, reward_processor = lambda x: np.sign(x) * np.log(1 + np.abs(x)))

# additional env specific parameters
frame_shape = env.reset().shape
frame_skip = 4
num_stacked_frames = 4
num_actions = env.action_space.n

# replay parameters
batch_size = 32
max_frame_num = 2**20
prioritized_replay = True
prio_coeff = 0.4
is_schedule = [0.6, 1.0, 1500000]
replay_epsilon = 0.001
expert_epsilon = 1.0
memory_restore_path = None

# network training parameters
dueling = True
double_q = True
lr_schedule = [[0.0000625, 0.0000625, 10000000]]
optimizer = Adam
discount_factor = 0.99
n_step = 10
one_step_weight = 1.0/3.0
n_step_weight = 1.0/3.0
expert_weight = 1.0/3.0
l2_weight = 0.00001
large_margin_coeff = 0.8
model_restore_path = None

# network architecture
conv_layers = {'filters': [32, 64, 64, 1024],
               'kernel_sizes': [8, 4, 3, 7],
               'strides': [4, 2, 1, 1],
               'paddings': ['valid' for _ in range(4)],
               'activations': ['relu' for _ in range(4)],
               'initializers': [initializers.VarianceScaling(scale = 2.0) for _ in range(4)],
               'names': ['conv_%i'%(i) for i in range(1,5)]}
dense_layers = None

# exploration parameters
eps_schedule = [[0.25, 0.1, 250000],
                [0.1, 0.01, 5000000],
                [0.01, 0.001, 5000000]]

# training session parameters
target_interval = 10000
warmup_steps = 50000
pretrain_steps = 500000
learning_interval = 4
num_steps = 12000000
num_episodes = 10000
max_steps_per_episode = 18000
output_freq = 1000
save_freq = 500
store_memory = True
save_path = "experiments/montezuma_standard_experiment/"

# create replay memory
memory = PrioritizedExperienceReplay(frame_shape = frame_shape,
                                     max_frame_num = max_frame_num,
                                     num_stacked_frames = num_stacked_frames,
                                     batch_size = batch_size,
                                     prio_coeff = prio_coeff,
                                     is_schedule = is_schedule,
                                     epsilon = replay_epsilon,
                                     restore_path = memory_restore_path)

# expert memory
data_loader = LoadAtariHeadData(game_name = game_name, frame_processor = frame_processor)
expert_memory = data_loader.demonstrations_to_per(max_frame_num = max_frame_num,
                                                  num_stacked_frames = num_stacked_frames,
                                                  frame_shape = frame_shape,
                                                  batch_size = batch_size,
                                                  prio_coeff = prio_coeff,
                                                  is_schedule = is_schedule,
                                                  epsilon = expert_epsilon,
                                                  recompute_demonstrations = True,
                                                  only_highscore = False,
                                                  frame_skip = frame_skip)

# create policy network
policy_network = DeepQNetwork(in_shape = (num_stacked_frames, *frame_shape),
                              conv_layers = conv_layers,
                              dense_layers = dense_layers,
                              num_actions = num_actions,
                              optimizer = optimizer,
                              lr_schedule = lr_schedule,
                              dueling = dueling,
                              one_step_weight = one_step_weight,
                              n_step_weight = n_step_weight,
                              expert_weight = expert_weight)

if model_restore_path is not None:
    policy_network.model.load_weights(model_restore_path, by_name = True)

# create target network
target_network = DeepQNetwork(in_shape = (num_stacked_frames, *frame_shape),
                              conv_layers = conv_layers,
                              dense_layers = dense_layers,
                              num_actions = num_actions,
                              optimizer = optimizer,
                              lr_schedule = lr_schedule,
                              dueling = dueling,
                              one_step_weight = one_step_weight,
                              n_step_weight = n_step_weight,
                              expert_weight = expert_weight)

if model_restore_path is not None:
    target_network.model.load_weights(model_restore_path, by_name = True)

# create agent
agent = EpsAnnDQNAgent(env = env,
                       memory = memory,
                       policy_network = policy_network,
                       target_network = target_network,
                       num_actions = num_actions,
                       frame_shape = frame_shape,
                       discount_factor = discount_factor,
                       save_path = save_path,
                       eps_schedule = eps_schedule,
                       double_q = double_q,
                       n_step = n_step,
                       expert_memory = expert_memory,
                       prioritized_replay = prioritized_replay)

agent.policy_network.model.save(save_path + "/trained_models/initial_model.h5")

# train the agent
agent.train(num_episodes = num_episodes,
            num_steps = num_steps,
            max_steps_per_episode = max_steps_per_episode,
            warmup_steps = warmup_steps,
            pretrain_steps = pretrain_steps,
            target_interval = target_interval,
            learning_interval = learning_interval,
            output_freq = output_freq,
            save_freq = save_freq,
            store_memory = store_memory)