# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.6.1
#     language: julia
#     name: julia-1.6
# ---

using Rocket
using ReactiveMP
using GraphPPL
using Distributions

@model function coin_model()
    
    a = datavar(Float64)
    b = datavar(Float64)
    y = datavar(Float64)
    
    θ ~ Beta(a, b)
    y ~ Bernoulli(θ)
    
    return y, a, b, θ
end

# +
N = 10000   # number of coin tosses
p = 0.5    # p parameter of the Bernoulli distribution

dataset = float.(rand(Bernoulli(p), N));
# -

function inference(data)
    model, (y, a, b, θ) = coin_model()
    
    fe = ScoreActor(Float64)
    θs = keep(Marginal)
    
    fe_sub = subscribe!(score(Float64, BetheFreeEnergy(), model), fe)
    θ_sub = subscribe!(getmarginal(θ), θs)
    
    prior_sub = subscribe!(getmarginal(θ), (posterior) -> begin
        posterior_a, posterior_b = params(posterior)
        update!(a, posterior_a)
        update!(b, posterior_b)
    end)
    
    update!(a, 1.0)
    update!(b, 1.0)
    
    for d in data
        update!(y, d)
    end
    
    unsubscribe!(θ_sub)
    unsubscribe!(prior_sub)
    unsubscribe!(fe_sub)
    
    return getvalues(θs), getvalues(fe)
end

est, fe = inference(dataset);

using Plots

# +
pest = est[1:10:end]
pfe = fe[1:10:end]

p1 = plot(mean.(pest), ribbon = var.(pest), label = "estimate")
p1 = plot!([ p ], seriestype = :hline, label = "real")
 
p2 = plot(pfe, label = "Free-Energy")
p2 = plot!([ -log(p) ], seriestype = :hline, label = "-log(evidence)")

plot(p1, p2, layout = @layout([ a b ]), size = (1000, 400))
# -



StarcoderdataJupyterScriptsDedupFiltered