import tensorflow as tf

import tensorflow_probability as tfp
from normalizing_flows.flows import Transform
from . import Parameterize

def gaussianize(x, mus, log_sigmas, inverse=tf.constant(False)):
    if inverse:
        z = tf.math.exp(log_sigmas)*x + mus
        ldj = tf.math.reduce_sum(log_sigmas, axis=[1,2,3])
    else:
        z = (x - mus)*tf.math.exp(-log_sigmas)
        ldj = -tf.math.reduce_sum(log_sigmas, axis=[1,2,3])
    return z, ldj
