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

__all__ = ['PolicyNetwork', 'uniform_initializer', 'CriticNetwork']

# Cell
import numpy as np
import sonnet as snt
import tensorflow as tf
from acme import specs
from acme.tf import utils as tf2_utils
from acme.tf import networks
from acme import types

uniform_initializer = tf.initializers.VarianceScaling(
    distribution='uniform', mode='fan_out', scale=0.333)

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),
            # snt.LayerNorm(axis=slice(
            #     1, None), create_scale=True, create_offset=True),
            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, 10*3], axis=1)
        laser = self.conv(laser)

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

        return x

# Cell
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),
            # snt.LayerNorm(axis=slice(
            #     1, None), create_scale=True, create_offset=True),
            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, 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