const BASELINE_VALUES = Dict(:r => 1.1, :γ => 0.1, :pref => 0.01, :g => 200.0)
const MAKIE_COLORS = Makie.wong_colors()[1:3]

# base theme settings
fontsize_theme = Theme(fontsize = 17, font = :bold_italic)
attributes = Attributes(
    Axis = (
        xgridvisible = false,
        ygridvisible = false,
        ytickalign = 1,
        xtickalign = 1,
    ),
)
set_theme!(merge(fontsize_theme, attributes))

function standard_deviation_bands!(cycle, trend, colors)
    # Compute residuals from cyclicality of the time series
    residuals = cycle - trend
    # Compute the standard deviation of the residuals
    sigma = std(residuals)
    # Calculate upper and lower bands
    upper_bound = trend .+ sigma
    lower_bound = trend .- sigma

    # Plot the standard deviations as dashed lines
    lines!(lower_bound; color = (colors, 0.3), linestyle = :dash)
    lines!(upper_bound; color = (colors, 0.3), linestyle = :dash)
end

function generic_plot(dataframes::Vector{DataFrame}, vars::Vector{Symbol}, labels::Vector{String}, param::Symbol; debt_to_gdp::Bool = false)
    fig =  Figure(size = (800, 400))

    if length(dataframes) > 1
        axes = tuple(collect((1, i) for i in 1:length(dataframes))...)
        for i in 1:length(dataframes)
            gdf = @pipe dataframes[i] |>
                groupby(_, param)
            ax = fig[axes[i]...] = Axis(fig, title = labels[i])

            for j in 1:length(gdf)
                cycle, trend = hp_filter(gdf[j][!, vars[1]], 1600)
                lines!(trend; label = "$(param) = $(only(unique(gdf[j][!, param])))", color = only(unique(gdf[j][:, param])) == BASELINE_VALUES[param] ? :black : MAKIE_COLORS[j])
                standard_deviation_bands!(cycle, trend, only(unique(gdf[j][:, param])) == BASELINE_VALUES[param] ? :black : MAKIE_COLORS[j])            
            end
            ax.xticks = 0:200:800
            fig.content[1].ylabel = "Volumes"
            ax.xlabel = "Steps"
        end
    else
        if length(vars) > 1 
            axes = tuple(collect((1, i) for i in 1:length(vars))...)
            for i in 1:length(vars)
                gdf = groupby(dataframes[1], param)
                ax = fig[axes[i]...] = Axis(fig, title = labels[i])

                for j in 1:length(gdf)
                    cycle, trend = hp_filter(gdf[j][!, vars[i]], 1600)
                    lines!(trend; label = "$(param) = $(only(unique(gdf[j][!, param])))", color = only(unique(gdf[j][:, param])) == BASELINE_VALUES[param] ? :black : MAKIE_COLORS[j])
                    standard_deviation_bands!(cycle, trend, only(unique(gdf[j][:, param])) == BASELINE_VALUES[param] ? :black : MAKIE_COLORS[j])            
                end
                ax.xticks = 0:200:800
                fig.content[1].ylabel = "Volumes"
                ax.xlabel = "Steps"
            end
        else
            for i in 1:length(vars)
                gdf = groupby(dataframes[1], param)
                ax = fig[1,1] = Axis(fig, title = labels[i])

                for j in 1:length(gdf)
                    cycle, trend = hp_filter(gdf[j][!, vars[1]], 1600)
                    lines!(trend; label = "$(param) = $(only(unique(gdf[j][!, param])))", color = only(unique(gdf[j][:, param])) == BASELINE_VALUES[param] ? :black : MAKIE_COLORS[j])
                    standard_deviation_bands!(cycle, trend, only(unique(gdf[j][:, param])) == BASELINE_VALUES[param] ? :black : MAKIE_COLORS[j])            
                end
                ax.xticks = 0:200:800
                fig.content[1].ylabel = debt_to_gdp ? "%" : "Volumes"
                ax.xlabel = "Steps"
            end
        end
    end

    length_custom = (debt_to_gdp && length(vars) == 1) ? 1 : 1:2
    fig[end + 1, length_custom] = Legend(fig, fig.content[1]; tellheight=true, tellwidth=false, orientation=:horizontal)
    return fig
end