"""
Read Omega_m and sigma8 from a DES Y6 chain and plot with ChainConsumer.

Works for chains obtained with
  i)  the Nautilus nested sampler, and
  ii) DES combined with external data via normalizing flows (tensiometer),
as long as the column layout below matches.
"""

import os

import numpy as np
import pandas as pd
from chainconsumer import Chain, ChainConsumer, PlotConfig

#######################
# i) NAUTILUS SAMPLER #
#######################
FILENAME = "32pt_wolensbin2_lin_lcdm_DV0723_R0723.txt"

# Column layout: col 0 = Omega_m, col -7 = sigma8, col -3 = log_weight
samples = np.loadtxt(os.path.join(PATH, FILENAME))

# Remove any samples containing NaNs or infs
samples = samples[np.isfinite(samples).all(axis=1)]

# Nautilus considers natural log weights, need to convert
log_weight = samples[:, -3]
weight = np.exp(log_weight[np.isfinite(log_weight)])

# remove samples with weight = 0.0
mask = weight > 0.0
samples = samples[mask, :]
weight = weight[mask]

df = pd.DataFrame({
    r"$\Omega_m$": samples[:, 0],
    r"$\sigma_8$": samples[:, -7],
    "weight": weight,
})

# Initialize empty chain
c = ChainConsumer()
c.add_chain(Chain(
    samples=df,
    name="DES Y6 3x2pt",
    kde=1.5,
    plot_point=False,
    plot_cloud=False,
    shade=True,
))

########################
# ii) NORMALIZING FLOWS #
########################
FILENAME = "chain_des_sn___des_bao_des_y6_3x2pt___des_y3_clusters_no_gg_chain_csislike.txt"

samples = np.loadtxt(os.path.join(PATH, FILENAME))

# Remove any samples containing NaNs or infs
samples = samples[np.isfinite(samples).all(axis=1)]
weight = np.ones(samples.shape[0])

df = pd.DataFrame({
    r"$\Omega_m$": samples[:, 0],
    r"$\sigma_8$": samples[:, 1],
    "weight": weight,
})

# Add "DES all probes" chain
c.add_chain(Chain(
    samples=df,
    name="DES All probes",
    kde=1.5,
    plot_point=False,
    plot_cloud=False,
    shade=True,
))

c.set_plot_config(PlotConfig(
    label_font_size=15,
    tick_font_size=8,
    legend_kwargs={"fontsize": 10},
    legend_color_text=True,
    summarise=False
))

fig = c.plotter.plot(columns=[r"$\Omega_m$", r"$\sigma_8$"])
fig.savefig("des_y6_omegam_sig8.png", dpi=300, bbox_inches="tight")