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

__all__ = [
    "thinfp_string",
    "thinfp_table",
    "get_thinfp",
    "get_tube_years",
    "load_noise",
]

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

# Cell

thinfp_string = """telescope	tube	thinfp
LAT	ULFPL	1
LAT	LFL	4
LAT	LFPL	4
LAT	MFL	16
LAT	HFL	16
LAT	MFPL	16
LAT	HFPL	16
SAT	LFS	1
SAT	MFLS	4
SAT	MFHS	4
SAT	HFS	8
"""

from astropy.io import ascii

thinfp_table = ascii.read(thinfp_string)

# Cell


def get_thinfp(channel):
    """Get the focal plane thinning factor for noise simulations

    Parameters
    ----------
    channel : str
        CMB-S4 channel tag e.g. HFL2

    Returns
    -------
    thinfp : int
        thinning factor
    """
    return (thinfp_table[thinfp_table["tube"] == channel[:-1]])["thinfp"][0]


# Cell


def get_tube_years(config, site, channel):
    """Compute the number of tube/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'
    """
    tube_years = 0
    for telescope_name, telescope_config in config["telescopes"][
        get_telescope(channel)
    ].items():
        if telescope_config["site"].lower() == site.lower():
            num_tubes = telescope_config.get(channel[:-1], 0)
            tube_years += num_tubes * telescope_config.get(
                "years", config["experiment"]["total_experiment_length_years"]
            )
    return tube_years


# Cell


def load_noise(config, site, channel, realization=0):
    """Load noise 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()}_noise_{telescope}_{channel_noP}_{mapmaking_naming[telescope]}"
    )
    log.info(f"Base folder: {base_folder}")
    log.info(f"Reading {map_filename}")
    output_map = hp.read_map(
        Path(base_folder) / map_filename, (0, 1, 2), dtype=None, verbose=False
    )
    output_map[output_map == hp.UNSEEN] = np.nan
    output_map[output_map == 0] = np.nan

    # input map is 10 days
    output_map *= np.sqrt(
        10
        * simulations_observing_efficiency[site.lower()][channel]
        / (
            365.25
            * get_observing_efficiency(
                config["experiment"]["observing_efficiency"], site, telescope, channel
            )
            * get_observing_efficiency(
                config["experiment"]["sensitivity_factor"], site, telescope, channel
            )
        )
    )
    output_map /= np.sqrt(
        get_tube_years(config, site, channel) / simulated_tubes[channel[:-1]]
    )
    # focal plane thinning factor of TOD simulations
    output_map /= np.sqrt(get_thinfp(channel))
    return output_map
