# pylint: disable=g-bad-file-header
# Copyright 2021 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Dummy model for testing."""

import abc
from typing import NamedTuple, Sequence, Union

from absl import logging
from dm_c19_modelling.modelling import definitions
from dm_c19_modelling.modelling.models import utils
import haiku as hk
import jax
import jax.numpy as jnp
import optax


class LossMinimizerHaikuModelState(NamedTuple):
  trainable_params: hk.Params
  non_trainable_state: hk.State
  optimizer_state: Union[optax.OptState, Sequence[optax.OptState]]
  rng: jnp.ndarray  # jax.random.PRNGKey


class LossMinimizerHaikuModel(definitions.TrainableModel):
  """Model based on Haiku that is trained by minimizing a loss.

  Child classes have to implement the "_loss_fn" and "_predict_fn" methods.

  This base class abstracts away all of the complexity around updating params
  and states, as well as `hk.transform`, `hk.init` and `hk.apply`.

  """

  def __init__(self, init_seed, training_seed, optimizer_kwargs,
               learning_rate_annealing_kwargs, jit_predict_fn=True,
               jit_update_fn=True, **parent_kwargs):
    super().__init__(**parent_kwargs)

    self._init_seed = init_seed
    self._training_seed = training_seed
    self._optimizer_kwargs = optimizer_kwargs
    self._learning_rate_annealing_kwargs = learning_rate_annealing_kwargs

    # Prepare the update fn.
    self._transformed_loss_fn = hk.transform_with_state(self._loss_fn)

    # `hk.transform_with_state(self._loss_fn)`, would have an output structured
    # as:
    #   ((loss, scalars_dict), updated_state)
    # so we change the packing to:
    #   (loss, updated_state, scalars_dict)
    def loss_fn(*args, **kwargs):
      ((loss, scalars_dict),
       updated_state) = self._transformed_loss_fn.apply(*args, **kwargs)
      return loss, updated_state, scalars_dict
    self._update_fn = utils.get_optimizer_params_update_step(
        loss_fn, self._get_optimizer)

    if jit_update_fn:
      self._update_fn = jax.jit(self._update_fn)

    # Prepare the predict fn.
    self._apply_predict_fn = hk.transform_with_state(self._predict_fn).apply
    if jit_predict_fn:
      self._apply_predict_fn = jax.jit(self._apply_predict_fn)

  @abc.abstractmethod
  def _loss_fn(self, batch):
    """Computes the loss to be minimized.

    Args:
      batch: batch of data as generated by `build_training_generator`.

    Returns:
      Tuple (loss, scalars_dict) with the scalar
      loss, and a dictionary of additional scalar metrics.

    """

  @abc.abstractmethod
  def _prepare_predict_fn_inputs(self, dataset):
    """Builds the inputs to be passed to the `predict_fn`."""

  @abc.abstractmethod
  def _predict_fn(self, inputs):
    """Makes a prediction using the inputs from `_prepare_predict_fn_inputs`."""

  def _get_optimizer(self, global_step):
    return utils.get_optimizer_with_learning_rate_annealing(
        global_step,
        self._optimizer_kwargs,
        self._learning_rate_annealing_kwargs)

  def _get_initial_state(self, sample_batch):

    # Initialize the model.
    rng = jax.random.PRNGKey(self._init_seed)
    trainable_params, non_trainable_state = self._transformed_loss_fn.init(
        rng, sample_batch)

    _log_param_types_and_shapes(trainable_params, non_trainable_state)

    # Initialize the optimizer.
    (optimizer_init_fn, _), _ = self._get_optimizer(global_step=0)
    optimizer_state = optimizer_init_fn(trainable_params)

    # Initialize a random seed.
    rng = jax.random.PRNGKey(self._training_seed)
    return LossMinimizerHaikuModelState(
        trainable_params=trainable_params,
        non_trainable_state=non_trainable_state,
        optimizer_state=optimizer_state,
        rng=rng)

  def training_update(self, previous_state, batch, global_step):
    """Updates the state (params, optimizer step, logic) of the model."""

    # Obtain a state if we do not have one.
    if previous_state is None:
      previous_state = self._get_initial_state(sample_batch=batch)

    next_rng, rng_to_use = jax.random.split(previous_state.rng)

    with jax.profiler.StepTraceContext(
        "training_update", step_num=global_step):
      (updated_optimizer_state, updated_params, updated_non_trainable_state,
       loss, aux_loss_outputs, aux_optimizer_outputs) = self._update_fn(
           global_step,
           previous_state.optimizer_state,
           previous_state.trainable_params,
           previous_state.non_trainable_state,
           rng_to_use,
           batch)

    # Build the next module state.
    updated_state = LossMinimizerHaikuModelState(
        trainable_params=updated_params,
        non_trainable_state=updated_non_trainable_state,
        optimizer_state=updated_optimizer_state,
        rng=next_rng)

    # Aggregate all of the summary scalars.
    scalars = {}
    scalars["loss"] = loss
    scalars.update(aux_optimizer_outputs)
    scalars.update(aux_loss_outputs)

    return updated_state, scalars

  def _evaluate(self, model_state, dataset):
    # This asumes that the evaluation inputs are already something simple
    # like a number array, if it is something non-jittable, we may need to
    # add an additional tranformation, here or change slighly how we do this.
    inputs = self._prepare_predict_fn_inputs(dataset)
    ((predictions, aux_data), _) = self._apply_predict_fn(
        model_state.trainable_params,
        model_state.non_trainable_state,
        model_state.rng, inputs)
    return predictions, aux_data

  def _prepare_features(self, features):
    """Prepares features for a neural network."""
    return utils.prepare_features(
        features,
        self._dataset_spec.feature_names,
        self._dataset_spec.feature_stats,
        categorical_features_dict={
            definitions.SITE_ID_INTEGER: len(self._dataset_spec.sites),
            definitions.WEEK_DAY_INTEGER: 7
        },
        features_with_missing_values=(
            self._dataset_spec.features_with_missing_values))


def _log_param_types_and_shapes(trainable_params, non_trainable_state):
  logging.info("Model trainable params:")
  _log_dict_recursively([], trainable_params)

  logging.info("Model non-trainable state:")
  _log_dict_recursively([], non_trainable_state)


def _log_dict_recursively(parent_names, param):
  """For now, we assume that they are just nested dicts or arrays."""
  if not hasattr(param, "dtype"):
    for name, child_param in param.items():
      _log_dict_recursively(parent_names + [name], child_param)
  else:
    logging.info("  %s: shape %s dtype %s",
                 "/".join(parent_names), param.shape, param.dtype)
