from typing import Tuple

import jax
import jax.numpy as jnp

from jax_rl.agents.common import InfoDict, Params
from jax_rl.agents.sac.state import SAC
from jax_rl.datasets import Batch


def target_update(sac: SAC) -> SAC:
    # A mix of soft and hard target updates.
    params = sac.critic.optimizer.target

    tau = sac.config.tau
    if tau < 1:
        new_target_params = jax.tree_multimap(
            lambda p, tp: p * tau + tp * (1 - tau), params,
            sac.target_critic.params)
    else:
        new_target_params = params

    step = sac.critic.optimizer.state.step
    target_update_period = sac.config.target_update_period
    if target_update_period > 1:
        new_target_params = jax.lax.cond(step % target_update_period == 0,
                                         lambda _: new_target_params,
                                         lambda _: sac.target_critic.params,
                                         operand=None)

    new_target_critic = sac.target_critic.replace(params=new_target_params)

    return sac.replace(target_critic=new_target_critic)


def update(sac: SAC, batch: Batch) -> Tuple[SAC, InfoDict]:
    dist = sac.actor(batch.next_observations)
    rng, key = jax.random.split(sac.rng)
    next_actions = dist.sample(seed=key)
    next_log_probs = dist.log_prob(next_actions)

    next_q1, next_q2 = sac.target_critic(batch.next_observations, next_actions)
    next_q = jnp.minimum(next_q1, next_q2)

    target_q = batch.rewards + sac.config.discount * batch.masks * (
        next_q - sac.temp() * next_log_probs)

    def critic_loss_fn(
        critic_params: Params
    ) -> Tuple[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]]:
        q1, q2 = sac.critic.apply({'params': critic_params},
                                  batch.observations, batch.actions)
        critic_loss = ((q1 - target_q)**2 + (q2 - target_q)**2).mean()
        return critic_loss, {
            'critic_loss': critic_loss,
            'q1': q1.mean(),
            'q2': q2.mean()
        }

    info, new_critic = sac.critic.apply_gradient(critic_loss_fn)

    new_sac = sac.replace(critic=new_critic, rng=rng)

    return new_sac, info
