"""
Copyright (C) since 2013 Calliope contributors listed in AUTHORS.
Licensed under the Apache 2.0 License (see LICENSE file).

generate_runs.py
~~~~~~~~~~~~~~~~

Generate scripts to run multiple versions of the same model
in parallel on a cluster or sequentially on any machine.

"""

import os

import pandas as pd
from calliope.core import AttrDict


def generate_runs(model_file, scenarios=None, additional_args=None, override_dict=None):
    """
    Returns a list of "calliope run" invocations.

    ``scenarios`` must be specified as either a semicolon-separated
    list of scenarios or a semicolon-separated list of comma-separated
    individual override combinations, such as:

        ``override1,override2;override1,override3;...``

    If ``scenarios`` is not given, use all scenarios in the model
    configuration, and if no scenarios given in the model configuration,
    uses all individual overrides, one by one.

    """
    if scenarios is None:
        config = AttrDict.from_yaml(model_file)
        if override_dict:
            override = AttrDict.from_yaml_string(override_dict)
            config.union(override, allow_override=True, allow_replacement=True)

        if "scenarios" in config:
            runs = config.scenarios.keys()
        else:
            runs = config.overrides.keys()
    else:
        runs = scenarios.split(";")

    commands = []

    # len(str(x)) gives us the number of digits in x, for padding
    i_string = "{:0>" + str(len(str(len(runs)))) + "d}"

    for i, run in enumerate(runs):
        cmd = (
            (
                "calliope run {model} --scenario {scenario} "
                "--save_netcdf out_{i}_{scenario}.nc "
                "--save_plots plots_{i}_{scenario}.html"
            )
            .format(
                i=i_string.format(i + 1),
                model=model_file,
                scenario=run,
                override_dict=override_dict,
            )
            .strip()
        )

        if override_dict:
            cmd = cmd + ' --override_dict="{}"'.format(override_dict)

        if additional_args:
            cmd = cmd + " " + additional_args

        commands.append(cmd)

    return commands


def generate_bash_script(
    out_file, model_file, scenarios, additional_args=None, override_dict=None, **kwargs
):

    commands = generate_runs(model_file, scenarios, additional_args, override_dict)

    base_string = "    {i}) {cmd} ;;\n"
    lines_start = [
        "#!/bin/sh",
        "",
        "function process_case () {",
        '    case "$1" in',
        "",
    ]
    lines_end = [
        "    esac",
        "}",
        "",
        "if [[ $# -eq 0 ]] ; then",
        "    echo No parameter given, running all runs sequentially...",
        "    for i in $(seq 1 {}); do process_case $i; done".format(len(commands)),
        "else",
        "    echo Running run $1",
        "    process_case $1",
        "fi",
        "",
    ]

    lines_all = (
        lines_start
        + [base_string.format(i=i + 1, cmd=cmd) for i, cmd in enumerate(commands)]
        + lines_end
    )

    with open(out_file, "wb") as f:
        f.write(bytes("\n".join(lines_all), "UTF-8"))

    os.chmod(out_file, 0o755)

    return commands


def generate_bsub_script(
    out_file,
    model_file,
    scenarios,
    additional_args,
    override_dict,
    cluster_mem,
    cluster_time,
    cluster_threads=1,
    **kwargs,
):

    # We also need to generate the bash script to run on the cluster
    bash_out_file = out_file + ".array.sh"
    bash_out_file_basename = os.path.basename(bash_out_file)
    commands = generate_bash_script(
        bash_out_file, model_file, scenarios, additional_args, override_dict
    )

    lines = [
        "#!/bin/sh",
        "#BSUB -J calliope[1-{}]".format(len(commands)),
        "#BSUB -n {}".format(cluster_threads),
        '#BSUB -R "rusage[mem={}]"'.format(cluster_mem),
        "#BSUB -W {}".format(cluster_time),
        "#BSUB -r",  # Automatically restart failed jobs
        "#BSUB -o log_%I.log",
        "",
        "./" + bash_out_file_basename + " ${LSB_JOBINDEX}",
        "",
    ]

    with open(out_file, "wb") as f:
        f.write(bytes("\n".join(lines), "UTF-8"))


def generate_sbatch_script(
    out_file,
    model_file,
    scenarios,
    additional_args,
    override_dict,
    cluster_mem,
    cluster_time,
    cluster_threads=1,
    **kwargs,
):
    """
    SBATCH (SLURM) script generator.
    """

    # We also need to generate the bash script to run on the cluster
    bash_out_file = out_file + ".array.sh"
    bash_out_file_basename = os.path.basename(bash_out_file)
    commands = generate_bash_script(
        bash_out_file, model_file, scenarios, additional_args, override_dict
    )

    if ":" not in cluster_time:
        # Assuming time given as minutes, so needs changing to %H:%M%S
        cluster_time = pd.to_datetime(cluster_time, unit="m").strftime("%H:%M:%S")

    lines = [
        "#!/bin/bash",
        "#SBATCH -J calliope",  # Name of the job
        "#SBATCH --array=1-{}".format(len(commands)),  # How many jobs there are
        "#SBATCH --ntasks={}".format(cluster_threads),
        "#SBATCH --mem={}".format(cluster_mem),
        "#SBATCH --time={}".format(
            cluster_time
        ),  # How much wallclock time will be required
        "#SBATCH -o log_%a.log",
        "",
        "#! Optional add-ins for SBATCH (uncomment and add info as necessary):",
        "##SBATCH -A project_name",  # Which project should be charged'
        "##SBATCH --nodes=X",  # X whole nodes should be allocated'
        "##SBATCH -p partition_name",
        "",
        "#! Insert module load commands after this line, if needed:",
        "#! (Note: you can load this in ~.bashrc if you want them loaded every time you log in)",
        "",
        "#! module load gurobi (or glpk/cplex)",
        "#! module load miniconda3",
        "#! module load /path/to/miniconda3/envs/your_env_name/",
        "",
        "cd $SLURM_SUBMIT_DIR",
        "",
        "./" + bash_out_file_basename + " ${SLURM_ARRAY_TASK_ID}",
    ]

    with open(out_file, "wb") as f:
        f.write(bytes("\n".join(lines), "UTF-8"))


def generate_windows_script(
    out_file, model_file, scenarios, additional_args=None, override_dict=None, **kwargs
):

    commands = generate_runs(model_file, scenarios, additional_args, override_dict)

    # \r\n are Windows line endings
    base_string = 'echo "Run {i}"\r\n{cmd}\r\n'
    lines_start = [
        "@echo off",
        "",
    ]

    lines_all = lines_start + [
        base_string.format(i=i + 1, cmd=cmd) for i, cmd in enumerate(commands)
    ]

    with open(out_file, "wb") as f:
        f.write(bytes("\r\n".join(lines_all), "UTF-8"))

    os.chmod(out_file, 0o755)

    return commands


_KINDS = {
    "bash": generate_bash_script,
    "bsub": generate_bsub_script,
    "windows": generate_windows_script,
    "sbatch": generate_sbatch_script,
}


def generate(kind, **kwargs):
    _KINDS[kind](**kwargs)
