import os
import tempfile
import unittest

import numpy as np

packages_path = os.environ.get("COBAYA_PACKAGES_PATH") or os.path.join(
    tempfile.gettempdir(), "SPT_packages"
)

cosmo_params = dict(
    cosmomc_theta=0.010411,
    As=1e-10 * np.exp(3.1),
    ombh2=0.0221,
    omch2=0.1200,
    ns=0.96,
    Alens=1.0,
    tau=0.09,
)

fg_params = dict(
    kappa=0.0,
    Dl_Poisson_90x90=0.1,
    Dl_Poisson_90x150=0.1,
    Dl_Poisson_90x220=0.1,
    Dl_Poisson_150x150=0.1,
    Dl_Poisson_150x220=0.1,
    Dl_Poisson_220x220=0.1,
    TDust=19.6,
    ADust_TE_150=0.1647,
    BetaDust_TE=1.59,
    AlphaDust_TE=-2.42,
    ADust_EE_150=0.0236,
    BetaDust_EE=1.59,
    AlphaDust_EE=-2.42,
    mapTcal90=1.0,
    mapTcal150=1.0,
    mapTcal220=1.0,
    mapPcal90=1.0,
    mapPcal150=1.0,
    mapPcal220=1.0,
)


class SPT3GTest(unittest.TestCase):
    def setUp(self):
        from cobaya.install import install

        install({"likelihood": {"spt3g_2020.TEEE": None}}, path=packages_path)

    def test_with_cosmomc_data(self):
        from spt3g_2020 import TEEE

        my_spt = TEEE(dict(packages_path=packages_path))

        # Read Dl generated by CosmoMC
        test_dir = os.path.dirname(__file__)
        dl_te, dl_ee = np.load(os.path.join(test_dir, "data", "dlte_ee_from_cosmomc.npy"))
        loglike = my_spt.loglike(dl_te, dl_ee, **fg_params)
        # self.assertAlmostEqual(-2 * loglike, 1160.1925138672325, 5)
        # New number from SPT3G v3 release
        self.assertAlmostEqual(-2 * loglike, 1142.999285, 5)

    def test_cobaya(self):
        """Test the Cobaya interface to the SPT3G likelihood."""
        from cobaya.model import get_model

        info = {
            "debug": True,
            "likelihood": {"spt3g_2020.TEEE": None},
            "theory": {"camb": {"extra_args": {"lens_potential_accuracy": 1}}},
            "params": {**cosmo_params, **fg_params},
            "packages_path": packages_path,
        }

        model = get_model(info)
        chi2 = -2 * model.loglike({})[0]
        self.assertAlmostEqual(chi2, 1143.009, 2)


if __name__ == "__main__":
    unittest.main()
