# AUTOGENERATED! DO NOT EDIT! File to edit: 01_foregrounds.ipynb (unless otherwise specified).

__all__ = ['emission_naming', 'load_sky_emission']

# Cell

import healpy as hp
import numpy as np

from pathlib import Path
import logging as log

from .core import get_telescope, base_folder

# Cell

emission_naming = dict(
    foreground_emission="foreground",
    CMB_unlensed="cmb-unlensed",
    CMB_lensing_signal="cmb-lensing",
    CMB_tensor_to_scalar_ratio="cmb-tensors",
)

# Cell


def load_sky_emission(config_sky_emission, site, channel):
    """Load foreground maps for a channel

    Parameters
    ----------
    config_sky_emission : dict
        CMB-S4 sky emission configuration,
        generally config['sky_emission'] for a configuration
        loaded from a TOML file
        dictionary with standard emission names and their weights
    site : str
        'Pole' or 'Chile', case doesn't matter
    channel : str
        Channel tag, e.g. 'MFHS1'

    Returns
    -------
    output_map : numpy array
        Output map with all emissions combined, uses nan for missing pixels
    """
    telescope = get_telescope(channel)
    nside = 512 if telescope == "SAT" else 4096
    npix = hp.nside2npix(nside)
    output_map = np.zeros((3, npix), dtype=np.float32)
    realization = 0  # foregrounds are deterministic
    for emission, weight in config_sky_emission.items():
        if weight == 0:
            log.info("Skip %s", emission)
            continue
        log.info("Processing %s", emission)

        emission_map = hp.read_map(
            Path(base_folder)
            / f"{realization:08d}"
            / f"{site.lower()}_{emission_naming[emission]}_{telescope}_{channel}_filtered_telescope_all_time_all_bmap.fits",
            (0, 1, 2),
            verbose=False,
        )
        emission_map[emission_map == hp.UNSEEN] = np.nan
        if emission == "CMB_tensor_to_scalar_ratio":
            weight /= 0.01  # r of the input maps
        output_map += emission_map * weight

    output_map /= 1e6  # uK -> K
    return output_map