
# # plotting functions ------------------------------------------------------------
function plot_omegahat(chain::Chains, s::NamedTuple; prior_only = false, legend = :bottomleft)
    w, w1, w2 = summarize_omega(chain)
    p1 = plot(s.timesteps, w.μ', ribbon = [w.lb', w.ub'], 
                fillalpha = 0.3, linewidth = 2, color = :blue,
                title = "Total velocity",
                label = "Estimated",
                xlabel = "Time (s)",
                ylabel = "Angular velocity (deg/s)")

    plot!(p1, s.timesteps, s.ω, 
            linestyle = :dash, 
            linewidth = 2, 
            color = :cornflowerblue,
            label = "True")
    if !prior_only
        scatter!(p1, s.timesteps, s.y,
            label = "Observations",
            markershape = :circle,
            markersize = 2,
            markerstrokealpha = 0.8,
            markeralpha = 0.2,
            markercolor = :black)
    end

    p2 = plot(s.timesteps, w1.μ', ribbon = [w1.lb', w1.ub'], 
            fillalpha = 0.3, linewidth = 2, color = :darkorange,
            title = "Active",
            label = "Estimated")
    plot!(p2, s.timesteps, s.ωᵤ,
            linestyle = :dash, 
            linewidth = 2, 
            color = :orangered,
            label = "True")

    p3 = plot(s.timesteps, w2.μ', ribbon = [w2.lb', w2.ub'], 
            fillalpha = 0.3, linewidth = 2, color = :palegreen,
            title = "Passive",
            label = "Estimated",
            xlabel = "Time (s)")
    plot!(p3, s.timesteps, s.ωₑ,
            linestyle = :dash, 
            linewidth = 2, 
            color = :darkgreen,
            label = "True")

    l = @layout [a [b; c]]
    p = plot(p1, p2, p3, layout = l, 
            link=:all,
            titleloc = :right,
            legend = legend,
            foreground_color_legend = nothing )
    return p
end


function plot_omegahat_simple(chain::Chains, s::NamedTuple; 
                                prior_only = false, 
                                spaghetti = false,
                                legend = :bottomleft)
    if spaghetti === true
        out = Array(group(chain, :ω))
        nrows, ncols = size(out)
        idx = sample(1:nrows, 100)
        samples = out[idx, :]

        p1 = plot(s.timesteps, samples', 
                        linewidth = 2, 
                    # color = :blue,
                    alpha = 0.1, color = :steelblue,
                    title = "Total velocity",
                    legend = legend,
                    label = "Estimated",
                    xlabel = "Time (s)",
                    ylabel = "Angular velocity (deg/s)")
    else
        w_array = Array(group(chain, :ω))
        w = compute_stats(w_array)
        p1 = plot(s.timesteps, w.μ', ribbon = [w.lb', w.ub'], 
                fillalpha = 0.3, 
                linewidth = 2, 
                color = :blue,
                title = "Total velocity",
                legend = legend,
                label = "Estimated",
                xlabel = "Time (s)",
                ylabel = "Angular velocity (deg/s)")
    end


    plot!(p1, s.timesteps, s.ω, 
            linestyle = :dash, 
            linewidth = 2, 
            color = :cornflowerblue,
            label = "True")
    if !prior_only
        scatter!(p1, s.timesteps, s.y,
            label = "Observations",
            markershape = :circle,
            markersize = 2,
            markerstrokealpha = 0.8,
            markeralpha = 0.2,
            markercolor = :black)
    end
    return p1
end


function plot_omegatrue(s::NamedTuple; 
                        displacement = false, 
                        legend = :bottomleft)
    plot(s.timesteps, [s.ωᵤ, s.ωₑ], label = ["self" "external"], 
        linewidth = 3, 
        linestyle = :dot, 
        linecolor = ["red" "steelblue"],
        legend = legend,
        foreground_color_legend = nothing)

    plot!(s.timesteps, s.ω, 
        label = "total velocity", 
        lw = 3, 
        color = "purple")

    if displacement===true
        plot!(s.timesteps, s.θ, 
            label = "total displacement", 
            lw = 3, 
            color = "orange")
    end

    plot!(s.timesteps, s.ω - s.c, 
        label = "lag (cupula dynamics)",
        lw = 3,
        color = :grey)

    scatter!(s.timesteps, s.y,
        xlabel = "Time (s)",
        ylabel = "Angular velocity (deg/s)",
        label = "observations",
        markershape = :circle,
        markersize = 3,
        markeralpha = 0.5,
        markercolor = :black)
end




function plot_omegahat_simulation(chain::Chains, s::NamedTuple; prior_only = false, legend = :topright)
    w_array = Array(group(chain, :ω))
    w = compute_stats(w_array)
    p1 = plot(s.timesteps, w.μ', ribbon = [w.lb', w.ub'], 
                fillalpha = 0.3, 
                linewidth = 2, 
                color = :blue,
                title = "Total velocity",
                legend = legend,
                label = "Estimated",
                xlabel = "Time (s)",
                ylabel = "Angular velocity (deg/s)")

    if !prior_only
        scatter!(p1, s.timesteps, s.y,
            label = "Observations",
            markershape = :circle,
            markersize = 2,
            markerstrokealpha = 0.8,
            markeralpha = 0.2,
            markercolor = :black)
    end
    return p1
end



function plot_nigmatullina_trial(motion)
    p = scatter(motion.timesteps, motion.y,
    xlabel = "Time (s)",
    ylabel = "Angular velocity (deg/s)",
    legend = :none,
    markershape = :circle,
    markersize = 2,
    markeralpha = 0.1,
    markercolor = :black)
    plot!(motion.timesteps, motion.α,
        linewidth = 4,
        linecolor = :steelblue)
    vline!([motion.onset], linewidth = 2, linestyle = :dot, linecolor = "grey")
    plot!(p, motion.timesteps, motion.ω,
        linewidth = 4,
        linecolor = :orange)
end