"""
Illustrative mu-distortion (PIXIE-like) Fisher block.

mu-type spectral distortions are sourced by Silk damping of acoustic
waves in the primordial power spectrum between z~5e4 and z~2e6, i.e.
they integrate the *dimensionless* power spectrum Delta^2_zeta(k) =
k^3 P(k) / (2 pi^2) against a window function J(k) that is significant
over roughly 1 <~ k <~ 1e4 Mpc^-1 (e.g. Chluba, Khatri & Sunyaev 2012;
Chluba 2016 review).

CAVEAT: the window function below is a SCHEMATIC double-Gaussian
bandpass (peaked ~few x 1e2 Mpc^-1, spanning ~1-1e4 Mpc^-1) chosen only
to reproduce the qualitative k-range and shape of the real window, NOT
a precise reproduction of the literature fitting function. For a
publication-grade forecast, use the actual Green's-function window from
a Boltzmann/thermalisation code (e.g. CosmoTherm, SZpack, or the
Chluba/Khatri-Sunyaev analytic fits) instead of this placeholder.

sigma(mu) ~ 1e-8 is the commonly quoted PIXIE-class target sensitivity
(order-of-magnitude, e.g. Kogut et al. 2011) -- again treat as a
ballpark literature number, not a rigorously reforecast value.
"""
import numpy as np
import cmbs4_forecast as base


def window_mu(k, k_lo=1.0, k_hi=1.0e4, k_peak=3.0e2):
    """Schematic bandpass: rises from k_lo, falls off by k_hi."""
    return np.exp(-(k / k_hi) ** 2) * (1 - np.exp(-(k / k_lo) ** 2)) \
        * np.exp(-0.5 * (np.log(k / k_peak) / 1.5) ** 2) / \
        np.exp(-0.5 * (np.log(k_peak / k_peak) / 1.5) ** 2)  # normalise peak to ~1 envelope


def compute_mu(params, coeff=2.2, kmin=0.3, kmax=1e4, n=4000):
    k = np.logspace(np.log10(kmin), np.log10(kmax), n)
    Pk = base.rlmt_power_spectrum(k, params["As"], params["ns"], params["beta_eff"],
                                   params["c1"], params["c2"], params["c3"],
                                   params["kstar"], params["kc"], params["eps3"])
    Delta2 = k ** 3 * Pk / (2 * np.pi ** 2)
    J = window_mu(k)
    integrand = Delta2 * J
    lnk = np.log(k)
    mu = coeff * np.trapezoid(integrand, lnk)
    return mu


def mu_distortion_fisher(fiducial, varied, step, sigma_mu=1.0e-8, mu_lcdm_ref=2.0e-8):
    """
    The schematic window above is only trustworthy for its *relative*
    k-dependence, not its absolute normalisation. We therefore calibrate
    the overall amplitude so that a pure power-law spectrum (kc->inf,
    eps3->0, i.e. RLMT reducing to LCDM) reproduces the standard quoted
    LCDM mu-distortion amplitude mu_lcdm_ref ~ 2e-8 (e.g. Chluba & Sunyaev
    2012 order of magnitude for ns~0.965). The RLMT-specific k-dependence
    (from kc, eps3) then rides on top of that calibrated amplitude.
    """
    ref_params = dict(fiducial); ref_params["kc"] = 1e8; ref_params["eps3"] = 0.0
    calib = mu_lcdm_ref / compute_mu(ref_params)

    def mu_phys(p):
        return calib * compute_mu(p)

    mu0 = mu_phys(fiducial)
    n = len(varied)
    grad = np.zeros(n)
    for i, name in enumerate(varied):
        pp = dict(fiducial); pp[name] += step[name]
        pm = dict(fiducial); pm[name] -= step[name]
        grad[i] = (mu_phys(pp) - mu_phys(pm)) / (2 * step[name])
    F = np.outer(grad, grad) / sigma_mu ** 2
    return F, mu0, grad


if __name__ == "__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}
    F, mu0, grad = mu_distortion_fisher(fid, varied, step)
    print("mu0 (fiducial) =", mu0)
    print("gradients:", dict(zip(varied, grad)))
    print("Single-probe (mu only) sigma:", np.sqrt(np.diag(np.linalg.inv(F + 1e-30*np.eye(len(varied))))) if np.linalg.matrix_rank(F)==len(varied) else "F is rank-1 (mu is a single number, cannot constrain 4 params alone)")