import torch


def bellman(qf, targ_qf, targ_pol, batch, gamma, continuous=True, deterministic=True, sampling=1, reduction='elementwise_mean'):
    """
    Bellman loss.
    Mean Squared Error of left hand side and right hand side of Bellman Equation.

    Parameters
    ----------
    qf : SAVfunction
    targ_qf : SAVfunction
    targ_pol : Pol
    batch : dict of torch.Tensor
    gamma : float
    continuous : bool
        action space is continuous or not
    sampling : int
        Number of sampling in calculating expectation.
    reduction : str
      This argument takes only elementwise, sum, and none.
      Loss shape is pytorch's manner.

    Returns
    -------
    bellman_loss : torch.Tensor
    """
    if continuous:
        obs = batch['obs']
        acs = batch['acs']
        rews = batch['rews']
        next_obs = batch['next_obs']
        dones = batch['dones']

        targ_pol.reset()
        _, _, pd_params = targ_pol(next_obs)
        pd = targ_pol.pd

        next_acs = pd.sample(pd_params, torch.Size([sampling]))
        next_obs = next_obs.expand([sampling] + list(next_obs.size()))
        targ_q, _ = targ_qf(next_obs, next_acs)
        next_q = torch.mean(targ_q, dim=0)

        targ = rews + gamma * next_q * (1 - dones)
        targ = targ.detach()
        q, _ = qf(obs, acs)

        ret = 0.5 * (q - targ)**2
        if reduction != 'none':
            ret = torch.mean(
                ret) if reduction == 'elementwise_mean' else torch.sum(ret)
        return ret
    else:
        raise NotImplementedError(
            "Only Q function with continuous action space is supported now.")