import getdist.plots as gdplt
from getdist import MCSamples, loadMCSamples
from pathlib import Path

# Load chain file
samples = loadMCSamples('chains/BK18_fig4_20/2024-04-15_20000_', settings={"ignore_rows": 0.3})
samples.paramNames.setLabelsFromParamNames('chains/BK18_fig4_20/2024-04-15_20000_.paramnames')
# Plot settings
plot_settings = {
    "smooth_scale_1D": 0.15,
    "smooth_scale_2D": 2,
    "num_bins": 100,
    "num_bins_2D": 40,
    "contours": [0.68, 0.95],
    "triangle_plot": True,
    "plot_2D_param": 0,
    "plot_2D_num": 1,
    "params": ["r", "BBdust", "BBsync"],
    "triangle_params": ["r", "BBdust", "BBsync"],
    "param_limits": {"r": [0, 0.18], "BBdust": [0, 11], "BBsync": [0, 8], "BBdustsynccorr": [-1, 1], "BBbetadust": [1, 2], 
                     "BBbetasync": [-4.2, -1.8], "BBalphadust": [-1, 0], "BBalphasync": [-1, 0]},
}

# Make triangle plot of figure 4 of BKXIII
gdplt.get_subplot_plotter()
gdplot = gdplt.getSubplotPlotter()
gdplot.settings.num_plot_contours = 2
gdplot.settings.alpha_filled_add = 0.6
gdplot.triangle_plot(samples, **plot_settings, filled=True)

# Make 1d posterior plots of figure 4 of BKXIII
plot_settings['params'] = ["BBbetadust", "BBbetasync", "BBdustsynccorr", "BBalphadust", "BBalphasync"]
plot_settings['xlims'] = [plot_settings['param_limits'][p] for p in plot_settings['params']]
gdplot.plots_1d(samples, **plot_settings)