#!python

# This file is part of rddl2tf.

# rddl2tf is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# rddl2tf is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with rddl2tf. If not, see <http://www.gnu.org/licenses/>.

import argparse
import os
import tensorflow as tf

import rddlgym

import rddl2tf
from rddl2tf.compilers import DefaultCompiler


def parse_args():
    description = 'rddl2tf ({}): RDDL2TensorFlow compiler in Python3.'.format(rddl2tf.__version__)
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument(
        'rddl',
        type=str,
        help='path to RDDL file or rddlgym problem id'
    )
    parser.add_argument(
        '-b', '--batch-size',
        type=int, default=256,
        help='number of fluents in a batch (default=256)'
    )
    parser.add_argument(
        '--logdir',
        type=str, default='/tmp/rddl2tf',
        help='log directory for tensorboard graph visualization (default=/tmp/rddl2tf)'
    )
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()

    # parse RDDL into an AST
    model = rddlgym.make(args.rddl, mode=rddlgym.AST)

    # create a RDDL-to-TF compiler
    compiler = DefaultCompiler(model, batch_size=256)
    compiler.init()

    # compile initial state and default action fluents
    state = compiler.initial_state()
    action = compiler.default_action()

    # compile state invariants and action preconditions
    invariants = compiler.state_invariants(state)
    preconditions = compiler.action_preconditions(state, action)

    # compile action bounds
    bounds = compiler.action_bound_constraints(state)

    # compile intermediate fluents and next state fluents
    interms, next_state = compiler.cpfs(state, action)

    # compile reward function
    reward = compiler.reward(state, action, next_state)

    # save and visualize computation graph
    logdir = os.path.join(args.logdir, model.domain.name, model.instance.name)
    file_writer = tf.summary.FileWriter(logdir, compiler.graph)
    print('tensorboard --logdir {}\n'.format(logdir))
