#!/usr/bin/env python3
# -*- coding: utf-8 -*-
r"""
make_fig_pk_pipeline.py
=======================
Generate fig_pk_pipeline.png for the DFD production P(k) pipeline appendix.

Reuses dfd_pk_pipeline.py end-to-end (same CAMB run, same window convolution,
same Patchy covariance, same best-fit nuisances) so the figure is a faithful
picture of the EXACT fit reported in the appendix -- no separate/parallel model.

The figure shows, for two representative bins (NGC z=0.38 with the full P0+P2
mock covariance, and NGC z=0.61 P0-only), the window-convolved best-fit
monopole P0 (and quadrupole where fit) for DFD and LCDM overlaid on the real
Beutler BOSS DR12 data with mock-covariance error bars, plus a residual panel
(data - model)/sigma.

Run:
  python3 make_fig_pk_pipeline.py
Writes:
  fig_pk_pipeline.png
"""
import os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

import dfd_pk_pipeline as P

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "fig_pk_pipeline.png")


def windowed_best(dat_model, fitres):
    """Window-convolved best-fit 200-vector [P0(100),P2(100)] for a model."""
    return P.windowed_model(dat_model["kh"], dat_model["Pk"], dat_model["W"],
                            fitres["b1"], fitres["f_fixed"], fitres["sigma_v"],
                            fitres["Pshot"])


def main():
    zs = sorted(set(P.ZEFF.values()))
    print("[fig] CAMB DFD ...")
    pk_dfd = P.camb_pk(P.DFD_COSMO, zs, sigma8_target=P.SIGMA8_DFD_TARGET,
                       nonlinear=True)
    print("[fig] CAMB LCDM ...")
    pk_lcdm = P.camb_pk(P.LCDM_COSMO, zs, sigma8_target=P.PLANCK["sigma8"],
                        nonlinear=True)

    # Two representative bins: NGC_z1 (full P0+P2 mock cov), NGC_z3 (P0-only).
    panels = [("NGC", "z1"), ("NGC", "z3")]

    fig = plt.figure(figsize=(11.5, 7.6))
    gs = fig.add_gridspec(2, 2, height_ratios=[2.4, 1.0], hspace=0.06, wspace=0.22)

    for col, (cap, zbin) in enumerate(panels):
        b = P.load_bin(cap, zbin)
        use_p2 = b["use_p2"]
        d_dfd = P.build_dat_for_model(b, pk_dfd)
        d_lcdm = P.build_dat_for_model(b, pk_lcdm)
        r_dfd = P.fit_bin(d_dfd)
        r_lcdm = P.fit_bin(d_lcdm)
        r_dfd["f_fixed"] = d_dfd["f_fixed"]
        r_lcdm["f_fixed"] = d_lcdm["f_fixed"]
        d_dfd.update(r_dfd)   # carry b1, sigma_v, Pshot, f_fixed
        d_lcdm.update(r_lcdm)

        m_dfd = windowed_best(d_dfd, d_dfd)
        m_lcdm = windowed_best(d_lcdm, d_lcdm)

        k = P.K_MEAS
        err = np.sqrt(np.diag(b["C"]))
        d = b["d"]
        sel0 = (k >= P.KMIN_FIT) & (k <= P.KMAX_FIT)

        ax = fig.add_subplot(gs[0, col])
        axr = fig.add_subplot(gs[1, col], sharex=ax)

        # ---- monopole P0 (always shown) ----
        kk = k[sel0]
        ax.errorbar(kk, (kk * d[:P.N_MEAS][sel0]), yerr=kk * err[:P.N_MEAS][sel0],
                    fmt="o", ms=4, color="k", capsize=2, lw=1, zorder=5,
                    label="BOSS DR12 data $P_0$")
        ax.plot(kk, kk * m_lcdm[:P.N_MEAS][sel0], "-", color="#1f77b4", lw=2.0,
                label=r"$\Lambda$CDM $P_0$")
        ax.plot(kk, kk * m_dfd[:P.N_MEAS][sel0], "--", color="#d62728", lw=2.0,
                label="DFD $P_0$")

        # ---- quadrupole P2 (only where it is in the fit) ----
        if use_p2:
            d2 = b["d"][P.N_MEAS:]
            e2 = err[P.N_MEAS:]
            ax.errorbar(kk, kk * d2[sel0], yerr=kk * e2[sel0], fmt="s", ms=4,
                        color="0.45", capsize=2, lw=1, zorder=4,
                        label="data $P_2$")
            ax.plot(kk, kk * m_lcdm[P.N_MEAS:][sel0], "-", color="#1f77b4",
                    lw=1.2, alpha=0.7)
            ax.plot(kk, kk * m_dfd[P.N_MEAS:][sel0], "--", color="#d62728",
                    lw=1.2, alpha=0.7)

        mult = "P0+P2" if use_p2 else "P0 only"
        ax.set_title(f"{cap} $z={b['z']:.2f}$  ({mult}, {b['cov_src'].split('(')[0].strip()})",
                     fontsize=10)
        ax.set_ylabel(r"$k\,P_\ell(k)\ [(\mathrm{Mpc}/h)^2]$", fontsize=10)
        ax.legend(fontsize=8, loc="upper right", ncol=1, framealpha=0.9)
        ax.grid(alpha=0.25)
        ax.tick_params(labelbottom=False)
        ax.text(0.03, 0.05,
                r"$\chi^2/\nu$: DFD %.2f, $\Lambda$CDM %.2f" %
                (r_dfd["chi2"] / r_dfd["dof"], r_lcdm["chi2"] / r_lcdm["dof"]),
                transform=ax.transAxes, fontsize=9,
                bbox=dict(boxstyle="round", fc="white", ec="0.6", alpha=0.9))

        # ---- residual panel (P0) ----
        resid_dfd = (d[:P.N_MEAS][sel0] - m_dfd[:P.N_MEAS][sel0]) / err[:P.N_MEAS][sel0]
        resid_lcdm = (d[:P.N_MEAS][sel0] - m_lcdm[:P.N_MEAS][sel0]) / err[:P.N_MEAS][sel0]
        axr.axhline(0, color="0.5", lw=0.8)
        axr.axhspan(-1, 1, color="0.85", alpha=0.5, zorder=0)
        axr.plot(kk, resid_lcdm, "-o", color="#1f77b4", ms=3, lw=1.0,
                 label=r"$\Lambda$CDM")
        axr.plot(kk, resid_dfd, "--s", color="#d62728", ms=3, lw=1.0, label="DFD")
        axr.set_ylim(-4, 4)
        axr.set_xlabel(r"$k\ [h/\mathrm{Mpc}]$", fontsize=10)
        axr.set_ylabel(r"$\Delta P_0/\sigma$", fontsize=9)
        axr.grid(alpha=0.25)
        if col == 0:
            axr.legend(fontsize=8, loc="upper right", ncol=2)

    fig.suptitle("DFD-native production galaxy $P(k)$ pipeline vs BOSS DR12 "
                 "(window-convolved, Patchy covariance, profiled nuisances)",
                 fontsize=12, y=0.985)
    fig.text(0.5, 0.005,
             "CAMB linear+HMcode  ->  chi-matter Q=1 growth  ->  Kaiser+FoG  ->  "
             "Beutler 200x2000 window  ->  Patchy-cov chi^2.   "
             "Combined: DFD chi2/dof=1.59, LCDM=1.09 (DFD acceptable, not superior).",
             ha="center", fontsize=8, color="0.3")
    fig.savefig(OUT, dpi=140, bbox_inches="tight")
    print(f"[fig] wrote {OUT}")


if __name__ == "__main__":
    main()
