import math


def wrap_command_with_sbatch(
        cmd,
        config,
        n_tasks,
):
    """
    Wrap a command around a call to sbatch

    Usage:
    ```
    config = SlurmConfig()
    cmd = "echo 'Hello, world!'"
    sbatch_cmd = wrap_command_with_sbatch(cmd, config)
    """
    cmd = cmd.replace("'", "\\'")
    num_cpus = n_tasks * config.n_cpus_per_task
    n_nodes = config.n_nodes or math.ceil(
        num_cpus / config.max_num_cores_per_node)
    if config.n_gpus > 0:
        full_cmd = (
            "sbatch -A {account_name} -p {partition} -t {time}"
            " -N {nodes} -n {n_tasks} --cpus-per-task={cpus_per_task}"
            " --gres=gpu:TITAN:{n_gpus} {extra_flags} --wrap=$'{cmd}'".format(
                account_name=config.account_name,
                partition=config.partition,
                time=config.time_in_mins,
                nodes=n_nodes,
                n_tasks=n_tasks,
                cpus_per_task=config.n_cpus_per_task,
                cmd=cmd,
                n_gpus=config.n_gpus,
                extra_flags=config.extra_flags,
            )
        )
    else:
        full_cmd = (
            "sbatch -A {account_name} -p {partition} -t {time}"
            " -N {nodes} -n {n_tasks} --cpus-per-task={cpus_per_task}"
            " {extra_flags} --wrap=$'{cmd}'".format(
                account_name=config.account_name,
                partition=config.partition,
                time=config.time_in_mins,
                nodes=n_nodes,
                n_tasks=n_tasks,
                cpus_per_task=config.n_cpus_per_task,
                cmd=cmd,
                extra_flags=config.extra_flags,
            )
        )
    return full_cmd