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

__all__ = ['md5sum_string', 'md5sum_file', 'parse_channels', 's4_channels', 'merge_dict', 'parse_config',
           'S4DesSimTool', 'command_line_script']

# Cell

import os
import toml
import healpy as hp
hp.disable_warnings()

import numpy as np
import h5py
from pathlib import Path
import logging as log
from datetime import date

from .core import get_telescope

from s4_design_sim_tool import __version__

from .foregrounds import load_sky_emission
from .atmosphere import load_atmosphere, get_telecope_years
from .noise import load_noise
from .hitmap_wcov import load_hitmap_wcov

# Cell

import hashlib


def md5sum_string(string):
    return hashlib.md5(string.encode("utf-8")).hexdigest()


def md5sum_file(filename):
    """Compute md5 checksum of the contents of a file"""

    return md5sum_string(open(filename, "r").read())

# Cell

s4_channels = {
    "LAT": [
        "ULFPL1",
        "LFL1",
        "LFPL1",
        "LFL2",
        "LFPL2",
        "MFPL1",
        "MFL1",
        "MFL2",
        "MFPL2",
        "HFL1",
        "HFPL1",
        "HFL2",
        "HFPL2",
    ],
    "SAT": ["LFS1", "LFS2", "MFLS1", "MFHS1", "MFLS2", "MFHS2", "HFS1", "HFS2"],
}


def parse_channels(channels):
    """Parse a comma separated list of channels or all or SAT/LAT into channel tag list"""
    if channels in ["SAT", "LAT"]:
        channels = s4_channels[channels]
    elif channels in ["all", None]:
        channels = s4_channels["SAT"] + s4_channels["LAT"]
    elif isinstance(channels, str):
        channels = channels.split(",")
    return channels

# Cell

import collections


def merge_dict(d1, d2):
    """
    Modifies d1 in-place to contain values from d2.  If any value
    in d1 is a dictionary (or dict-like), *and* the corresponding
    value in d2 is also a dictionary, then merge them in-place.
    """
    for k, v2 in d2.items():
        v1 = d1.get(k)  # returns None if v1 has no value for this key
        if isinstance(v1, collections.Mapping) and isinstance(v2, collections.Mapping):
            merge_dict(v1, v2)
        else:
            d1[k] = v2


def parse_config(*config_files):
    """Parse TOML configuration files

    Later TOML configuration files override the previous ones,
    dictionaries at the same level are merged.

    Parameters
    ----------
    config_files : one or more str
        paths to TOML configuration files

    Returns
    -------
    config : dict
        parsed dictionary"""
    config = toml.load(config_files[0])
    for conf in config_files[1:]:
        merge_dict(config, toml.load(conf))
    return config

# Cell


class S4DesSimTool:
    def __init__(self, config_filename, output_folder="output"):
        """Simulate CMB-S4 maps based on the experiment configuration

        Parameters
        ----------
        config : str or Path or List
            CMB-S4 configuration stored in a TOML file
            see for example s4_reference_design.toml in the repository
            It also supports multiple TOML files as a List, in this case
            later files override configuration files of the earlier files.
            check the `config` attribute to verify that the parsing behaved
            as expected.
        output_folder : str or Path
            Output path
        """
        self.config_filename = (
            [config_filename]
            if isinstance(config_filename, (str, Path))
            else config_filename
        )

        self.config = parse_config(*self.config_filename)
        self.output_filename_template = "cmbs4_KCMB_{telescope}-{band}_{site}_nside{nside}_{split}_of_{nsplits}.fits"
        self.output_folder = Path(output_folder)
        self.output_folder.mkdir(parents=True, exist_ok=True)

    def run(self, channels="all", sites=["Pole", "Chile"], overwrite=False):
        """Run the simulation

        Parameters
        ----------
        channels : str or list[str]
            list of channel tags, e.g.
            * ["LFS1", "LFS2"] or
            * "SAT" or "LAT"
            * "all" (default)
        site : list[str]
            ['Pole'] or ['Chile'], by default ["Pole", "Chile"]
        """
        nsplits = self.config["experiment"].get("number_of_splits", 0)
        if nsplits == 1:
            nsplits = 0
        assert (
            nsplits < 8
        ), "We currently only have 7 independent realizations of atmosphere and noise"
        conf_md5 = ",".join(map(md5sum_file, self.config_filename))
        for site in sites:
            for channel in parse_channels(channels):

                if get_telecope_years(self.config, site, channel) == 0:
                    continue
                telescope = get_telescope(channel)
                subfolder = self.output_folder / f"{telescope}-{channel}_{site.lower()}"
                subfolder.mkdir(parents=True, exist_ok=True)
                log.info("Created output folder %s", str(subfolder))


                for split in range(nsplits + 1):
                    nside = 512 if telescope == "SAT" else 4096
                    output_filename = self.output_filename_template.format(
                        nside=nside,
                        telescope=telescope,
                        band=channel,
                        site=site.lower(),
                        split=max(1, split),  # split=0 is full mission and we want 1
                        nsplits=1 if split == 0 else nsplits,
                    )
                    if os.path.exists(subfolder / output_filename) and not overwrite:
                        log.info("File %s already exists, SKIP", output_filename)
                        continue

                    if split == 0:
                        log.info(f"Simulate channel {channel} at {site}")
                        sky_emission = load_sky_emission(
                            self.config["sky_emission"], site, channel
                        )

                    output_map = np.zeros_like(sky_emission)
                    if self.config["experiment"].get("include_atmosphere", True):
                        output_map += load_atmosphere(
                            self.config, site, channel, realization=split
                        )
                    else:
                        log.info("Skip the atmosphere noise")

                    if self.config["experiment"].get("include_noise", True):
                        output_map += load_noise(
                            self.config, site, channel, realization=split
                        )

                    else:
                        log.info("Skip the instrument noise")

                    if split > 0:
                        output_map *= np.sqrt(nsplits)
                    output_map += sky_emission
                    # Use UNSEEN instead of nan for missing pixels
                    output_map[np.isnan(output_map)] = hp.UNSEEN

                    log.info(f"Writing {output_filename}")
                    noise_version = "1.0"
                    hp.write_map(
                        subfolder / output_filename,
                        output_map,
                        column_units="K_CMB",
                        extra_header=[
                            ("SOFTWARE", "s4_design_sim_tool"),
                            ("SW_VERS", __version__),
                            ("SKY_VERS", "1.0"),
                            ("ATM_VERS", "1.0"),
                            ("NOI_VERS", noise_version),
                            ("SITE", site),
                            ("SPLIT", split),
                            ("NSPLITS", nsplits),
                            ("CHANNEL", channel),
                            ("DATE", str(date.today())),
                            ("CONFMD5", conf_md5),
                        ],
                        coord="Q",
                        overwrite=True,
                    )
                    # only run of full mission and the first split
                    if split in [0, 1] and self.config["experiment"].get(
                        "include_noise", True
                    ):

                        log.info(f"Loading hitmap and white noise covariance matrix")
                        if split == 0:
                            hitmap, wcov = load_hitmap_wcov(
                                self.config, site, channel, realization=0
                            )
                        else:
                            hitmap = np.round(hitmap / nsplits).astype(np.int64)
                            wcov = hp.ma(wcov) * nsplits

                        hitmap_filename = output_filename.replace("KCMB", "hitmap")
                        log.info(f"Writing {hitmap_filename}")
                        hp.write_map(
                            subfolder / hitmap_filename,
                            hitmap,
                            column_units="hits",
                            extra_header=[
                                ("SOFTWARE", "s4_design_sim_tool"),
                                ("SW_VERS", __version__),
                                ("NOI_VERS", noise_version),
                                ("SITE", site),
                                ("SPLIT", split),
                                ("NSPLITS", nsplits),
                                ("CHANNEL", channel),
                                ("DATE", str(date.today())),
                                ("CONFMD5", conf_md5),
                            ],
                            coord="Q",
                            overwrite=True,
                        )

                        wcov_filename = output_filename.replace("KCMB", "wcov")
                        log.info(f"Writing {wcov_filename}")

                        hp.write_map(
                            subfolder / wcov_filename,
                            wcov,
                            column_units="K_CMB**2",
                            extra_header=[
                                ("SOFTWARE", "s4_design_sim_tool"),
                                ("SW_VERS", __version__),
                                ("NOI_VERS", noise_version),
                                ("SITE", site),
                                ("SPLIT", split),
                                ("NSPLITS", nsplits),
                                ("CHANNEL", channel),
                                ("DATE", str(date.today())),
                                ("CONFMD5", conf_md5),
                            ],
                            coord="Q",
                            overwrite=True,
                        )
                        if split == 1:
                            del hitmap, wcov

# Cell


def command_line_script(args=None):

    import logging as log

    log.basicConfig(level=log.INFO)

    import argparse

    parser = argparse.ArgumentParser(description="Run s4_design_sim_tool")
    parser.add_argument("config", type=str, nargs="*", help="TOML Configuration files")
    parser.add_argument(
        "--channels",
        type=str,
        help="Channels e.g. all, SAT, LAT, LFL1 or comma separated list of channels",
        required=False,
        default="all",
    )
    parser.add_argument(
        "--site",
        type=str,
        help="Pole, Chile or all, default all",
        required=False,
        default="all",
    )
    parser.add_argument(
        "--output_folder",
        type=str,
        help="Output folder, optional",
        required=False,
        default="output",
    )
    res = parser.parse_args(args)
    if res.site == "all":
        sites = ["Chile", "Pole"]
    else:
        sites = [res.site]
    sim = S4DefSimTool(res.config, output_folder=res.output_folder)
    sim.run(channels=res.channels, sites=sites)