"""
Job Submitter
-------------


The ``job_submitter`` allows to execute a parametric study using a script mask and a `dictionary` of parameters to replace in this mask, from the command line.
These parameters must be present in the given mask in the ``%(PARAMETER)s`` format (other types apart from string are also allowed).

The type of script and executable is freely choosable, but defaults to ``madx``, for which this submitter was originally written.
When submitting to ``HTCondor``, data to be transferred back to the working directory must be written in a sub-folder defined by ``job_output_directory`` which defaults to **Outputdata**.

This script also allows to check if all ``HTCondor`` jobs finished successfully, for resubmissions with a different parameter grid, and for local execution.
A **Jobs.tfs** file is created in the working directory containing the Job Id, parameter per job
and job directory for further post processing.

For additional information and guides, see the `Job Submitter page
<https://pylhc.github.io/packages/pylhcsubmitter/job_submitter/>`_ in the ``OMC`` documentation site.

*--Required--*

- **mask** *(str)*: Program mask to use

- **replace_dict** *(DictAsString)*: Dict containing the str to replace as
  keys and values a list of parameters to replace

- **working_directory** *(str)*: Directory where data should be put


*--Optional--*

- **append_jobs**: Flag to rerun job with finer/wider grid,
  already existing points will not be reexecuted.

  Action: ``store_true``
- **check_files** *(str)*: List of files/file-name-masks expected to be in the
  'job_output_dir' after a successful job (for appending/resuming). Uses the 'glob'
  function, so unix-wildcards (*) are allowed. If not given, only the presence of the folder itself is checked.
- **dryrun**: Flag to only prepare folders and scripts,
  but does not start/submit jobs.
  Together with `resume_jobs` this can be use to check which jobs succeeded and which failed.

  Action: ``store_true``
- **executable** *(str)*: Path to executable or job-type (of ['madx', 'python3', 'python2']) to use.

- **htc_arguments** *(DictAsString)*: Additional arguments for htcondor, as Dict-String.
  For AccountingGroup please use 'accounting_group'. 'max_retries' and 'notification' have defaults (if not given).
  Others are just passed on.

  Default: ``{}``
- **job_output_dir** *(str)*: The name of the output dir of the job. (Make sure your script puts its data there!)

  Default: ``Outputdata``
- **jobflavour** *(str)*: Jobflavour to give rough estimate of runtime of one job

  Choices: ``('espresso', 'microcentury', 'longlunch', 'workday', 'tomorrow', 'testmatch', 'nextweek')``
  Default: ``workday``
- **jobid_mask** *(str)*: Mask to name jobs from replace_dict

- **num_processes** *(int)*: Number of processes to be used if run locally

  Default: ``4``
- **resume_jobs**: Only do jobs that did not work.

  Action: ``store_true``
- **run_local**: Flag to run the jobs on the local machine. Not suggested.

  Action: ``store_true``
- **script_arguments** *(DictAsString)*: Additional arguments to pass to the script,
  as dict in key-value pairs ('--' need to be included in the keys).

  Default: ``{}``
- **script_extension** *(str)*: New extension for the scripts created from the masks.
  This is inferred automatically for ['madx', 'python3', 'python2']. Otherwise not changed.

- **ssh** *(str)*: Run htcondor from this machine via ssh (needs access to the `working_directory`)


:author: mihofer, jdilly, fesoubel
"""
import itertools
import logging
import multiprocessing
import subprocess
import sys
from pathlib import Path

import numpy as np
import tfs
from generic_parser import EntryPointParameters, entrypoint
from generic_parser.entry_datatypes import DictAsString
from generic_parser.tools import print_dict_tree

import pylhc_submitter.htc.utils as htcutils
from pylhc_submitter.htc.mask import (
    check_percentage_signs_in_mask,
    create_jobs_from_mask,
    find_named_variables_in_mask,
    generate_jobdf_index,
)
from pylhc_submitter.htc.utils import (
    COLUMN_JOB_DIRECTORY,
    COLUMN_SHELL_SCRIPT,
    EXECUTEABLEPATH,
    HTCONDOR_JOBLIMIT,
    JOBFLAVOURS,
)
from pylhc_submitter.utils.environment_tools import on_windows
from pylhc_submitter.utils.iotools import PathOrStr, save_config, make_replace_entries_iterable, keys_to_path
from pylhc_submitter.utils.logging_tools import log_setup

JOBSUMMARY_FILE = "Jobs.tfs"
JOBDIRECTORY_PREFIX = "Job"
COLUMN_JOBID = "JobId"
CONFIG_FILE = "config.ini"

SCRIPT_EXTENSIONS = {
    "madx": ".madx",
    "python3": ".py",
    "python2": ".py",
}

LOG = logging.getLogger(__name__)


try:
    import htcondor
    HAS_HTCONDOR = True
except ImportError:
    platform = "macOS" if sys.platform == "darwin" else "windows"
    LOG.warning(
        f"htcondor python bindings are linux-only. You can still use job_submitter on {platform}, "
        "but only for local runs."
    )
    HAS_HTCONDOR = False


def get_params():
    params = EntryPointParameters()
    params.add_parameter(
        name="mask",
        type=PathOrStr,
        required=True,
        help="Program mask to use",
    )
    params.add_parameter(
        name="working_directory",
        type=PathOrStr,
        required=True,
        help="Directory where data should be put",
    )
    params.add_parameter(
        name="executable",
        default="madx",
        type=PathOrStr,
        help=(
            "Path to executable or job-type " f"(of {str(list(EXECUTEABLEPATH.keys()))}) to use."
        ),
    )
    params.add_parameter(
        name="jobflavour",
        type=str,
        choices=JOBFLAVOURS,
        default="workday",
        help="Jobflavour to give rough estimate of runtime of one job ",
    )
    params.add_parameter(
        name="run_local",
        action="store_true",
        help="Flag to run the jobs on the local machine. Not suggested.",
    )
    params.add_parameter(
        name="resume_jobs",
        action="store_true",
        help="Only do jobs that did not work.",
    )
    params.add_parameter(
        name="append_jobs",
        action="store_true",
        help=(
            "Flag to rerun job with finer/wider grid, already existing points will not be "
            "reexecuted."
        ),
    )
    params.add_parameter(
        name="dryrun",
        action="store_true",
        help=(
            "Flag to only prepare folders and scripts, but does not start/submit jobs. "
            "Together with `resume_jobs` this can be use to check which jobs "
            "succeeded and which failed."
        ),
    )
    params.add_parameter(
        name="replace_dict",
        help=(
            "Dict containing the str to replace as keys and values a list of parameters to "
            "replace"
        ),
        type=DictAsString,
        required=True,
    )
    params.add_parameter(
        name="script_arguments",
        help=(
            "Additional arguments to pass to the script, as dict in key-value pairs "
            "('--' need to be included in the keys)."
        ),
        type=DictAsString,
        default={},
    )
    params.add_parameter(
        name="script_extension",
        help=(
            "New extension for the scripts created from the masks. This is inferred "
            f"automatically for {str(list(SCRIPT_EXTENSIONS.keys()))}. Otherwise not changed."
        ),
        type=str,
    )
    params.add_parameter(
        name="num_processes",
        help="Number of processes to be used if run locally",
        type=int,
        default=4,
    )
    params.add_parameter(
        name="check_files",
        help=(
            "List of files/file-name-masks expected to be in the "
            "'job_output_dir' after a successful job "
            "(for appending/resuming). Uses the 'glob' function, so "
            "unix-wildcards (*) are allowed. If not given, only the "
            "presence of the folder itself is checked."
        ),
        type=str,
        nargs="+",
    )
    params.add_parameter(
        name="jobid_mask",
        help="Mask to name jobs from replace_dict",
        type=str,
    )
    params.add_parameter(
        name="job_output_dir",
        help="The name of the output dir of the job. (Make sure your script puts its data there!)",
        type=str,
        default="Outputdata",
    )
    params.add_parameter(
        name="htc_arguments",
        help=(
            "Additional arguments for htcondor, as Dict-String. "
            "For AccountingGroup please use 'accounting_group'. "
            "'max_retries' and 'notification' have defaults (if not given). "
            "Others are just passed on. "
        ),
        type=DictAsString,
        default={},
    )
    params.add_parameter(
        name="ssh",
        help="Run htcondor from this machine via ssh (needs access to the `working_directory`)",
        type=str,
    )

    return params


@entrypoint(get_params(), strict=True)
def main(opt):
    if not opt.run_local:
        LOG.info("Starting HTCondor Job-submitter.")
        _check_htcondor_presence()
    else:
        LOG.info("Starting Job-submitter.")

    opt = _check_opts(opt)
    save_config(opt.working_directory, opt, "job_submitter")

    job_df = _create_jobs(
        opt.working_directory,
        opt.mask,
        opt.jobid_mask,
        opt.replace_dict,
        opt.job_output_dir,
        opt.append_jobs,
        opt.executable,
        opt.script_arguments,
        opt.script_extension,
    )
    job_df, dropped_jobs = _drop_already_ran_jobs(
        job_df, opt.resume_jobs or opt.append_jobs, opt.job_output_dir, opt.check_files
    )

    if opt.run_local and not opt.dryrun:
        _run_local(job_df, opt.num_processes)
    else:
        _run_htc(
            job_df,
            opt.working_directory,
            opt.job_output_dir,
            opt.jobflavour,
            opt.ssh,
            opt.dryrun,
            opt.htc_arguments,
        )
    if opt.dryrun:
        _print_stats(job_df.index, dropped_jobs)


# Main Functions ---------------------------------------------------------------


def _create_jobs(
    cwd,
    mask_path_or_string,
    jobid_mask,
    replace_dict,
    output_dir,
    append_jobs,
    executable,
    script_args,
    script_extension,
) -> tfs.TfsDataFrame:
    LOG.debug("Creating Jobs.")
    values_grid = np.array(list(itertools.product(*replace_dict.values())), dtype=object)

    if append_jobs:
        jobfile_path = cwd / JOBSUMMARY_FILE
        try:
            job_df = tfs.read(str(jobfile_path.absolute()), index=COLUMN_JOBID)
        except FileNotFoundError as filerror:
            raise FileNotFoundError(
                "Cannot append jobs, as no previous jobfile was found at " f"'{jobfile_path}'"
            ) from filerror
        mask = [elem not in job_df[replace_dict.keys()].values for elem in values_grid]
        njobs = mask.count(True)
        values_grid = values_grid[mask]
    else:
        njobs = len(values_grid)
        job_df = tfs.TfsDataFrame()

    if njobs == 0:
        raise ValueError(f"No (new) jobs found!")
    if njobs > HTCONDOR_JOBLIMIT:
        LOG.warning(
            f"You are attempting to submit an important number of jobs ({njobs})."
            "This can be a high stress on your system, make sure you know what you are doing."
        )

    LOG.debug(f"Initial number of jobs: {njobs:d}")
    data_df = tfs.TfsDataFrame(
        index=generate_jobdf_index(job_df, jobid_mask, replace_dict.keys(), values_grid),
        columns=list(replace_dict.keys()),
        data=values_grid,
    )
    job_df = job_df.append(data_df, sort=False, how_headers='left')
    job_df = _setup_folders(job_df, cwd)

    if htcutils.is_mask_file(mask_path_or_string):
        LOG.debug("Creating all jobs from mask.")
        script_extension = _get_script_extension(script_extension, executable, mask_path_or_string)
        job_df = create_jobs_from_mask(
            job_df, mask_path_or_string, replace_dict.keys(), script_extension
        )

    LOG.debug("Creating shell scripts for submission.")
    job_df = htcutils.write_bash(
        job_df,
        output_dir,
        executable=executable,
        cmdline_arguments=script_args,
        mask=mask_path_or_string,
    )

    job_df[COLUMN_JOB_DIRECTORY] = job_df[COLUMN_JOB_DIRECTORY].apply(str)
    tfs.write(str(cwd / JOBSUMMARY_FILE), job_df, save_index=COLUMN_JOBID)
    return job_df


def _drop_already_ran_jobs(
    job_df: tfs.TfsDataFrame, drop_jobs: bool, output_dir: str, check_files: str
):
    LOG.debug("Dropping already finished jobs, if necessary.")
    finished_jobs = []
    if drop_jobs:
        finished_jobs = [
            idx
            for idx, row in job_df.iterrows()
            if _job_was_successful(row, output_dir, check_files)
        ]
        LOG.info(
            f"{len(finished_jobs):d} of {len(job_df.index):d}"
            " Jobs have already finished and will be skipped."
        )
        job_df = job_df.drop(index=finished_jobs)
    return job_df, finished_jobs


def _run_local(job_df: tfs.TfsDataFrame, num_processes: int) -> None:
    LOG.info(f"Running {len(job_df.index)} jobs locally in {num_processes:d} processes.")
    pool = multiprocessing.Pool(processes=num_processes)
    res = pool.map(_execute_shell, job_df.iterrows())
    if any(res):
        LOG.error("At least one job has failed.")
        raise RuntimeError("At least one job has failed. Check output logs!")


def _run_htc(
    job_df: tfs.TfsDataFrame,
    cwd: str,
    output_dir: str,
    flavour: str,
    ssh: str,
    dryrun: bool,
    additional_htc_arguments: DictAsString,
) -> None:
    LOG.info(f"Submitting {len(job_df.index)} jobs on htcondor, flavour '{flavour}'.")
    LOG.debug("Creating htcondor subfile.")
    subfile = htcutils.make_subfile(
        cwd, job_df, output_dir=output_dir, duration=flavour, **additional_htc_arguments
    )
    if not dryrun:
        LOG.debug("Submitting jobs to htcondor.")
        htcutils.submit_jobfile(subfile, ssh)


def _get_script_extension(script_extension: str, executable: PathOrStr, mask: PathOrStr) -> str:
    if script_extension is not None:
        return script_extension
    return SCRIPT_EXTENSIONS.get(executable, mask.suffix)


# Sub Functions ----------------------------------------------------------------


def _check_htcondor_presence() -> None:
    """Checks the ``HAS_HTCONDOR`` variable and raises EnvironmentError if it is ``False``."""
    if not HAS_HTCONDOR:
        raise EnvironmentError("htcondor bindings are necessary to run this module.")


def _setup_folders(job_df: tfs.TfsDataFrame, working_directory: PathOrStr) -> tfs.TfsDataFrame:
    def _return_job_dir(job_id):
        return working_directory / f"{JOBDIRECTORY_PREFIX}.{job_id}"

    LOG.debug("Setting up folders: ")
    job_df[COLUMN_JOB_DIRECTORY] = [_return_job_dir(id_) for id_ in job_df.index]

    for job_dir in job_df[COLUMN_JOB_DIRECTORY]:
        try:
            job_dir.mkdir()
        except IOError:
            LOG.debug(f"   failed '{job_dir}' (might already exist).")
        else:
            LOG.debug(f"   created '{job_dir}'.")
    return job_df


def _job_was_successful(job_row, output_dir, files) -> bool:
    output_dir = Path(job_row[COLUMN_JOB_DIRECTORY], output_dir)
    success = output_dir.is_dir() and any(output_dir.iterdir())
    if success and files is not None and len(files):
        for f in files:
            success &= len(list(output_dir.glob(f))) > 0
    return success


def _execute_shell(df_row) -> int:
    idx, column = df_row
    cmd = [] if on_windows() else ["sh"]

    with Path(column[COLUMN_JOB_DIRECTORY], "log.tmp").open("w") as logfile:
        process = subprocess.Popen(
            cmd + [column[COLUMN_SHELL_SCRIPT]],
            shell=on_windows(),
            stdout=logfile,
            stderr=subprocess.STDOUT,
            cwd=column[COLUMN_JOB_DIRECTORY],
        )
    return process.wait()


def _check_opts(opt):
    LOG.debug("Checking options.")
    if opt.resume_jobs and opt.append_jobs:
        raise ValueError("Select either Resume jobs or Append jobs")

    # Paths ---
    opt = keys_to_path(opt, "working_directory", "executable")

    if str(opt.executable) in EXECUTEABLEPATH.keys():
        opt.executable = str(opt.executable)

    if htcutils.is_mask_file(opt.mask):
        mask = Path(opt.mask).read_text()  # checks that mask and dir are there
        opt["mask"] = Path(opt["mask"])
    else:
        mask = opt.mask

    # Replace dict ---
    dict_keys = set(opt.replace_dict.keys())
    mask_keys = find_named_variables_in_mask(mask)
    not_in_mask = dict_keys - mask_keys
    not_in_dict = mask_keys - dict_keys

    if len(not_in_dict):
        raise KeyError(
            "The following keys in the mask were not found in the given replace_dict: "
            f"{str(not_in_dict).strip('{}')}"
        )

    if len(not_in_mask):
        LOG.warning(
            "The following replace_dict keys were not found in the given mask: "
            f"{str(not_in_mask).strip('{}')}"
        )

        # remove all keys which are not present in mask (otherwise unnecessary jobs)
        [opt.replace_dict.pop(key) for key in not_in_mask]
        if len(opt.replace_dict) == 0:
            raise KeyError("Empty replace-dictionary")
    check_percentage_signs_in_mask(mask)

    print_dict_tree(opt, name="Input parameter", print_fun=LOG.debug)
    opt.replace_dict = make_replace_entries_iterable(opt.replace_dict)
    return opt


def _print_stats(new_jobs, finished_jobs):
    """Print some quick statistics."""
    LOG.info("------------- QUICK STATS ----------------")
    LOG.info(f"Jobs total:{len(new_jobs) + len(finished_jobs):d}")
    LOG.info(f"Jobs to run: {len(new_jobs):d}")
    LOG.info(f"Jobs already finished: {len(finished_jobs):d}")
    LOG.info("---------- JOBS TO RUN: NAMES -------------")
    for job_name in new_jobs:
        LOG.info(job_name)
    LOG.info("--------- JOBS FINISHED: NAMES ------------")
    for job_name in finished_jobs:
        LOG.info(job_name)


# Script Mode ------------------------------------------------------------------


if __name__ == "__main__":
    log_setup()
    main()
