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

__all__ = ['md5sum_string', 'md5sum_file', 'S4RefSimTool', 'command_line_script']

# Cell

import os
import toml
import healpy as hp
import numpy as np
import h5py
from pathlib import Path
import logging as log
from datetime import date

from .core import get_telescope, parse_channels

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


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

        Parameters
        ----------
        config : filename
            CMB-S4 configuration stored in a TOML file
            see for example s4_reference_design.toml in the repository
        output_folder : str or Path
            Output path
        """
        self.config_filename = config_filename
        self.config = toml.load(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"]):
        """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"
        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))

                log.info(f"Simulate channel {channel} at {site}")
                sky_emission = load_sky_emission(
                    self.config["sky_emission"], site, channel
                )
                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):
                        log.info("File %s already exists, SKIP", output_filename)
                        continue

                    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", md5sum_file(self.config_filename)),
                        ],
                        coord="Q",
                        overwrite=True,
                    )
                    # only run of full mission and the first split
                    if split in [0, 1]:

                        if split == 0:
                            log.info(
                                f"Loading hitmap and white noise covariance matrix"
                            )
                            hitmap, wcov = load_hitmap_wcov(
                                self.config, site, channel, realization=0
                            )
                        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", md5sum_file(self.config_filename)),
                            ],
                            coord="Q",
                            overwrite=True,
                        )
                        hitmap = np.round(hitmap / nsplits).astype(np.int64)

                        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", md5sum_file(self.config_filename)),
                            ],
                            coord="Q",
                            overwrite=True,
                        )
                        wcov = hp.ma(wcov) * nsplits
                        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, help="Configuration file")
    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 = S4RefSimTool(res.config, output_folder=res.output_folder)
    sim.run(channels=res.channels, sites=sites)