#only run this once
include("RateDistortionDecisionMaking.jl")

using RateDistortionDecisionMaking, Distances, DataFrames, Colors, Gadfly, Distributions, Interact, Reactive

#make the default plot size a bit larger
set_default_plot_size(15cm, 12cm)

#set up medical diagnosis system example
include("MedicalSystemExample.jl")
w_values, w_strings, a_values, a_strings, p_w_uni, U = setup_medical_example()


numa = length(a_strings)
a_vec = collect(1:numa)

numw = length(w_strings)
w_vec = collect(1:numw)

#pre-compute utility
U_pre, Umax = setuputilityarrays(a_values,w_values,U)

#visualize utility

plt_utility = visualizeMatrix(U_pre, w_values, a_values, w_strings, a_strings, xlabel="Disease type w",
                              ylabel="Treatment a", legendlabel="U(a,w)")


#nonuniform p_w
p_w_nonuni = medical_nonuniform_pw(numw)

plt_pw_uni = visualizeBAmarginal(p_w_uni, w_vec, w_strings, alabel="Disease type w", legendlabel="p1(w)")
plt_pw_nonuni = visualizeBAmarginal(p_w_nonuni, w_vec, w_strings, alabel="Disease type w", legendlabel="p2(w)")

display(hstack(plt_pw_uni, plt_pw_nonuni))

#inverse temperatures
#β1: model selection -> inv. price for automated diagnosis (I(X;W))
#β2=∞: high-level decision -> inv. price for I(A;X)
#β3: low-level decision -> inv. price for manual diagnosis by specialist (I(A;W|X))

#parameters for case 1 - hierarchical, uniform environment
β1 = 1#2.5
β2 = Inf  
β3 = 0.9#1
pw = p_w_uni
#cardinality of x - number of different kinds of specialists
numx = 3


#paremters for case 2 - hierarchical, nonuniform environment
β1_n = 1#2.5
β2_n = Inf  
β3_n = 0.9#1
pw_n = p_w_nonuni
#cardinality of x - number of different kinds of specialists
numx_n = 3

#paremters for case 2 - sequential, uniform environment
#β1_n = 10
#β2_n = 10  
#β3_n = 0
#pw_n = p_w_uni


#parameters for iteration of self-consistent equations
ε = 0.0001 #convergence critetion for BAiterations
maxiter = 10000 #maximum number of BA iterations

x_vec = collect(1:numx)
x_strings = map((x)->"x"*string(x), x_vec)


#This function performs Blahut-Arimoto iterations for the three-variable general case
px, pa, pxgw, pagx, pagxw, pagw, performance_df = threevarBAiterations(numx, β1, β2, β3, U_pre, pw, ε, maxiter,
                                                  compute_performance=true, performance_as_dataframe=true,
                                                  init_pogw_sparse = true, init_pogw_uniformly = false,
                                                  init_pagow_uniformly = true);

#call the routine that creates plots for all the probability distributions
plt_px, plt_pa, plt_pxgw, plt_pagx, plt_pagw, dpdown, plt_pagxw_vis = visualize_three_var_BAsolution(px, pa,
                                                                      pxgw, pagx, pagxw, pagw,
                                                                      x_vec, a_vec, w_vec,
                                                                      x_strings, a_strings, w_strings, 
                                                                      olabel_string="x", alabel_string="a", wlabel_string="w")
; #suppress output

x_vec_n = collect(1:numx_n)
x_strings_n = map((x)->"x"*string(x), x_vec_n)

#This function performs Blahut-Arimoto iterations for the three-variable general case
px_n, pa_n, pxgw_n, pagx_n, pagxw_n, pagw_n, performance_df_n = threevarBAiterations(numx_n, β1_n, β2_n, β3_n, U_pre, pw_n, ε, maxiter,
                                                                                     compute_performance=true, performance_as_dataframe=true,
                                                                                     init_pogw_sparse = true, init_pogw_uniformly = false,
                                                                                     init_pagow_uniformly = true);


#call the routine that creates plots for all the probability distributions
plt_px_n, plt_pa_n, plt_pxgw_n, plt_pagx_n, plt_pagw_n, dpdown_n, plt_pagxw_vis_n = visualize_three_var_BAsolution(px_n, pa_n,
                                                                                    pxgw_n, pagx_n, pagxw_n, pagw_n,
                                                                                    x_vec_n, a_vec, w_vec,
                                                                                    x_strings_n, a_strings, w_strings, 
                                                                                    olabel_string="x", alabel_string="a",
                                                                                    wlabel_string="w")

; #suppress output

#visualize p(x|w)
display(hstack(plt_pxgw,plt_pxgw_n))
#visualize p(a|x)
display(hstack(plt_pagx,plt_pagx_n))
#visualize p(a|w)
display(hstack(plt_pagw,plt_pagw_n))

#visualize p(a|x,w)
display(dpdown)
display(plt_pagxw_vis)

#visualize p(a|x,w)
display(dpdown_n)
display(plt_pagxw_vis_n)

p_MI, p_composed, p_perf = plot_three_var_performancemeasures(performance_df, maximum(U_pre), β1, β2, β3)
p_MI_n, p_composed_n, p_perf_n = plot_three_var_performancemeasures(performance_df_n, maximum(U_pre), β1_n, β2_n, β3_n)

display(hstack(p_MI, p_MI_n))
display(hstack(p_composed, p_composed_n))
display(hstack(p_perf, p_perf_n))

#for the paper visualize p(w) as a histogram (bar-plot)
plt_pw_uni_bar = plot(x=w_strings,y=p_w_uni,Geom.bar, Scale.x_discrete, Scale.y_continuous(minvalue=0, maxvalue=1), 
Guide.xlabel("Disease type w", orientation=:horizontal),
Guide.ylabel("p1(w)",orientation=:vertical),BAtheme(default_color=colorant"darkred"))

plt_pw_nonuni_bar = plot(x=w_strings,y=p_w_nonuni,Geom.bar, Scale.x_discrete, Scale.y_continuous(minvalue=0, maxvalue=1), 
Guide.xlabel("Disease type w", orientation=:horizontal),
Guide.ylabel("p2(w)",orientation=:vertical),BAtheme(default_color=colorant"darkred"));

#generate plots for comparison

plt_setup = hstack(plt_utility, vstack(plt_pw_uni_bar, plt_pw_nonuni_bar))

plt_final = vstack(hstack(plt_pxgw, plt_pagx, plt_pagw), hstack(plt_pxgw_n, plt_pagx_n, plt_pagw_n))

display(plt_setup)
display(plt_final)



w = 18cm
h = 12cm
#draw(SVG("Figures/MedicalExampleSetup.svg", w, h), plt_setup) #uncomment to store figure
#draw(SVG("Figures/HierarchicalComparison.svg", w, h), plt_final) #uncomment to store figure#