# 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, mapmaking_naming, read_instrument_model

# Cell

emission_naming = dict(
    foreground_emission="combined_foregrounds",
    CMB_unlensed="cmb_unlensed_solardipole",
    CMB_lensing_signal="cmb_lensing_signal",
    CMB_tensor_to_scalar_ratio="cmb_tensor",
)

# 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
    """
    log.info("Configuration %s", str(config_sky_emission))
    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
    channel_noP = channel.replace("P", "")
    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_noP}_{mapmaking_naming[telescope]}",
            (0, 1, 2),
            dtype=None,
            verbose=False,
        )
        emission_map[emission_map == 0] = np.nan
        emission_map[emission_map == hp.UNSEEN] = np.nan

        if emission == "CMB_tensor_to_scalar_ratio":
            # input maps are simulated with r=3e-3
            # tensor-to-scalar ratio is defined in the power spectrum, we need to weight the maps with the square root
            weight = np.sqrt(weight / 3e-3)
        output_map += emission_map * weight

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