using LaTeXStrings
using Plots
const color_palette_bright = palette(:seaborn_bright) #bright color theme for the plots

# default(dpi = 300)
# default(framestyle = :box)
# pgfplotsx()
distribution_keys = [Poisson]
function plot_all()
    fnames = ["hh","ws","rest"]
    map(plot_estimate,fnames)
end
function plot_estimate(fname)
    data = deserialize(joinpath(PACKAGE_FOLDER,"intervals_model_output/simulation_output","$fname.dat"))
    plot_dists("$fname",distribution_keys,collect(values(data)))
    for (k,v) in zip(keys(data),values(data)) 
    plot_posteriors("$(k)_$fname",v)
    end
end
function plot_dists(fname,dist_constructors,data)

    p_estimate_as_arrays = map(d -> get_params(d.P),data)
    p_matrix = map(x -> plot(),p_estimate_as_arrays[1])
    ymo = ["Y","M","O"]
    x_range = 0.0:144.0
    for i in YOUNG:OLD, j in YOUNG:OLD
        for (k,(dist_constructor,p_estimate)) in collect(enumerate(zip(dist_constructors,p_estimate_as_arrays)))[1:end]
            dists =  map(p -> p.particles, p_estimate[i,j]) |>
                        t -> zip(t...) |>
                        l -> map(t -> dist_constructor(t...),l) 
        
            dist_pts = [pdf.(dists,i) for i in x_range]
            mean_dat = median.(dist_pts)
            
            hasnans = any(any.(map(l -> isnan.(l),dist_pts)))
            err_down = hasnans ? 0 : quantile.(dist_pts,0.05)
            err_up = hasnans ? 0 : quantile.(dist_pts,0.95)
            plot!(p_matrix[i,j] ,x_range,mean_dat;
                ribbon = ( mean_dat .- err_down,err_up .- mean_dat),
                legend = false,
                label = string(dist_constructor),
                seriescolor = color_palette_bright[k]
            )
        end
        annotate!(p_matrix[i,j],compute_x_pos(p_matrix[i,j]),compute_y_pos(p_matrix[i,j]), Plots.text("$(ymo[i])→$(ymo[j])", :left, 10))
    end
    plot!(p_matrix[end,1]; legend = true)
    p = plot(p_matrix..., size = (800,600))
    savefig(p,joinpath(PACKAGE_FOLDER,"plots", "$fname.pdf"))
end
using Statistics
compute_x_pos(p) = xlims(p)[1] + 0.02*((xlims(p)[2]  - xlims(p)[1]))
compute_y_pos(p) = ylims(p)[2] - 0.11*((ylims(p)[2]  - ylims(p)[1]))

function plot_posteriors(fname,data)
    sym_data = as_symmetric_matrix(data.P)
    p_matrix = map(x -> plot(),sym_data)

    ymo = ["Y","M","O"]
    for  i in YOUNG:OLD, j in YOUNG:OLD
        # display(data.P[i].particles)
        hist = StatsBase.fit(Histogram,sym_data[i,j].particles; nbins = 40)
        plot!(p_matrix[i,j],hist;legend = false, xlabel = L"\lambda_{%$i,%$j}", color = color_palette[1]  )
        # display(kernel_data)
        # vline!(p_list[i],[argmax(kernel_data)]; seriescolor = color_palette[2]
        annotate!(p_matrix[i,j],compute_x_pos(p_matrix[i,j]),compute_y_pos(p_matrix[i,j]), Plots.text("$(ymo[i])→$(ymo[j])", :left, 10))
    end   
    plot!(p_matrix[1]; ylabel = "No. of particles")
    plot!(p_matrix[4]; ylabel = "No. of particles")
    plot!(p_matrix[7]; ylabel = "No. of particles")

    p = plot(p_matrix...; size = (800,600), seriescolor = color_palette[1])
    savefig(p,joinpath(PACKAGE_FOLDER,"plots","$fname.pdf"))
end
function add_subplot_letters!(plot_list; pos = :top)
    for (i,sp) in enumerate(plot_list)
        letter = string(Char(i+96))
        if pos == :top
            annotate!(sp,xlims(sp)[1] + 0.02*((xlims(sp)[2]  - xlims(sp)[1])),ylims(sp)[2] - 0.11*((ylims(sp)[2]  - ylims(sp)[1])), Plots.text("$letter)", :left, 18))
        elseif pos == :bottom
            annotate!(sp,xlims(sp)[1] + 0.02*((xlims(sp)[2]  - xlims(sp)[1])),ylims(sp)[1] + 0.11*((ylims(sp)[2]  - ylims(sp)[1])), Plots.text("$letter)", :left, 18))
        end
    end
end