from copy import deepcopy
from typing import Any

import numpy as np
import pytest

from cobaya.likelihood import Likelihood, LikelihoodInterface
from cobaya.model import get_model
from cobaya.theory import Theory
from cobaya.typing import InputDict, ParamDict, ParamsDict, empty_dict

from .common import process_packages_path
from .conftest import install_test_wrapper
from .test_cosmo_camb import get_camb

# Test separating out the BBN consistency constraint into separate theory code,
# using CAMB's BBN interpolator class. Tests dependencies/multi-theory with one
# agnostic theory


class BBN(Theory):
    bbn: Any

    def get_requirements(self):
        return {"ombh2", "nnu"}

    def calculate(self, state, want_derived=True, **params_values_dict):
        yhe = self.bbn.Y_He(
            self.provider.get_param("ombh2"), self.provider.get_param("nnu") - 3.044
        )
        state["derived"] = {"YHe": yhe}

    def get_can_provide_params(self):
        return ["YHe"]


class BBN2(Theory):
    params = {"ombh2": None, "nnu": None, "YHe": {"derived": True}}
    bbn: Any

    def calculate(self, state, want_derived=True, **params_values_dict):
        if want_derived:
            state["derived"] = {
                "YHe": self.bbn.Y_He(
                    params_values_dict["ombh2"], params_values_dict["nnu"] - 3.044
                )
            }


def cmb_likelihood(_self):
    results = _self.provider.get_CAMBdata()
    derived = {"check": results.Params.YHe}
    return results.Params.YHe, derived


cmb_likelihood_info = {
    "external": cmb_likelihood,
    "output_params": "check",
    "requires": {"Hubble": {"z": [0.5]}, "CAMBdata": None},
}

camb_params: ParamsDict = {
    "ombh2": 0.022274,
    "omch2": 0.11913,
    "cosmomc_theta": 0.01040867,
    "As": 0.2132755716e-8,
    "ns": 0.96597,
    "tau": 0.0639,
    "nnu": 3.044,
}

bbn_table = "PRIMAT_Yp_DH_Error.dat"
debug = True
info: InputDict = {
    "likelihood": {"cmb": cmb_likelihood_info},
    "theory": {
        "camb": {
            "extra_args": {"lens_potential_accuracy": 1},
            "requires": ["YHe", "ombh2"],
        },
        "bbn": {"external": BBN, "provides": ["YHe"]},
    },
    "params": camb_params,
    "debug": debug,
    "stop_at_error": True,
}

info2: InputDict = {
    "likelihood": {"cmb": cmb_likelihood_info},
    "theory": {"camb": {"requires": ["YHe", "ombh2"]}, "bbn": BBN2},
    "params": camb_params,
    "debug": debug,
}


def test_bbn_yhe(packages_path, skip_not_installed):
    packages_path = process_packages_path(packages_path)
    camb = install_test_wrapper(skip_not_installed, get_camb, packages_path)

    BBN.bbn = camb.bbn.BBN_table_interpolator(bbn_table)
    BBN2.bbn = BBN.bbn

    info["params"]["check"] = {"derived": True}

    for inf in (info, info2):
        inf["packages_path"] = packages_path
        for order in [1, -1]:
            for explicit_derived in [None, None, {"derived": True}]:
                print(inf, order, explicit_derived)
                model = get_model(inf)
                _, derived = model.loglikes({})
                vals = set([BBN.bbn.Y_He(camb_params["ombh2"])] + derived)
                assert len(vals) == 1, f"wrong Yhe value: {vals}"
                inf["params"]["YHe"] = explicit_derived
            inf["params"].pop("YHe")
            inf["theory"] = {p: v for p, v in reversed(list(inf["theory"].items()))}


# Not inherit from BBN to derive likelihoods that account for the theory error


class BBN_likelihood(BBN2, LikelihoodInterface):
    """
    Sample YHe and just calculate a direct theory likelihood
    """

    params = dict.fromkeys(["ombh2", "nnu", "YHe"])

    def calculate(self, state, want_derived=True, **params_values_dict):
        ombh2 = params_values_dict["ombh2"]
        delta_neff = params_values_dict["nnu"] - 3.044
        yhemean = self.bbn.Y_He(ombh2, delta_neff)
        error = self.bbn.get("sig(Yp^BBN)", ombh2, delta_neff)
        state["logp"] = -((params_values_dict["YHe"] - yhemean) ** 2) / (2 * error**2)

    def get_can_provide_params(self):
        return {}


class BBN_with_theory_errors(BBN, LikelihoodInterface):
    """
    The BBN theory prediction has an error. So we can derive a likelihood for
    sampling over a separate error parameter to account for it if we approximate
    the distribution as Gaussian.
    """

    params = {
        "BBN_delta": {
            "prior": {"min": 0, "max": 1},
            "ref": {"dist": "norm", "loc": 0, "scale": 1},
        }
    }

    def calculate(self, state, want_derived=True, **params_values_dict):
        ombh2, nnu = self.provider.get_param(["ombh2", "nnu"])
        delta_neff = nnu - 3.044
        yhemean = self.bbn.Y_He(ombh2, delta_neff)
        # get theory error (neglecting difference between Y_He and Yp^BBN)
        error = self.bbn.get("sig(Yp^BBN)", ombh2, delta_neff)
        yhe = yhemean + error * params_values_dict["BBN_delta"]
        if want_derived:
            state["derived"] = {"YHe": yhe}
        # take the error to be Gaussian
        state["logp"] = -(params_values_dict["BBN_delta"] ** 2) / 2


info_error: InputDict = {
    "likelihood": {"cmb": cmb_likelihood_info, "BBN": BBN_likelihood},
    "theory": {"camb": {"requires": ["YHe", "ombh2"]}},
    "params": dict(YHe={"prior": {"min": 0, "max": 1}}, **camb_params),
    "debug": debug,
}

info_error2: InputDict = {
    "likelihood": {
        "cmb": cmb_likelihood_info,
        "BBN": {"external": BBN_with_theory_errors, "provides": "YHe"},
    },
    "theory": {
        "camb": {
            "requires": ["YHe", "ombh2"],
            "extra_args": {"bbn_predictor": "PArthENoPE_880.2_standard.dat"},
        }
    },
    "params": dict(BBN_delta={"prior": {"min": -5, "max": 5}}, **camb_params),
    "debug": debug,
}


def test_bbn_likelihood(packages_path, skip_not_installed):
    packages_path = process_packages_path(packages_path)
    camb = install_test_wrapper(skip_not_installed, get_camb, packages_path)

    BBN_likelihood.bbn = camb.bbn.BBN_table_interpolator(bbn_table)
    model = get_model(info_error, packages_path=packages_path)
    assert np.allclose(model.loglikes({"YHe": 0.246})[0], [0.246, -0.8696], rtol=1e-3), (
        f"Failed BBN likelihood with {info_error}"
    )

    # second case, BBN likelihood has to be calculated before CAMB
    BBN_with_theory_errors.bbn = BBN_likelihood.bbn
    model = get_model(info_error2, packages_path=packages_path)
    assert np.allclose(
        model.loglikes({"BBN_delta": 1.0})[0], [0.24594834, -0.5], rtol=1e-4
    )


class ExamplePrimordialPk(Theory):
    def initialize(self):
        # need to provide valid results at wide k range, any that might be used
        self.ks = np.logspace(-5.5, 2, 1000)

    def calculate(self, state, want_derived=True, **params_values_dict):
        pivot_scalar = 0.05
        pk = (self.ks / pivot_scalar) ** (
            params_values_dict["testns"] - 1
        ) * params_values_dict["testAs"]
        state["primordial_scalar_pk"] = {
            "kmin": self.ks[0],
            "kmax": self.ks[-1],
            "Pk": pk,
            "log_regular": True,
        }

    def get_primordial_scalar_pk(self):
        return self.current_state["primordial_scalar_pk"]

    def get_can_support_params(self):
        return ["testAs", "testns"]


testAs = 1.8e-9
testns = 0.8


class Pklike(Likelihood):
    def logp(self, **params_values):
        results = self.provider.get_CAMBdata()
        print(results.Params.scalar_power(1.1), testAs * (1.1 / 0.05) ** (testns - 1))
        assert np.allclose(
            results.Params.scalar_power(1.1),
            testAs * (1.1 / 0.05) ** (testns - 1),
            rtol=1e-3,
            atol=1e-20,
        )
        return 0

    def get_requirements(self):
        return {"Cl": {"tt": 1000}, "CAMBdata": None}


info_pk: InputDict = {
    "likelihood": {"cmb": Pklike},
    "theory": {"camb": {"external_primordial_pk": True}, "my_pk": ExamplePrimordialPk},
    "params": {
        "ombh2": 0.022274,
        "omch2": 0.11913,
        "cosmomc_theta": 0.01040867,
        "tau": 0.0639,
        "nnu": 3.044,
        "testAs": {"prior": {"min": 1e-9, "max": 1e-8}},
        "testns": {"prior": {"min": 0.8, "max": 1.2}},
    },
    "stop_at_error": True,
    "debug": debug,
}


def test_primordial_pk(packages_path, skip_not_installed):
    packages_path = process_packages_path(packages_path)
    info_pk["packages_path"] = packages_path
    model = install_test_wrapper(skip_not_installed, get_model, info_pk)
    model.loglikes({"testAs": testAs, "testns": testns})


ratio_amp = 2.0


class TrivialNonLinearRatio(Theory):
    def get_requirements(self):
        return "CAMB_transfers"

    def get_non_linear_ratio(self, results):
        assert np.isclose(results.Params.InitPower.ns, 0.9667)
        kh = results.Params.Transfer.kmax
        k_h = np.logspace(-4, np.log10(kh), 200)
        z = np.array(
            results.Params.Transfer.PK_redshifts[
                : results.Params.Transfer.PK_num_redshifts
            ]
        )
        z = np.sort(z)
        ratio = ratio_amp * np.ones((len(z), len(k_h)))
        return {"k_h": k_h, "z": z, "ratio": ratio}


class NonLinearRatioLike(Likelihood):
    def get_requirements(self):
        return {"Pk_grid": {"z": [0, 0.5, 1.0], "k_max": 10, "nonlinear": [False, True]}}

    def logp(self, **params_values):
        k_lin, z_lin, pk_lin = self.provider.get_Pk_grid(nonlinear=False)
        k_nonlin, z_nonlin, pk_nonlin = self.provider.get_Pk_grid(nonlinear=True)
        np.testing.assert_allclose(k_nonlin, k_lin)
        np.testing.assert_allclose(z_nonlin, z_lin)
        np.testing.assert_allclose(pk_nonlin, ratio_amp**2 * pk_lin, rtol=1e-4, atol=0)
        return 0


info_non_linear_ratio: InputDict = {
    "likelihood": {"like": NonLinearRatioLike},
    "theory": {
        "camb": {"use_non_linear_ratio": True},
        "my_nonlin": TrivialNonLinearRatio,
    },
    "params": {
        "ombh2": 0.022274,
        "omch2": 0.11913,
        "cosmomc_theta": 0.01040867,
        "tau": 0.0639,
        "ns": 0.9667,
        "As": 2.105e-9,
    },
    "stop_at_error": True,
    "debug": debug,
}


def test_trivial_non_linear_ratio(packages_path, skip_not_installed):
    packages_path = process_packages_path(packages_path)
    camb = install_test_wrapper(skip_not_installed, get_camb, packages_path)

    if not hasattr(camb.nonlinear, "ExternalNonLinearRatio"):
        pytest.skip("CAMB < 1.6.6 does not have ExternalNonLinearRatio")

    info_non_linear_ratio["packages_path"] = packages_path
    model = get_model(info_non_linear_ratio)
    model.loglikes({})


class BinnedPk(Theory):
    # example splined power spectrum exp(-2 tau)P(k) based on bin values.
    # Can pass dense sampling to CAMB to reproduce any function, here bins are directly
    # cubic spline values use by CAMB.
    # Note need to  have wide k bounds or latest CAMB (which takes value beyond start
    # and end bins to be equal to the end bins).
    nbins: int = 20
    k_min_bin: float = np.log10(0.001)
    k_max_bin: float = np.log10(0.35)
    scale: float = 1e-9
    bin_par: ParamDict = {"prior": {"min": 0, "max": 100}}

    def initialize(self):
        self.ks = np.logspace(self.k_min_bin, self.k_max_bin, self.nbins)

    def get_requirements(self):
        return {"tau"}

    def calculate(self, state, want_derived=True, **params_values_dict):
        pk = np.zeros_like(self.ks)
        for b in range(self.nbins):
            pk[b] = params_values_dict["b%s" % (b + 1)]
        pk *= self.scale * np.exp(2 * self.provider.get_param("tau"))
        # should use log_regular: True for speed if the binning is log regular
        # here test the non-regular option for coverage
        state["primordial_scalar_pk"] = {"k": self.ks, "Pk": pk, "log_regular": False}
        state["primordial_tensor_pk"] = {
            "k": self.ks,
            "Pk": pk * 0.1,
            "log_regular": False,
        }

    def get_primordial_scalar_pk(self):
        return self.current_state["primordial_scalar_pk"]

    def get_primordial_tensor_pk(self):
        return self.current_state["primordial_tensor_pk"]

    @classmethod
    def get_class_options(cls, input_options=empty_dict):
        # dynamically generate defaults for params based on nbins
        options = super().get_class_options().copy()
        nbins = input_options["nbins"]
        bin_par = input_options.get("bin_par", cls.bin_par)
        params = {}
        for b in range(nbins):
            par = deepcopy(bin_par)
            par["latex"] = "b_%s" % (b + 1)
            params["b%s" % (b + 1)] = par
        options["params"] = params
        return options


def test_pk_binning(packages_path, skip_not_installed):
    # reproduce power law by sending in spline point values
    # has to be fine sampling to get to 1e-3 precision in test.
    nbins = 40
    tau = 0.05
    k_min_bin = -5.5
    k_max_bin = 2

    _info: InputDict = {
        "packages_path": process_packages_path(packages_path),
        "likelihood": {"cmb": Pklike},
        "theory": {
            "camb": {"external_primordial_pk": True, "extra_args": {"WantTensors": True}},
            "my_pk": {
                "external": BinnedPk,
                "nbins": nbins,
                "k_min_bin": k_min_bin,
                "k_max_bin": k_max_bin,
            },
        },
        "params": {
            "ombh2": 0.022274,
            "omch2": 0.11913,
            "cosmomc_theta": 0.01040867,
            "tau": tau,
            "nnu": 3.044,
        },
        "stop_at_error": True,
        "debug": debug,
    }
    scale = 1e-9
    ks = np.logspace(k_min_bin, k_max_bin, nbins)

    def pk_test(k):
        return testAs * (k / 0.05) ** (testns - 1) / scale * np.exp(-2 * tau)

    pars = {"b%s" % (b + 1): pk_test(ks[b]) for b in range(nbins)}
    model = install_test_wrapper(skip_not_installed, get_model, _info)
    model.loglikes(pars)
