"""
Copyright (c) Meta Platforms, Inc. and affiliates.
All rights reserved.

This source code is licensed under the license found in the
LICENSE file in the root directory of this source tree.
"""

from functools import partial
import jax
import jax.numpy as jnp
import chex


class RollingStats:
    """
    This class tracks episodic stats, such as final returns
    and env complexity metrics. Works on a per-env basis.
    """

    def __init__(self, names, step_metrics_names=[], window=None):
        self.names = names
        self.step_metric_names = step_metrics_names
        self.window = window

    @partial(jax.jit, static_argnums=(0, 1))
    def reset_stats(self, batch_shape=(1,)):
        stats = {
            "n_episodes": jnp.zeros((*(batch_shape), 1), dtype=jnp.uint32),
            "n_steps": jnp.zeros((*(batch_shape), 1), dtype=jnp.uint32),
        }
        stats.update({name: jnp.zeros((*(batch_shape), 1)) for name in self.names})

        if self.window is not None:
            # Average over window
            stats.update(
                {
                    f"{name}_buffer": jnp.zeros((*(batch_shape), self.window))
                    for name in self.names
                }
            )

        return stats

    @partial(jax.jit, static_argnums=(0,))
    def update_stats(self, stats, done, info, max_episodes=jnp.inf):
        n_eps = stats["n_episodes"]
        n_steps = stats["n_steps"]

        print("update names", self.names, "step metric names", self.step_metric_names)

        for name in self.names:
            # Update stat
            if name not in info:
                continue

            new_val = info[name]

            # Only record first max_episode episodes
            done = done * (n_eps < max_episodes)
            if name in self.step_metric_names:  # NOTE don't understand this
                n_incr_prev = n_steps
                n_incr_total = n_steps + 1
                _metric_done = True
            else:
                n_incr_prev = n_eps
                n_incr_total = n_eps + done
                _metric_done = done

            if self.window is None:
                mean = stats[name]
                new_mean = self._update_stat_mean(
                    new_val, mean, n_incr_total, _metric_done
                )
            else:
                buffer_key = f"{name}_buffer"
                buffer = stats[buffer_key]
                new_mean, buffer = self._update_stat_window(
                    new_val, buffer, n_incr_total, n_incr_prev, _metric_done
                )
                stats.update({buffer_key: buffer})

            stats.update(
                {
                    name: new_mean,
                }
            )

            # Only update n_episodes based on real episodes
            if name in self.step_metric_names:
                stats.update({"n_steps": n_incr_total})
            else:
                stats.update({"n_episodes": n_incr_total})

        return stats

    @partial(jax.jit, static_argnums=(0,))
    def _update_stat_mean(self, new_val, mean, n_eps_total, done):
        z = 1 / jnp.maximum(1, n_eps_total)
        new_mean = done * (mean * (1 - z) + new_val * z) + (1 - done) * mean

        return new_mean

    @partial(jax.jit, static_argnums=(0,))
    def _update_stat_window(self, new_val, buffer, n_eps_total, n_eps_prev, done):
        cur_val = buffer[n_eps_prev % self.window]
        new_val = done * new_val + (1 - done) * cur_val
        buffer = buffer.at[n_eps_prev % self.window].set(new_val)
        new_mean = buffer.sum() / jnp.maximum(jnp.minimum(self.window, n_eps_total), 1)

        return new_mean, buffer


""" 
log wrapper: want to count certain values throughout an episode
"""


class LogEpisodicStats:

    def __init__(self, names):
        self.names = names

    @partial(jax.jit, static_argnums=(0, 1))
    def reset_stats(self, batch_shape=(1,)):
        stats = {
            "n_episodes": jnp.zeros((*(batch_shape), 1), dtype=jnp.uint32),
            "episode_lenghts": jnp.zeros((*(batch_shape), 1), dtype=jnp.uint32),
            "returned_episode_lengths": jnp.zeros(
                (*(batch_shape), 1), dtype=jnp.uint32
            ),
        }
        stats.update({name: jnp.zeros((*(batch_shape), 1)) for name in self.names})
        stats.update(
            {"ret_ep_" + name: jnp.zeros((*(batch_shape), 1)) for name in self.names}
        )

        return stats

    @partial(jax.jit, static_argnums=(0,))
    def update_stats(self, stats, ep_done: chex.Array, info, max_episodes=jnp.inf):
        n_eps = stats["n_episodes"]

        ep_done = ep_done * (n_eps < max_episodes)

        # Update stats
        new_episode_length = stats["episode_lenghts"] + 1
        stats.update(
            {
                "n_episodes": stats["n_episodes"] + ep_done,
                "episode_lenghts": new_episode_length * (1 - ep_done),
                "returned_episode_lengths": new_episode_length * ep_done
                + stats["returned_episode_lengths"] * (1 - ep_done),
            }
        )

        for name in self.names:
            if name not in info:
                continue

            new_val = info[name] + stats[name]
            stats.update(
                {
                    name: new_val * (1 - ep_done),
                    "ret_ep_"
                    + name: stats["ret_ep_" + name] * (1 - ep_done)
                    + new_val * ep_done,
                }
            )

        return stats
