"""
lss_shape_fisher_v2.py
------------------------
Extends cmbs4_forecast_extended.lss_shape_fisher with an optional
k-dependent suppression weight, to give a (still simplified, but less
optimistic) 21cm intensity-mapping forecast that accounts for:

  (1) Foreground wedge / removal of large-scale line-of-sight modes:
      approximated isotropically as a smooth high-pass filter
          W_fg(k) = 1 - exp(-(k/k_fg)^2)
      with k_fg ~ 0.02-0.05 Mpc^-1 a rough proxy for "foreground-dominated
      radial modes below this scale are unusable". This is a crude
      isotropic stand-in for the anisotropic foreground wedge in
      (k_perp, k_par) space -- a real forecast needs the full 2D wedge.

  (2) Nonlinear/small-scale information loss: approximated as a smooth
      low-pass filter
          W_nl(k) = exp(-(k/k_NL)^2)
      with k_NL ~ 0.2-0.3 Mpc^-1 (z-dependent nonlinear scale), which
      down-weights (rather than hard-cuts) modes approaching and beyond
      the nonlinear scale where linear Tegmark-Veff Fisher is no longer
      trustworthy.

Both are illustrative smooth proxies, not a substitute for a real
foreground-wedge + nonlinear-P(k)-covariance forecast, but they capture
the right qualitative behaviour: information is suppressed at both very
large and very small scales, leaving an effective usable window.
"""

import numpy as np
import camb
from camb import initialpower
import cmbs4_forecast as base


def suppression_weight(k, k_fg=None, k_nl=None):
    w = np.ones_like(k)
    if k_fg is not None:
        w *= (1 - np.exp(-(k / k_fg) ** 2))
    if k_nl is not None:
        w *= np.exp(-(k / k_nl) ** 2)
    return w


def lss_shape_fisher_weighted(fiducial, varied, step, kh_min=0.005, kh_max=0.5,
                               npoints=200, z=1.0, V_survey=20e9, n_g=5e-2,
                               k_fg=None, k_nl=None):
    def get_pk(params):
        pk = base.rlmt_power_spectrum(base.KGRID, params["As"], params["ns"],
                                       params["beta_eff"], params["c1"],
                                       params["c2"], params["c3"],
                                       params["kstar"], params["kc"],
                                       params["eps3"])
        pars = camb.CAMBparams()
        pars.set_cosmology(H0=params["H0"], ombh2=params["ombh2"],
                            omch2=params["omch2"], tau=params["tau"])
        pk_ini = initialpower.SplinedInitialPower()
        pk_ini.set_scalar_table(base.KGRID, pk)
        pk_ini.effective_ns_for_nonlinear = params["ns"]
        pars.InitPower = pk_ini
        pars.set_matter_power(redshifts=[z], kmax=kh_max * 2)
        pars.NonLinear = camb.model.NonLinear_none
        results = camb.get_results(pars)
        kh, _, pkm = results.get_matter_power_spectrum(minkh=kh_min, maxkh=kh_max,
                                                         npoints=npoints)
        return kh, pkm[0]

    kh, P0 = get_pk(fiducial)
    derivs = {}
    for name in varied:
        pp = dict(fiducial); pp[name] += step[name]
        pm = dict(fiducial); pm[name] -= step[name]
        _, Pp = get_pk(pp)
        _, Pm = get_pk(pm)
        derivs[name] = (Pp - Pm) / (2 * step[name])

    W = suppression_weight(kh, k_fg=k_fg, k_nl=k_nl)

    n = len(varied)
    F = np.zeros((n, n))
    dk = kh[1] - kh[0]
    for idx, k in enumerate(kh):
        Pk = P0[idx]
        Veff = (n_g * Pk / (1 + n_g * Pk)) ** 2 * V_survey
        prefac = k ** 2 * dk / (4 * np.pi ** 2) * Veff / 2 * W[idx] ** 2
        for i, ni in enumerate(varied):
            dlnPi = derivs[ni][idx] / Pk
            for j, nj in enumerate(varied):
                dlnPj = derivs[nj][idx] / Pk
                F[i, j] += prefac * dlnPi * dlnPj
    return F, kh, W