"""
future_probes_forecast.py
---------------------------
Combines CMB-S4 (TT+TE+EE, amplitude nuisance externally fixed) with three
probes offering genuinely different k-leverage on the RLMT eps3/kc
degeneracy: mu-distortions (PIXIE-like), Lyman-alpha forest P(k), and
21cm intensity mapping P(k).

Requires cmbs4_forecast.py, cmbs4_forecast_extended.py (for
lss_shape_fisher) and mu_distortion.py in the same directory.

Headline result (fiducial kc=0.30, eps3=0.02, threshold from the paper's
Outlook, amplitude coefficients beta_eff/c1/c2/c3 externally fixed):

  combo                        sigma(kc)   kc/sig   sigma(eps3)  eps3/sig
  CMB-S4 only                   0.0126      23.9      0.0731       0.27
  + mu-distortion (PIXIE)       0.0125      24.1      0.0725       0.28
  + Lyman-alpha forest          0.0081      37.1      0.0482       0.42
  + 21cm IM (optimistic)        0.0013     236.       0.0157       1.27
  + Lya + 21cm                  0.0013     239.       0.0156       1.28

Interpretation:
  - mu-distortions probe a very different k-range (~1-1e4 Mpc^-1) but
    carry only ONE number's worth of statistical weight (sigma_mu~1e-8)
    -- utterly outweighed by CMB-S4's thousands of TT/TE/EE multipoles.
    Different k-dependence alone is not sufficient; the new probe also
    needs comparable statistical weight (mode count) to actually move a
    joint fit dominated by CMB data.
  - Lyman-alpha forest (kmax~1 Mpc^-1) and especially 21cm intensity
    mapping (kmax~2 Mpc^-1, optimistic/futuristic survey assumptions)
    have enough independent k-modes to meaningfully shrink sigma(eps3),
    pushing eps3/sigma from 0.27 to ~1.3 at this optimistic 21cm
    configuration -- i.e. finally into marginally-interesting territory,
    though still well short of a clean detection.
  - The 21cm number in particular assumes linear-theory Tegmark Veff all
    the way to kmax=2 Mpc^-1, which is optimistic; realistic foreground
    removal / nonlinearity limits typically cap usable kmax around
    0.3-0.5 Mpc^-1, which would degrade these numbers substantially.
    Treat this as a best-case ceiling, not a realistic near-term forecast.
"""

import numpy as np
import cmbs4_forecast as base
import cmbs4_forecast_extended as ext
import mu_distortion as mu


def main():
    fid = dict(base.FIDUCIAL)
    fid["kc"], fid["eps3"] = 0.30, 0.02
    varied = ["kc", "eps3", "ns", "As"]
    step = {"kc": 0.009, "eps3": 0.005, "ns": 0.004828, "As": 2.102e-11}

    print("Building CMB-S4 (TT+TE+EE) baseline, amplitude coefficients fixed...")
    varied_full = ["kc", "eps3", "ns", "As", "beta_eff", "c1", "c2", "c3"]
    step_full = {**step, "beta_eff": 0.069, "c1": 0.0459, "c2": 0.0462, "c3": 0.075}
    orig = base.FIDUCIAL
    base.FIDUCIAL = fid
    F_full = base.build_fisher(varied_full, lmin=30, lmax_T=3000, lmax_P=5000, fsky=0.4)
    base.FIDUCIAL = orig
    keep = [varied_full.index(n) for n in varied]
    F_cmb = F_full[np.ix_(keep, keep)]

    print("Computing mu-distortion Fisher (PIXIE-like, sigma_mu=1e-8)...")
    F_mu, mu0, _ = mu.mu_distortion_fisher(fid, varied, step)
    print(f"  fiducial mu = {mu0:.3e}")

    print("Computing Lyman-alpha forest shape Fisher (kmax=1.0 Mpc^-1, z=2.5)...")
    F_lya = ext.lss_shape_fisher(fid, varied, step, kh_min=0.05, kh_max=1.0,
                                  npoints=150, z=2.5, V_survey=5e9, n_g=1e-2)

    print("Computing 21cm intensity-mapping shape Fisher "
          "(kmax=2.0 Mpc^-1, z=1.0, optimistic survey)...")
    F_21 = ext.lss_shape_fisher(fid, varied, step, kh_min=0.01, kh_max=2.0,
                                 npoints=200, z=1.0, V_survey=20e9, n_g=5e-2)

    combos = [
        ("CMB-S4 only", F_cmb),
        ("+ mu-distortion (PIXIE)", F_cmb + F_mu),
        ("+ Lyman-alpha forest", F_cmb + F_lya),
        ("+ 21cm IM (optimistic)", F_cmb + F_21),
        ("+ Lya + 21cm", F_cmb + F_lya + F_21),
        ("+ mu + Lya + 21cm (all)", F_cmb + F_mu + F_lya + F_21),
    ]
    print(f"\n{'combo':28s}  sigma(kc)   kc/sig   sigma(eps3)  eps3/sig")
    for label, F in combos:
        s = np.sqrt(np.diag(np.linalg.inv(F)))
        print(f"{label:28s}  {s[0]:.5f}   {fid['kc']/s[0]:6.1f}   "
              f"{s[1]:.5f}     {fid['eps3']/s[1]:.2f}")


if __name__ == "__main__":
    main()