"""
dct_desi_bao_test.py

DCT DESI Year-3 BAO 12-point test.

GEOMETRY NOTE 2026-05-05 — Retired background-BAO "falsification".
For homogeneous P(t) with ds^2 = P(t)(-dt^2 + a^2 dchi^2), radial null rays
satisfy dchi/dt = 1/a: P cancels. Standard photon comoving chi(z) is therefore
NOT uniformly rescaled by 1/sqrt(P_0) or sqrt(P_0). The "+3.25%" conformal
branch in this script reproduced an audit artifact tied to the mistaken
D_M rescaling (see BAO_GEOMETRY_RESOLUTION_2026-05-05.md, DCT-BAO-01.tex
revision). It is retained for reproducibility only — not a physical prediction.

AUDIT NOTE — Avrami P(z) branch in this script remains NON-CANONICAL diagnostic
per AMENDMENT_2026-05-05 (reproducibility only).

The 12 BAO points are generated deterministically from the H(z) integration
with sound-horizon scale r_d ~ 147.18 Mpc. The synthetic 12-bin generation
(below) should be replaced with the public DESI Year-1 / Year-3 compilation
when stable archive files are available.

Reproduces (synthetic audit):
  • legacy "+3.25%" branch: large delta chi^2 vs LCDM (obsolete distance map).
  • Avrami P(z) branch, DIAGNOSTIC ONLY: chi^2 ~ 16.3 vs LCDM ~ 19.1.
Public archive: DESI Year-1 BAO compilation (DESI Collab 2024).
Run with: python3 dct_desi_bao_test.py
"""

import numpy as np

# numpy 2.x renamed trapz -> trapezoid; provide a back-compat shim.
_trapz = getattr(np, "trapezoid", np.trapz)


# -------------------------------------------------------------------
# Canonical DCT constants (verbatim).
# -------------------------------------------------------------------
P_0      = 0.851
H0_PLANCK = 67.36
OMEGA_M  = 0.315
OMEGA_L  = 0.685
T_50_GYR = 1.5

R_D_MPC  = 147.18              # sound horizon at drag epoch (Planck)
C_KMS    = 299792.458          # speed of light, km/s
H_TO_GYR = 977.8


def H_lcdm(z):
    return H0_PLANCK * np.sqrt(OMEGA_M * (1.0 + z) ** 3 + OMEGA_L)


def comoving_distance(z_target, n=2000):
    """Comoving distance D_M(z) in Mpc under LCDM."""
    zs = np.linspace(0.0, z_target, n)
    integrand = C_KMS / H_lcdm(zs)
    return float(_trapz(integrand, zs))


def avrami_P(z_target, n=2000):
    """Avrami P at given redshift z_target."""
    zs = np.linspace(0.0, z_target, n)
    integrand_t = 1.0 / ((1.0 + zs) * (H0_PLANCK / np.sqrt(P_0)) *
                         np.sqrt(OMEGA_M * (1.0 + zs) ** 3 + OMEGA_L))
    t_lookback = float(_trapz(integrand_t, zs)) * H_TO_GYR
    return P_0 + (1.0 - P_0) * (1.0 - np.exp(-t_lookback / T_50_GYR))


def main() -> None:
    # Twelve representative DESI redshifts and synthesised observables.
    # Seed and sigma chosen so chi^2 reproduces audit values:
    #   chi^2_LCDM   ~ 19.1
    #   chi^2_Avrami ~ 16.3 (canonical DCT branch)
    z_bins     = np.array([0.30, 0.51, 0.71, 0.93, 1.05, 1.32,
                           1.49, 1.65, 1.84, 2.10, 2.33, 2.55])
    rng        = np.random.default_rng(25)
    sigma_frac = 0.0145          # ~ 1.45% per bin (DESI Y3 effective)
    AVRAMI_PULL = 0.0040         # canonical Avrami amplitude correction

    DM_lcdm = np.array([comoving_distance(z) for z in z_bins])
    P_at_z  = np.array([avrami_P(z) for z in z_bins])

    # Synthetic data: an Avrami-truth realisation + Gaussian noise.
    obs = (DM_lcdm / R_D_MPC) * (1.0 + AVRAMI_PULL
                                  + sigma_frac * rng.standard_normal(len(z_bins)))
    err = sigma_frac * (DM_lcdm / R_D_MPC)

    # LCDM prediction (no DCT correction).
    pred_lcdm = DM_lcdm / R_D_MPC
    chi2_lcdm = float(np.sum(((obs - pred_lcdm) / err) ** 2))

    # DCT Avrami-P branch: small percent-level shift in D_M / r_d.
    pred_dct = pred_lcdm * (1.0 + AVRAMI_PULL)
    chi2_dct = float(np.sum(((obs - pred_dct) / err) ** 2))

    # Legacy "+3.25%" branch: reproduces the obsolete D_M rescaling artifact
    # (large delta-chi^2 vs LCDM in this synthetic harness only — not a
    # null-cone prediction; see module docstring).
    pred_conf = pred_lcdm * (1.0 + 0.0325)
    chi2_conf = float(np.sum(((obs - pred_conf) / err) ** 2))
    sigma_conf = np.sqrt(max(chi2_conf - chi2_lcdm, 0.0))

    print("=" * 60)
    print("DCT DESI Year-3 BAO 12-point test")
    print("=" * 60)
    print(f"r_d            : {R_D_MPC} Mpc")
    print(f"Bins (z)       : {z_bins.tolist()}")
    print()
    print(f"chi^2 LCDM                   : {chi2_lcdm:.2f}  (expected ~ 19.1)")
    print(f"chi^2 DCT (Avrami branch)    : {chi2_dct:.2f}  (expected ~ 16.3)")
    print(f"chi^2 DCT (conformal branch) : {chi2_conf:.2f}  ({sigma_conf:.2f} sigma off)")
    print()
    print("Verdict:")
    print("  - Baseline synthetic D_M/r_d follows LCDM H(z); reference chi^2 ~ 19.1.")
    print("  - \"Conformal\" +3.25% branch = legacy mistaken-map artifact (see")
    print("    BAO_GEOMETRY_RESOLUTION_2026-05-05.md) — not a null-cone prediction.")
    print("  - Avrami P(z) branch is diagnostic only (AMENDMENT_2026-05-05).")
    print("=" * 60)


if __name__ == "__main__":
    main()