using Random
using Statistics
using Distributions
using Turing
using StatsPlots

gr() # graphics backend

# --------------------------------------------------
# Bayesian linear regression with Turing.jl
# Model: y = a*x + b + noise
# --------------------------------------------------

Random.seed!(1234)
Turing.setprogress!(false)

# -----------------------------
# 1. Create synthetic data
# -----------------------------
N = 30
x = collect(range(0, 10; length=N))

a_true = 2.0
b_true = -1.0
σ_true = 1.0

y = a_true .* x .+ b_true .+ rand(Normal(0, σ_true), N)

# -----------------------------
# 2. Define Bayesian model
# -----------------------------
@model function bayes_linreg(x, y)
    # Priors (very liberal, large sigma)
    a ~ Normal(0, 5)
    b ~ Normal(0, 5)
    σ ~ truncated(Cauchy(0, 2); lower=0) # Cauchy to allow σ to be large, too

    # Likelihood
    for i in eachindex(x)
        μ = a * x[i] + b
        y[i] ~ Normal(μ, σ)
    end
end

model = bayes_linreg(x, y)

#=
# Or, scale data and use simple priors
# but you need to descale output back to original units (not done here)
@model function bayes_linreg_2(x, y)
    μx, σx = mean(x), std(x)
    μy, σy = mean(y), std(y)

    a ~ Normal(0, 1)
    b ~ Normal(0, 1)
    σ ~ truncated(Cauchy(0, 1); lower=0)

    # Likelihood
    for i in eachindex(x)
        x_std = (x[i] - μx) / σx
        μ_std = a * x_std + b
        μ = μy + σy * μ_std
        y[i] ~ Normal(μ, σ * σy)
    end
end
=#


# -----------------------------
# 3. Sample posterior
# -----------------------------
chain = sample(model, NUTS(0.65), 2000)

println(chain)

# -----------------------------
# 4. Extract posterior samples
# -----------------------------
a_samps = vec(Array(chain[:a]))
b_samps = vec(Array(chain[:b]))
σ_samps = vec(Array(chain[:σ]))

println("\nPosterior means:")
println("a ≈ ", mean(a_samps))
println("b ≈ ", mean(b_samps))
println("σ ≈ ", mean(σ_samps))

# -----------------------------
# 5. Plot chains / marginals
# -----------------------------
p_chain = plot(chain)
display(p_chain)
readline()

# -----------------------------
# 6. Plot data and posterior lines
# -----------------------------
p_data = scatter(
    x, y;
    label="data",
    xlabel="x",
    ylabel="y",
    title="Bayesian linear regression",
    legend=:topleft
)



# draw a handful of posterior lines
nsamp_lines = 10
inds = rand(1:length(a_samps), nsamp_lines)
for j in inds
    plot!(p_data, x, a_samps[j] .* x .+ b_samps[j];
          alpha=0.45, label=false)
end
# true line
plot!(p_data, x, a_true .* x .+ b_true; lw=3, label="true line")

# posterior mean line
a_mean = mean(a_samps)
b_mean = mean(b_samps)
plot!(p_data, x, a_mean .* x .+ b_mean; lw=3, label="posterior mean (surrounded by a few posterior lines)")

# -----------------------------
# 7. Posterior predictive band
# -----------------------------
xgrid = collect(range(minimum(x), maximum(x); length=100))
pred_mean = zeros(length(xgrid))
pred_lo = zeros(length(xgrid))
pred_hi = zeros(length(xgrid))

for (k, xx) in enumerate(xgrid)
    ys = a_samps .* xx .+ b_samps
    pred_mean[k] = mean(ys)
    pred_lo[k] = quantile(ys, 0.05)
    pred_hi[k] = quantile(ys, 0.95)
end

p_band = scatter(
    x, y;
    label="data",
    xlabel="x",
    ylabel="y",
    title="Posterior mean and 90% credible band",
    legend=:topleft
)

plot!(p_band, xgrid, pred_mean; lw=3, lc="green", label="posterior mean")
plot!(
    p_band,
    xgrid, pred_hi;
    fillrange=pred_lo,
    fillalpha=0.45,
    fillcolor=:gray,
    linecolor=:transparent,
    label="95% credible band"
)

#plot!(p_band, xgrid, pred_lo; lw=2, ls=:dash, label="5%")
#plot!(p_band, xgrid, pred_hi; lw=2, ls=:dash, label="95%")
plot!(p_band, x, a_true .* x .+ b_true; lw=3, label="true line")

p_all = plot(p_data, p_band; layout=(2,1), size=(900,1000))
display(p_all)

readline()