### A Pluto.jl notebook ###
# v0.20.21

using Markdown
using InteractiveUtils

# ╔═╡ a106ff56-59d8-11f1-35fd-6d8c01cb271e
# ╠═╡ show_logs = false
begin
	using Pkg
	Pkg.activate(mktempdir())
	Pkg.add([
		Pkg.PackageSpec("NumericalDistributions"),
		Pkg.PackageSpec("HadronicLineshapes"),
		Pkg.PackageSpec(; url="https://github.com/JuliaHEP/DistributionsHEP.jl.git", rev="codex/extended-mixture-model"),
		Pkg.PackageSpec("Distributions"),
		Pkg.PackageSpec("Plots"),
	])
	# 
	using NumericalDistributions
	using HadronicLineshapes
	using DistributionsHEP
	using Distributions
	using Plots
end

# ╔═╡ f46bec1c-b892-40c0-a2c4-8a082138c5fa
md"""
## Example 2: 2D Extended Model
"""

# ╔═╡ 3e27d54d-65da-431a-96ea-c797628421e6
theme(:boxed)

# ╔═╡ 193c0c34-6be2-4bb4-971a-d588d46531f9
begin
	# this code should go to NumericalDistributions
	# 
	struct Abs2 <: AbstractFlexFunc
	    F
	end
	(o::Abs2)(x) = abs2(o.F(x))
	 
	struct ofAbs2 <: AbstractFlexFunc
	    F
	end
	(o::ofAbs2)(x) = o.F(x^2)	
end

# ╔═╡ 11ae6a03-5b08-4058-b200-bdf81ac1e434
lb, ub = 0.9, 1.1

# ╔═╡ 06b1917b-c80d-4201-82e1-2cf479f6db52
model = let
    y_phiphi = 10_000
	y_mixed = 1_500
	y_kkkk = 3_000

    signal = DoubleCrystalBall(1.0, 0.01, 1.5, 4.0, 1.5, 4.0)
    background = Chebyshev([1.0,0.5], lb, ub)

    phiphi = product_distribution([signal, signal])
    mixed = MixtureModel(
        [
            product_distribution([signal, background]),
            product_distribution([background, signal]),
        ],
        [0.5, 0.5],
    )
    kkkk = product_distribution([background, background])

    ExtendedMixtureModel([phiphi, mixed, kkkk], [y_phiphi, y_mixed, y_kkkk])	
end

# ╔═╡ b1ac3354-e104-4c3b-843f-df8ff33f1f00
let
	heatmap(
		range(lb, ub, 100)[2:end-1],
		range(lb, ub, 100)[2:end-1],
		(x,y)->model([x,y]))
end

# ╔═╡ aad1df39-309a-436e-a9fd-196709dc13e5
begin
	proj1 = DistributionsHEP.marginalize(model, 1)
	proj2 = DistributionsHEP.marginalize(model, 2)
end;

# ╔═╡ 9b56f491-d709-4a37-b4cd-d11bcd09a593
let
	plot(x->proj1(x), lb, ub, fill=0, fillalpha=0.3)
end

# ╔═╡ Cell order:
# ╟─f46bec1c-b892-40c0-a2c4-8a082138c5fa
# ╠═a106ff56-59d8-11f1-35fd-6d8c01cb271e
# ╠═3e27d54d-65da-431a-96ea-c797628421e6
# ╠═193c0c34-6be2-4bb4-971a-d588d46531f9
# ╠═11ae6a03-5b08-4058-b200-bdf81ac1e434
# ╠═06b1917b-c80d-4201-82e1-2cf479f6db52
# ╠═b1ac3354-e104-4c3b-843f-df8ff33f1f00
# ╠═aad1df39-309a-436e-a9fd-196709dc13e5
# ╠═9b56f491-d709-4a37-b4cd-d11bcd09a593