import os
from posixpath import isabs
from functools import partial
from typing import Tuple
import typing

import chex
import jax
import jax.numpy as jnp
from flax.training.train_state import TrainState

from gymnax import EnvParams, EnvState
from gymnax.environments.environment import Environment
from kinetix.models.actor_critic import ScannedRNN

Observation = typing.Any


def evaluate_rnn(  # from jaxued
    rng: chex.PRNGKey,
    env: Environment,
    env_params: EnvParams,
    train_state: TrainState,
    init_hstate: chex.ArrayTree,
    init_obs: Observation,
    init_env_state: EnvState,
    max_episode_length: int,
    keep_states=True,
    return_trajectories=False,
) -> Tuple[chex.Array, chex.Array, chex.Array]:
    """This runs the RNN on the environment, given an initial state and observation, and returns (states, rewards, episode_lengths)

    Args:
        rng (chex.PRNGKey):
        env (Environment):
        env_params (EnvParams):
        train_state (TrainState):
        init_hstate (chex.ArrayTree): Shape (num_levels, )
        init_obs (Observation): Shape (num_levels, )
        init_env_state (EnvState): Shape (num_levels, )
        max_episode_length (int):

    Returns:
        Tuple[chex.Array, chex.Array, chex.Array]: (States, rewards, episode lengths) ((NUM_STEPS, NUM_LEVELS), (NUM_STEPS, NUM_LEVELS), (NUM_LEVELS,)
    """
    num_levels = jax.tree_util.tree_flatten(init_obs)[0][0].shape[0]

    def step(carry, _):
        rng, hstate, obs, state, done, mask, episode_length = carry
        rng, rng_action, rng_step = jax.random.split(rng, 3)

        x = jax.tree.map(lambda x: x[None, ...], (obs, done))
        hstate, pi, _ = train_state.apply_fn(train_state.params, hstate, x)
        action = pi.sample(seed=rng_action).squeeze(0)

        obs, next_state, reward, done, info = jax.vmap(env.step, in_axes=(0, 0, 0, None, 0))(
            jax.random.split(rng_step, num_levels),
            state,
            action,
            env_params,
            init_env_state,  # use this to reset to the init env state, keyword arguments are not easily vmappable
        )

        next_mask = mask & ~done
        episode_length += mask

        if keep_states:
            return (rng, hstate, obs, next_state, done, next_mask, episode_length), (state, reward, done, info)
        else:
            return (rng, hstate, obs, next_state, done, next_mask, episode_length), (None, reward, done, info)

    (_, _, _, _, _, _, episode_lengths), (states, rewards, dones, infos) = jax.lax.scan(
        step,
        (
            rng,
            init_hstate,
            init_obs,
            init_env_state,
            jnp.zeros(num_levels, dtype=bool),
            jnp.ones(num_levels, dtype=bool),
            jnp.zeros(num_levels, dtype=jnp.int32),
        ),
        None,
        length=max_episode_length,
    )
    done_idx = jnp.argmax(dones, axis=0)

    to_return = (states, rewards, done_idx, episode_lengths, infos)
    if return_trajectories:
        return to_return, (dones, rewards)
    return to_return


def general_eval(
    rng: chex.PRNGKey,
    eval_env: Environment,
    env_params: EnvParams,
    train_state: TrainState,
    levels: EnvState,
    num_eval_steps: int,
    num_levels: int,
    keep_states=True,
    return_trajectories=False,
    initialize_carry_fn=ScannedRNN.initialize_carry,
):
    """
    This evaluates the current policy on the set of evaluation levels
    It returns (states, cum_rewards, episode_lengths), with shapes (num_steps, num_eval_levels, ...), (num_eval_levels,), (num_eval_levels,)
    """
    rng, rng_reset = jax.random.split(rng)
    init_obs, init_env_state = jax.vmap(eval_env.reset, (0, None, 0))(
        jax.random.split(rng_reset, num_levels), env_params, levels
    )
    init_hstate = initialize_carry_fn(num_levels)
    (states, rewards, done_idx, episode_lengths, infos), (dones, reward) = evaluate_rnn(
        rng,
        eval_env,
        env_params,
        train_state,
        init_hstate,
        init_obs,
        init_env_state,
        num_eval_steps,
        keep_states=keep_states,
        return_trajectories=True,
    )
    mask = jnp.arange(num_eval_steps)[..., None] < episode_lengths
    cum_rewards = (rewards * mask).sum(axis=0)
    to_return = (
        states,
        cum_rewards,
        done_idx,
        episode_lengths,
        infos,
    )  # (num_steps, num_eval_levels, ...), (num_eval_levels,), (num_eval_levels,)

    if return_trajectories:
        return to_return, (dones, reward)
    return to_return


def compute_gae(
    gamma: float,
    lambd: float,
    last_value: chex.Array,
    values: chex.Array,
    rewards: chex.Array,
    dones: chex.Array,
) -> Tuple[chex.Array, chex.Array]:
    """This takes in arrays of shape (NUM_STEPS, NUM_ENVS) and returns the advantages and targets.

    Args:
        gamma (float):
        lambd (float):
        last_value (chex.Array):  Shape (NUM_ENVS)
        values (chex.Array): Shape (NUM_STEPS, NUM_ENVS)
        rewards (chex.Array): Shape (NUM_STEPS, NUM_ENVS)
        dones (chex.Array): Shape (NUM_STEPS, NUM_ENVS)

    Returns:
        Tuple[chex.Array, chex.Array]: advantages, targets; each of shape (NUM_STEPS, NUM_ENVS)
    """

    def compute_gae_at_timestep(carry, x):
        gae, next_value = carry
        value, reward, done = x
        delta = reward + gamma * next_value * (1 - done) - value
        gae = delta + gamma * lambd * (1 - done) * gae
        return (gae, value), gae

    _, advantages = jax.lax.scan(
        compute_gae_at_timestep,
        (jnp.zeros_like(last_value), last_value),
        (values, rewards, dones),
        reverse=True,
        unroll=16,
    )
    return advantages, advantages + values


def sample_trajectories_rnn(
    rng: chex.PRNGKey,
    env: Environment,
    env_params: EnvParams,
    train_state: TrainState,
    init_hstate: chex.ArrayTree,
    init_obs: Observation,
    init_env_state: EnvState,
    num_envs: int,
    max_episode_length: int,
    return_states: bool = False,
) -> Tuple[
    Tuple[chex.PRNGKey, TrainState, chex.ArrayTree, Observation, EnvState, chex.Array],
    Tuple[Observation, chex.Array, chex.Array, chex.Array, chex.Array, chex.Array, dict],
]:
    """This samples trajectories from the environment using the agent specified by the `train_state`.

    Args:

        rng (chex.PRNGKey): Singleton
        env (Environment):
        env_params (EnvParams):
        train_state (TrainState): Singleton
        init_hstate (chex.ArrayTree): This is the init RNN hidden state, has to have shape (NUM_ENVS, ...)
        init_obs (Observation): The initial observation, shape (NUM_ENVS, ...)
        init_env_state (EnvState): The initial env state (NUM_ENVS, ...)
        num_envs (int): The number of envs that are vmapped over.
        max_episode_length (int): The maximum episode length, i.e., the number of steps to do the rollouts for.

    Returns:
        Tuple[Tuple[chex.PRNGKey, TrainState, chex.ArrayTree, Observation, EnvState, chex.Array], Tuple[Observation, chex.Array, chex.Array, chex.Array, chex.Array, chex.Array, dict]]: (rng, train_state, hstate, last_obs, last_env_state, last_value), traj, where traj is (obs, action, reward, done, log_prob, value, info). The first element in the tuple consists of arrays that have shapes (NUM_ENVS, ...) (except `rng` and and `train_state` which are singleton). The second element in the tuple is of shape (NUM_STEPS, NUM_ENVS, ...), and it contains the trajectory.
    """

    def sample_step(carry, _):
        rng, train_state, hstate, obs, env_state, last_done = carry
        prev_state = env_state
        rng, rng_action, rng_step = jax.random.split(rng, 3)

        x = jax.tree.map(lambda x: x[None, ...], (obs, last_done))
        hstate, pi, value = train_state.apply_fn(train_state.params, hstate, x)
        action = pi.sample(seed=rng_action)
        log_prob = pi.log_prob(action)
        value, action, log_prob = jax.tree.map(lambda x: x.squeeze(0), (value, action, log_prob))

        next_obs, env_state, reward, done, info = jax.vmap(env.step, in_axes=(0, 0, 0, None, 0))(
            jax.random.split(rng_step, num_envs), env_state, action, env_params, init_env_state
        )

        carry = (rng, train_state, hstate, next_obs, env_state, done)
        step = (obs, action, reward, done, log_prob, value, info)
        if return_states:
            step += (prev_state,)
        return carry, step

    (rng, train_state, hstate, last_obs, last_env_state, last_done), traj = jax.lax.scan(
        sample_step,
        (
            rng,
            train_state,
            init_hstate,
            init_obs,
            init_env_state,
            jnp.zeros(num_envs, dtype=bool),
        ),
        None,
        length=max_episode_length,
    )

    x = jax.tree.map(lambda x: x[None, ...], (last_obs, last_done))
    _, _, last_value = train_state.apply_fn(train_state.params, hstate, x)

    my_obs = traj[0]
    rew = traj[2]

    return (rng, train_state, hstate, last_obs, last_env_state, last_value.squeeze(0)), traj


def update_actor_critic_rnn(
    rng: chex.PRNGKey,
    train_state: TrainState,
    init_hstate: chex.ArrayTree,
    batch: chex.ArrayTree,
    num_envs: int,
    n_steps: int,
    n_minibatch: int,
    n_epochs: int,
    clip_eps: float,
    entropy_coeff: float,
    critic_coeff: float,
    update_grad: bool = True,
) -> Tuple[Tuple[chex.PRNGKey, TrainState], chex.ArrayTree]:
    """This function takes in a rollout, and PPO hyperparameters, and updates the train state.

    Args:
        rng (chex.PRNGKey):
        train_state (TrainState):
        init_hstate (chex.ArrayTree):
        batch (chex.ArrayTree): obs, actions, dones, log_probs, values, targets, advantages
        num_envs (int):
        n_steps (int):
        n_minibatch (int):
        n_epochs (int):
        clip_eps (float):
        entropy_coeff (float):
        critic_coeff (float):
        update_grad (bool, optional): If False, the train state does not actually get updated. Defaults to True.

    Returns:
        Tuple[Tuple[chex.PRNGKey, TrainState], chex.ArrayTree]: It returns a new rng, the updated train_state, and the losses. The losses have structure (loss, (l_vf, l_clip, entropy))
    """
    obs, actions, dones, log_probs, values, targets, advantages = batch
    last_dones = jnp.roll(dones, 1, axis=0).at[0].set(False)
    batch = obs, actions, last_dones, log_probs, values, targets, advantages

    def update_epoch(carry, _):
        def update_minibatch(train_state, minibatch):
            init_hstate, obs, actions, last_dones, log_probs, values, targets, advantages = minibatch

            def loss_fn(params):
                _, pi, values_pred = train_state.apply_fn(params, init_hstate, (obs, last_dones))
                log_probs_pred = pi.log_prob(actions)
                entropy = pi.entropy().mean()

                ratio = jnp.exp(log_probs_pred - log_probs)
                A = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
                l_clip = (-jnp.minimum(ratio * A, jnp.clip(ratio, 1 - clip_eps, 1 + clip_eps) * A)).mean()

                values_pred_clipped = values + (values_pred - values).clip(-clip_eps, clip_eps)
                l_vf = 0.5 * jnp.maximum((values_pred - targets) ** 2, (values_pred_clipped - targets) ** 2).mean()

                loss = l_clip + critic_coeff * l_vf - entropy_coeff * entropy

                return loss, (l_vf, l_clip, entropy)

            grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
            loss, grads = grad_fn(train_state.params)
            if update_grad:
                train_state = train_state.apply_gradients(grads=grads)
            grad_norm = jnp.linalg.norm(
                jnp.concatenate(jax.tree.map(lambda x: x.flatten(), jax.tree_util.tree_flatten(grads)[0]))
            )
            return train_state, (loss, grad_norm)

        rng, train_state = carry
        rng, rng_perm = jax.random.split(rng)
        permutation = jax.random.permutation(rng_perm, num_envs)
        minibatches = (
            jax.tree.map(
                lambda x: jnp.take(x, permutation, axis=0).reshape(n_minibatch, -1, *x.shape[1:]),
                init_hstate,
            ),
            *jax.tree.map(
                lambda x: jnp.take(x, permutation, axis=1)
                .reshape(x.shape[0], n_minibatch, -1, *x.shape[2:])
                .swapaxes(0, 1),
                batch,
            ),
        )
        train_state, (losses, grads) = jax.lax.scan(update_minibatch, train_state, minibatches)
        return (rng, train_state), (losses, grads)

    return jax.lax.scan(update_epoch, (rng, train_state), None, n_epochs)


# Cannot jit outside otherwise jax errors with trying to hash a tracer.
# @partial(jax.jit, static_argnums=(0, 2, 8, 9))
def sample_trajectories_and_learn(
    env: Environment,
    env_params: EnvParams,
    config: dict,
    rng: chex.PRNGKey,
    train_state: TrainState,
    init_hstate: chex.Array,
    init_obs: Observation,
    init_env_state: EnvState,
    update_grad: bool = True,
    return_states: bool = False,
) -> Tuple[
    Tuple[chex.PRNGKey, TrainState, Observation, EnvState],
    Tuple[
        Observation,
        chex.Array,
        chex.Array,
        chex.Array,
        chex.Array,
        chex.Array,
        dict,
        chex.Array,
        chex.Array,
        chex.ArrayTree,
        chex.Array,
    ],
]:
    """This function loops the following:
        - rollout for config['num_steps']
        - learn / update policy

    And it loops it for config['outer_rollout_steps'].
    What is returns is a new carry (rng, train_state, init_obs, init_env_state), and concatenated rollouts. The shape of the rollouts are config['num_steps'] * config['outer_rollout_steps']. In other words, the trajectories returned by this function are the same as if we ran rollouts for config['num_steps'] * config['outer_rollout_steps'] steps, but the agent does perform PPO updates in between.

    Args:
        env (Environment):
        env_params (EnvParams):
        config (dict):
        rng (chex.PRNGKey):
        train_state (TrainState):
        init_obs (Observation):
        init_env_state (EnvState):
        update_grad (bool, optional): Defaults to True.

    Returns:
        Tuple[Tuple[chex.PRNGKey, TrainState, Observation, EnvState], Tuple[Observation, chex.Array, chex.Array, chex.Array, chex.Array, chex.Array, dict, chex.Array, chex.Array, chex.ArrayTree, chex.Array]]: This returns a tuple:
        (
            (rng, train_state, init_obs, init_env_state),
            (obs, actions, rewards, dones, log_probs, values, info, advantages, targets, losses, grads)
        )
    """

    @jax.jit
    def single_step(carry, _):
        rng, train_state, init_hstate, init_obs, init_env_state = carry
        ((rng, train_state, new_hstate, last_obs, last_env_state, last_value), traj,) = sample_trajectories_rnn(
            rng,
            env,
            env_params,
            train_state,
            init_hstate,
            init_obs,
            init_env_state,
            config["num_train_envs"],
            config["num_steps"],
            return_states=return_states,
        )
        if return_states:
            states = traj[-1]
            traj = traj[:-1]

        (obs, actions, rewards, dones, log_probs, values, info) = traj
        advantages, targets = compute_gae(config["gamma"], config["gae_lambda"], last_value, values, rewards, dones)

        # Update the policy using trajectories collected from replay levels
        (rng, train_state), (losses, grads) = update_actor_critic_rnn(
            rng,
            train_state,
            init_hstate,
            (obs, actions, dones, log_probs, values, targets, advantages),
            config["num_train_envs"],
            config["num_steps"],
            config["num_minibatches"],
            config["update_epochs"],
            config["clip_eps"],
            config["ent_coef"],
            config["vf_coef"],
            update_grad=update_grad,
        )
        new_carry = (rng, train_state, new_hstate, last_obs, last_env_state)
        step = (obs, actions, rewards, dones, log_probs, values, info, advantages, targets, losses, grads)
        if return_states:
            step += (states,)
        return new_carry, step

    @jax.jit
    def inner(rng, train_state, init_hstate, init_obs, init_env_state):
        carry = (rng, train_state, init_hstate, init_obs, init_env_state)
        new_carry, all_rollouts = jax.lax.scan(single_step, carry, None, length=config["outer_rollout_steps"])

        all_rollouts = jax.tree.map(lambda x: jnp.concatenate(x, axis=0), all_rollouts)
        return new_carry, all_rollouts

    return inner(rng, train_state, init_hstate, init_obs, init_env_state)


def no_op_rollout(
    env: Environment,
    env_params: EnvParams,
    rng: chex.PRNGKey,
    init_obs: Observation,
    init_env_state: EnvState,
    num_envs: int,
    max_episode_length: int,
    do_random=False,
):

    noop = jnp.array(env.action_type.noop_action())
    zero_action = jnp.repeat(noop[None, ...], num_envs, axis=0)
    SHAPE = zero_action.shape

    def sample_step(carry, _):
        rng, obs, env_state, last_done = carry
        rng, rng_step, _rng = jax.random.split(rng, 3)
        if do_random:
            action = jax.vmap(env.action_space(env_params).sample)(jax.random.split(_rng, num_envs))
        else:
            action = zero_action

        next_obs, env_state, reward, done, info = jax.vmap(env.step, in_axes=(0, 0, 0, None, 0))(
            jax.random.split(rng_step, num_envs), env_state, action, env_params, init_env_state
        )

        carry = (rng, next_obs, env_state, done)
        return carry, (obs, action, reward, done, info)

    (rng, last_obs, last_env_state, last_done), traj = jax.lax.scan(
        sample_step,
        (
            rng,
            init_obs,
            init_env_state,
            jnp.zeros(num_envs, dtype=bool),
        ),
        None,
        length=max_episode_length,
    )

    info = traj[-1]
    dones = traj[-2]

    returns_per_env = (info["returned_episode_returns"] * dones).sum(axis=0) / jnp.maximum(1, dones.sum(axis=0))
    lens_per_env = (info["returned_episode_lengths"] * dones).sum(axis=0) / jnp.maximum(1, dones.sum(axis=0))
    success_per_env = (info["returned_episode_solved"] * dones).sum(axis=0) / jnp.maximum(1, dones.sum(axis=0))
    return returns_per_env, lens_per_env, success_per_env


def no_op_and_random_rollout(
    env: Environment,
    env_params: EnvParams,
    rng: chex.PRNGKey,
    init_obs: Observation,
    init_env_state: EnvState,
    num_envs: int,
    max_episode_length: int,
):
    returns_noop, lens_noop, success_noop = no_op_rollout(
        env, env_params, rng, init_obs, init_env_state, num_envs, max_episode_length, do_random=False
    )
    returns_random, lens_random, success_random = no_op_rollout(
        env, env_params, rng, init_obs, init_env_state, num_envs, max_episode_length, do_random=True
    )
    return returns_noop, lens_noop, success_noop, returns_random, lens_random, success_random
