"""
Figure for Appendix LS: the DFD late-time LOS optical-screen amplitude.

LEFT  panel: amplitude vs LOS depth -- the cumulative screen Delta_psi(<chi) built up by the
             halo-model 1-halo + 2-halo contributions (Strand 1, los_screen_2halo.py), with the
             target 0.030 and data-preferred 0.039 marked.  Shows the 1h/2h decomposition.
RIGHT panel: the HONEST sensitivity -- amplitude as a function of the UNDERIVED MOND external-field
             (EFE) floor x_ext (Strand 3, halo_los_2halo_efe.py).  The two corpus-named floors
             (cosmological deep-MOND x_ext->0 ; galactic/cluster Hubble-EFE x_ext=5.85) bracket a
             ~10x swing; the in-range corner where 0.030/0.039 are met is shaded.

No tuning: every curve is the direct output of the strand scripts at their pre-registered settings.
"""
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

import los_screen_2halo as S1            # halo model (1h+2h), fiducial
import halo_los_2halo_efe as S3          # EFE-floor sensitivity

TARGET = 0.030
DATA_PREF = 0.039
BAND = (0.023, 0.045)

# ----------------------------------------------------------------------------- LEFT: depth build-up
# Reconstruct the cumulative screen vs LOS depth from the Strand-1 fiducial.
# Delta_psi = kappa_struct * Q ; kappa_struct = kappa_1h + kappa_2h.  We distribute the amplitude
# along the LOS using the lensing efficiency kernel W(chi)=chi(chi_s-chi)/chi_s of the fiducial run,
# so the curve is the genuine cumulative C_l-weighted convergence the halo model produces, split 1h/2h.
fid = S1.halo_model_amplitude(verbose=False)
Dpsi_fid, kappa_struct, Qb, bg, f1h, xbar = fid
kap1h = kappa_struct * f1h
kap2h = kappa_struct * (1 - f1h)
amp_1h = kap1h * Qb
amp_2h = kap2h * Qb

chis = S1.chi_of_z(S1.Z)
chi = np.linspace(1.0, chis - 1.0, 200)
W = chi * (chis - chi) / chis                 # lensing efficiency along the slab
w = W / np.trapezoid(W, chi)                  # normalized build-up weight
cum = np.array([np.trapezoid(w[:i + 1], chi[:i + 1]) for i in range(len(chi))])
cum /= cum[-1]
cum_1h = cum * amp_1h
cum_2h = cum * amp_2h
cum_tot = cum_1h + cum_2h

# ----------------------------------------------------------------------------- RIGHT: EFE-floor scan
xext = np.array([0.0, 0.05, 0.1, 0.15, 0.2, 0.24, 0.3, 0.5, 1.0, 2.0, S3.X_HUBBLE])
kappa_gal = 1.8 * 1.0 * 0.0029                 # Strand-3 measured kappa_gal
amp_efe = np.array([kappa_gal * S3.toy_Q(xe)[0] for xe in xext])

fig, (axL, axR) = plt.subplots(1, 2, figsize=(11.6, 4.7))

# ----- LEFT -----
axL.fill_between(chi, 0, cum_1h, color="#2a6f97", alpha=0.85, label="1-halo")
axL.fill_between(chi, cum_1h, cum_tot, color="#89c2d9", alpha=0.9, label="2-halo")
axL.plot(chi, cum_tot, color="#013a63", lw=2.0)
axL.axhline(TARGET, color="#d62828", ls="--", lw=1.6, label=f"target 0.030")
axL.axhline(DATA_PREF, color="#7209b7", ls=":", lw=1.8, label=f"data-preferred 0.039")
axL.axhspan(BAND[0], BAND[1], color="grey", alpha=0.12)
axL.text(chi[3], Dpsi_fid + 0.0015,
         f"halo-model fiducial = {Dpsi_fid:.3f}\n(x$_{{\\rm ext}}$=0 limit)",
         fontsize=8.5, color="#013a63")
axL.set_xlabel("LOS comoving depth  $\\chi$  [Mpc]")
axL.set_ylabel("cumulative screen  $\\Delta\\psi_{\\rm BAO}(<\\chi)$")
axL.set_title("(a) Halo-model LOS build-up (1h+2h), $z\\simeq0.5$", fontsize=10.5)
axL.set_ylim(0, max(Dpsi_fid, DATA_PREF) * 1.18)
axL.set_xlim(chi[0], chi[-1])
axL.legend(fontsize=8, loc="lower right", framealpha=0.9)
axL.grid(alpha=0.25)

# ----- RIGHT -----
axR.plot(xext, amp_efe, "o-", color="#013a63", lw=2.0, ms=5, label="halo-model amplitude")
axR.axhline(TARGET, color="#d62828", ls="--", lw=1.6, label="target 0.030")
axR.axhline(DATA_PREF, color="#7209b7", ls=":", lw=1.8, label="data-preferred 0.039")
# in-range corner: x_ext where amplitude is between target and ~data-pref
axR.axvspan(0.031, 0.107, color="#90be6d", alpha=0.22)
axR.text(0.058, DATA_PREF * 0.62, "in-range\ncorner", fontsize=8, color="#386641", rotation=0)
# corpus-named floors
axR.axvline(S3.X_HUBBLE, color="grey", ls="-.", lw=1.3)
axR.text(S3.X_HUBBLE * 0.5, 0.0065, "corpus Hubble-EFE\n$x_{\\rm ext}=5.85$ (galactic)",
         fontsize=7.8, color="grey")
axR.text(0.0008, amp_efe[0] * 0.93, "cosmological\ndeep-MOND\n$x_{\\rm ext}\\to0$",
         fontsize=7.8, color="#013a63")
axR.set_xscale("symlog", linthresh=0.05)
axR.set_xlabel("MOND external-field floor  $x_{\\rm ext}=g_{\\rm ext}/a_0$  (UNDERIVED)")
axR.set_ylabel("screen amplitude  $\\Delta\\psi_{\\rm BAO}$")
axR.set_title("(b) Amplitude vs the undetermined EFE floor", fontsize=10.5)
axR.set_ylim(0, max(amp_efe[1:]) * 1.18)
axR.legend(fontsize=8, loc="upper right", framealpha=0.9)
axR.grid(alpha=0.25, which="both")

fig.suptitle("DFD late-time line-of-sight optical screen: amplitude, structure decomposition, "
             "and EFE-floor sensitivity", fontsize=11.5, y=1.005)
fig.tight_layout()
out = "fig_los_screen.png"
fig.savefig(out, dpi=160, bbox_inches="tight")
print(f"wrote {out}")
print(f"  fiducial(x_ext=0)  Delta_psi = {Dpsi_fid:.4f}   [1h {amp_1h:.4f} + 2h {amp_2h:.4f}]")
print(f"  EFE-floor scan      amp(x=0)={amp_efe[0]:.4f}  amp(x=0.24)={kappa_gal*S3.toy_Q(0.24)[0]:.4f}  "
      f"amp(x=5.85)={amp_efe[-1]:.4f}")
