"""
analyze_results.py  (v2 - getdist version-safe)
"""
import os, sys, glob, numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CHAIN    = os.path.join(BASE_DIR, "chains", "rlmt_lite_run")
OUT_DIR  = os.path.join(BASE_DIR, "results")
os.makedirs(OUT_DIR, exist_ok=True)

# ── getdist version check ────────────────────────────────────────────
import getdist
print(f"getdist version: {getdist.__version__}")
from getdist.mcsamples import loadMCSamples, MCSamplesError
from getdist import plots

# ── 1. チェーンファイル確認 ─────────────────────────────────────────
print("\n=== Chain files ===")
chain_files = sorted(glob.glob(CHAIN + "*.txt"))
print(f"  Found: {chain_files}")
if not chain_files:
    print("  ERROR: No chain files found.")
    print(f"  Looking in: {os.path.join(BASE_DIR, 'chains')}")
    sys.exit(1)

# ── 2. チェーン読み込み（バージョン対応） ───────────────────────────
print("\n=== Loading chain ===")
try:
    # getdist 1.x
    s = loadMCSamples(CHAIN, settings={"ignore_rows": 0.3})
except TypeError:
    try:
        # getdist 0.x
        s = loadMCSamples(CHAIN, ignore_rows=0.3)
    except TypeError:
        # フォールバック: settingsなしで読み込み
        s = loadMCSamples(CHAIN)

print(f"  Rows loaded  : {s.numrows}")
params = s.getParamNames().list()
print(f"  Parameters   : {params}")

# ── 3. パラメータ統計 ────────────────────────────────────────────────
print("\n=== Parameter constraints (mean ± 1-sigma) ===")
means = s.getMeans()
stds  = s.getVars()**0.5

RLMT_PARAMS = ["As","ns","beta_eff","c1","c2","c3","k_star","kc","eps3"]
BG_PARAMS   = ["H0","ombh2","omch2","tau"]
NUISANCE    = ["A_planck"]

lines = []
for grp, label in [(BG_PARAMS,"Background"),
                   (RLMT_PARAMS,"RLMT primordial"),
                   (NUISANCE,"Nuisance")]:
    print(f"\n  [{label}]")
    lines.append(f"\n[{label}]")
    for pname in grp:
        if pname in params:
            i = params.index(pname)
            line = f"  {pname:12s} = {means[i]:.5g}  +/- {stds[i]:.5g}"
            print(line)
            lines.append(line)

with open(os.path.join(OUT_DIR,"parameter_constraints.txt"),"w") as f:
    f.write("\n".join(lines))
print(f"\n  Saved: results/parameter_constraints.txt")

# ── 4. Best-fit chi^2 (チェーンファイルから直接読む) ─────────────────
print("\n=== Best-fit statistics ===")
try:
    data = np.loadtxt(chain_files[0])
    # Cobaya format: col0=weight, col1=-logpost, col2..=params
    neg_logpost = data[:, 1]
    best_neg_logpost = float(np.min(neg_logpost))
    # -logpost ≈ -loglike (with flat prior), so chi^2 ≈ 2*min(-logpost)
    best_chi2_approx = 2.0 * best_neg_logpost
    n_data  = data.shape[0]
    n_rlmt  = 15   # total varied params in RLMT run
    n_lcdm  = 6    # standard LCDM (H0,ombh2,omch2,tau,As,ns) + A_planck = 7
    n_extra = n_rlmt - n_lcdm
    print(f"  Chains rows          : {n_data}")
    print(f"  Best -logpost        : {best_neg_logpost:.4f}")
    print(f"  Best chi^2 (approx)  : {best_chi2_approx:.4f}")
    print(f"  Extra RLMT params    : {n_extra}")
    print(f"  AIC penalty (2k)     : +{2*n_extra}")
    print(f"  BIC note: need N_data for BIC; AIC sufficient for model comparison")
    print(f"\n  *** Run lcdm_baseline to get Delta_chi^2 ***")
    lines2 = [
        f"\n[Best-fit statistics]",
        f"  best -logpost       = {best_neg_logpost:.4f}",
        f"  best chi^2 (approx) = {best_chi2_approx:.4f}",
        f"  extra RLMT params   = {n_extra}",
        f"  AIC penalty         = +{2*n_extra}",
    ]
    with open(os.path.join(OUT_DIR,"parameter_constraints.txt"),"a") as f:
        f.write("\n".join(lines2))
except Exception as e:
    print(f"  Could not read chi^2: {e}")

# ── 5. コーナープロット ──────────────────────────────────────────────
print("\n=== Corner plot (RLMT primordial params) ===")
plot_params = [p for p in RLMT_PARAMS if p in params]
if len(plot_params) >= 2:
    try:
        g = plots.get_subplot_plotter(width_inch=12)
        g.triangle_plot(s, plot_params, filled=True, title_limit=1)
        for ext in ["pdf","png"]:
            out = os.path.join(OUT_DIR, f"corner_rlmt_primordial.{ext}")
            g.export(out)
            print(f"  Saved: {out}")
    except Exception as e:
        print(f"  Corner plot error: {e}")

# ── 6. 1D事後分布 ──────────────────────────────────────────────────
print("\n=== 1D posterior plots ===")
fig, axes = plt.subplots(3, 3, figsize=(12, 9))
axes = axes.flatten()
for i, pname in enumerate(RLMT_PARAMS):
    if pname in params:
        idx = params.index(pname)
        col = s.samples[:, idx]
        axes[i].hist(col, bins=60, density=True,
                     color="steelblue", alpha=0.7)
        axes[i].axvline(means[idx], color="red",
                        lw=1.5, label=f"mean={means[idx]:.3g}")
        axes[i].set_xlabel(pname, fontsize=10)
        axes[i].legend(fontsize=7)
plt.suptitle("RLMT primordial parameter posteriors\n"
             "(Planck 2018 TT+TE+EE+lensing lite)", fontsize=11)
plt.tight_layout()
for ext in ["pdf","png"]:
    out = os.path.join(OUT_DIR, f"posteriors_1d.{ext}")
    plt.savefig(out, dpi=150)
    print(f"  Saved: {out}")

print("\n=== Done. Check the results/ folder. ===")