#!/usr/bin/env python
# coding: utf-8

# In[1]:


import os
os.environ["ZFIT_DISABLE_TF_WARNINGS"] = "1"

import zfit
from zfit import z
import numpy as np


# In[2]:


print(zfit.pdf.__all__)


# In[3]:


obs = zfit.Space('x', limits=(4800, 6000))

# Creating the parameters for the crystal ball
mu = zfit.Parameter("mu", 5279, 5100, 5300)
sigma = zfit.Parameter("sigma", 20, 0, 50)
a = zfit.Parameter("a", 1, 0, 10)
n = zfit.Parameter("n", 1, 0, 10)

# Single crystal Ball
model_cb = zfit.pdf.CrystalBall(obs=obs, mu=mu, sigma=sigma, alpha=a, n=n)


# In[4]:


# Get the probabilities of some random generated events
probs = model_cb.pdf(x=np.random.random(10))
print(probs)


# In[5]:


# A different normalisation range can be provided as an argument
probs = model_cb.pdf(x=np.random.random(10), norm=(5000, 5250))


# In[6]:


# Calculate the integral between 5000 and 5250 over the PDF normalized
integral_norm = model_cb.integrate(limits=(5000, 5250))
print(f"Integral={integral_norm}")


# In[7]:


# New tail parameters for the second CB
a2 = zfit.Parameter("a2", -1, -10, 0)
n2 = zfit.Parameter("n2", 1, 0, 10)

# New crystal Ball function defined in the same observable range
model_cb2 = zfit.pdf.CrystalBall(obs=obs, mu=mu, sigma=sigma, alpha=a2, n=n2)


# In[8]:


# or via the class API
frac = 0.3  # can also be a Parameter
double_cb_class = zfit.pdf.SumPDF(pdfs=[model_cb, model_cb2], fracs=frac)


# In[9]:


# Defining two Gaussians in two different observables
mu1 = zfit.Parameter("mu1", 1.)
sigma1 = zfit.Parameter("sigma1", 1.)
gauss1 = zfit.pdf.Gauss(obs=obs, mu=mu1, sigma=sigma1)

obs2 = zfit.Space('y', limits=(5, 11))

mu2 = zfit.Parameter("mu2", 1.)
sigma2 = zfit.Parameter("sigma2", 1.)
gauss2 = zfit.pdf.Gauss(obs=obs2, mu=mu2, sigma=sigma2)

# Producing the product of two PDFs
prod_gauss = gauss1 * gauss2
# Or alternatively
prod_gauss_class = zfit.pdf.ProductPDF(pdfs=[gauss2, gauss1])  # notice the different order or the pdf


# In[10]:


print("python syntax product obs", prod_gauss.obs)


# In[11]:


print("class API product obs", prod_gauss_class.obs)


# In[12]:


# Create a parameter for the number of events
yield_gauss = zfit.Parameter("yield_gauss", 100, 0, 1000)

# Extended PDF using a predefined method
extended_gauss = gauss1.create_extended(yield_gauss)


# In[13]:


print(f"Gauss is extended: {gauss1.is_extended}")
gauss1.set_yield(yield_gauss)
print(f"Gauss is extended: {gauss1.is_extended}")


# In[14]:


class MyGauss(zfit.pdf.ZPDF):
    """Simple implementation of a Gaussian similar to zfit.pdf.Gauss class"""
    _N_OBS = 1  # dimension, can be omitted
    _PARAMS = ['mean', 'std']  # name of the parameters

    def _unnormalized_pdf(self, x):
       x = z.unstack_x(x)
       mean = self.params['mean']
       std  = self.params['std']
       return z.exp(- ((x - mean) / std) ** 2)


# In[15]:


obs_own = zfit.Space("my obs", limits=(-4, 4))

mean = zfit.Parameter("mean", 1.)
std  = zfit.Parameter("std", 1.)
my_gauss = MyGauss(obs=obs_own, mean=mean, std=std)


# For instance sampling, integral and probabilities
data     = my_gauss.sample(15)
integral = my_gauss.integrate(limits=(-1, 2))
probs    = my_gauss.pdf(data,norm=(-3, 4))
print(f"Probs: {probs} and integral: {integral}")


# In[16]:


def gauss_integral_from_any_to_any(limits, params, model):
   (lower,), (upper,) = limits.limits
   mean = params['mean']
   std = params['std']
   # Write you integral
   return 42. # Dummy value

# Register the integral
limits = zfit.Space(axes=0, limits=(zfit.Space.ANY_LOWER, zfit.Space.ANY_UPPER))
MyGauss.register_analytic_integral(func=gauss_integral_from_any_to_any, limits=limits)


# In[17]:


# pdf creation
obs = zfit.Space("x", limits=(-10, 10))
mu = zfit.Parameter("mu", 1.0, -4, 6)
sigma = zfit.Parameter("sigma", 1.0, 0.1, 10)
lambd = zfit.Parameter("lambda", -1.0, -5.0, 0)
frac = zfit.Parameter("fraction", 0.5, 0.0, 1.0)

gauss = zfit.pdf.Gauss(mu=mu, sigma=sigma, obs=obs)
exponential = zfit.pdf.Exponential(lambd, obs=obs)

# make exponential pdf cacheable
cached_exponential = zfit.pdf.CachedPDF(exponential)

# create SumPDF with cacheable exponential pdf
sum_pdf = zfit.pdf.SumPDF([gauss, cached_exponential], fracs=frac)