"""
DFD-CLASS — Optical-psi distance screen module (module 2, distance part).

DFD posits that observed angular/luminosity distances carry a line-of-sight optical
bias from the GRIN metric n=e^psi:
    D_L^DFD = D_L^FRW · e^{Dpsi(z)},   D_A^DFD = D_A^FRW · e^{Dpsi(z)}
(section_cosmology.tex:65,73). The radial rate D_H = c/H(z) is NOT a propagation
distance, so the screen acts on the TRANSVERSE D_M = (1+z)D_A, not on D_H.

CORPUS HONESTY (corpus_extract.md sec 1.4): the screen FORM is derived (optical
metric + Etherington reciprocity), but the PROFILE Dpsi(z) is reconstructed AS the
LCDM distance ratio (appendix_AU:205) -> w_eff=-1 by construction. So a screen that
matches data is a TIE with LCDM, not an independent zero-parameter win.

THIS MODULE asks the decisive, honest question: does a PHYSICALLY-ADMISSIBLE screen
(monotonic in z, -> 0 as z->0 so the local H0=72.09 is preserved) reconcile DFD's
derived H0=72.09 + STANDARD r_drag with BAO and theta*? We compute the screen each
datum REQUIRES and report whether one smooth profile fits all. No fudge.
"""

import numpy as np
from scipy.optimize import minimize_scalar
import dfd_params as P
import background as B


def screened_predict(z, kind, rd, dpsi):
    """BAO prediction with screen dpsi applied to transverse D_M (and D_V's D_M part)."""
    DM = B.D_M(z) * np.exp(dpsi)        # transverse distance screened
    DH = B.D_H(z)                       # radial rate NOT screened
    if kind == "DM":
        return DM / rd
    if kind == "DH":
        return DH / rd
    if kind == "DV":
        return (z * DM**2 * DH)**(1.0 / 3.0) / rd
    raise ValueError(kind)


def needed_dpsi_per_DM():
    """For each transverse (DM/DV) datum, the screen it alone requires: ln[obs/pred_noscreen]."""
    rd = B.r_drag()
    rows = []
    for b in B.BAO:
        if b["kind"] == "DH":
            continue  # radial, screen-insensitive
        pred0 = B.bao_predict(b["z"], b["kind"], rd)
        dpsi_need = np.log(b["val"] / pred0)
        rows.append((b["name"], b["z"], dpsi_need))
    # theta* requirement
    rs_star = B.sound_horizon(P.Z_STAR)
    DM_star = B.D_M(P.Z_STAR)
    th_obs = P.PLANCK["theta_star"]
    # theta = rs/(DM e^dpsi)  ->  dpsi = ln(rs/(DM*th_obs))
    dpsi_star = np.log(rs_star / (DM_star * th_obs))
    rows.append(("theta* (z*)", P.Z_STAR, dpsi_star))
    return rows


def bao_chi2_screen(dpsi_fn, rd=None):
    rd = rd if rd is not None else B.r_drag()
    chi2 = 0.0
    rows = []
    for b in B.BAO:
        dp = dpsi_fn(b["z"])
        pred = screened_predict(b["z"], b["kind"], rd, dp)
        r = (pred - b["val"]) / b["err"]
        chi2 += r * r
        rows.append((b["name"], b["z"], b["val"], pred, r))
    return chi2, rows


def main():
    print("=" * 66)
    print("DFD optical-psi distance screen — does it reconcile H0=72.09 + BAO/theta*?")
    print("=" * 66)

    # 1) The screen each transverse datum requires on its own
    print("\n[1] Screen Dpsi(z) REQUIRED by each transverse datum (and theta*):")
    print(f"{'datum':14s} {'z':>7s} {'Dpsi_needed':>12s}")
    needs = needed_dpsi_per_DM()
    for name, z, dp in needs:
        print(f"{name:14s} {z:7.2f} {dp:+12.4f}")
    dvals = [dp for _, z, dp in needs if z < 3]  # exclude z* for the low-z spread
    print(f"  low-z (BAO) spread: mean {np.mean(dvals):+.4f}, std {np.std(dvals):.4f}, "
          f"range [{min(dvals):+.4f}, {max(dvals):+.4f}]")
    print(f"  theta* requires Dpsi(z*) = {needs[-1][2]:+.4f}")
    print("  -> MONOTONIC test: a line-of-sight screen must grow with z and ->0 at z=0.")

    # 2) No screen
    chi2_0, _ = B.bao_chi2()
    print(f"\n[2] No screen:           BAO chi^2 = {chi2_0:.2f}  (chi^2/N = {chi2_0/12:.2f})")

    # 3) Best-fit CONSTANT screen (1 param) on transverse distances
    def chi2_const(a):
        c, _ = bao_chi2_screen(lambda z: a)
        return c
    res = minimize_scalar(chi2_const, bounds=(-0.05, 0.10), method="bounded")
    a_best = res.x
    chi2_c = res.fun
    print(f"[3] Best constant screen Dpsi={a_best:+.4f}: BAO chi^2 = {chi2_c:.2f} "
          f"(chi^2/N = {chi2_c/12:.2f})")

    # 4) Best-fit SMOOTH monotonic screen Dpsi(z)=A*(1-e^{-z/z0}), ->0 at z=0, ->A plateau
    #    (physically-admissible LOS shape). 2 params (A, z0).
    def make_smooth(A, z0):
        return lambda z: A * (1.0 - np.exp(-z / z0))
    def chi2_smooth(p):
        A, z0 = p
        if z0 <= 0:
            return 1e9
        c, _ = bao_chi2_screen(make_smooth(A, z0))
        return c
    from scipy.optimize import minimize
    r2 = minimize(chi2_smooth, x0=[0.04, 0.5], method="Nelder-Mead")
    A_s, z0_s = r2.x
    chi2_s = r2.fun
    print(f"[4] Best smooth monotonic Dpsi=A(1-e^-z/z0), A={A_s:+.4f}, z0={z0_s:.3f}: "
          f"BAO chi^2 = {chi2_s:.2f} (chi^2/N = {chi2_s/12:.2f})")
    # what does this smooth screen give at z*?
    dpsi_star_smooth = make_smooth(A_s, z0_s)(P.Z_STAR)
    th_screened = B.sound_horizon(P.Z_STAR) / (B.D_M(P.Z_STAR) * np.exp(dpsi_star_smooth))
    print(f"     -> at z*, this screen gives Dpsi={dpsi_star_smooth:+.4f}, "
          f"100*theta*={100*th_screened:.5f} (Planck {100*P.PLANCK['theta_star']:.5f}, "
          f"need Dpsi(z*)={needs[-1][2]:+.4f})")

    chi2_l, _ = B.bao_chi2(par=_lcdm_par(), rd=P.PLANCK["r_drag"])
    print(f"[5] LambdaCDM reference:  BAO chi^2 = {chi2_l:.2f}  (chi^2/N = {chi2_l/12:.2f})")

    # 6) DDR unit test: screen multiplies D_L and D_A equally -> Etherington preserved
    z_t = 1.0
    DA = B.D_A(z_t) * np.exp(0.3)
    DL = (1 + z_t)**2 * B.D_A(z_t) * np.exp(0.3)   # both screened by same e^dpsi
    ddr = DL / ((1 + z_t)**2 * DA)
    print(f"\n[6] DDR (Etherington) unit test: D_L/((1+z)^2 D_A) = {ddr:.10f}  "
          f"(screen cancels, eta=1) {'PASS' if abs(ddr-1)<1e-9 else 'FAIL'}")

    # honest verdict — must check BOTH BAO chi^2 AND whether the SAME screen fixes theta*
    print("\n" + "-" * 66)
    print("HONEST VERDICT:")
    dpsi_bao = a_best                       # screen the BAO data prefer
    dpsi_th_need = needs[-1][2]             # screen theta* needs
    th_with_bao_screen = B.sound_horizon(P.Z_STAR) / (B.D_M(P.Z_STAR) * np.exp(dpsi_bao))
    th_resid = 100 * (th_with_bao_screen - P.PLANCK["theta_star"]) / P.PLANCK["theta_star"]
    consistent = np.std(dvals) < 0.5 * abs(np.mean(dvals))
    compat = abs(dpsi_bao - dpsi_th_need) < 0.010    # do BAO and theta* want the same screen?
    print(f"  BAO prefers screen Dpsi~{dpsi_bao:+.3f}; theta* needs Dpsi(z*)~{dpsi_th_need:+.3f}"
          f"  -> {'COMPATIBLE' if compat else 'INCOMPATIBLE (different amplitudes)'}")
    print(f"  applying the BAO screen to z* gives 100*theta*={100*th_with_bao_screen:.4f} "
          f"(Planck {100*P.PLANCK['theta_star']:.4f}, resid {th_resid:+.2f}%)")
    print(f"  low-z BAO screen mean {np.mean(dvals):+.3f} +/- {np.std(dvals):.3f} "
          f"({'consistent' if consistent else 'scattered/sign-changing across z'})")
    print(f"  best screen drops BAO chi^2/N {chi2_0/12:.2f} -> {chi2_c/12:.2f} (LCDM {chi2_l/12:.2f})")
    if compat and chi2_c/12 <= 1.3 * chi2_l/12:
        print("  => one physical screen fits BAO AND theta*: a genuine TIE with LCDM.")
    else:
        print("  => PARTIAL: the screen helps BAO a lot but the BAO-screen OVER-corrects theta*")
        print("     (BAO and the CMB acoustic scale want different screen amplitudes). The high")
        print("     derived H0 + standard r_drag leaves a residual distance-SHAPE tension -- the")
        print("     same wall every high-H0/standard-r_drag model hits. NOT yet a clean tie; OPEN.")
        print("     Genuine wins stand: H0 derived (resolves the H0 tension axis), n_s derived,")
        print("     chi mass derived, fewer free params. Next lever: z-dependent screen or the")
        print("     perturbation-level G_eff (module 3), not a single constant bias.")
    return dict(chi2_noscreen=chi2_0, chi2_const=chi2_c, chi2_smooth=chi2_s, chi2_lcdm=chi2_l,
                dpsi_bao=dpsi_bao, dpsi_theta=dpsi_th_need, theta_resid_pct=th_resid)


def _lcdm_par():
    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 - P.PLANCK["Omega_m"] - OMEGA_R
        OMEGA_B_H2 = P.OMEGA_B_H2; OMEGA_G_H2 = P.OMEGA_G_H2
    return _PL


if __name__ == "__main__":
    main()
