"""Implementations of algorithms for continuous control."""

import functools
from typing import Optional, Sequence, Tuple

import jax
import jax.numpy as jnp
import numpy as np
import optax

from jaxrl.agents.sac import temperature
from jaxrl.agents.sac.actor import update as update_actor
from jaxrl.agents.sac.critic import target_update
from jaxrl.agents.sac_v1.critic import update_q, update_v
from jaxrl.datasets import Batch
from jaxrl.networks import critic_net, policies
from jaxrl.networks.common import InfoDict, Model, PRNGKey


@functools.partial(jax.jit, static_argnames=('update_target'))
def _update_jit(
    rng: PRNGKey, actor: Model, critic: Model, value: Model,
    target_value: Model, temp: Model, batch: Batch, discount: float,
    tau: float, target_entropy: float, update_target: bool
) -> Tuple[PRNGKey, Model, Model, Model, Model, Model, InfoDict]:

    new_critic, critic_info = update_q(critic, target_value, batch, discount)

    rng, key = jax.random.split(rng)
    new_actor, actor_info = update_actor(key, actor, new_critic, temp, batch)

    rng, key = jax.random.split(rng)
    new_value, value_info = update_v(key, new_actor, new_critic, value, temp,
                                     batch, True)

    if update_target:
        new_target_value = target_update(new_value, target_value, tau)
    else:
        new_target_value = target_value

    new_temp, alpha_info = temperature.update(temp, actor_info['entropy'],
                                              target_entropy)

    return rng, new_actor, new_critic, new_value, new_target_value, new_temp, {
        **critic_info,
        **value_info,
        **actor_info,
        **alpha_info
    }


class SACV1Learner(object):

    def __init__(self,
                 seed: int,
                 observations: jnp.ndarray,
                 actions: jnp.ndarray,
                 actor_lr: float = 3e-4,
                 value_lr: float = 3e-4,
                 critic_lr: float = 3e-4,
                 temp_lr: float = 3e-4,
                 hidden_dims: Sequence[int] = (256, 256),
                 discount: float = 0.99,
                 tau: float = 0.005,
                 target_update_period: int = 1,
                 target_entropy: Optional[float] = None,
                 init_temperature: float = 1.0):
        """
        An implementation of the version of Soft-Actor-Critic described in https://arxiv.org/abs/1801.01290
        """

        action_dim = actions.shape[-1]

        if target_entropy is None:
            self.target_entropy = -action_dim / 2
        else:
            self.target_entropy = target_entropy

        self.tau = tau
        self.target_update_period = target_update_period
        self.discount = discount

        rng = jax.random.PRNGKey(seed)
        rng, actor_key, critic_key, temp_key = jax.random.split(rng, 4)

        actor_def = policies.NormalTanhPolicy(hidden_dims, action_dim)
        actor = Model.create(actor_def,
                             inputs=[actor_key, observations],
                             tx=optax.adam(learning_rate=actor_lr))

        critic_def = critic_net.DoubleCritic(hidden_dims)
        critic = Model.create(critic_def,
                              inputs=[critic_key, observations, actions],
                              tx=optax.adam(learning_rate=critic_lr))

        value_def = critic_net.ValueCritic(hidden_dims)
        value = Model.create(value_def,
                             inputs=[critic_key, observations],
                             tx=optax.adam(learning_rate=value_lr))

        target_value = Model.create(value_def,
                                    inputs=[critic_key, observations])

        temp = Model.create(temperature.Temperature(init_temperature),
                            inputs=[temp_key],
                            tx=optax.adam(learning_rate=temp_lr))

        self.actor = actor
        self.critic = critic
        self.value = value
        self.target_value = target_value
        self.temp = temp
        self.rng = rng

        self.step = 1

    def sample_actions(self,
                       observations: np.ndarray,
                       temperature: float = 1.0) -> jnp.ndarray:
        rng, actions = policies.sample_actions(self.rng, self.actor.apply_fn,
                                               self.actor.params, observations,
                                               temperature)
        self.rng = rng

        actions = np.asarray(actions)
        return np.clip(actions, -1, 1)

    def update(self, batch: Batch) -> InfoDict:
        self.step += 1

        new_rng, new_actor, new_critic, new_value, new_target_value, new_temp, info = _update_jit(
            self.rng, self.actor, self.critic, self.value, self.target_value,
            self.temp, batch, self.discount, self.tau, self.target_entropy,
            self.step % self.target_update_period == 0)

        self.rng = new_rng
        self.actor = new_actor
        self.critic = new_critic
        self.value = new_value
        self.target_value = new_target_value
        self.temp = new_temp

        return info

[((20, 1, 20, 62), 'functools.partial', 'functools.partial', (), '', False, 'import functools\n'), ((27, 30, 27, 77), 'jaxrl.agents.sac_v1.critic.update_q', 'update_q', ({(27, 39, 27, 45): 'critic', (27, 47, 27, 59): 'target_value', (27, 61, 27, 66): 'batch', (27, 68, 27, 76): 'discount'}, {}), '(critic, target_value, batch, discount)', False, 'from jaxrl.agents.sac_v1.critic import update_q, update_v\n'), ((29, 15, 29, 36), 'jax.random.split', 'jax.random.split', ({(29, 32, 29, 35): 'rng'}, {}), '(rng)', False, 'import jax\n'), ((30, 28, 30, 77), 'jaxrl.agents.sac.actor.update', 'update_actor', ({(30, 41, 30, 44): 'key', (30, 46, 30, 51): 'actor', (30, 53, 30, 63): 'new_critic', (30, 65, 30, 69): 'temp', (30, 71, 30, 76): 'batch'}, {}), '(key, actor, new_critic, temp, batch)', True, 'from jaxrl.agents.sac.actor import update as update_actor\n'), ((32, 15, 32, 36), 'jax.random.split', 'jax.random.split', ({(32, 32, 32, 35): 'rng'}, {}), '(rng)', False, 'import jax\n'), ((33, 28, 34, 49), 'jaxrl.agents.sac_v1.critic.update_v', 'update_v', ({(33, 37, 33, 40): 'key', (33, 42, 33, 51): 'new_actor', (33, 53, 33, 63): 'new_critic', (33, 65, 33, 70): 'value', (33, 72, 33, 76): 'temp', (34, 37, 34, 42): 'batch', (34, 44, 34, 48): 'True'}, {}), '(key, new_actor, new_critic, value, temp, batch, True)', False, 'from jaxrl.agents.sac_v1.critic import update_q, update_v\n'), ((41, 27, 42, 61), 'jaxrl.agents.sac.temperature.update', 'temperature.update', ({(41, 46, 41, 50): 'temp', (41, 52, 41, 73): "actor_info['entropy']", (42, 46, 42, 60): 'target_entropy'}, {}), "(temp, actor_info['entropy'], target_entropy)", False, 'from jaxrl.agents.sac import temperature\n'), ((37, 27, 37, 70), 'jaxrl.agents.sac.critic.target_update', 'target_update', ({(37, 41, 37, 50): 'new_value', (37, 52, 37, 64): 'target_value', (37, 66, 37, 69): 'tau'}, {}), '(new_value, target_value, tau)', False, 'from jaxrl.agents.sac.critic import target_update\n'), ((83, 14, 83, 38), 'jax.random.PRNGKey', 'jax.random.PRNGKey', ({(83, 33, 83, 37): 'seed'}, {}), '(seed)', False, 'import jax\n'), ((84, 47, 84, 71), 'jax.random.split', 'jax.random.split', ({(84, 64, 84, 67): 'rng', (84, 69, 84, 70): '4'}, {}), '(rng, 4)', False, 'import jax\n'), ((86, 20, 86, 70), 'jaxrl.networks.policies.NormalTanhPolicy', 'policies.NormalTanhPolicy', ({(86, 46, 86, 57): 'hidden_dims', (86, 59, 86, 69): 'action_dim'}, {}), '(hidden_dims, action_dim)', False, 'from jaxrl.networks import critic_net, policies\n'), ((91, 21, 91, 57), 'jaxrl.networks.critic_net.DoubleCritic', 'critic_net.DoubleCritic', ({(91, 45, 91, 56): 'hidden_dims'}, {}), '(hidden_dims)', False, 'from jaxrl.networks import critic_net, policies\n'), ((96, 20, 96, 55), 'jaxrl.networks.critic_net.ValueCritic', 'critic_net.ValueCritic', ({(96, 43, 96, 54): 'hidden_dims'}, {}), '(hidden_dims)', False, 'from jaxrl.networks import critic_net, policies\n'), ((101, 23, 102, 70), 'jaxrl.networks.common.Model.create', 'Model.create', (), '', False, 'from jaxrl.networks.common import InfoDict, Model, PRNGKey\n'), ((120, 23, 122, 59), 'jaxrl.networks.policies.sample_actions', 'policies.sample_actions', ({(120, 47, 120, 55): 'self.rng', (120, 57, 120, 76): 'self.actor.apply_fn', (121, 47, 121, 64): 'self.actor.params', (121, 66, 121, 78): 'observations', (122, 47, 122, 58): 'temperature'}, {}), '(self.rng, self.actor.apply_fn, self.actor.params,\n    observations, temperature)', False, 'from jaxrl.networks import critic_net, policies\n'), ((125, 18, 125, 37), 'numpy.asarray', 'np.asarray', ({(125, 29, 125, 36): 'actions'}, {}), '(actions)', True, 'import numpy as np\n'), ((126, 15, 126, 38), 'numpy.clip', 'np.clip', ({(126, 23, 126, 30): 'actions', (126, 32, 126, 34): '(-1)', (126, 36, 126, 37): '(1)'}, {}), '(actions, -1, 1)', True, 'import numpy as np\n'), ((104, 28, 104, 69), 'jaxrl.agents.sac.temperature.Temperature', 'temperature.Temperature', ({(104, 52, 104, 68): 'init_temperature'}, {}), '(init_temperature)', False, 'from jaxrl.agents.sac import temperature\n'), ((89, 32, 89, 66), 'optax.adam', 'optax.adam', (), '', False, 'import optax\n'), ((94, 33, 94, 68), 'optax.adam', 'optax.adam', (), '', False, 'import optax\n'), ((99, 32, 99, 66), 'optax.adam', 'optax.adam', (), '', False, 'import optax\n'), ((106, 31, 106, 64), 'optax.adam', 'optax.adam', (), '', False, 'import optax\n')]