import dill
import haiku as hk
import jax
from jax.experimental import optix
import jax.numpy as jnp

from dataset import load_data


MINERL_ENV = 'MineRLTreechopVectorObf-v0'
PARAMS_FILENAME = 'bc_params_treechop.pkl'


class PovStack(hk.Module):
    """ PovStack is a module for processing the point-of-view image data that
    comes from the agent's viewport. This input is in NHWC format for a shape
    of (N, 64, 64, 3).

    This model is inspired from
    https://github.com/minerllabs/baselines/blob/master/general/chainerrl/baselines/behavioral_cloning.py
    """
    def __init__(self, name=None):
        super().__init__(name=name)
        conv_0 = hk.Conv2D(output_channels=32,
                           kernel_shape=(8, 8),
                           stride=4,
                           padding='SAME',
                           name='conv_0')
        layer_0 = (conv_0, jax.nn.relu)

        conv_1 = hk.Conv2D(output_channels=64,
                           kernel_shape=(4, 4),
                           stride=2,
                           padding='SAME',
                           name='conv_1')
        layer_1 = (conv_1, jax.nn.relu)

        conv_2 = hk.Conv2D(output_channels=64,
                           kernel_shape=(3, 3),
                           stride=1,
                           padding='SAME',
                           name='conv_2')
        layer_2 = (conv_2, jax.nn.relu)

        layer_3 = (hk.Flatten(),
                   hk.Linear(512, name='fc_0'), jax.nn.relu,
                   hk.Linear(128, name='fc_1'), jax.nn.relu)

        self.layers = layer_0 + layer_1 + layer_2 + layer_3

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


class VectorStack(hk.Module):
    """ VectorStack is a module for processing the obfuscated "vector" data that
    is included in the agent's observation. This is a densely encoded form of
    the discrete information regarding the state of the agent other than the
    viewport, e.g. current inventory. The input is of shape (N, 64)
    """
    def __init__(self, name=None):
        super().__init__(name=name)
        layer_0 = (hk.Linear(32, name='fc_0'), jax.nn.relu)

        self.layers = layer_0

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


def behavioral_cloning(batch):
    """ The full forward model definition """
    x_0 = PovStack(name='pov_stack')(batch[0])
    x_1 = VectorStack(name='vector_stack')(batch[1])
    x = jnp.concatenate((x_0, x_1), axis=1)
    return jnp.tanh(hk.Linear(64)(x))


@jax.jit
def mse_loss(logits, labels):
    """ Mean Squared Error loss """
    return jnp.mean(jnp.power(logits - labels, 2))


def main():
    net = hk.transform(behavioral_cloning)
    opt = optix.adam(0.001)

    @jax.jit
    def loss(params, batch):
        """ The loss criterion for our model """
        logits = net.apply(params, None, batch)
        return mse_loss(logits, batch[2])

    @jax.jit
    def update(opt_state, params, batch):
        grads = jax.grad(loss)(params, batch)
        updates, opt_state = opt.update(grads, opt_state)
        params = optix.apply_updates(params, updates)
        return params, opt_state

    @jax.jit
    def accuracy(params, batch):
        """ Simply report the loss for the current batch """
        logits = net.apply(params, None, batch)
        return mse_loss(logits, batch[2])

    train_dataset, val_dataset = load_data(MINERL_ENV,
                                           batch_size=32, epochs=100)

    rng = jax.random.PRNGKey(2020)
    batch = next(train_dataset)
    params = net.init(rng, batch)
    opt_state = opt.init(params)

    for i, batch in enumerate(train_dataset):
        params, opt_state = update(opt_state, params, batch)
        if i % 1000 == 0:
            print(accuracy(params, val_dataset))
        
        if i % 10000 == 0:
            with open(PARAMS_FILENAME, 'wb') as fh:
                dill.dump(params, fh)

    with open(PARAMS_FILENAME, 'wb') as fh:
        dill.dump(params, fh)



if __name__ == '__main__':
    main()

[((79, 8, 79, 43), 'jax.numpy.concatenate', 'jnp.concatenate', (), '', True, 'import jax.numpy as jnp\n'), ((90, 10, 90, 42), 'haiku.transform', 'hk.transform', ({(90, 23, 90, 41): 'behavioral_cloning'}, {}), '(behavioral_cloning)', True, 'import haiku as hk\n'), ((91, 10, 91, 27), 'jax.experimental.optix.adam', 'optix.adam', ({(91, 21, 91, 26): '0.001'}, {}), '(0.001)', False, 'from jax.experimental import optix\n'), ((112, 33, 113, 69), 'dataset.load_data', 'load_data', (), '', False, 'from dataset import load_data\n'), ((115, 10, 115, 34), 'jax.random.PRNGKey', 'jax.random.PRNGKey', ({(115, 29, 115, 33): '2020'}, {}), '(2020)', False, 'import jax\n'), ((24, 17, 28, 41), 'haiku.Conv2D', 'hk.Conv2D', (), '', True, 'import haiku as hk\n'), ((31, 17, 35, 41), 'haiku.Conv2D', 'hk.Conv2D', (), '', True, 'import haiku as hk\n'), ((38, 17, 42, 41), 'haiku.Conv2D', 'hk.Conv2D', (), '', True, 'import haiku as hk\n'), ((86, 20, 86, 49), 'jax.numpy.power', 'jnp.power', ({(86, 30, 86, 45): '(logits - labels)', (86, 47, 86, 48): '(2)'}, {}), '(logits - labels, 2)', True, 'import jax.numpy as jnp\n'), ((103, 17, 103, 53), 'jax.experimental.optix.apply_updates', 'optix.apply_updates', ({(103, 37, 103, 43): 'params', (103, 45, 103, 52): 'updates'}, {}), '(params, updates)', False, 'from jax.experimental import optix\n'), ((130, 8, 130, 29), 'dill.dump', 'dill.dump', ({(130, 18, 130, 24): 'params', (130, 26, 130, 28): 'fh'}, {}), '(params, fh)', False, 'import dill\n'), ((45, 19, 45, 31), 'haiku.Flatten', 'hk.Flatten', ({}, {}), '()', True, 'import haiku as hk\n'), ((46, 19, 46, 46), 'haiku.Linear', 'hk.Linear', (), '', True, 'import haiku as hk\n'), ((47, 19, 47, 46), 'haiku.Linear', 'hk.Linear', (), '', True, 'import haiku as hk\n'), ((65, 19, 65, 45), 'haiku.Linear', 'hk.Linear', (), '', True, 'import haiku as hk\n'), ((80, 20, 80, 33), 'haiku.Linear', 'hk.Linear', ({(80, 30, 80, 32): '(64)'}, {}), '(64)', True, 'import haiku as hk\n'), ((101, 16, 101, 30), 'jax.grad', 'jax.grad', ({(101, 25, 101, 29): 'loss'}, {}), '(loss)', False, 'import jax\n'), ((127, 16, 127, 37), 'dill.dump', 'dill.dump', ({(127, 26, 127, 32): 'params', (127, 34, 127, 36): 'fh'}, {}), '(params, fh)', False, 'import dill\n')]