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

using RateDistortionDecisionMaking
using Gadfly, Distances

#create n discrete actions a
a = collect(1:4)
na = length(a)
a_strings = ["a$aval" for aval in a] #string representation for plotting


#--------- choose either one of the p0s (by commenting/uncommenting) ---------
#uniform p0(a)
p0 = ones(na)/na

#skewed p0(a)
#p0 = ones(na)
#p0[1]=2
#p0 = p0/sum(p0)
#---------------------


#U(a,o) - utility with two peaks that are almost identical in value
U = ones(na)
U[1]=3
U[2]=0.02
U[end]=3.1


#create (bar) plots for utility and prior
plt_u = plot(x=a_strings,y=U,Geom.bar, Scale.x_discrete, 
             Guide.xlabel("a", orientation=:horizontal),Guide.ylabel("U(w,a)",orientation=:vertical),
             Guide.title("Utility"),BAtheme())

plt_p0 = plot(x=a_strings,y=p0,Geom.bar, Scale.x_discrete, Scale.y_continuous(minvalue=0, maxvalue=1), 
              Guide.xlabel("a", orientation=:horizontal),Guide.ylabel("p0(a)",orientation=:vertical),
              Guide.title("Prior"),BAtheme())

#show both plots stacked horizontally
hstack(plt_u,plt_p0)

#compute the posterior p(a|w), using the inverse temperature β as a parameter
#using a slider to interactively change the inverse temperature
using Reactive, Interact

println("Drag the slider and see how the inverse temperature influences the posterior.")
sl_β = slider(0:0.1:50, label="inverse temp. β")    
display(sl_β) #displays the slider at the output of the current cell

#tie the slider with the plotting routine (this will also show the figure, since it's the last command in this cell)
plt_pa = lift(β->begin
                    pagw = boltzmanndist(p0,β,U) #compute posterior
                    plt_pagw = plot(x=a_strings,y=pagw,Geom.bar, Scale.x_discrete, Scale.y_continuous(minvalue=0, maxvalue=1), 
                                    Guide.xlabel("a", orientation=:horizontal),Guide.ylabel("p*(a|w)",orientation=:vertical),
                                    Guide.title("Posterior, β=$β"),BAtheme())
               end,sl_β)


#Sweep through β values and compute expected utility and KL-divergence

β_sweep = collect(0:0.1:50)
nβ = length(β_sweep)

EU = zeros(nβ)
DKL = zeros(nβ)

#inverse temp. sweep
for i=1:nβ
    #compute posterior p(a|w) using the current β
    pagw_s = boltzmanndist(p0,β_sweep[i],U)
    #compute expected utility = ∑ p(a)U(a)
    EUvec = (pagw_s') * U  #this will create a Float vector with only one element 
                           #(which is still different from a scalar in Julia)
    EU[i] = EUvec[1]  #this will actually extract the Float scalar
    DKL[i] = kl_divergence(pagw_s,p0)/log(2) #divide by log(2) for bits
end

#plot expected utility and KL divergence as a function of β
plt_EU = plot(x=β_sweep,y=EU,Geom.line,
              Guide.xlabel("β", orientation=:horizontal),Guide.ylabel("E[U]",orientation=:vertical),
              Guide.title("Expected utility"),BAtheme())

plt_DKL = plot(x=β_sweep,y=DKL,Geom.line,
               Guide.xlabel("β", orientation=:horizontal),Guide.ylabel("D<sub>KL</sub>(p||p0) [bits]",
               orientation=:vertical),Guide.title("KL-divergence"),BAtheme())

hstack(plt_EU,plt_DKL)

#for final figure: compute two different posteriors (low and high inverse temp.)

β_lo = 1  #inv. temp. for posterior with low computational resources (will stick close to prior)
β_hi = 100   #inv. temp. posterior with high computational resources (can deviate a lot from prior)

#compute posteriors
pagw_lo = boltzmanndist(p0,β_lo,U)
pagw_hi = boltzmanndist(p0,β_hi,U)


#create plots
plt_pagw_lo = plot(x=a_strings,y=pagw_lo,Geom.bar, Scale.x_discrete, Scale.y_continuous(minvalue=0, maxvalue=1), 
                   Guide.xlabel("a", orientation=:horizontal),
                   Guide.ylabel("p*(a|w)",orientation=:vertical),Guide.title("Posterior, β=$β_lo"),BAtheme())
plt_pagw_hi = plot(x=a_strings,y=pagw_hi,Geom.bar, Scale.x_discrete, Scale.y_continuous(minvalue=0, maxvalue=1), 
                   Guide.xlabel("a", orientation=:horizontal),Guide.ylabel("p*(a|w)",orientation=:vertical),
                   Guide.title("Posterior, β=$β_hi"),BAtheme())


hstack(plt_pagw_lo, plt_pagw_hi)

#compose final figure - stacking all the plots
plt_final = vstack(hstack(plt_u,plt_p0,plt_EU),hstack(plt_pagw_lo, plt_pagw_hi,plt_DKL))

#show figure
display(plt_final)

#store figure (make sure that the folder 'Figures' exists or change path)
w = 18cm
h = 12cm
#draw(SVG("Figures/FEOptimization.svg", w, h), plt_final) #uncomment to store figure