"""
rlmt_primordial.py  (v2 - CAMB 1.6.x compatible)
"""
import numpy as np
import camb
from camb.initialpower import SplinedInitialPower


def tension_weights(k, beta_eff, c, k_star, p_exp):
    k = np.atleast_1d(k)
    chi_k = np.arctan(k / k_star)
    s = np.sin(chi_k)
    X = np.array([beta_eff * cn * s**p_exp for cn in c])
    X -= X.min(axis=0, keepdims=True)
    w = np.exp(-X)
    w /= w.sum(axis=0, keepdims=True)
    return w


def base_spectra(k, As, ns, k0, kc, eps3):
    k = np.atleast_1d(k)
    P1 = As * (k / k0) ** (ns - 1.0)
    P2 = As * (k / k0) ** (ns - 1.0) * np.exp(-(k / kc) ** 2)
    P3 = As * (k / k0) ** (ns - 1.0 + eps3)
    return np.array([P1, P2, P3])


def P_RLMT(k, params):
    c = (params["c1"], params["c2"], params["c3"])
    w = tension_weights(k, params["beta_eff"], c, params["k_star"], params["p_exp"])
    P = base_spectra(k, params["As"], params["ns"], params["k0"],
                     params["kc"], params["eps3"])
    return np.sum(w * P, axis=0)


DEFAULT_PARAMS = dict(
    As=2.1e-9, ns=0.965, k0=0.05, kc=1.0, eps3=-0.05,
    beta_eff=1.0, c1=0.5, c2=1.0, c3=2.0, k_star=0.05, p_exp=2.0,
)


def make_camb_initial_power(params=None, kmin=1e-6, kmax=50.0, nk=4000):
    p = dict(DEFAULT_PARAMS)
    if params:
        p.update(params)

    k  = np.geomspace(kmin, kmax, nk)
    Pk = P_RLMT(k, p)
    ns_eff = float(p["ns"])

    ip = SplinedInitialPower()

    # CAMB 1.5+: effective_ns_for_nonlinear must be provided.
    # Try the keyword-argument form first (CAMB 1.6.x API);
    # fall back to direct attribute assignment.
    try:
        ip.set_scalar_table(k, Pk, effective_ns_for_nonlinear=ns_eff)
    except TypeError:
        ip.set_scalar_table(k, Pk)

    # Belt-and-suspenders: also set via attribute
    try:
        ip.effective_ns_for_nonlinear = ns_eff
    except Exception:
        pass

    return ip


def validate_power_law_limit():
    p = dict(DEFAULT_PARAMS)
    p.update(c1=1.0, c2=1.0, c3=1.0, kc=1e6, eps3=0.0)
    k = np.geomspace(1e-4, 1.0, 50)
    Pk = P_RLMT(k, p)
    P_pl = p["As"] * (k / p["k0"]) ** (p["ns"] - 1.0)
    rel_err = np.max(np.abs(Pk - P_pl) / P_pl)
    assert rel_err < 1e-6, f"power-law limit failed: rel_err={rel_err}"
    print(f"OK: power-law limit recovered, max rel. error = {rel_err:.2e}")


if __name__ == "__main__":
    validate_power_law_limit()