"""
FIGURE bao_hubble — BAO distance ladder for DFD-CLASS.

Plots DFD predictions D_V/r_d, D_M/r_d, D_H/r_d as smooth curves vs redshift,
overlaid with the 12 consensus BAO data points (with error bars) from
background.BAO, using the DERIVED r_drag (~147 Mpc). LCDM (Planck) curves are
shown for contrast. The BAO chi^2/N is computed honestly and annotated.

Run:
  cd . && \
  PYTHONPATH=. \
  python3 figures/make_bao_hubble.py
"""
import sys
sys.path.insert(0, ".")

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

import dfd_params as P
import background as bg

OUT = "./figures/bao_hubble.png"

# ---- DFD background ----------------------------------------------------------
rd_dfd = bg.r_drag()                       # DERIVED drag-epoch sound horizon [Mpc]
chi2_dfd, rows = bg.bao_chi2()             # honest chi^2 vs the 12 consensus points
ndat = len(rows)

# ---- LCDM (Planck) reference background --------------------------------------
class _PL:
    H0 = P.PLANCK["H0"]; OMEGA_M = P.PLANCK["Omega_m"]
    OMEGA_R = P.OMEGA_R_H2 / P.PLANCK["h"] ** 2
    OMEGA_LAMBDA = 1.0 - P.PLANCK["Omega_m"] - OMEGA_R
    OMEGA_B_H2 = P.OMEGA_B_H2; OMEGA_G_H2 = P.OMEGA_G_H2
rd_lcdm = P.PLANCK["r_drag"]
chi2_lcdm, _ = bg.bao_chi2(par=_PL, rd=rd_lcdm)

# ---- smooth model curves -----------------------------------------------------
zz = np.linspace(0.02, 2.6, 400)
DV_dfd  = np.array([bg.D_V(z) / rd_dfd for z in zz])
DM_dfd  = np.array([bg.D_M(z) / rd_dfd for z in zz])
DH_dfd  = np.array([bg.D_H(z) / rd_dfd for z in zz])
DV_lcdm = np.array([bg.D_V(z, par=_PL) / rd_lcdm for z in zz])
DM_lcdm = np.array([bg.D_M(z, par=_PL) / rd_lcdm for z in zz])
DH_lcdm = np.array([bg.D_H(z, par=_PL) / rd_lcdm for z in zz])

# ---- split BAO data by observable type ---------------------------------------
def collect(kind):
    z, v, e = [], [], []
    for b in bg.BAO:
        if b["kind"] == kind:
            z.append(b["z"]); v.append(b["val"]); e.append(b["err"])
    return np.array(z), np.array(v), np.array(e)

zDV, vDV, eDV = collect("DV")
zDM, vDM, eDM = collect("DM")
zDH, vDH, eDH = collect("DH")

# ---- plot --------------------------------------------------------------------
plt.rcParams.update({"font.size": 11, "axes.grid": True,
                     "grid.alpha": 0.25, "figure.dpi": 150})
fig, ax = plt.subplots(figsize=(9.2, 6.4))

cDM, cDH, cDV = "#1f5fb4", "#c1392b", "#2e8b57"

# LCDM dashed reference
ax.plot(zz, DM_lcdm, "--", color=cDM, lw=1.3, alpha=0.55)
ax.plot(zz, DH_lcdm, "--", color=cDH, lw=1.3, alpha=0.55)
ax.plot(zz, DV_lcdm, "--", color=cDV, lw=1.3, alpha=0.55,
        label=r"LCDM (Planck), $r_d=%.2f$ Mpc" % rd_lcdm)

# DFD solid curves
ax.plot(zz, DM_dfd, "-", color=cDM, lw=2.2, label=r"DFD  $D_M/r_d$")
ax.plot(zz, DH_dfd, "-", color=cDH, lw=2.2, label=r"DFD  $D_H/r_d$")
ax.plot(zz, DV_dfd, "-", color=cDV, lw=2.2, label=r"DFD  $D_V/r_d$")

# data points
ax.errorbar(zDM, vDM, yerr=eDM, fmt="o", ms=7, color=cDM, mfc="white",
            mew=1.6, capsize=3, zorder=5, label=r"BAO $D_M/r_d$")
ax.errorbar(zDH, vDH, yerr=eDH, fmt="s", ms=7, color=cDH, mfc="white",
            mew=1.6, capsize=3, zorder=5, label=r"BAO $D_H/r_d$")
ax.errorbar(zDV, vDV, yerr=eDV, fmt="D", ms=7, color=cDV, mfc="white",
            mew=1.6, capsize=3, zorder=5, label=r"BAO $D_V/r_d$")

ax.set_xlabel("redshift  $z$")
ax.set_ylabel(r"distance / sound horizon   $D_X / r_d$")
ax.set_title("DFD BAO distance ladder vs consensus ruler "
             "(6dFGS, MGS, BOSS DR12, eBOSS QSO, Ly$\\alpha$)")
ax.set_xlim(0, 2.6)
ax.set_ylim(0, 41)

# annotation box: honest chi^2 ledger
txt = (r"$r_d^{\rm DFD}=%.2f$ Mpc  (Planck %.2f)" % (rd_dfd, rd_lcdm) + "\n"
       r"DFD  $\chi^2=%.1f$ / %d  $\Rightarrow\ \chi^2/N=%.2f$" % (chi2_dfd, ndat, chi2_dfd / ndat) + "\n"
       r"LCDM $\chi^2=%.1f$ / %d  $\Rightarrow\ \chi^2/N=%.2f$" % (chi2_lcdm, ndat, chi2_lcdm / ndat) + "\n"
       r"(DFD distances run $\sim$2% short: mild ruler offset)")
ax.text(0.025, 0.975, txt, transform=ax.transAxes, va="top", ha="left",
        fontsize=9.5, family="monospace",
        bbox=dict(boxstyle="round,pad=0.5", fc="#fffbe6", ec="#999", alpha=0.95))

ax.legend(loc="lower right", fontsize=8.5, ncol=2, framealpha=0.92)
fig.tight_layout()
fig.savefig(OUT, dpi=150)
print("saved", OUT)
print("rd_dfd = %.3f Mpc | chi2_dfd = %.2f (chi2/N=%.3f) | chi2_lcdm = %.2f (chi2/N=%.3f)"
      % (rd_dfd, chi2_dfd, chi2_dfd / ndat, chi2_lcdm, chi2_lcdm / ndat))
