# Copyright 2023 DeepMind Technologies Limited
#
# 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.
# ==============================================================================

"""A Jaxline script for training and evaluating TAPNet."""

import sys
from typing import Callable, Dict, Iterable, Iterator, Mapping, Optional, Tuple

from absl import app
from absl import flags
from absl import logging
import chex
import haiku as hk
import jax
import jax.numpy as jnp

from jaxline import experiment
from jaxline import platform
from jaxline import utils
from kubric.challenges.point_tracking import dataset
from ml_collections import config_dict

import numpy as np
import optax
import tensorflow as tf
import tensorflow_datasets as tfds

from juxtapose.trackers.tapnet import supervised_point_prediction
from juxtapose.trackers.tapnet import tapir_model
from juxtapose.trackers.tapnet import tapnet_model
from juxtapose.trackers.tapnet import task
from juxtapose.trackers.tapnet.utils import experiment_utils as exputils


class Experiment(experiment.AbstractExperiment):
    """TAPNet experiment.

    This constructs a set of tasks which are used to compute gradients, as well as
    a set of datasets which are passed into every task.  After the tasks compute
    gradients using the data, this experiment will aggregate and apply those
    gradients using an optimizer.

    This experiment is organized around the needs of Haiku's multi transform:
    each task specifies its own forward function, but different tasks need to
    share parameters.  Therefore, the design of this experiment is that each
    task has a forward function, but the global Experiment class is responsible
    for creating the multi transform out of these.  The Experiment class keeps
    track of the variables.  Each task doesn't directly call its own forward
    function.  Instead receives a "wrapped" forward function which contains
    the Haiku state, which injected (via closure) by the Experiment class.

    To do its job, multi transform needs two things: a function
    which initializes all variables, and the set of forward functions that need
    to be called separately, one for each task.  The  latter is just the set of
    forward functions for each task.  The initialization function is
    just a function that calls every forward function sequentially.
    """

    # A map from object properties that will be checkpointed to their name
    # in a checkpoint. Currently we assume that these are all sharded
    # device arrays.
    CHECKPOINT_ATTRS = {
        "_params": "params",
        "_state": "state",
        "_opt_state": "opt_state",
    }

    def __init__(
        self,
        mode: str,
        init_rng: jnp.ndarray,
        config: config_dict.ConfigDict,
    ):
        """Initializes the experiment.

        This includes constructing all of the requested tasks, creating the Haiku
        transforms, and pmapping the update function that will be used throughout
        training.  Currently the two supported training tasks are 'selfsup_tracks'
        and 'kubric'.

        Args:
          mode: either 'train' (for training) or one of the recognized eval modes
            (see Experiment.evaluate).
          init_rng: jax.random.PRNGKey for random number generation.
          config: config options.  See configs/tapnet_config.py for an example.
        """

        super().__init__(mode=mode, init_rng=init_rng)

        self.mode = mode
        self.init_rng = init_rng
        self.config = config

        # Checkpointed experiment state.
        self._params = None
        self._state = None
        self._opt_state = None

        self._optimizer = None

        # Input pipelines.
        self._train_input = None
        self._eval_input = None

        self.point_prediction = supervised_point_prediction.SupervisedPointPrediction(
            config, **config.supervised_point_prediction_kwargs
        )

        def forward(*args, is_training=True, **kwargs):
            shared_modules = self._construct_shared_modules()
            return self.point_prediction.forward_fn(
                *args,
                shared_modules=shared_modules,
                is_training=is_training,
                **kwargs,
            )

        self._transform = hk.transform_with_state(forward)

        # NOTE: We "donate" the `params, state, opt_state` arguments which allows
        # JAX (on some backends) to reuse the device memory associated with these
        # inputs to store the outputs of our function (which also start with
        # `params, state, opt_state`).
        self._update_func = jax.pmap(
            self._update_func, axis_name="i", donate_argnums=(0, 1, 2)
        )

    def _construct_shared_modules(self) -> Mapping[str, task.SharedModule]:
        """Constructs the TAPNet module which is used for all tasks.

        More generally, these are Haiku modules that are passed to all tasks so that
        weights are shared across tasks.

        Returns:
          A dict with a single key 'tapnet_model' containing the tapnet model.
        """
        shared_module_constructors = {
            "tapnet_model": tapnet_model.TAPNet,
            "tapir_model": tapir_model.TAPIR,
        }
        shared_modules = {}

        for shared_mod_name in self.config.shared_modules.shared_module_names:
            ctor = shared_module_constructors[shared_mod_name]
            kwargs = self.config.shared_modules[shared_mod_name + "_kwargs"]
            shared_modules[shared_mod_name] = ctor(**kwargs)
        return shared_modules

    #  _             _
    # | |_ _ __ __ _(_)_ __
    # | __| '__/ _` | | '_ \
    # | |_| | | (_| | | | | |
    #  \__|_|  \__,_|_|_| |_|
    #
    def step(  # pytype: disable=signature-mismatch  # numpy-scalars
        self,
        global_step: chex.Array,
        rng: chex.PRNGKey,
        *unused_args,
        **unused_kwargs,
    ) -> Dict[str, chex.Array]:
        """See base class."""

        if self._train_input is None:
            self._initialize_train()

        inputs = next(self._train_input)

        self._params, self._state, self._opt_state, scalars = self._update_func(
            self._params,
            self._state,
            self._opt_state,
            inputs,
            rng,
            global_step,
        )

        scalars = utils.get_first(scalars)

        if ((utils.get_first(global_step) + 1) % self.config.evaluate_every) == 0:
            for mode in self.config.eval_modes:
                eval_scalars = self.evaluate(global_step, rng=rng, mode=mode)
                scalars.update(eval_scalars)

        return scalars

    def _initialize_train(self):
        """Initializes the training parameters using the first training input.

        self._multi_transform.init will call the forward_fn, which allows
        Haiku to identify the params that need to be initialized.
        """
        self._train_input = utils.py_prefetch(self._build_train_input)

        total_steps = self.config.training.n_training_steps
        # Scale by the (negative) learning rate.
        # Check we haven't already restored params
        if self._params is None:
            logging.info("Initializing parameters.")

            inputs = next(self._train_input)

            init_net = jax.pmap(
                lambda *a: self._transform.init(*a, is_training=True),
                axis_name="i",
            )

            # Init uses the same RNG key on all hosts+devices to ensure everyone
            # computes the same initial state.
            init_rng = utils.bcast_local_devices(self.init_rng)
            self._params, self._state = init_net(init_rng, inputs)

        self._lr_schedule = exputils.get_lr_schedule(
            total_steps,
            self.config.optimizer,
        )

        self._optimizer = exputils.make_optimizer(
            self.config.optimizer,
            self._lr_schedule,
        )

        init_opt = jax.pmap(self._optimizer.init, axis_name="i")

        if self._opt_state is None:
            self._opt_state = init_opt(self._params)

    def _build_train_input(
        self,
    ) -> Iterable[Mapping[str, Mapping[str, np.ndarray]]]:
        """Builds the training input.

        For each dataset specified in the config, this will call the appropriate
        constructor for the dataset and then yield elements from each dataset.
        Currently supported datasets are 'selfsup_kinetics_tracks' and 'kubric'.

        Yields:
          Elements of a joint dataset, where each item is a dict.  The keys are
          the dataset names, and the associated values are items generated by
          those dataset classes.
        """

        # Note: a dataset constructor is a function which takes only kwargs from
        # the experiment config, and returns a dataset.
        #
        # A dataset consists of a generator function (calling next() on it will
        # produce a data value)
        dataset_constructors = {
            "kubric": dataset.create_point_tracking_dataset,
        }
        dataset_generators = {}
        for dset_name in self.config.datasets.dataset_names:
            ds_generator = self.create_dataset_generator(
                dataset_constructors,
                dset_name,
                color_augmentation=True,
            )

            dataset_generators[dset_name] = ds_generator

        while True:
            combined_dset = {}
            for dset_name in self.config.datasets.dataset_names:
                next_data = next(dataset_generators[dset_name])
                combined_dset[dset_name] = next_data
            yield combined_dset

    def create_dataset_generator(
        self,
        dataset_constructors: Mapping[str, Callable[..., tf.data.Dataset]],
        dset_name: str,
        color_augmentation: bool = False,
    ) -> Iterator[Mapping[str, np.ndarray]]:
        # Batch data on available devices.
        # Number of devices is unknown when an interpreter reads a config file.
        # Here we re-define batch dims to guide jax for right batching.
        batch_dims = [
            jax.local_device_count(),
            self.config.datasets[dset_name + "_kwargs"].batch_dims,
        ]
        dset_kwargs = dict(self.config.datasets[dset_name + "_kwargs"])
        dset_kwargs["batch_dims"] = []
        ds = dataset_constructors[dset_name](**dset_kwargs)
        if color_augmentation:
            ds = exputils.add_default_data_augmentation(ds)

        for dim in batch_dims[::-1]:
            ds = ds.batch(dim)
        np_ds = tfds.as_numpy(ds)
        return iter(np_ds)

    def _update_func(
        self,
        params: chex.ArrayTree,
        state: chex.ArrayTree,
        opt_state: chex.ArrayTree,
        inputs: chex.ArrayTree,
        rng: chex.PRNGKey,
        global_step: chex.Array,
    ) -> Tuple[chex.Array, chex.Array, chex.Array, Mapping[str, chex.Numeric]]:
        """Applies an update to parameters and returns new state."""

        updates = None
        grads, state, scalars = self.point_prediction.get_gradients(
            params,
            state,
            inputs,
            rng,
            global_step,
            self._transform.apply,
            is_training=True,
        )
        scalars = {**scalars}  # Mutable copy.
        if grads is not None:
            grads = jax.lax.psum(grads, axis_name="i")
            task_updates, opt_state = self._optimizer.update(
                grads,
                opt_state,
                params,
            )

            # We accumulate task_updates across tasks; if this is the first task
            # processed, then updates is None and we just initialize with the
            # updates from the first task.
            if updates is None:
                updates = task_updates
            else:
                updates = jax.tree_map(jnp.add, updates, task_updates)
            scalars["gradient_norm"] = optax.global_norm(grads)

        # Grab the learning rate to log before performing the step.
        learning_rate = self._lr_schedule(global_step)

        # Make the BYOL predictor train faster than everything else.
        def multiply_lr(variable_name, variable_update):
            need_fast_lr = False
            for substr in self.config.fast_variables:
                need_fast_lr = need_fast_lr or any([substr in x for x in variable_name])
            if need_fast_lr:
                logging.info("Boosting learning rate for: %s", "/".join(variable_name))
                return variable_update * 10.0
            return variable_update

        mapping_type = type(updates)

        def map_fn(nested_dict, prefix=None):
            prefix = prefix or []
            result = {}
            for k, v in nested_dict.items():
                if isinstance(v, mapping_type):
                    result[k] = map_fn(v, prefix + [k])
                else:
                    result[k] = multiply_lr(prefix + [k], v)

            return mapping_type(**result)

        updates = map_fn(updates)

        params = optax.apply_updates(params, updates)

        n_params = 0
        for k in params.keys():  # pytype: disable=attribute-error  # numpy-scalars
            for l in params[k]:
                n_params = n_params + np.prod(
                    params[k][l].shape
                )  # pytype: disable=attribute-error  # numpy-scalars

        # Scalars to log (note: we log the mean across all hosts/devices).
        scalars.update(
            {
                "learning_rate": learning_rate,
                "n_params (M)": float(n_params / 1e6),
            }
        )
        scalars = jax.lax.pmean(scalars, axis_name="i")

        return params, state, opt_state, scalars

    #                  _
    #   _____   ____ _| |
    #  / _ \ \ / / _` | |
    # |  __/\ V / (_| | |
    #  \___| \_/ \__,_|_|
    #
    def evaluate(  # pytype: disable=signature-mismatch  # numpy-scalars
        self,
        global_step: chex.Array,
        rng: chex.PRNGKey,
        mode: Optional[str] = None,
        **unused_args,
    ) -> Dict[str, chex.Array]:
        mode = mode or self.mode
        point_prediction_task = self.point_prediction
        forward_fn = self._transform.apply
        eval_scalars = point_prediction_task.evaluate(
            global_step=global_step,
            params=self._params,
            state=self._state,
            rng=rng,
            wrapped_forward_fn=forward_fn,
            mode=mode,
        )
        return {f"eval/{mode}/{key}": value for key, value in eval_scalars.items()}


def main(_):
    flags.mark_flag_as_required("config")
    # Keep TF off the GPU; otherwise it hogs all the memory and leaves none for
    # JAX.
    tf.config.experimental.set_visible_devices([], "GPU")
    tf.config.experimental.set_visible_devices([], "TPU")
    platform.main(
        Experiment,
        sys.argv[1:],
        checkpointer_factory=exputils.NumpyFileCheckpointer.create,
    )


if __name__ == "__main__":
    app.run(main)
