using LinearAlgebra, Random, Statistics
include("conjugate.jl"); include("model.jl"); include("accounting.jl")

ols(y, cols...) = hcat(ones(length(y)), cols...) \ y

# The usual single-equation MMMs, then the joint model, on the same data
function fit(d)
    n1 = ols(d.sales, d.tv, d.social)
    n2 = ols(d.sales, d.tv, d.social, d.brand)
    n3 = ols(d.sales, d.tv, d.social, d.search)
    n4 = ols(d.sales, d.tv)
    j  = account_mean(posteriors(d))
    [n1[2], n1[3], n2[2], n2[3], n2[4], n3[2], n3[3], n4[2], values(j)...]
end

labels = ["no search: TV", "no search: social", "brand spend: TV", "brand spend: social",
          "brand spend: brand", "search: TV", "search: social", "TV only",
          string.("joint ", keys(truth()))...]
println("truth: ", map(x -> round(x, digits = 3), truth()))
Random.seed!(1)
for sd in (1.0, 0.5)
    R = reduce(hcat, [fit(simulate(156; sd_social = sd)) for _ in 1:2000])
    println("sd_social = $sd")
    for (l, m, s) in zip(labels, mean(R, dims = 2), std(R, dims = 2))
        println("  ", rpad(l, 24), round(m, digits = 2), "  sd ", round(s, digits = 2))
    end
end