#!/usr/bin/env python3

"""
Benchmarking LSTM training times for different LSTM implementations, both on GPU and CPU.

Some runtimes:

GeForce GTX 680:
  Settings:
  {'batch_size': 2000,
   'chunking': '50:25',
   'max_seqs': 40,
   'n_hidden': 500,
   'num_layers': 5,
   'num_seqs': 500}
  Final results:
    GPU:CudnnLSTM: 0:00:12.3727
    GPU:NativeLSTM: 0:00:19.0306
    GPU:LSTMBlockFused: 0:00:41.5129
    GPU:LSTMBlock: 0:00:52.6296
    GPU:StandardLSTM: 0:00:53.2629
    GPU:BasicLSTM: 0:00:53.6150
    CPU:NativeLSTM: 0:02:37.1753
    CPU:LSTMBlockFused: 0:03:17.5448
    CPU:BasicLSTM: 0:03:28.9882
    CPU:StandardLSTM: 0:03:29.3730
    CPU:LSTMBlock: 0:03:33.7039

  Settings:
  {'batch_size': 2000,
   'chunking': '0',
   'max_seqs': 50,
   'n_hidden': 500,
   'num_layers': 5,
   'num_seqs': 500}
  Final results:
    GPU:CudnnLSTM: 0:00:12.8325
    GPU:NativeLSTM: 0:00:15.8518
    GPU:LSTMBlockFused: 0:01:05.2157
    GPU:StandardLSTM: 0:01:20.6564
    GPU:BasicLSTM: 0:01:20.9896
    GPU:LSTMBlock: 0:01:21.4046
    CPU:NativeLSTM: 0:02:12.5029
    CPU:LSTMBlockFused: 0:03:39.0270
    CPU:LSTMBlock: 0:03:51.9667
    CPU:StandardLSTM: 0:03:56.6404
    CPU:BasicLSTM: 0:03:58.1545
"""

from __future__ import annotations
import sys
import time
from argparse import ArgumentParser
from pprint import pprint

import _setup_returnn_env  # noqa
from returnn.util import better_exchook
from returnn.log import log
from returnn.config import Config
from returnn.util.basic import hms_fraction, describe_returnn_version, describe_tensorflow_version
from returnn.tf.engine import Engine
from returnn.tf.util.basic import setup_tf_thread_pools, is_gpu_available, print_available_devices
from returnn.datasets import init_dataset, Dataset


LstmCellTypes = [
    "BasicLSTM",
    "StandardLSTM",
    "LSTMBlock",
    "LSTMBlockFused",
    "NativeLSTM",
    "NativeLstm2",
    "NativeLstmLowMem",
    "CudnnLSTM",
]

GpuOnlyCellTypes = ["CudnnLSTM"]

# Fixed by dataset. See make_config_dict().
_input_dim = 9
_output_dim = 2

# You can play around with these. E.g. use "chunking=0", "max_seqs=20" as command-line args.
base_settings = {
    "num_layers": 5,  # number of LSTM layers
    "n_hidden": 500,  # per direction, per layer, the number of units
    "num_seqs": 500,  # for the dataset generation
    "batch_size": 2000,  # upper limit for n_seqs * max_seq_len in a batch
    "max_seqs": 40,  # upper limit for n_seqs in a batch
    "chunking": "50:25",  # set to "0" to disable
}


def make_config_dict(lstm_unit, use_gpu):
    """
    :param str lstm_unit: "NativeLSTM", "LSTMBlock", "LSTMBlockFused", "CudnnLSTM", etc, one of LstmCellTypes
    :param bool use_gpu:
    :return: config dict
    :rtype: dict[str]
    """
    num_layers = base_settings["num_layers"]
    network = {}
    for i in range(num_layers):
        for direction in [-1, 1]:
            dir_str = {-1: "bwd", 1: "fwd"}[direction]
            layer = {"class": "rec", "unit": lstm_unit, "n_out": base_settings["n_hidden"], "direction": direction}
            if i > 0:
                layer["from"] = ["lstm%i_fwd" % i, "lstm%i_bwd" % i]
            network["lstm%i_%s" % (i + 1, dir_str)] = layer
    network["output"] = {
        "class": "softmax",
        "loss": "ce",
        "target": "classes",
        "from": ["lstm%i_fwd" % num_layers, "lstm%i_bwd" % num_layers],
    }
    return {
        "device": "gpu" if use_gpu else "cpu",
        "train": {"class": "Task12AXDataset", "num_seqs": base_settings["num_seqs"]},
        "num_inputs": _input_dim,
        "num_outputs": _output_dim,
        "num_epochs": 1,
        "model": None,  # don't save
        "tf_log_dir": None,  # no TF logs
        "tf_log_memory_usage": True,
        "network": network,
        # batching
        "batch_size": base_settings["batch_size"],
        "max_seqs": base_settings["max_seqs"],
        "chunking": base_settings["chunking"],
        # optimization
        "optimizer": {"class": "adam"},
        "learning_rate": 0.01,
    }


def benchmark(lstm_unit, use_gpu):
    """
    :param str lstm_unit: e.g. "LSTMBlock", one of LstmCellTypes
    :param bool use_gpu:
    :return: runtime in seconds of the training itself, excluding initialization
    :rtype: float
    """
    device = {True: "GPU", False: "CPU"}[use_gpu]
    key = "%s:%s" % (device, lstm_unit)
    print(">>> Start benchmark for %s." % key)
    config = Config()
    config.update(make_config_dict(lstm_unit=lstm_unit, use_gpu=use_gpu))
    dataset_kwargs = config.typed_value("train")
    Dataset.kwargs_update_from_config(config, dataset_kwargs)
    dataset = init_dataset(dataset_kwargs)
    engine = Engine(config=config)
    engine.init_train_from_config(config=config, train_data=dataset)
    print(">>> Start training now for %s." % key)
    start_time = time.time()
    engine.train()
    runtime = time.time() - start_time
    print(">>> Runtime of %s: %s" % (key, hms_fraction(runtime)))
    engine.finalize()
    return runtime


def main():
    """
    Main entry.
    """
    global LstmCellTypes
    print("Benchmarking LSTMs.")
    better_exchook.install()
    print("Args:", " ".join(sys.argv))
    arg_parser = ArgumentParser()
    arg_parser.add_argument("cfg", nargs="*", help="opt=value, opt in %r" % sorted(base_settings.keys()))
    arg_parser.add_argument("--no-cpu", action="store_true")
    arg_parser.add_argument("--no-gpu", action="store_true")
    arg_parser.add_argument("--selected", help="comma-separated list from %r" % LstmCellTypes)
    arg_parser.add_argument("--no-setup-tf-thread-pools", action="store_true")
    args = arg_parser.parse_args()
    for opt in args.cfg:
        key, value = opt.split("=", 1)
        assert key in base_settings
        value_type = type(base_settings[key])
        base_settings[key] = value_type(value)
    print("Settings:")
    pprint(base_settings)

    log.initialize(verbosity=[4])
    print("Returnn:", describe_returnn_version(), file=log.v3)
    print("TensorFlow:", describe_tensorflow_version(), file=log.v3)
    print("Python:", sys.version.replace("\n", ""), sys.platform)
    if not args.no_setup_tf_thread_pools:
        setup_tf_thread_pools(log_file=log.v2)
    else:
        print("Not setting up the TF thread pools. Will be done automatically by TF to number of CPU cores.")
    if args.no_gpu:
        print("GPU will not be used.")
    else:
        print("GPU available: %r" % is_gpu_available())
    print_available_devices()

    if args.selected:
        LstmCellTypes = args.selected.split(",")
    benchmarks = {}
    if not args.no_gpu and is_gpu_available():
        for lstm_unit in LstmCellTypes:
            benchmarks["GPU:" + lstm_unit] = benchmark(lstm_unit=lstm_unit, use_gpu=True)
    if not args.no_cpu:
        for lstm_unit in LstmCellTypes:
            if lstm_unit in GpuOnlyCellTypes:
                continue
            benchmarks["CPU:" + lstm_unit] = benchmark(lstm_unit=lstm_unit, use_gpu=False)

    print("-" * 20)
    print("Settings:")
    pprint(base_settings)
    print("Final results:")
    for t, lstm_unit in sorted([(t, lstm_unit) for (lstm_unit, t) in sorted(benchmarks.items())]):
        print("  %s: %s" % (lstm_unit, hms_fraction(t)))
    print("Done.")


if __name__ == "__main__":
    main()
