from typing import Tuple
import flax.linen as nn
from flax.core.scope import VariableDict as Params
from jax import Array
import jax.numpy as jnp
import optax
import rlax

from helx_base.mdp import TERMINATION, Timestep


def flatten_timesteps(timesteps: Timestep, discount: float = 1.0) -> Tuple[Array, Array, Array, Array, Array, Array]:
    s_tm1 = timesteps.observation[:-1]
    s_t = timesteps.observation[1:]
    a_tm1 = timesteps.action[:-1][0]  # [0] because scalar
    r_t = timesteps.reward[:-1][0]  # [0] because scalar
    terminal_tm1 = timesteps.step_type[:-1] != TERMINATION
    discount_t = discount ** timesteps.t[:-1][0]  # [0] because scalar
    return s_tm1, s_t, a_tm1, r_t, terminal_tm1, discount_t

def dqn_loss(
    timesteps: Timestep,
    critic: nn.Module,
    params: Params,
    params_target: Params,
    discount: float = 0.99,
) -> Array:
    s_tm1, s_t, a_tm1, r_t, terminal_tm1, discount_t = flatten_timesteps(timesteps, discount)

    q_tm1 = jnp.asarray(critic.apply(params, s_tm1))
    q_t = jnp.asarray(critic.apply(params_target, s_t)) * terminal_tm1

    td_error = rlax.q_learning(
        q_tm1, a_tm1, r_t, discount_t, q_t, stop_target_gradients=True
    )
    td_error = jnp.asarray(td_error)
    td_loss = jnp.mean(0.5 * td_error**2)
    return td_loss


def double_dqn_loss(
    timesteps: Timestep,
    critic: nn.Module,
    params: Params,
    params_target: Params,
    discount: float = 0.99,
):
    s_tm1, s_t, a_tm1, r_t, terminal_tm1, discount_t = flatten_timesteps(timesteps, discount)

    q_tm1 = jnp.asarray(critic.apply(params, s_tm1))
    a_t = jnp.argmax(jnp.asarray(critic.apply(params, s_t)) * terminal_tm1)
    q_t = critic.apply(params_target, s_t)
    q_target = r_t + discount_t * q_t[a_t]

    td_loss = optax.l2_loss(q_tm1[a_tm1] - q_target)
    return jnp.asarray(td_loss)


def soft_q_loss(
    timesteps: Timestep,
    critic: nn.Module,
    params: Params,
    params_target: Params,
    entropy: Array,
    discount: float = 0.99,
):
    s_tm1, s_t, a_tm1, r_t, terminal_tm1, discount_t = flatten_timesteps(timesteps, discount)


    raise NotImplementedError()
