# 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,
    get_observing_efficiency,
    base_folder,
    simulations_observing_efficiency,
    mapmaking_naming,
    read_instrument_model,
)

from .noise import get_thinfp

# 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 telescope_name, telescope_config in config["telescopes"][
        get_telescope(channel)
    ].items():
        if telescope_config["site"].lower() == site.lower():
            has_band = telescope_config.get(channel[:-1], 0) > 0
            telescope_years += has_band * telescope_config.get(
                "years", config["experiment"]["total_experiment_length_years"]
            )
    return telescope_years


# Cell


def load_atmosphere(config, site, channel, realization=0, raw=False):
    """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)
    channel_noP = channel.replace("P", "")

    map_filename = (
        Path(f"{realization:08d}")
        / f"{site.lower()}_atmosphere_{telescope}_{channel_noP}_{mapmaking_naming[telescope]}"
    )
    log.info(f"Reading {map_filename}")
    atmosphere_map = hp.read_map(
        Path(base_folder) / map_filename, (0, 1, 2), dtype=None, verbose=False
    )
    if raw:
        atmosphere_map[atmosphere_map == 0] = hp.UNSEEN
        return atmosphere_map

    atmosphere_map[atmosphere_map == hp.UNSEEN] = np.nan
    atmosphere_map[atmosphere_map == 0] = np.nan

    # input map is 10 days at 100% efficiency
    atmosphere_map *= np.sqrt(
        10
        * simulations_observing_efficiency[site.lower()][channel]
        / (
            365.25
            * get_observing_efficiency(
                config["experiment"]["observing_efficiency"], site, telescope, channel
            )
        )
    )
    atmosphere_map /= np.sqrt(get_telecope_years(config, site, channel))

    atmosphere_map[0] *= config["experiment"].get("atmosphere_scaling_T", 1)
    atmosphere_map[1:] *= config["experiment"].get("atmosphere_scaling_P", 1)
    atmosphere_map[1:] /= np.sqrt(get_thinfp(channel))

    return atmosphere_map
