"""
external_ns_prior.py
-----------------------
Instead of adding a new k-range probe, this operationalises what DESI and
SKA actually do in practice (per DESI 2024 VII, arXiv:2411.12022 -- "a
loose prior on the scalar spectral index ns"; and the SKA Red Book 2018,
arXiv:1811.02743 -- cosmological/shape parameters held fixed to Planck in
their Fisher forecasts): import an EXTERNAL, independently-obtained ns
constraint rather than trying to measure ns and eps3 jointly from the
same k-range.

Physical reason this works where a same-k-range probe (LSS/21cm) did not:
the bottleneck was not raw statistical weight, it was the INTERNAL
eps3-ns degeneracy within CMB-S4's own joint fit. An ns value measured
under an independent assumption (e.g. Planck's own tilt measurement,
which does not carry the RLMT eps3-ns degeneracy because Planck's
baseline analysis assumes a pure power law) breaks that degeneracy from
outside, in a way that adding more same-shape data at a different k
cannot.

Caveat: this is only as trustworthy as the assumption that an externally
measured ns (from an analysis that assumes a featureless power law) can
be legitimately imported as a prior on the RLMT model's ns parameter.
If eps3 != 0, the "true" ns inferred by an external power-law fit is
itself a biased/projected version of the RLMT ns -- so this prior is
approximate, not rigorous. A fully consistent treatment would need the
external experiment's own likelihood re-fit under the RLMT model.
"""

import numpy as np

def add_gaussian_prior(F, param_index, sigma_prior):
    Fp = F.copy()
    Fp[param_index, param_index] += 1.0 / sigma_prior**2
    return Fp


if __name__ == "__main__":
    F_full = np.load("F_threshold.npy")
    names8 = ["kc", "eps3", "ns", "As", "beta_eff", "c1", "c2", "c3"]
    varied = ["kc", "eps3", "ns", "As"]
    keep = [names8.index(n) for n in varied]
    F_cmb4 = F_full[np.ix_(keep, keep)]

    idx_ns = varied.index("ns")
    idx_eps3 = varied.index("eps3")

    print("CMB-S4 (TT+TE+EE) joint fit only, amplitude coeffs fixed:")
    s0 = np.sqrt(np.diag(np.linalg.inv(F_cmb4)))
    print(f"  sigma(ns)={s0[idx_ns]:.4f}  sigma(eps3)={s0[idx_eps3]:.4f}  "
          f"eps3/sigma={0.02/s0[idx_eps3]:.2f}")

    print("\nWith an external ns prior (mimicking DESI/SKA practice of "
          "importing ns rather than self-measuring it):")
    for label, sig in [("Planck 2018 (TT+TE+EE+lowE+lensing), sigma(ns)=0.0042", 0.0042),
                        ("Illustrative DESI-quoted 'loose' prior, sigma(ns)~0.01", 0.01),
                        ("Idealized single-parameter CMB-S4 ns prior, sigma(ns)=0.002", 0.002)]:
        Fp = add_gaussian_prior(F_cmb4, idx_ns, sig)
        s = np.sqrt(np.diag(np.linalg.inv(Fp)))
        print(f"  {label}")
        print(f"    -> sigma(ns)={s[idx_ns]:.4f}  sigma(eps3)={s[idx_eps3]:.4f}  "
              f"eps3/sigma={0.02/s[idx_eps3]:.2f}")