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

__all__ = [
    "simulations_sampling_frequency_scaling",
    "load_hitmap_wcov",
    "hitmaps_naming",
    "wcov_naming",
]

# 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,
)
from .noise import get_thinfp, get_tube_years

# Cell

simulations_sampling_frequency_scaling = dict(SAT=5, LAT=2)

# Cell

hitmaps_naming = {
    "SAT": "telescope_all_time_all_hits.fits.gz",
    "LAT": "filtered_telescope_all_time_all_hmap.fits",
}
wcov_naming = {
    "SAT": "telescope_all_time_all_wcov.fits.gz",
    "LAT": "filtered_telescope_all_time_all_wcov.fits",
}


def load_hitmap_wcov(config, site, channel, realization=0, raw_hitmap=False):
    """Load hitmaps and white noise covariance matrices for a channel

    This loads the simulated hitmaps and white noise covariance matrices
    and scales them properly to the experiment configuration and duration
    as defined in the input config file.
    Hitmaps assumes a sampling frequency of 100 Hz for SAT and 400 Hz for
    LAT.

    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
    -------
    hitmap : numpy array
        Hitmap for all channels all tubes combined
    wcov : numpy array
        White noise covariance matrix, rows are:
        "II", "IQ", "IU", "QQ", "QU", "UU", units are K^2
    """

    # it is the same scaling for hitmap and white noise covariance matrix,
    # which is the same as noise except squared

    telescope = get_telescope(channel)
    tube = channel[:-1]

    channel_noP = channel.replace("P", "")

    map_filename = (
        Path(f"{realization:08d}")
        / f"{site.lower()}_noise_{telescope}_{channel_noP}_{hitmaps_naming[telescope]}"
    )
    wcov_filename = (
        Path(f"{realization:08d}")
        / f"{site.lower()}_noise_{telescope}_{channel_noP}_{wcov_naming[telescope]}"
    )
    log.info(f"Base folder: {base_folder}")
    log.info(f"Reading {map_filename}")
    hitmap = hp.read_map(Path(base_folder) / map_filename, 0, dtype=None, verbose=False)

    hitmap[hitmap == hp.UNSEEN] = 0
    if raw_hitmap:
        return hitmap
    wcov = hp.read_map(
        Path(base_folder) / wcov_filename, (0, 1, 2, 3, 4, 5), dtype=None, verbose=False
    )
    scaling = (
        365.25
        * get_observing_efficiency(
            config["experiment"]["observing_efficiency"], site, telescope, channel
        )
        * get_observing_efficiency(
            config["experiment"]["sensitivity_factor"], site, telescope, channel
        )
    ) / (10 * simulations_observing_efficiency[site.lower()][channel])
    # focal plane thinning factor of TOD simulations
    scaling *= get_thinfp(channel)
    scaling *= get_tube_years(config, site, channel) / simulated_tubes[tube]
    hitmap = np.round(hitmap * scaling).astype(np.int64)
    hitmap *= simulations_sampling_frequency_scaling[telescope]
    wcov /= scaling
    wcov[:, hitmap == 0] = hp.UNSEEN
    return hitmap, wcov
