#!/usr/bin/env python3.6

""" Front-end script for training a Snake agent. """

import json
import sys

import tensorflow as tf

from agent import DeepQNetworkAgent
from gameplay.environment import Environment
from utils.argparser import HelpOnFailArgumentParser

Sequential = tf.keras.models.Sequential
Conv2D = tf.keras.layers.Conv2D
Activation = tf.keras.layers.Activation
Dense = tf.keras.layers.Dense
Flatten = tf.keras.layers.Flatten
SGD = tf.keras.optimizers.SGD


def parse_command_line_args(args):
    """
    Parse command-line arguments and organize them into a
    single structured object.
    """

    parser = HelpOnFailArgumentParser(
        description='Snake AI training client.',
        epilog='Example: train.py --level base_10x10.json --num-episodes 30000'
    )

    parser.add_argument(
        '--level',
        required=True,
        type=str,
        help='JSON file containing a level definition.',
    )
    parser.add_argument(
        '--num-episodes',
        required=True,
        type=int,
        default=30000,
        help='The number of episodes to run consecutively.',
    )

    return parser.parse_args(args)


def create_snake_environment(level_filename):
    """ Create a new Snake environment from the config file. """

    with open(level_filename) as cfg:
        env_config = json.load(cfg)

    return Environment(config=env_config, verbose=1)


def create_dqn_model(env, num_last_frames):
    """
    Build a new DQN model to be used for training.

    Args:
        env: an instance of Snake environment.
        num_last_frames: the number of last frames the agent considers as state.

    Returns:
        A compiled DQN model.
    """

    model = Sequential()
    model.add(Conv2D(
        16,
        kernel_size=(3, 3),
        strides=(1, 1),
        padding='SAME',
        input_shape=(num_last_frames,) + env.observation_shape
    ))
    model.add(Activation('selu'))
    model.add(Conv2D(
        32,
        kernel_size=(3, 3),
        strides=(1, 1),
        padding='SAME'
    ))
    model.add(Activation('selu'))

    # Dense layers.
    model.add(Flatten())
    model.add(Dense(256))
    model.add(Activation('selu'))
    model.add(Dense(env.num_actions))
    model.add(Activation('softmax'))

    model.summary()
    model.compile(SGD(lr=0.001), 'MSE')

    return model


def main():
    parsed_args = parse_command_line_args(sys.argv[1:])

    env = create_snake_environment(parsed_args.level)
    model = create_dqn_model(env, num_last_frames=4)

    agent = DeepQNetworkAgent(
        model=model,
        memory_size=-1,
        num_last_frames=model.input_shape[1]
    )
    agent.train(
        env,
        batch_size=64,
        num_episodes=parsed_args.num_episodes,
        checkpoint_freq=parsed_args.num_episodes // 10,
        discount_factor=0.95
    )


if __name__ == '__main__':
    main()