# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.7.1
#     language: julia
#     name: julia-1.7
# ---

# # Homework 2 Exercise 3
# Now suppose the causal association between age and weight might be dif-
# ferent for boys and girls. Use a single linear regression, with a categorical
# variable for sex, to estimate the total causal effect of age on weight separately
# for boys and girls. How do girls and boys differ? Provide one or more pos-
# terior contrasts as a summary.

using StatsPlots
using StatisticalRethinking: sr_datadir, PI
using StatisticalRethinkingCommon
using Statistics, StatsBase
import CSV
using DataFrames
using Gen, Distributions

@time data = CSV.read(sr_datadir("Howell1.csv"), DataFrame)

children = data[data.age .< 13 , :]
boys = children[children.male .== 1, :]
girls = children[children.male .== 0, :];

@df boys plot(:age, :weight, seriestype=:scatter, title="Age/Weight (Children)", xlabel="Age (years)", ylabel="weight(kg)", labels="boys", legend=:bottomright)
@df girls plot!(:age, :weight, seriestype=:scatter, labels="girls")

dt = standardize_column!(children, :age, scale=false)
StatsBase.transform!(dt,boys.age)
StatsBase.transform!(dt,girls.age)
dt

@gen function children_age_weight_model(ages, males)
    a0 ~ normal(12., 7.5)
    b0 ~ gamma(0.25, 4.)
    noise0 ~ gamma(1., 1.)
    a1 ~ normal(12., 7.5)
    b1 ~ gamma(0.25, 4.)
    noise1 ~ gamma(1., 1.)
    a = [a0 a1]
    b = [b0 b1]
    noise = [noise0 noise1]
    
    function f(age,male)
        return a[male+1] + b[male+1] * age
    end
    
    for (i,age) in enumerate(ages)
        male = Int(males[i])
        {(:y, i)} ~ normal(f(age, male), noise[male+1])
    end
    return f
end

@df boys plot(:age, :weight, seriestype=:scatter, title="Age/Weight (Children)", xlabel="Age (years)", ylabel="weight(kg)", labels="boys", legend=:bottomright)
@df girls plot!(:age, :weight, seriestype=:scatter, labels="girls")
priors = [children_age_weight_model((), ()) for _ in 1:15]
test_xs = Vector(range(minimum(children.age)-1, maximum(children.age)+1, length=1000))
plot!(test_xs, [f.(test_xs, zeros(Int8,1000)) for f in priors[1:7]], labels=nothing)
plot!(test_xs, [f.(test_xs, ones(Int8,1000)) for f in priors[8:end]], labels=nothing)
plot!(ylim=(minimum(children.weight) - 10, maximum(children.weight) + 10))

# +
observations = Gen.choicemap()
for (i,weight) in enumerate(children.weight)
    observations[(:y, i)] = weight
end

@time traces, = mcmc(children_age_weight_model, (children.age, children.male), observations, warmup=1_500, steps=15_000);

# +
params = (:a0, :a1, :b0, :b1, :noise0, :noise1)

@time inferred,logweight = infer(traces, children_age_weight_model, (children.age,children.male), params)
a0 = inferred[:a0]
b0 = inferred[:b0]
noise0 = inferred[:noise0]
a1 = inferred[:a1]
b1 = inferred[:b1]
noise1 = inferred[:noise1]
@info "Parameters:" a0 b0 noise0 a1 b1 noise1
# -

@df boys plot(:age, :weight, seriestype=:scatter, title="Age/Weight (Children)", xlabel="Age (years)", ylabel="weight(kg)", labels="boys", legend=:bottomright)
@df girls plot!(:age, :weight, seriestype=:scatter, labels="girls")
plot!(test_xs, (@. a0 + b0 * test_xs), linewidth=2., ribbon=noise0, labels="posterior (girls)")
plot!(test_xs, (@. a1 + b1 * test_xs), linewidth=2., ribbon=noise1, labels="posterior (boys)")

@df boys plot(:age, :weight, seriestype=:scatter, title="Age/Weight (Children)", xlabel="Age (years)", ylabel="weight(kg)", labels="boys", legend=:bottomright)
@df girls plot!(:age, :weight, seriestype=:scatter, labels="girls")
plot!(test_xs, (@. a0 + b0 * test_xs), labels="posterior (girls)", linewidth=2.)
plot!(test_xs, (@. a1 + b1 * test_xs), labels="posterior (boys)", linewidth=2.)
test_predictions_girls = @. rand(Normal(a0 + b0 * girls.age, noise0))
test_predictions_boys = @. rand(Normal(a1 + b1 * girls.age, noise1))
plot!(children.age, test_predictions_girls, seriestype=:scatter, alpha=1/2, label="model (girls)")
plot!(children.age, test_predictions_boys, seriestype=:scatter, alpha=1/2, label="model (boys)")



StarcoderdataJupyterScriptsDedupFiltered