
"""
Provides the learning rate scheduling logic.
The base class is :class:`LearningRateControl`.
"""

from __future__ import print_function

import os
import typing
from returnn.util.basic import better_repr, simple_obj_repr, ObjAsDict, unicode
from returnn.log import log
import numpy


class LearningRateControl(object):
  """
  Base class for learning rate control / scheduling.
  """

  need_error_info = True

  class EpochData:
    """
    Encapsulates all relevant information for one epoch,
    needed to perform learning rate scheduling,
    such as the individual scores (cv or train; cross-entropy or frame-error or whatever).
    """

    # Need to keep the non-PEP8 name for compatibility, because we store the repr of the object.
    # noinspection PyPep8Naming
    def __init__(self, learningRate, error=None):
      """
      :type learningRate: float
      :type error: dict[str,float] | None
      """
      self.learning_rate = learningRate
      if isinstance(error, float):  # Old format.
        error = {"old_format_score": error}
      if error is None:
        error = {}
      self.error = error

    def __repr__(self):
      # This is being used for serialization, and we want some forward/backward compatibility,
      # so we should try to keep this consistent.
      return "EpochData(learningRate=%s, error=%s)" % (
        better_repr(self.learning_rate), better_repr(self.error))

  @classmethod
  def load_initial_kwargs_from_config(cls, config):
    """
    :type config: returnn.config.Config
    :rtype: dict[str]
    """
    return {
      "default_learning_rate": config.float('learning_rate', 1.0),
      "min_learning_rate": config.float('min_learning_rate', 0.0),
      "default_learning_rates": config.typed_value('learning_rates') or config.float_list('learning_rates'),
      "error_measure_key": (
        config.typed_value('learning_rate_control_error_measure')
        or config.value('learning_rate_control_error_measure', None)),
      "relative_error_also_relative_to_learning_rate": (
        config.bool('learning_rate_control_relative_error_relative_lr', False)),
      "min_num_epochs_per_new_learning_rate": config.int("learning_rate_control_min_num_epochs_per_new_lr", 0),
      "relative_error_div_by_old": config.bool('newbob_relative_error_div_by_old', False),
      "learning_rate_decay": config.typed_value(
        'learning_rate_decay', config.opt_typed_value('newbob_learning_rate_decay', 0.5)),
      "learning_rate_growth": config.typed_value(
        'learning_rate_growth', config.opt_typed_value('newbob_learning_rate_growth', 1.0)),
      "filename": config.value('learning_rate_file', None),
    }

  @classmethod
  def load_initial_from_config(cls, config):
    """
    :type config: returnn.config.Config
    :rtype: LearningRateControl
    """
    kwargs = cls.load_initial_kwargs_from_config(config)
    return cls(**kwargs)

  def __init__(self, default_learning_rate, min_learning_rate=0.0, default_learning_rates=None,
               error_measure_key=None,
               relative_error_also_relative_to_learning_rate=False,
               min_num_epochs_per_new_learning_rate=0,
               relative_error_div_by_old=False,
               learning_rate_decay=1.0,
               learning_rate_growth=1.0,
               filename=None):
    """
    :param float default_learning_rate: default learning rate. usually for epoch 1
    :param list[float] | dict[int,float] default_learning_rates: learning rates
    :param str|list[str]|None error_measure_key: for get_epoch_error_value() the key for EpochData.error which is a dict
    :param int min_num_epochs_per_new_learning_rate: if the lr was recently updated, use it for at least N epochs
    :param bool relative_error_div_by_old: if True, compute relative error as (new - old) / old.
    :param float|(float)->float learning_rate_decay:
    :param float|(float)->float learning_rate_growth:
    :param str filename: load from and save to file
    """
    self.epoch_data = {}  # type: typing.Dict[int,LearningRateControl.EpochData]
    self.filename = filename
    if filename:
      if os.path.exists(filename):
        print("Learning-rate-control: loading file %s" % filename, file=log.v4)
        # Load now, such that default_learning_rates is correctly handled.
        self.load()
      else:
        print("Learning-rate-control: file %s does not exist yet" % filename, file=log.v4)
    else:
      print("Learning-rate-control: no file specified, not saving history (no proper restart possible)", file=log.v4)
    self.default_learning_rate = default_learning_rate
    self.min_learning_rate = min_learning_rate
    if default_learning_rates:
      if isinstance(default_learning_rates, list):
        default_learning_rates = {i + 1: v for (i, v) in enumerate(default_learning_rates)}
      if isinstance(default_learning_rates, (str, unicode)):
        default_learning_rates = eval(default_learning_rates)
      assert isinstance(default_learning_rates, dict)
      for epoch, v in default_learning_rates.items():
        self.set_default_learning_rate_for_epoch(epoch, v)
    self.default_learning_rates = default_learning_rates
    self.error_measure_key = error_measure_key
    self.relative_error_also_relative_to_learning_rate = relative_error_also_relative_to_learning_rate
    self.min_num_epochs_per_new_learning_rate = min_num_epochs_per_new_learning_rate
    self.relative_error_div_by_old = relative_error_div_by_old
    self.learning_rate_decay = learning_rate_decay
    self.learning_rate_growth = learning_rate_growth

  __repr__ = simple_obj_repr

  def __str__(self):
    epochs = sorted(self.epoch_data.keys())
    if len(epochs) > 6:
      epoch_str = ", ".join(
        ["%i: %s" % (epoch, self.epoch_data[epoch]) for epoch in epochs[:3]] +
        ["..."] +
        ["%i: %s" % (epoch, self.epoch_data[epoch]) for epoch in epochs[-3:]])
    else:
      epoch_str = ", ".join(["%i: %s" % (epoch, self.epoch_data[epoch]) for epoch in epochs])
    return "%r, epoch data: %s, error key: %s" % (self, epoch_str, self.get_error_key(epoch=1))

  @staticmethod
  def _calc_learning_rate_update(learning_rate, update):
    """
    :param float learning_rate:
    :param None|float|(float)->float update: factor, or generic func
    :return: lr with update applied (e.g. decay factor)
    :rtype: float
    """
    if update is None:
      return learning_rate
    if isinstance(update, float):
      return learning_rate * update
    assert callable(update)
    learning_rate = update(learning_rate)
    assert isinstance(learning_rate, float)
    return learning_rate

  def _calc_learning_rate_decay(self, learning_rate):
    """
    :param float learning_rate:
    :return: lr with decay applied
    :rtype: float
    """
    return self._calc_learning_rate_update(learning_rate, update=self.learning_rate_decay)

  def _calc_learning_rate_growth(self, learning_rate):
    """
    :param float learning_rate:
    :return: lr with growth applied
    :rtype: float
    """
    return self._calc_learning_rate_update(learning_rate, update=self.learning_rate_growth)

  def calc_learning_rate_decay_or_grow(self, learning_rate, decay, grow=None):
    """
    :param float learning_rate:
    :param bool decay:
    :param bool|None grow: default is not decay
    :return: lr with decay or growth applied
    :rtype: float
    """
    assert isinstance(decay, bool)
    if grow is None:
      grow = not decay
    assert isinstance(grow, bool)
    assert not (grow and decay)  # not sure if this makes sense...
    if decay:
      learning_rate = self._calc_learning_rate_decay(learning_rate)
      if learning_rate < self.min_learning_rate:
        learning_rate = self.min_learning_rate
    if grow:
      learning_rate = self._calc_learning_rate_growth(learning_rate)
    return learning_rate

  def calc_learning_rate_for_epoch(self, epoch):
    """
    :type epoch: int
    :returns learning rate
    :rtype: float
    """
    raise NotImplementedError

  def calc_new_learning_rate_for_epoch(self, epoch):
    """
    :param int epoch:
    :return: new learning rate for this epoch
    :rtype: float
    """
    if self.min_num_epochs_per_new_learning_rate > 1:
      last_lrs = [self.epoch_data[e].learning_rate
                  for e in self._last_epochs_for_epoch(epoch, num_epochs=self.min_num_epochs_per_new_learning_rate)]
      if len(set(last_lrs)) >= 2 or 0 < len(last_lrs) < self.min_num_epochs_per_new_learning_rate:
        return last_lrs[-1]
    learning_rate = self.calc_learning_rate_for_epoch(epoch)
    return learning_rate

  def _last_epochs_for_epoch(self, epoch, num_epochs):
    """
    :param int epoch:
    :param int num_epochs:
    :return: last N epochs where we have some epoch data
    :rtype: list[int]
    """
    last_epochs = sorted([e for e in self.epoch_data.keys() if e < epoch])
    if not last_epochs:
      return []
    last_epochs = last_epochs[-num_epochs:]
    return last_epochs

  def get_learning_rate_for_epoch(self, epoch):
    """
    :type epoch: int
    :rtype: float
    """
    assert epoch >= 1
    if epoch in self.epoch_data:
      return self.epoch_data[epoch].learning_rate
    learning_rate = self.calc_new_learning_rate_for_epoch(epoch)
    self.set_default_learning_rate_for_epoch(epoch, learning_rate)
    return learning_rate

  def set_default_learning_rate_for_epoch(self, epoch, learning_rate):
    """
    :type epoch: int
    :type learning_rate: float
    """
    if epoch in self.epoch_data:
      if not self.epoch_data[epoch].learning_rate:
        self.epoch_data[epoch].learning_rate = learning_rate
    else:
      self.epoch_data[epoch] = self.EpochData(learning_rate)

  def get_last_epoch(self, epoch):
    """
    :param int epoch:
    :return: last epoch before ``epoch`` where we have some epoch data
    :rtype: int
    """
    epochs = sorted([e for e in self.epoch_data.keys() if e < epoch])
    if not epochs:
      return None
    return epochs[-1]

  def get_most_recent_learning_rate(self, epoch, exclude_current=True):
    """
    :param int epoch:
    :param bool exclude_current:
    :return: most learning rate before or including ``epoch``
    :rtype: float
    """
    for e, data in reversed(sorted(self.epoch_data.items())):
      assert isinstance(data, LearningRateControl.EpochData)
      if e > epoch:
        continue
      if exclude_current and e == epoch:
        continue
      if data.learning_rate is None:
        continue
      return data.learning_rate
    return self.default_learning_rate

  def calc_relative_error(self, old_epoch, new_epoch):
    """
    :param int old_epoch:
    :param int new_epoch:
    :return: relative error between old epoch and new epoch
    :rtype: float
    """
    old_key, old_error = self.get_epoch_error_key_value(old_epoch)
    new_key, new_error = self.get_epoch_error_key_value(new_epoch)
    if old_error is None or new_error is None:
      return None
    if old_key != new_key:
      return None
    if self.relative_error_div_by_old:
      relative_error = (new_error - old_error) / abs(old_error)
    else:
      relative_error = (new_error - old_error) / abs(new_error)
    if self.relative_error_also_relative_to_learning_rate:
      learning_rate = self.get_most_recent_learning_rate(new_epoch, exclude_current=False)
      if learning_rate > 0:
        # If the learning rate is lower than the initial learning rate,
        # the relative error is also expected to be lower, so correct for that here.
        relative_error /= learning_rate / self.default_learning_rate
    return relative_error

  def set_epoch_error(self, epoch, error):
    """
    :type epoch: int
    :type error: dict[str,float|dict[str,float]]
    """
    if epoch not in self.epoch_data:
      print("Learning rate not set for epoch %i. Assuming default." % epoch, file=log.v4)
      self.get_learning_rate_for_epoch(epoch)  # This will set it.
    assert isinstance(error, dict)
    error = error.copy()
    for k, v in list(error.items()):
      if isinstance(v, dict):  # like error = {"dev_score": {"cost:output1": .., "cost:output2": ...}, ...}
        del error[k]
        if len(v) == 1:
          error[k] = list(v.values())[0]
          continue
        for k1, v1 in v.items():
          if ":" in k1:
            k1 = k1[k1.index(":") + 1:]
          error[k + "_" + k1] = v1
    for v in error.values():
      assert isinstance(v, float)
    self.epoch_data[epoch].error.update(error)
    if epoch == 1:
      print("Learning-rate-control: error key %r from %r" % (self.get_error_key(epoch), error), file=log.v4)

  def get_error_key(self, epoch):
    """
    :param int epoch:
    :return: key which we should look in scores/errors, for this epoch
    :rtype: str
    """
    if epoch not in self.epoch_data:
      if isinstance(self.error_measure_key, list):
        return self.error_measure_key[0]
      assert isinstance(self.error_measure_key, (str, type(None)))
      return self.error_measure_key
    epoch_data = self.epoch_data[epoch]
    if not epoch_data.error:
      return None
    if len(epoch_data.error) == 1 and "old_format_score" in epoch_data.error:
      return "old_format_score"
    keys = []
    if isinstance(self.error_measure_key, list):
      for key in self.error_measure_key:
        keys += [key, key + "_output"]  # for multiple outputs, try default output
    elif isinstance(self.error_measure_key, str):
      keys += [self.error_measure_key, self.error_measure_key + "_output"]
    else:
      assert self.error_measure_key is None
    keys += ["dev_score", "dev_score_output"]
    for key in keys:
      if key in epoch_data.error:
        return key
    for key in sorted(epoch_data.error.keys()):
      if key == "dev_score_output/output" or key.startswith("dev_score_output/output_"):
        return key
    for key in sorted(epoch_data.error.keys()):
      if key.startswith("dev_score_output/"):
        return key
    for key in sorted(epoch_data.error.keys()):
      if key.startswith("dev_"):
        return key
    for key in ["train_score", "train_score_output"]:
      if key in epoch_data.error:
        return key
    return min(epoch_data.error.keys())

  def get_epoch_error_dict(self, epoch):
    """
    :param int epoch:
    :rtype: dict[str,float]
    """
    if epoch not in self.epoch_data:
      return {}
    return self.epoch_data[epoch].error

  def get_epoch_error_value(self, epoch):
    """
    :param int epoch:
    :return: error/score for the specific epoch, given the error-key, see :func:`get_error_key`
    :rtype: float
    """
    error = self.get_epoch_error_dict(epoch)
    if not error:
      return None
    key = self.get_error_key(epoch)
    assert key
    assert key in error, (
      "%r not in %r. fix %r in config. set it to %r or so." % (
        key, error, 'learning_rate_control_error_measure', 'dev_error'))
    return error[key]

  def get_epoch_error_key_value(self, epoch):
    """
    :param int epoch:
    :return: key, error
    :rtype: (str, float)
    """
    error = self.get_epoch_error_dict(epoch)
    if not error:
      return None, None
    key = self.get_error_key(epoch)
    assert key
    assert key in error, (
      "%r not in %r. fix %r in config. set it to %r or so." %
      (key, error, 'learning_rate_control_error_measure', 'dev_error'))
    return key, error[key]

  def get_last_best_epoch(self, last_epoch, first_epoch=1, only_last_epochs=None,
                          filter_score=float("inf"),
                          only_last_n=-1,
                          min_score_dist=0.0):
    """
    :param int first_epoch: will check all epochs >= first_epoch
    :param int last_epoch: inclusive. will check all epochs <= last_epoch
    :param int|None only_last_epochs: if set, will only check the last N epochs, inclusive
    :param float filter_score: all epochs which values over this score are not considered
    :param int only_last_n: if set (>=1), *from the resulting list*, we consider only the last only_last_n
    :param float min_score_dist: filter out epochs where the diff to the most recent is not big enough
    :return: the last best epoch. to get the details then, you might want to use getEpochErrorDict.
    :rtype: int|None
    """
    if only_last_epochs:
      first_epoch = max(first_epoch, last_epoch - only_last_epochs + 1)
    if first_epoch > last_epoch:
      return None
    values = [(self.get_epoch_error_key_value(ep), ep) for ep in range(first_epoch, last_epoch + 1)]
    # Note that the order of the checks here is a bit arbitrary but I had some thoughts on it.
    # Changing the order will also slightly change the behavior, so be sure it make sense.
    values = [((key, v), ep) for ((key, v), ep) in values if v is not None]
    if not values:
      return None
    last_key, latest_score = values[-1][0]
    values = [(v, ep) for ((key, v), ep) in values if key == last_key]  # only same key
    values = [(v, ep) for (v, ep) in values if v <= filter_score]
    if not values:
      return None
    if only_last_n >= 1:
      values = values[-only_last_n:]
    values = [(v, ep) for (v, ep) in values if v + min_score_dist < latest_score]
    if not values:
      return None
    return min(values)[1]

  def save(self):
    """
    Save the current epoch data to file (self.filename).
    """
    if not self.filename:
      return
    # First write to a temp-file, to be sure that the write happens without errors.
    # Otherwise, it could happen that we delete the old existing file, then
    # some error happens (e.g. disk quota), and we loose the newbob data.
    # Loosing that data is very bad because it basically means that we have to redo all the training.
    tmp_filename = self.filename + ".new_tmp"
    f = open(tmp_filename, "w")
    f.write(better_repr(self.epoch_data))
    f.write("\n")
    f.close()
    os.rename(tmp_filename, self.filename)

  def load(self):
    """
    Loads the saved epoch data from file (self.filename).
    """
    s = open(self.filename).read()
    self.epoch_data = eval(s, {"nan": float("nan"), "inf": float("inf")}, ObjAsDict(self))


class ConstantLearningRate(LearningRateControl):
  """
  Just a constant learning rate.
  """

  need_error_info = False

  def calc_learning_rate_for_epoch(self, epoch):
    """
    Dummy constant learning rate. Returns initial learning rate.
    :type epoch: int
    :returns learning rate
    :rtype: float
    """
    while True:
      last_epoch = self.get_last_epoch(epoch)
      if last_epoch is None:
        return self.default_learning_rate
      learning_rate = self.epoch_data[last_epoch].learning_rate
      if learning_rate is None:
        epoch = last_epoch
        continue
      return learning_rate


class NewbobRelative(LearningRateControl):
  """
  If relative diff between old and new error is over some threshold, decay learning rate.
  """

  @classmethod
  def load_initial_kwargs_from_config(cls, config):
    """
    :type config: returnn.config.Config
    :rtype: dict[str]
    """
    kwargs = super(NewbobRelative, cls).load_initial_kwargs_from_config(config)
    kwargs.update({
      "relative_error_threshold": config.float('newbob_relative_error_threshold', -0.01),
    })
    return kwargs

  def __init__(self, relative_error_threshold, **kwargs):
    """
    :type relative_error_threshold: float
    """
    super(NewbobRelative, self).__init__(**kwargs)
    self.relative_error_threshold = relative_error_threshold

  def calc_learning_rate_for_epoch(self, epoch):
    """
    Newbob+ on train data.
    :type epoch: int
    :returns learning rate
    :rtype: float
    """
    last_epoch = self.get_last_epoch(epoch)
    if last_epoch is None:
      return self.default_learning_rate
    learning_rate = self.epoch_data[last_epoch].learning_rate
    if learning_rate is None:
      return self.default_learning_rate
    last2_epoch = self.get_last_epoch(last_epoch)
    if last2_epoch is None:
      return learning_rate
    relative_error = self.calc_relative_error(last2_epoch, last_epoch)
    if relative_error is None:
      return learning_rate
    learning_rate = self.calc_learning_rate_decay_or_grow(
      learning_rate, decay=relative_error > self.relative_error_threshold)
    return learning_rate


class NewbobAbs(LearningRateControl):
  """
  If absolute diff between old and new error is over some threshold, decay learning rate.
  """

  @classmethod
  def load_initial_kwargs_from_config(cls, config):
    """
    :type config: returnn.config.Config
    :rtype: dict[str]
    """
    kwargs = super(NewbobAbs, cls).load_initial_kwargs_from_config(config)
    kwargs.update({
      "error_threshold": config.float('newbob_error_threshold', -0.01),
    })
    return kwargs

  def __init__(self, error_threshold, **kwargs):
    """
    :type error_threshold: float
    """
    super(NewbobAbs, self).__init__(**kwargs)
    self.error_threshold = error_threshold

  def calc_learning_rate_for_epoch(self, epoch):
    """
    Newbob+ on train data.

    :type epoch: int
    :returns learning rate
    :rtype: float
    """
    last_epoch = self.get_last_epoch(epoch)
    if last_epoch is None:
      return self.default_learning_rate
    learning_rate = self.epoch_data[last_epoch].learning_rate
    if learning_rate is None:
      return self.default_learning_rate
    last2_epoch = self.get_last_epoch(last_epoch)
    if last2_epoch is None:
      return learning_rate
    old_key, old_error = self.get_epoch_error_key_value(last2_epoch)
    new_key, new_error = self.get_epoch_error_key_value(last_epoch)
    if old_error is None or new_error is None:
      return learning_rate
    if old_key != new_key:
      return learning_rate
    error_diff = new_error - old_error
    learning_rate = self.calc_learning_rate_decay_or_grow(learning_rate, decay=error_diff > self.error_threshold)
    return learning_rate


class NewbobMultiEpoch(LearningRateControl):
  """
  Like :class:`NewbobRelative`, but looks at the average relative error over multiple epochs.
  This is useful together with ``partition_epoch`` from :class:`Dataset`.
  """

  @classmethod
  def load_initial_kwargs_from_config(cls, config):
    """
    :type config: returnn.config.Config
    :rtype: dict[str]
    """
    kwargs = super(NewbobMultiEpoch, cls).load_initial_kwargs_from_config(config)
    kwargs.update({
      "num_epochs": config.int("newbob_multi_num_epochs", 5),
      "update_interval": config.int("newbob_multi_update_interval", config.int("newbob_multi_num_epochs", 5)),
      "relative_error_threshold": config.float('newbob_relative_error_threshold', -0.01),
      "relative_error_grow_threshold": config.float(
        'newbob_relative_error_grow_threshold', config.float('newbob_relative_error_threshold', -0.01)),
    })
    return kwargs

  def __init__(self, num_epochs, update_interval, relative_error_threshold, relative_error_grow_threshold, **kwargs):
    """
    :param int num_epochs:
    :param int update_interval:
    :param float relative_error_threshold:
    :param float relative_error_grow_threshold:
    """
    super(NewbobMultiEpoch, self).__init__(**kwargs)
    self.num_epochs = num_epochs
    assert self.num_epochs >= 1
    self.update_interval = update_interval
    assert self.update_interval >= 1
    self.relative_error_threshold = relative_error_threshold
    self.relative_error_grow_threshold = relative_error_grow_threshold

  def _calc_mean_relative_error(self, epochs):
    """
    :param list[int] epochs:
    :return: mean of relative errors
    :rtype: float|None
    """
    assert len(epochs) >= 2
    errors = [self.calc_relative_error(epochs[i], epochs[i + 1]) for i in range(len(epochs) - 1)]
    if any([e is None for e in errors]):
      return None
    return float(numpy.mean(errors))

  def _calc_recent_mean_relative_error(self, epoch):
    """
    :param int epoch:
    :return: recent mean of relative errors
    :rtype: float|None
    """
    # Take one more than numEpochs because we are looking at the diffs.
    last_epochs = self._last_epochs_for_epoch(epoch, num_epochs=self.num_epochs + 1)
    if not last_epochs:
      return None
    # We could also use the self.numEpochs limit here. But maybe this is better.
    if len(last_epochs) <= 1:
      return None
    return self._calc_mean_relative_error(last_epochs)

  def calc_learning_rate_for_epoch(self, epoch):
    """
    Newbob+ on train data.
    :type epoch: int
    :returns learning rate
    :rtype: float
    """
    learning_rate = self.get_most_recent_learning_rate(epoch)
    # We start counting epochs at 1.
    if self.update_interval > 1 and epoch % self.update_interval != 1:
      return learning_rate
    mean_relative_error = self._calc_recent_mean_relative_error(epoch)
    if mean_relative_error is None:
      return learning_rate
    learning_rate = self.calc_learning_rate_decay_or_grow(
      learning_rate,
      decay=mean_relative_error > self.relative_error_threshold,
      grow=mean_relative_error < self.relative_error_grow_threshold)
    return learning_rate


def learning_rate_control_type(type_name):
  """
  :param str type_name:
  :rtype: type[LearningRateControl]|LearningRateControl
  """
  if type_name == "constant":
    return ConstantLearningRate
  elif type_name in ("newbob", "newbob_rel", "newbob_relative"):  # Old setups expect the relative version.
    return NewbobRelative
  elif type_name == "newbob_abs":
    return NewbobAbs
  elif type_name == "newbob_multi_epoch":
    return NewbobMultiEpoch
  else:
    assert False, "unknown learning-rate-control type %s" % type_name


def load_learning_rate_control_from_config(config):
  """
  :type config: returnn.config.Config
  :rtype: LearningRateControl
  """
  control_type = config.value("learning_rate_control", "constant")
  cls = learning_rate_control_type(control_type)
  return cls.load_initial_from_config(config)


def demo():
  """
  Demo run. Given some learning rate file (with scores / existing lrs), will calculate how lrs would have been set,
  given some config.
  """
  from returnn.util import better_exchook
  better_exchook.install()
  import returnn.__main__ as rnn
  import sys
  if len(sys.argv) <= 1:
    print("usage: python %s [config] [other options] [++check_learning_rates 1]" % __file__)
    print(
      ("example usage: "
       "python %s ++learning_rate_control newbob ++learning_rate_file newbob.data ++learning_rate 0.001") % __file__)
  rnn.init_config(command_line_options=sys.argv[1:])
  # noinspection PyProtectedMember
  rnn.config._hack_value_reading_debug()
  rnn.config.update({"log": []})
  rnn.init_log()
  rnn.init_backend_engine()
  check_lr = rnn.config.bool("check_learning_rates", False)
  from returnn.pretrain import pretrain_from_config
  pretrain = pretrain_from_config(rnn.config)
  first_non_pretrain_epoch = 1
  pretrain_learning_rate = None
  if pretrain:
    first_non_pretrain_epoch = pretrain.get_train_num_epochs() + 1
  log.initialize(verbosity=[5])
  control = load_learning_rate_control_from_config(rnn.config)
  print("LearningRateControl: %r" % control)
  if not control.epoch_data:
    print("No epoch data so far.")
    return
  first_epoch = min(control.epoch_data.keys())
  if first_epoch != 1:
    print("Strange, first epoch from epoch data is %i." % first_epoch)
  print("Error key: %s from %r" % (control.get_error_key(epoch=first_epoch), control.epoch_data[first_epoch].error))
  if pretrain:
    pretrain_learning_rate = rnn.config.float('pretrain_learning_rate', control.default_learning_rate)
  max_epoch = max(control.epoch_data.keys())
  for epoch in range(1, max_epoch + 2):  # all epochs [1..max_epoch+1]
    old_learning_rate = None
    if epoch in control.epoch_data:
      old_learning_rate = control.epoch_data[epoch].learning_rate
    if epoch < first_non_pretrain_epoch:
      learning_rate = pretrain_learning_rate
      s = "Pretrain epoch %i, fixed learning rate: %s (was: %s)" % (epoch, learning_rate, old_learning_rate)
    elif 1 < first_non_pretrain_epoch == epoch:
      learning_rate = control.default_learning_rate
      s = "First epoch after pretrain, epoch %i, fixed learning rate: %s (was %s)" % (
        epoch, learning_rate, old_learning_rate)
    else:
      learning_rate = control.calc_new_learning_rate_for_epoch(epoch)
      s = "Calculated learning rate for epoch %i: %s (was: %s)" % (epoch, learning_rate, old_learning_rate)
    if learning_rate < control.min_learning_rate:
      learning_rate = control.min_learning_rate
      s += ", clipped to %s" % learning_rate
    s += ", previous relative error: %s" % control.calc_relative_error(epoch - 2, epoch - 1)
    if hasattr(control, "_calc_recent_mean_relative_error"):
      # noinspection PyProtectedMember
      s += ", previous mean relative error: %s" % control._calc_recent_mean_relative_error(epoch)
    print(s)
    if check_lr and old_learning_rate is not None:
      if old_learning_rate != learning_rate:
        print("Learning rate is different in epoch %i!" % epoch)
        sys.exit(1)
    # Overwrite new learning rate so that the calculation for further learning rates stays consistent.
    if epoch in control.epoch_data:
      control.epoch_data[epoch].learning_rate = learning_rate
    else:
      control.epoch_data[epoch] = control.EpochData(learningRate=learning_rate)
  print("Finished, last stored epoch was %i." % max_epoch)


if __name__ == "__main__":
  demo()
