#!/usr/bin/env python

"""
Dumps the network topology as JSON on stdout.
"""

from __future__ import print_function

import sys
import argparse
import json
import typing

import _setup_returnn_env  # noqa
import returnn.__main__ as rnn
from returnn.log import log
from returnn.pretrain import pretrain_from_config
from returnn.config import network_json_from_config


config = None  # type: typing.Optional["returnn.config.Config"]


def init(config_filename, command_line_options):
  """
  :param str config_filename:
  :param list[str] command_line_options:
  """
  rnn.init_better_exchook()
  rnn.init_config(config_filename, command_line_options)
  global config
  config = rnn.config
  config.set("log", [])
  rnn.init_log()
  print("RETURNN dump-dataset starting up.", file=log.v3)
  rnn.init_config_json_network()


def main(argv):
  """
  Main entry.
  """
  arg_parser = argparse.ArgumentParser(description='Dump network as JSON.')
  arg_parser.add_argument('returnn_config_file')
  arg_parser.add_argument('--epoch', default=1, type=int)
  arg_parser.add_argument('--out', default="/dev/stdout")
  args = arg_parser.parse_args(argv[1:])
  init(config_filename=args.returnn_config_file, command_line_options=[])

  pretrain = pretrain_from_config(config)
  if pretrain:
    json_data = pretrain.get_network_json_for_epoch(args.epoch)
  else:
    json_data = network_json_from_config(config)

  f = open(args.out, 'w')
  print(json.dumps(json_data, indent=2, sort_keys=True), file=f)
  f.close()

  rnn.finalize()


if __name__ == '__main__':
  main(sys.argv)
