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

__all__ = ['get_telecope_years', 'load_atmosphere']

# Cell

import healpy as hp
import numpy as np
from pathlib import Path
import logging as log

from .core import get_telescope, base_folder, simulations_observing_efficiency

# Cell


def get_telecope_years(config, site, channel):
    """Compute the number of telescope/years in the CMB-S4 configuration

    config_telescopes : dict
        CMB-S4 telescopes configuration,
        generally loaded from a TOML file
    site : str
        'Pole' or 'Chile', case doesn't matter
    channel : str
        Channel tag, e.g. 'MFHS1'
    """
    telescope_years = 0
    for tube, bands in config["telescopes"][get_telescope(channel)].items():
        if bands["site"].lower() == site.lower():
            has_band = bands.get(channel[:-1], 0) > 0
            telescope_years += has_band * bands.get(
                "years", config["experiment"]["total_experiment_length_years"]
            )
    return telescope_years

# Cell


def load_atmosphere(config, site, channel, realization=0):
    """Load foreground maps for a channel

    Parameters
    ----------
    config : dict
        CMB-S4 configuration,
        generally loaded from a TOML file
    site : str
        'Pole' or 'Chile', case doesn't matter
    channel : str
        Channel tag, e.g. 'MFHS1'
    realization : int
        Choose one of the available 8 realizations

    Returns
    -------
    output_map : numpy array
        Output map with all emissions combined, uses nan for missing pixels
    """

    telescope = get_telescope(channel)
    map_filename = (
        Path(f"{realization:08d}")
        / f"{site.lower()}_atmosphere_{telescope}_{channel}_filtered_telescope_all_time_all_bmap.fits"
    )
    log.info(f"Reading {map_filename}")
    atmosphere_map = hp.read_map(
        Path(base_folder) / map_filename, (0, 1, 2), dtype=np.float32, verbose=False
    )
    atmosphere_map[atmosphere_map == hp.UNSEEN] = np.nan
    # input map is 10 days at 100% efficiency
    atmosphere_map *= np.sqrt(10) / np.sqrt(
        365.25 * config["experiment"]["observing_efficiency"]
    )
    # pole sims have lower efficiency
    atmosphere_map *= np.sqrt(
        simulations_observing_efficiency[site.lower()].get(telescope, 1)
    )
    atmosphere_map /= np.sqrt(get_telecope_years(config, site, channel))
    return atmosphere_map