from typing import List, Union
from ray.rllib.utils.framework import try_import_torch

torch, nn = try_import_torch()


def make_time_major(
    tensor: Union["torch.Tensor", List["torch.Tensor"]],
    *,
    trajectory_len: int = None,
    recurrent_seq_len: int = None,
):
    """Swaps batch and trajectory axis.

    Args:
        tensor: A tensor or list of tensors to swap the axis of.
            NOTE: Each tensor must have the shape [B * T] where B is the batch size and
            T is the trajectory length.
        trajectory_len: The length of each trajectory being transformed.
            If None then `recurrent_seq_len` must be set.
        recurrent_seq_len: Sequence lengths if recurrent.
            If None then `trajectory_len` must be set.

    Returns:
        res: A tensor with swapped axes or a list of tensors with
        swapped axes.
    """
    if isinstance(tensor, (list, tuple)):
        return [
            make_time_major(_tensor, trajectory_len, recurrent_seq_len)
            for _tensor in tensor
        ]

    assert (trajectory_len != recurrent_seq_len) and (
        trajectory_len is None or recurrent_seq_len is None
    ), "Either trajectory_len or recurrent_seq_len must be set."

    if recurrent_seq_len:
        B = recurrent_seq_len.shape[0]
        T = tensor.shape[0] // B
    else:
        # Important: chop the tensor into batches at known episode cut
        # boundaries.
        # TODO: (sven) this is kind of a hack and won't work for
        #  batch_mode=complete_episodes.
        T = trajectory_len
        B = tensor.shape[0] // T
    rs = torch.reshape(tensor, [B, T] + list(tensor.shape[1:]))

    # Swap B and T axes.
    res = torch.transpose(rs, 1, 0)

    return res


def vtrace_torch(
    *,
    target_action_log_probs: "torch.Tensor",
    behaviour_action_log_probs: "torch.Tensor",
    discounts: "torch.Tensor",
    rewards: "torch.Tensor",
    values: "torch.Tensor",
    bootstrap_value: "torch.Tensor",
    clip_rho_threshold: Union[float, "torch.Tensor"] = 1.0,
    clip_pg_rho_threshold: Union[float, "torch.Tensor"] = 1.0,
):
    """V-trace for softmax policies implemented with torch.

    Calculates V-trace actor critic targets for softmax polices as described in
    "IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner
    Architectures" by Espeholt, Soyer, Munos et al. (https://arxiv.org/abs/1802.01561)

    The V-trace implementation used here closely resembles the one found in the
    scalable-agent repository by Google DeepMind, available at
    https://github.com/deepmind/scalable_agent. This version has been optimized to
    minimize the number of floating-point operations required per V-Trace
    calculation, achieved through the use of dynamic programming techniques. It's
    important to note that the mathematical expressions used in this implementation
    may appear quite different from those presented in the IMPALA paper.

    The following terminology applies:
        - `target policy` refers to the policy we are interested in improving.
        - `behaviour policy` refers to the policy that generated the given
            rewards and actions.
        - `T` refers to the time dimension. This is usually either the length of the
            trajectory or the length of the sequence if recurrent.
        - `B` refers to the batch size.

    Args:
        target_action_log_probs: Action log probs from the target policy. A float32
            tensor of shape [T, B].
        behaviour_action_log_probs: Action log probs from the behaviour policy. A
            float32 tensor of shape [T, B].
        discounts: A float32 tensor of shape [T, B] with the discount encountered when
            following the behaviour policy. This will be 0 for terminal timesteps
            (done=True) and gamma (the discount factor) otherwise.
        rewards: A float32 tensor of shape [T, B] with the rewards generated by
            following the behaviour policy.
        values: A float32 tensor of shape [T, B] with the value function estimates
            wrt. the target policy.
        bootstrap_value: A float32 of shape [B] with the value function estimate at
            time T.
        clip_rho_threshold: A scalar float32 tensor with the clipping threshold for
            importance weights (rho) when calculating the baseline targets (vs).
            rho^bar in the paper.
        clip_pg_rho_threshold: A scalar float32 tensor with the clipping threshold
            on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)).
    """
    log_rhos = target_action_log_probs - behaviour_action_log_probs

    rhos = torch.exp(log_rhos)
    if clip_rho_threshold is not None:
        clipped_rhos = torch.clamp(rhos, max=clip_rho_threshold)
    else:
        clipped_rhos = rhos

    cs = torch.clamp(rhos, max=1.0)
    # Append bootstrapped value to get [v1, ..., v_t+1]
    values_t_plus_1 = torch.cat(
        [values[1:], torch.unsqueeze(bootstrap_value, 0)], axis=0
    )

    deltas = clipped_rhos * (rewards + discounts * values_t_plus_1 - values)

    # Only move the for-loop to CPU.
    discounts_cpu = discounts.to("cpu")
    cs_cpu = cs.to("cpu")
    deltas_cpu = deltas.to("cpu")
    vs_minus_v_xs_cpu = [torch.zeros_like(bootstrap_value, device="cpu")]
    for i in reversed(range(len(discounts_cpu))):
        discount_t, c_t, delta_t = discounts_cpu[i], cs_cpu[i], deltas_cpu[i]
        vs_minus_v_xs_cpu.append(delta_t + discount_t * c_t * vs_minus_v_xs_cpu[-1])
    vs_minus_v_xs_cpu = torch.stack(vs_minus_v_xs_cpu[1:])
    # Move results back to GPU - if applicable.
    vs_minus_v_xs = vs_minus_v_xs_cpu.to(deltas.device)

    # Reverse the results back to original order.
    vs_minus_v_xs = torch.flip(vs_minus_v_xs, dims=[0])

    # Add V(x_s) to get v_s.
    vs = torch.add(vs_minus_v_xs, values)

    # Advantage for policy gradient.
    vs_t_plus_1 = torch.cat([vs[1:], torch.unsqueeze(bootstrap_value, 0)], axis=0)
    if clip_pg_rho_threshold is not None:
        clipped_pg_rhos = torch.clamp(rhos, max=clip_pg_rho_threshold)
    else:
        clipped_pg_rhos = rhos
    pg_advantages = clipped_pg_rhos * (rewards + discounts * vs_t_plus_1 - values)

    # Make sure no gradients backpropagated through the returned values.
    return torch.detach(vs), torch.detach(pg_advantages)
