### A Pluto.jl notebook ###
# v0.20.5

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    #! format: off
    return quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
    #! format: on
end

# ╔═╡ 3eabffcc-6c90-11f0-379c-c77420b3c6e6
begin
    using Pkg
    Pkg.add([
        PackageSpec(name="Distributions"),
        PackageSpec(name="Plots"),
        PackageSpec(url="https://github.com/mmikhasenko/NumericalDistributions.jl"),
        PackageSpec(url="https://github.com/RUB-EP1/HighEnergyTools.jl"),
        PackageSpec(url="https://github.com/JuliaHEP/DistributionsHEP.jl")
    ])
    using Distributions
    using NumericalDistributions
    using HighEnergyTools
    using DistributionsHEP
    using Plots
	using Random
	using ComponentArrays
    using PlutoUI
    Random.seed!(42)
end

# ╔═╡ d694faeb-048a-4742-8e29-c74506967e24
TableOfContents()

# ╔═╡ 85fd0ac9-a0be-407f-9a86-ad8ce82d34d7
md"""
# New tools for analysis in high-energy physics with Julia
"""

# ╔═╡ d4b1b450-c4b6-4bed-92d9-c17ebe028477
md"""
## 0. Setup
"""

# ╔═╡ 47f06fcd-9c05-45ca-aa67-624b7bc1f38e
begin
	support=[0.0,5.0]
	nbins=50
	step=(support[2] - support[1]) / nbins
	bins=support[1]:step:support[2]
	nevents=10000
end

# ╔═╡ 0055fd61-7508-499e-a8ee-250ec89c68b6
md"""
## 1. NumericalDistributions.jl
"""

# ╔═╡ eb75ce49-2497-4f3c-98d7-5c2927dbeafe
let
    f(x) = 0.5*x^2 * exp(-x*2)
    dist_num = NumericallyIntegrable(f, (support[1], support[2]))

    samples = rand(dist_num, nevents)

    stephist(samples, bins=bins, normalize=true, lw=2, label="Sampled")
    plot!(x -> pdf(dist_num, x), bins, lw=2, label="PDF")
end

# ╔═╡ d4af90d2-d342-4575-ba35-6a780f70c1c8
md"""
## 2. DistributionsHEP.jl
"""

# ╔═╡ bd679832-aa9f-4406-9cc1-66044ab73522
begin
    cb = CrystalBall(0.0, 1.0, 1.5, 2.0)
	cb_data = rand(cb, nevents)

	stephist(cb_data, bins=-support[2]:step:support[2], normalize=true, lw=2, label="Crystal Ball Samples")
    plot!(x -> pdf(cb, x), -support[2]:step:support[2], lw=2, label="Crystal Ball PDF")
end

# ╔═╡ 8dee29b6-c0bd-4888-848c-def321175d7f
md"""
## 3. HighEnergyTools.jl
"""

# ╔═╡ 3178cc13-380b-458b-bb88-3cf6f56c807e
md"""
### 3.1. Sampling
"""

# ╔═╡ 2f602f81-2e87-4744-a5a2-86b56e984f4c
let
	f(x) = 0.5*x^2 * exp(-x*2)
    dist_num = NumericallyIntegrable(f, (support[1], support[2]))
	
	samples_inv = sample_inversion(f, nevents, (support[1], support[2]))
	samples_rej = sample_rejection(f, nevents, (support[1], support[2]))
	
    stephist(samples_inv, bins=bins, normalize=true, lw=2, label="Sampled (Inv)")
	stephist!(samples_rej, bins=bins, normalize=true, lw=2, label="Sampled (Rej)")
    plot!(x -> pdf(dist_num, x), bins, lw=2, label="PDF")
end

# ╔═╡ f2877ff3-6f27-42b1-8af8-c75c57ea61ba
md"""
### 3.2. Models
"""

# ╔═╡ d02b2990-89a7-4057-8fb5-ec3c5cc66bb8
md"""

Anka Model parameters:

``μ_{1}`` = $(@bind μ1 Slider(0.0:0.1:5.0, default=2.2, show_value=true))

``σ_{1}`` = $(@bind σ1 Slider(0.0:0.01:1.0, default=0.35, show_value=true))

``bg^{(1)}`` = $(@bind bg1 Slider(0.0:0.1:5.0, default=1.5, show_value=true))

``bg^{(2)}`` = $(@bind bg2 Slider(0.0:0.1:5.0, default=1.1, show_value=true))

``logfB`` = $(@bind logfB Slider(-5.0:0.1:5.0, default=0.0, show_value=true))

Frida Model parameters:

``μ_{2}`` = $(@bind μ2 Slider(0.0:0.1:5.0, default=3.1, show_value=true))

``σ_{2}`` = $(@bind σ2 Slider(0.0:0.01:1.0, default=0.20, show_value=true))

``logfS1`` = $(@bind logfS1 Slider(-5.0:0.1:5.0, default=-1.0, show_value=true))

``logfS2`` = $(@bind logfS2 Slider(-5.0:0.1:5.0, default=-2.0, show_value=true))

"""

# ╔═╡ 69ecd1b0-fe4c-477c-8585-08abf3112e10
begin
    anka = Anka(support[1], support[2])
	anka_pars = (
		sig = (μ = μ1, σ = σ1), 
		bgd = (coeffs = [bg1, bg2],), 
		logfB = logfB)
	anka_model = build_model(anka, anka_pars)
	anka_data = rand(anka_model, nevents)

	stephist(anka_data, bins=bins, normalize=true, lw=2, label="Anka Samples")
    plot!(x -> pdf(anka_model, x), bins, lw=2, label="Anka PDF")
end

# ╔═╡ 126a9dc4-38d8-4bd8-bf94-2ee56e37bb86
begin
    frida = Frida(support[1], support[2])
	frida_pars = (
		sig1 = (μ = μ1, σ = σ2),
    	sig2 = (μ = μ2, σ = σ2),
    	bgd = (coeffs = [bg1, bg2],),
    	logfS1 = logfS1,
    	logfS2 = logfS2,
	)
	frida_model = build_model(frida, frida_pars)
	frida_data = rand(frida_model, nevents)

	stephist(frida_data, bins=bins, normalize=true, lw=2, label="Frida Samples")
    plot!(x -> pdf(frida_model, x), bins, lw=2, label="Frida PDF")
end

# ╔═╡ 1dc83e0b-7848-4af4-afcb-3caa4bc08132
md"""
### 3.3. Fitting
"""

# ╔═╡ 7350c897-4780-4402-8108-0472ae13734b
fit_result = let
	init_pars = ComponentArray(sig = (μ = 2.2, σ = 0.06), bgd = (coeffs = [1.5, 1.1],), logfB = 0.0)
	fit_res = fit_nll(anka_data, init_pars) do p
    	build_model(anka, p)
	end
	fit_res
end

# ╔═╡ 1548356a-148e-44d3-a7d2-bc4b975982f2
best_pars = fit_result.minimizer

# ╔═╡ bcee59cb-098e-47fb-ba14-9ee0893bd7c2
let
	best_model = build_model(anka, best_pars)
	stephist(anka_data, bins=bins, normalize=true, lw=2, label="Anka Samples")
    plot!(x -> pdf(anka_model, x), bins, lw=2, label="Fit PDF")
end

# ╔═╡ f5dee3b0-7b17-4cea-8084-85482b6eaf2c
md"""
### 3.4. sPlot in Julia
"""

# ╔═╡ 1b317d94-8e19-4ef6-a56c-2986a7551a68
let
	sP = sPlot(anka_model)
	sW = sWeights(sP, anka_data)
	wS, wB = sW |> eachcol
	fS(x) = sWeights(sP, [x])[1,1]
	fB(x) = sWeights(sP, [x])[1,2]
	plot(bins, x->fS(x); lab="Signal", lw=2)
	plot!(bins, x->fB(x); lab="Background", lw=2)
end

# ╔═╡ bb8ee67f-14ff-4e47-8680-d74a9dc5c9d3
let
	sP = sPlot(frida_model)
	sW = sWeights(sP, frida_data)
	wS, wB = sW |> eachcol
	fS1(x) = sWeights(sP, [x])[1,1]
	fS2(x) = sWeights(sP, [x])[1,2]
	fB(x) = sWeights(sP, [x])[1,3]
	plot(bins, x->fS1(x); lab="Signal 1", lw=2)
	plot!(bins, x->fS2(x); lab="Signal 2", lw=2)
	plot!(bins, x->fB(x); lab="Background", lw=2)
end

# ╔═╡ Cell order:
# ╟─d694faeb-048a-4742-8e29-c74506967e24
# ╟─85fd0ac9-a0be-407f-9a86-ad8ce82d34d7
# ╟─d4b1b450-c4b6-4bed-92d9-c17ebe028477
# ╠═3eabffcc-6c90-11f0-379c-c77420b3c6e6
# ╠═47f06fcd-9c05-45ca-aa67-624b7bc1f38e
# ╟─0055fd61-7508-499e-a8ee-250ec89c68b6
# ╠═eb75ce49-2497-4f3c-98d7-5c2927dbeafe
# ╟─d4af90d2-d342-4575-ba35-6a780f70c1c8
# ╠═bd679832-aa9f-4406-9cc1-66044ab73522
# ╟─8dee29b6-c0bd-4888-848c-def321175d7f
# ╟─3178cc13-380b-458b-bb88-3cf6f56c807e
# ╠═2f602f81-2e87-4744-a5a2-86b56e984f4c
# ╟─f2877ff3-6f27-42b1-8af8-c75c57ea61ba
# ╟─d02b2990-89a7-4057-8fb5-ec3c5cc66bb8
# ╠═69ecd1b0-fe4c-477c-8585-08abf3112e10
# ╠═126a9dc4-38d8-4bd8-bf94-2ee56e37bb86
# ╟─1dc83e0b-7848-4af4-afcb-3caa4bc08132
# ╠═7350c897-4780-4402-8108-0472ae13734b
# ╠═1548356a-148e-44d3-a7d2-bc4b975982f2
# ╠═bcee59cb-098e-47fb-ba14-9ee0893bd7c2
# ╟─f5dee3b0-7b17-4cea-8084-85482b6eaf2c
# ╠═1b317d94-8e19-4ef6-a56c-2986a7551a68
# ╠═bb8ee67f-14ff-4e47-8680-d74a9dc5c9d3