# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/04_cls.ipynb (unless otherwise specified).

__all__ = ['cls_maker', 'pyhf_cls_maker']

# Cell
import jax
import jax.numpy as jnp
from jax import config
from jax.experimental import stax
import pyhf

# avoid those precision errors!
config.update("jax_enable_x64", True)

pyhf.set_backend(pyhf.tensor.jax_backend())

from .fit import get_solvers
from .transforms import *
from .models import *

# Cell

def cls_maker(nn_model_maker, solver_kwargs):
    '''
    Wraps the construction of the `cls_jax` method.

    Args:
            nn_model_maker: Function that returns a Model object using the nn parameters.

    Returns:
            cls_jax: A callable function that takes the parameters of the observable as argument,
            and returns an expected CLs value from testing the background-only model against the
            nominal signal hypothesis (or whatever the value of 'test_mu' is)
    '''
    @jax.jit
    def cls_jax(nn_params, test_mu):
        g_fitter, c_fitter = get_solvers(nn_model_maker, **solver_kwargs)

        m, bonlypars = nn_model_maker(nn_params)
        exp_data = m.expected_data(bonlypars)
        #print(f'exp_data: {exp_data}')
        bounds = m.config.suggested_bounds()

        # map these
        initval = jnp.asarray([test_mu, 1.0])
        transforms = solver_kwargs.get("pdf_transform", False)
        if transforms:
            initval = to_inf_vec(initval, bounds)

        # the constrained fit

        #print('fitting constrained with init val %s setup %s', initval,[test_mu, nn_params])

        numerator = (
            to_bounded_vec(c_fitter(initval, [test_mu, nn_params]), bounds)
            if transforms
            else c_fitter(initval, [test_mu, nn_params])
        )

        # don't have to fit these -- we know them for expected limits!
        denominator = bonlypars  # to_bounded_vec(g_fitter(initval, [test_mu, nn_params]), bounds) if transforms else g_fitter(initval, [test_mu, nn_params])

        # print(f"constrained fit: {numerator}")
        # print(f"global fit: {denominator}")

        # compute test statistic (lambda(µ))
        profile_likelihood = -2 * (
            m.logpdf(numerator, exp_data)[0] - m.logpdf(denominator, exp_data)[0]
        )

        # in exclusion fit zero out test stat if best fit µ^ is larger than test µ
        muhat = denominator[0]
        sqrtqmu = jnp.sqrt(
            jnp.where(muhat < test_mu, profile_likelihood, 0.0)
        )
        # print(f"sqrt(q(mu)): {sqrtqmu}")
        # compute CLs
        nullval = sqrtqmu
        altval = 0
        CLsb = 1 - pyhf.tensorlib.normal_cdf(nullval)
        CLb = 1 - pyhf.tensorlib.normal_cdf(altval)
        CLs = CLsb / CLb
        return CLs

    return cls_jax

# Cell

def pyhf_cls_maker(nn_model_maker, solver_kwargs):
#     @jax.jit
    def cls_jax(nn_params, test_mu):
        g_fitter, c_fitter = get_solvers(nn_model_maker, **solver_kwargs)

        m, bonlypars = nn_model_maker(nn_params)
        exp_data = m.expected_data(bonlypars)
        #print(f'exp_data: {exp_data}')
#         bounds = m.config.suggested_bounds()[0]

        names = m.config.par_order
        bounds = m.config.suggested_bounds()

        # map these
        initval = jnp.asarray([test_mu, 1.0])
        transforms = solver_kwargs.get("pdf_transform", False)
        if transforms:
            initval = to_inf_vec(initval, bounds)

        # the constrained fit

        #print('fitting constrained with init val %s setup %s', initval,[test_mu, nn_params])

        numerator = (
            to_bounded_vec(c_fitter(initval, [test_mu, nn_params]), bounds)
            if transforms
            else c_fitter(initval, [test_mu, nn_params])
        )

        denominator = bonlypars  # to_bounded_vec(g_fitter(initval, [test_mu, nn_params]), bounds) if transforms else g_fitter(initval, [test_mu, nn_params])

        # print(f"constrained fit: {numerator}")
        # print(f"global fit: {denominator}")

        # compute test statistic (lambda(µ))
        profile_likelihood = -2 * (
            m.logpdf(numerator, exp_data)[0] - m.logpdf(denominator, exp_data)[0]
        )

        # in exclusion fit zero out test stat if best fit µ^ is larger than test µ
        muhat = denominator[0]
        sqrtqmu = jnp.sqrt(
            jnp.where(muhat < test_mu, profile_likelihood, 0.0)
        )
        # print(f"sqrt(q(mu)): {sqrtqmu}")
        # compute CLs
        nullval = sqrtqmu
        altval = 0
        CLsb = 1 - pyhf.tensorlib.normal_cdf(nullval)
        CLb = 1 - pyhf.tensorlib.normal_cdf(altval)
        CLs = CLsb / CLb
        return CLs

    return cls_jax
# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/04_cls.ipynb (unless otherwise specified).

__all__ = ['cls_maker', 'pyhf_cls_maker']

# Cell
import jax
import jax.numpy as jnp
from jax import config
from jax.experimental import stax
import pyhf

# avoid those precision errors!
config.update("jax_enable_x64", True)

pyhf.set_backend(pyhf.tensor.jax_backend())

from .fit import get_solvers
from .transforms import *
from .models import *

# Cell

def cls_maker(nn_model_maker, solver_kwargs):
    '''
    Wraps the construction of the `cls_jax` method.

    Args:
            nn_model_maker: Function that returns a Model object using the nn parameters.

    Returns:
            cls_jax: A callable function that takes the parameters of the observable as argument,
            and returns an expected CLs value from testing the background-only model against the
            nominal signal hypothesis (or whatever the value of 'test_mu' is)
    '''
    @jax.jit
    def cls_jax(nn_params, test_mu):
        g_fitter, c_fitter = get_solvers(nn_model_maker, **solver_kwargs)

        m, bonlypars = nn_model_maker(nn_params)
        exp_data = m.expected_data(bonlypars)
        #print(f'exp_data: {exp_data}')
        bounds = m.config.suggested_bounds()

        # map these
        initval = jnp.asarray([test_mu, 1.0])
        transforms = solver_kwargs.get("pdf_transform", False)
        if transforms:
            initval = to_inf_vec(initval, bounds)

        # the constrained fit

        #print('fitting constrained with init val %s setup %s', initval,[test_mu, nn_params])

        numerator = (
            to_bounded_vec(c_fitter(initval, [test_mu, nn_params]), bounds)
            if transforms
            else c_fitter(initval, [test_mu, nn_params])
        )

        # don't have to fit these -- we know them for expected limits!
        denominator = bonlypars  # to_bounded_vec(g_fitter(initval, [test_mu, nn_params]), bounds) if transforms else g_fitter(initval, [test_mu, nn_params])

        # print(f"constrained fit: {numerator}")
        # print(f"global fit: {denominator}")

        # compute test statistic (lambda(µ))
        profile_likelihood = -2 * (
            m.logpdf(numerator, exp_data)[0] - m.logpdf(denominator, exp_data)[0]
        )

        # in exclusion fit zero out test stat if best fit µ^ is larger than test µ
        muhat = denominator[0]
        sqrtqmu = jnp.sqrt(
            jnp.where(muhat < test_mu, profile_likelihood, 0.0)
        )
        # print(f"sqrt(q(mu)): {sqrtqmu}")
        # compute CLs
        nullval = sqrtqmu
        altval = 0
        CLsb = 1 - pyhf.tensorlib.normal_cdf(nullval)
        CLb = 1 - pyhf.tensorlib.normal_cdf(altval)
        CLs = CLsb / CLb
        return CLs

    return cls_jax

# Cell

def pyhf_cls_maker(nn_model_maker, solver_kwargs):
#     @jax.jit
    def cls_jax(nn_params, test_mu):
        g_fitter, c_fitter = get_solvers(nn_model_maker, **solver_kwargs)

        m, bonlypars = nn_model_maker(nn_params)
        exp_data = m.expected_data(bonlypars)
        #print(f'exp_data: {exp_data}')
#         bounds = m.config.suggested_bounds()[0]

        names = m.config.par_order
        bounds = m.config.suggested_bounds()

        # map these
        initval = jnp.asarray([test_mu, 1.0])
        transforms = solver_kwargs.get("pdf_transform", False)
        if transforms:
            initval = to_inf_vec(initval, bounds)

        # the constrained fit

        #print('fitting constrained with init val %s setup %s', initval,[test_mu, nn_params])

        numerator = (
            to_bounded_vec(c_fitter(initval, [test_mu, nn_params]), bounds)
            if transforms
            else c_fitter(initval, [test_mu, nn_params])
        )

        denominator = bonlypars  # to_bounded_vec(g_fitter(initval, [test_mu, nn_params]), bounds) if transforms else g_fitter(initval, [test_mu, nn_params])

        # print(f"constrained fit: {numerator}")
        # print(f"global fit: {denominator}")

        # compute test statistic (lambda(µ))
        profile_likelihood = -2 * (
            m.logpdf(numerator, exp_data)[0] - m.logpdf(denominator, exp_data)[0]
        )

        # in exclusion fit zero out test stat if best fit µ^ is larger than test µ
        muhat = denominator[0]
        sqrtqmu = jnp.sqrt(
            jnp.where(muhat < test_mu, profile_likelihood, 0.0)
        )
        # print(f"sqrt(q(mu)): {sqrtqmu}")
        # compute CLs
        nullval = sqrtqmu
        altval = 0
        CLsb = 1 - pyhf.tensorlib.normal_cdf(nullval)
        CLb = 1 - pyhf.tensorlib.normal_cdf(altval)
        CLs = CLsb / CLb
        return CLs

    return cls_jax
# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/04_cls.ipynb (unless otherwise specified). # Cell # avoid those precision errors! # Cell Wraps the construction of the `cls_jax` method.      Args:             nn_model_maker: Function that returns a Model object using the nn parameters.      Returns:             cls_jax: A callable function that takes the parameters of the observable as argument,             and returns an expected CLs value from testing the background-only model against the             nominal signal hypothesis (or whatever the value of 'test_mu' is) #print(f'exp_data: {exp_data}') # map these # the constrained fit #print('fitting constrained with init val %s setup %s', initval,[test_mu, nn_params]) # don't have to fit these -- we know them for expected limits! # to_bounded_vec(g_fitter(initval, [test_mu, nn_params]), bounds) if transforms else g_fitter(initval, [test_mu, nn_params]) # print(f"constrained fit: {numerator}") # print(f"global fit: {denominator}") # compute test statistic (lambda(µ)) # in exclusion fit zero out test stat if best fit µ^ is larger than test µ # print(f"sqrt(q(mu)): {sqrtqmu}") # compute CLs # Cell #     @jax.jit #print(f'exp_data: {exp_data}') #         bounds = m.config.suggested_bounds()[0] # map these # the constrained fit #print('fitting constrained with init val %s setup %s', initval,[test_mu, nn_params]) # to_bounded_vec(g_fitter(initval, [test_mu, nn_params]), bounds) if transforms else g_fitter(initval, [test_mu, nn_params]) # print(f"constrained fit: {numerator}") # print(f"global fit: {denominator}") # compute test statistic (lambda(µ)) # in exclusion fit zero out test stat if best fit µ^ is larger than test µ # print(f"sqrt(q(mu)): {sqrtqmu}") # compute CLs