# AUTOGENERATED! DO NOT EDIT! File to edit: 02_pokingbot_training.ipynb (unless otherwise specified).

__all__ = ['PolicyNetwork', 'CriticNetwork', 'make_environment', 'make_networks', 'train_loop']

# Cell
from typing import Mapping, Sequence
from absl import app
from absl import flags
import acme
from acme import specs
from acme import types
from acme import wrappers
from acme.agents.tf import actors
from acme.agents.tf import d4pg
from acme.tf import networks
from acme.tf import utils as tf2_utils
from acme.tf import utils as tf2_utils
from acme.tf import networks
from acme import types
import dm_env
import gym
import numpy as np
import sonnet as snt
import sonnet as snt
import tensorflow as tf

# Cell
class PolicyNetwork(snt.Module):
    def __init__(self, action_spec):
        super(PolicyNetwork, self).__init__()

        # Get total number of action dimensions from action spec.
        output_dim = np.prod(action_spec.shape, dtype=int)

        self.conv = snt.Sequential([
            snt.Reshape([-1, 4]),
            snt.Conv1D(output_channels=32, kernel_shape=3, stride=1),
            tf.nn.relu,
            snt.Conv1D(output_channels=32, kernel_shape=3, stride=1),
            tf.nn.relu,
            snt.Flatten(),
        ])

        self.hidden = snt.Sequential([
            snt.Linear(512, w_init=uniform_initializer),
            tf.nn.elu,
            snt.Linear(512, w_init=uniform_initializer),
            tf.nn.elu,
            snt.Linear(256, w_init=uniform_initializer),
            tf.nn.elu,
            snt.Linear(256, w_init=uniform_initializer),
            tf.nn.elu,
        ])

        self.out = snt.Sequential([
            snt.Linear(
                output_dim, w_init=tf.initializers.VarianceScaling(1e-4)),
            networks.TanhToSpec(action_spec),
        ])

    def __call__(self, x):
        x = tf2_utils.batch_concat(x)

        laser, g= tf.split(x, [4*(241*2), 10*3], axis=1)
        laser = self.conv(laser)

        x = tf.concat([g, laser], 1)
        x = self.hidden(x)
        x = self.out(x)

        return x


class CriticNetwork(snt.Module):
    def __init__(self, vmin, vmax, num_atoms):
        super(CriticNetwork, self).__init__()

        self.conv = snt.Sequential([
            snt.Reshape([-1, 4]),
            snt.Conv1D(output_channels=32, kernel_shape=3, stride=1),
            tf.nn.relu,
            snt.Conv1D(output_channels=32, kernel_shape=3, stride=1),
            tf.nn.relu,
            snt.Flatten(),
        ])

        self.hidden = snt.Sequential([
            snt.Linear(512, w_init=uniform_initializer),
            tf.nn.elu,
            snt.Linear(512, w_init=uniform_initializer),
            tf.nn.elu,
            snt.Linear(256, w_init=uniform_initializer),
            tf.nn.elu,
            snt.Linear(256, w_init=uniform_initializer),
            tf.nn.elu,
        ])

        self.head = networks.DiscreteValuedHead(vmin, vmax, num_atoms)

    def __call__(self,
                 observation: types.NestedTensor,
                 action: types.NestedTensor) -> tf.Tensor:

        if hasattr(observation, 'dtype') and hasattr(action, 'dtype'):
            if observation.dtype != action.dtype:
                # Observation and action must be the same type for concat to work
                action = tf.cast(action, observation.dtype)

        x = tf2_utils.batch_concat(observation)
        action = tf2_utils.batch_concat(action)

        laser, g= tf.split(x, [4*(241*2), 10*3], axis=1)
        laser = self.conv(laser)

        x = tf.concat([g, laser, action], 1)
        x = self.hidden(x)
        x = self.head(x)

        return x


# Cell
def make_environment(
        task: str = 'MountainCarContinuous-v0') -> dm_env.Environment:
    """
    Creates an OpenAI Gym environment.
    Args:
        gym_id (String): The custom OpenAI gym environment have been registered.
    Returns:
        environment (env class)
    """
    environment = gym.make(task)
    environment = wrappers.GymWrapper(environment)
    environment = wrappers.SinglePrecisionWrapper(environment)

    return environment




def make_networks(
    action_spec: specs.BoundedArray,
    vmin: float = -40.,
    vmax: float = 130.,
    num_atoms: int = 51,
) -> Mapping[str, types.TensorTransformation]:
    """
    Creates the networks used by the agent.
    vmin/vmax and num_atoms hyperparameters should be set to
    give the distributional critic a good dynamic range over possible discounted
    returns. Note that this is very different than the scale of immediate rewards.

    The vmin|vmax hyperparameters of the distributional critic may need tuning
    depending on your environment's rewards. A good rule of thumb is to
    set vmax to the discounted sum of the maximum instantaneous rewards
    for the maximum episode length;
    then set vmin to -vmax.

    Args:
        action_spec (specs.BoundedArray)
        vmin (float)
        vmax (float)
        num_atoms (int)
    Returns:
        policy_network (snt.Module)
        critic_network (snt.Module)
        observation_network (tf2_utils.batch_concat)
    """

    observation_network = tf2_utils.batch_concat
    policy_network = PolicyNetwork(action_spec)
    critic_network = CriticNetwork(vmin, vmax, num_atoms)

    return {
        'policy': policy_network,
        'critic': critic_network,
        'observation': observation_network,
    }


def train_loop( gym_id,
                specified_path=None,
                num_episodes=350000000,
                discount=0.99,
                batch_size=512,
                prefetch_size=4,
                target_update_period=100,
                min_replay_size=10000,
                max_replay_size=200000,
                samples_per_insert=64.0,
                n_step=5,
                sigma=1.0,
                clipping=True):
    """
    Start running the train loop.
    Args:
        gym_id (String): The custom OpenAI gym environment have been registered. i.e. 'gym_poking:door-and-ball-v3'

        specified_path (String): While using a pretrained weight, the path of the specific weight, i.e.

                                 log_home = '/home/argsubt/pokingbot-rl/acme_logs/'
                                 log_id = '355a222e-7e90-11ec-9e19-04d4c401c495/'
                                 snapshots_path = 'snapshots/policy'
                                 checkpoints_path = 'checkpoints/d4pg_learner/'
                                 ckpt_id = 'ckpt-409'
                                 specified_path = log_home + log_id + checkpoints_path + ckpt_id


        Others hyperparameter for training.

    Returns:
        None
    """

    FLAGS = flags.FLAGS
    flags.DEFINE_integer('num_episodes', num_episodes,
                         'Number of training episodes to run for.')

    environment = make_environment(gym_id)
    environment_spec = specs.make_environment_spec(environment)
    agent_networks = make_networks(environment_spec.actions)

    # change batch_size together with samples_per_insert
    # batch_size/samples_per_insert = 8
    # depends on the performance of workstation
    if specified_path==None:
        agent = d4pg.D4PG(
            environment_spec=environment_spec,
            policy_network=agent_networks['policy'],
            critic_network=agent_networks['critic'],
            observation_network=agent_networks['observation'],
            discount=discount,
            batch_size=batch_size,
            prefetch_size=prefetch_size,
            target_update_period=target_update_period,
            min_replay_size=min_replay_size,
            max_replay_size=max_replay_size,
            samples_per_insert=samples_per_insert,
            n_step=n_step,
            sigma=sigma,
            clipping=clipping
        )
    else:
        agent = d4pg.D4PG(
            environment_spec=environment_spec,
            policy_network=agent_networks['policy'],
            critic_network=agent_networks['critic'],
            observation_network=agent_networks['observation'],
            discount=discount,
            batch_size=batch_size,
            prefetch_size=prefetch_size,
            target_update_period=target_update_period,
            min_replay_size=min_replay_size,
            max_replay_size=max_replay_size,
            samples_per_insert=samples_per_insert,
            n_step=n_step,
            sigma=sigma,
            clipping=clipping,
            specified_path=specified_path
        )


    train_loop = acme.EnvironmentLoop(environment, agent, label='train_loop')
    train_loop.run(num_episodes=FLAGS.num_episodes)
