import jax.numpy as jnp
import numpy as np

def sum_except_batch(x, num_dims=1):
    return x.reshape(*x.shape[:num_dims], -1).sum(-1)

def mean_except_batch(x, num_dims=1):
    return x.reshape(*x.shape[:num_dims], -1).mean(-1)

def bce_w_logits(x, y, weight=None, average=True):
    """
    Binary Cross Entropy Loss
    Should be numerically stable, built based on: https://github.com/pytorch/pytorch/issues/751
    :param x: Input tensor
    :param y: Target tensor
    :param weight: Vector of example weights
    :param average: Boolean to average resulting loss vector
    :return: Scalar value
    """
    max_val = jnp.clip(-x, 0, None)
    loss = x - x * y + max_val + jnp.log(jnp.exp(-max_val) + jnp.exp((-x - max_val)))

    if weight is not None:
        loss = loss * weight
    # print("here", x, y)
    if average:
        return loss.mean()
    else:
        return loss.sum()

def params_count(params):
    _params = []
    def flatten(_params, frozen_dict):
        for k in frozen_dict:
            if type(frozen_dict[k]) == type(frozen_dict):
                flatten(_params, frozen_dict[k])
            else:
                _params.append(frozen_dict[k])
    flatten(_params,params)
    m = 0
    for p in _params:
        if type(p) != type(None):
            m += np.array(p.shape).prod()
    return m
