"""
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
from enum import Enum
from collections import namedtuple

import numpy as np
import jax
import jax.numpy as jnp
import chex


class UEDScore(Enum):
    RELATIVE_REGRET = 1
    MEAN_RELATIVE_REGRET = 2
    POPULATION_REGRET = 3
    RETURN = 4
    NEG_RETURN = 5
    L1_VALUE_LOSS = 6
    POSITIVE_VALUE_LOSS = 7
    MAX_MC = 8
    VALUE_DISAGREEMENT = 9


@partial(jax.jit, static_argnums=(2, 3))
def compute_episodic_stats(
    metrics,
    dones,
    time_average=False,
    partial_metrics=0,
    partial_steps=0,
    return_partial=False,
):
    env_batch_shape = dones.shape[1:]
    n_episodes = jnp.zeros(env_batch_shape, dtype=jnp.uint32)
    sum_ep_metrics = jnp.zeros(env_batch_shape, dtype=jnp.float32)
    partial_metrics = jnp.zeros(env_batch_shape, dtype=jnp.float32)
    max_metrics = jnp.zeros(env_batch_shape, dtype=jnp.float32)
    steps = jnp.zeros(env_batch_shape, dtype=jnp.float32)
    partial_steps = jnp.zeros(env_batch_shape, dtype=jnp.float32)

    def _compute_metrics(carry, step):
        (n_episodes, sum_ep_metrics, max_metrics, partial_metrics, partial_steps) = (
            carry
        )

        _metrics, _dones = step

        partial_metrics += _metrics
        partial_steps += 1

        if time_average:
            ep_metric = partial_metrics / partial_steps
        else:
            ep_metric = partial_metrics

        sum_ep_metrics += _dones * ep_metric
        max_metrics = (
            _dones * jnp.maximum(max_metrics, ep_metric) + (1 - _dones) * max_metrics
        )

        n_episodes += _dones

        partial_metrics = (1 - _dones) * partial_metrics
        partial_steps = (1 - _dones) * partial_steps

        return (
            n_episodes,
            sum_ep_metrics,
            max_metrics,
            partial_metrics,
            partial_steps,
        ), None

    # @todo: Track and return partial means + partial n_steps
    (n_episodes, sum_ep_metrics, max_metrics, partial_metrics, partial_steps), _ = (
        jax.lax.scan(
            _compute_metrics,
            (n_episodes, sum_ep_metrics, max_metrics, partial_metrics, partial_steps),
            (metrics, dones),
            length=len(metrics),
        )
    )

    # Take mean over eval dimension
    total_metrics_per_env = sum_ep_metrics.sum(-1)
    n_episodes_per_env = n_episodes.sum(-1)
    n_episodes_per_env = jnp.maximum(n_episodes_per_env, 1)

    # Take max over eval dimension
    max_metrics_per_env = max_metrics.max(-1)

    return total_metrics_per_env / n_episodes_per_env, max_metrics_per_env


@partial(jax.jit, static_argnums=(3, 4))
def compute_episodic_stats_mask(
    metrics,
    dones,
    mask,
    time_average=False,
    partial_metrics=0,
    partial_steps=0,
    return_partial=False,
):

    env_batch_shape = dones.shape[1:]
    n_episodes = jnp.zeros(env_batch_shape, dtype=jnp.uint32)
    sum_ep_metrics = jnp.zeros(env_batch_shape, dtype=jnp.float32)
    partial_metrics = jnp.zeros(env_batch_shape, dtype=jnp.float32)
    max_metrics = jnp.zeros(env_batch_shape, dtype=jnp.float32)
    steps = jnp.zeros(env_batch_shape, dtype=jnp.float32)
    partial_steps = jnp.zeros(env_batch_shape, dtype=jnp.float32)

    def _compute_metrics(carry, step):
        (n_episodes, sum_ep_metrics, max_metrics, partial_metrics, partial_steps) = (
            carry
        )

        _metrics, _dones, _mask = step  # mask is 1 for valid, 0 for invalid
        just_done = _dones * _mask
        partial_metrics += _metrics * _mask
        partial_steps += 1

        if time_average:
            ep_metric = partial_metrics / partial_steps
        else:
            ep_metric = partial_metrics
        sum_ep_metrics += just_done * ep_metric
        max_metrics = (
            just_done * jnp.maximum(max_metrics, ep_metric)
            + (1 - just_done) * max_metrics
        )
        n_episodes += just_done

        partial_metrics = (1 - just_done) * partial_metrics
        partial_steps = (1 - just_done) * partial_steps

        return (
            n_episodes,
            sum_ep_metrics,
            max_metrics,
            partial_metrics,
            partial_steps,
        ), None

    # @todo: Track and return partial means + partial n_steps
    (n_episodes, sum_ep_metrics, max_metrics, partial_metrics, partial_steps), _ = (
        jax.lax.scan(
            _compute_metrics,
            (n_episodes, sum_ep_metrics, max_metrics, partial_metrics, partial_steps),
            (metrics, dones, mask),
            length=len(metrics),
        )
    )

    # Take mean over eval dimension
    total_metrics_per_env = sum_ep_metrics.sum(-1)
    n_episodes_per_env = n_episodes.sum(-1)
    n_episodes_per_env = jnp.maximum(n_episodes_per_env, 1)

    # Take max over eval dimension
    max_metrics_per_env = max_metrics.max(-1)

    return total_metrics_per_env / n_episodes_per_env, max_metrics_per_env


@partial(jax.jit, static_argnums=(0,))
def _compute_ued_scores(score_name: UEDScore, batch: namedtuple, info=None):
    """
    Compute UED score from a rollout batch.
    Individual score functions return a tuple of mean_scores and max_scores,
    where each is of dimension n_agents x n_envs.
    """
    print("score name", score_name)
    if score_name in [
        UEDScore.RELATIVE_REGRET,
        UEDScore.MEAN_RELATIVE_REGRET,
        UEDScore.POPULATION_REGRET,
    ]:
        mean_scores, max_scores, score_info = compute_return(batch)

    elif score_name == UEDScore.RETURN:
        mean_scores, max_scores, score_info = compute_return(batch)

    elif score_name == UEDScore.NEG_RETURN:
        batch = batch._replace(rewards=-batch.rewards)
        mean_scores, max_scores, score_info = compute_return(batch)

    elif score_name == UEDScore.MAX_MC:
        mean_scores, max_scores, score_info = compute_max_mc(batch, info)

    elif score_name == UEDScore.L1_VALUE_LOSS:
        mean_scores, max_scores, score_info = compute_l1_value_loss(batch)

    elif score_name == UEDScore.POSITIVE_VALUE_LOSS:
        mean_scores, max_scores, score_info = compute_positive_value_loss(batch)

    elif score_name == UEDScore.VALUE_DISAGREEMENT:
        mean_scores, max_scores, score_info = compute_value_disagreement(batch)

    return mean_scores, max_scores, score_info


@partial(jax.jit, static_argnums=(0, 2, 4, 5))
def compute_ued_scores(
    score_name: UEDScore,
    batch: namedtuple,
    n_eval: int,
    info: dict = None,
    ignore_val=None,
    per_agent=False,
):
    batch = jax.tree_util.tree_map(lambda x: x[None], batch)
    n_agents, n_steps, flat_batch_size = batch.dones.shape  # need to swap 0 and 1 axes
    # flat_batch_size = 1
    pop_batch_shape = (n_agents, n_steps, flat_batch_size // n_eval, n_eval)
    batch = jax.tree_util.tree_map(
        lambda x: x.reshape(*pop_batch_shape, *x.shape[3:]), batch
    )

    mean_env_returns_per_agent, max_env_returns_per_agent, score_info = jax.vmap(
        _compute_ued_scores, in_axes=(None, 0, 0)
    )(score_name, batch, info)

    if score_name in [UEDScore.RELATIVE_REGRET, UEDScore.MEAN_RELATIVE_REGRET]:
        assert (
            len(mean_env_returns_per_agent) == 2
        ), "Standard PAIRED requires exactly 2 agents."

    if score_name == UEDScore.RELATIVE_REGRET:
        scores = jnp.clip(
            max_env_returns_per_agent[1] - mean_env_returns_per_agent[0], 0
        )

    elif score_name == UEDScore.MEAN_RELATIVE_REGRET:
        scores = jnp.clip(
            mean_env_returns_per_agent[1] - mean_env_returns_per_agent[0], 0
        )

    elif score_name == UEDScore.POPULATION_REGRET:
        max_env_returns = max_env_returns_per_agent.max(0)
        mean_env_returns = mean_env_returns_per_agent.mean(0)
        scores = max_env_returns - mean_env_returns
    else:
        if per_agent:
            scores = mean_env_returns_per_agent
            max_scores = max_env_returns_per_agent
        else:
            scores = mean_env_returns_per_agent.mean(0)
            max_scores = max_env_returns_per_agent.max(0)

    if ignore_val is not None:
        if per_agent:
            incomplete_idxs = batch.dones.sum(axis=(1, -1)) == 0
        else:
            incomplete_idxs = batch.dones.sum(axis=(0, 1, -1)) == 0

        scores = jnp.where(incomplete_idxs, ignore_val, scores)

    return scores, score_info


# ======== UED score computations ========
def compute_return(batch):
    mean_scores, max_scores = compute_episodic_stats(
        batch.rewards, batch.dones, time_average=False
    )

    return mean_scores, max_scores, None


def compute_l1_value_loss(batch):
    mean_scores, max_scores = compute_episodic_stats(
        jnp.abs(batch.advantages), batch.dones, time_average=True
    )

    return mean_scores, max_scores, None


def compute_positive_value_loss(batch):
    mean_scores, max_scores = compute_episodic_stats_mask(
        jnp.clip(batch.advantages, 0), batch.dones, batch.mask, time_average=True
    )

    return mean_scores, max_scores, None


def compute_max_mc(batch, info):
    _, max_env_returns_per_agent = compute_episodic_stats(
        batch.rewards, batch.dones, time_average=False
    )

    max_returns = jnp.maximum(max_env_returns_per_agent, info["max_returns"])
    mean_scores, max_scores = compute_episodic_stats(
        max_returns[jnp.newaxis, :, jnp.newaxis] - batch.values,  # Can be negative
        batch.dones,
        time_average=True,
    )

    score_info = {"max_returns": max_env_returns_per_agent}

    return mean_scores, max_scores, score_info


def compute_value_disagreement(batch):
    mean_scores, max_scores = compute_episodic_stats(
        batch.values.std(-1), batch.dones, time_average=True
    )

    return mean_scores, max_scores, None
