from functools import partial

import jax
import jax.numpy as jnp

from .functional import chunk_encode, cache_grad, unchunk_args


def cache_train_step(loss_fn, state, ss, tt, axis='device'):
    def encode_with_params(params, **kwargs):
        return state.apply_fn(params=params, **kwargs)

    encode_fn = chunk_encode(partial(encode_with_params, state.params))
    grad_fn = cache_grad(encode_with_params)

    s_reps = encode_fn(**ss)
    t_reps = encode_fn(**tt)

    @unchunk_args(axis=0, argnums=(0, 1))
    def grad_cache_fn(xx, yy):
        return jnp.mean(loss_fn(xx, yy, axis=axis))
    loss, (s_grads, t_grads) = jax.value_and_grad(grad_cache_fn, argnums=(0, 1))(s_reps, t_reps)

    grads = jax.tree_map(lambda v: jnp.zeros_like(v), state.params)
    grads = grad_fn(state.params, grads, s_grads, **ss)
    grads = grad_fn(state.params, grads, t_grads, **tt)

    loss, grads = jax.lax.pmean([loss, grads], axis)
    new_state = state.apply_gradients(grads=grads)
    return loss, new_state