"""
rlmt_cobaya_theory.py  (v5 - cache-safe, NonLinear disabled)
"""
import numpy as np
from cobaya.theory import Theory
from cobaya.log import LoggedError
import camb
from rlmt_primordial import make_camb_initial_power, DEFAULT_PARAMS

_ALL_PARAMS = [
    "H0", "ombh2", "omch2", "tau",
    "As", "ns", "beta_eff", "c1", "c2", "c3",
    "k_star", "kc", "eps3", "k0", "p_exp",
]
_DEFAULTS = {
    "H0": 67.36, "ombh2": 0.02237, "omch2": 0.1200, "tau": 0.0544,
    "As": 2.1e-9, "ns": 0.965, "beta_eff": 1.0,
    "c1": 0.5, "c2": 1.0, "c3": 2.0,
    "k_star": 0.05, "kc": 1.0, "eps3": -0.05,
    "k0": 0.05, "p_exp": 2.0,
}


class RLMTCAMB(Theory):
    params = {k: None for k in _ALL_PARAMS}  # no type annotation
    lmax: int = 2600
    lens_potential_accuracy: int = 1

    def initialize(self):
        self._lmax = int(self.lmax)

    def get_requirements(self):
        return {}

    def must_provide(self, **requirements):
        return {}

    def get_can_provide(self):
        return {"Cl": {"tt": self._lmax, "te": self._lmax,
                       "ee": self._lmax, "bb": self._lmax,
                       "pp": self._lmax}}

    def get_can_provide_params(self):
        return []

    def calculate(self, state, want_derived=True, **params_values_dict):
        try:
            p = dict(_DEFAULTS)
            p.update({k: v for k, v in params_values_dict.items()
                      if k in _ALL_PARAMS})
            rlmt_params = dict(DEFAULT_PARAMS)
            rlmt_params.update(p)

            pars = camb.CAMBparams()
            pars.set_cosmology(
                H0=p["H0"], ombh2=p["ombh2"],
                omch2=p["omch2"], tau=p["tau"],
            )
            # Disable nonlinear: primary CMB does not need Halofit,
            # and SplinedInitialPower requires effective_ns_for_nonlinear
            # only when NonLinear != NonLinear_none.
            pars.NonLinear = camb.model.NonLinear_none

            ip = make_camb_initial_power(rlmt_params)
            pars.InitPower = ip

            # set_for_lmax AFTER NonLinear and InitPower
            pars.set_for_lmax(self._lmax,
                              lens_potential_accuracy=self.lens_potential_accuracy)
            # Re-apply NonLinear after set_for_lmax (it may reset settings)
            pars.NonLinear = camb.model.NonLinear_none

            results = camb.get_results(pars)
            powers  = results.get_cmb_power_spectra(
                pars, CMB_unit="muK", raw_cl=False)

            total = powers["total"]
            lp    = powers.get("lens_potential", None)
            ells  = np.arange(total.shape[0])

            state["Cl"] = {
                "ell": ells,
                "tt":  total[:, 0],
                "ee":  total[:, 1],
                "bb":  total[:, 2],
                "te":  total[:, 3],
                "pp":  (lp[:, 0] if lp is not None
                        else np.zeros_like(ells, dtype=float)),
            }
            return True

        except Exception as e:
            raise LoggedError(self.log, f"RLMTCAMB.calculate failed: {e}")

    def get_Cl(self, ell_factor=False, units="muK2"):
        return self.current_state["Cl"]


assert isinstance(RLMTCAMB.params, dict)
assert "H0" in RLMTCAMB.params