#!/usr/bin/env python

# Copyright (C) 2009 Martin Sandve Alnes
#
# This file is part of DOLFIN.
#
# DOLFIN is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# DOLFIN is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with DOLFIN. If not, see <http://www.gnu.org/licenses/>.

# Modified Johan Hake 2011

import torque
import sge
import os
import subprocess

def check_mem_argument(mem):
    allowed_mem_suffices = ["kb","mb","gb"]
    if not mem is None:
        assert isinstance(mem,str)
        if len(mem) < 3:
            return False
        if mem[-2:] not in allowed_mem_suffices:
            return False
        try:
            memory = int(mem[:-2])
            if memory < 0:
                return False
        except:
            return False
    return True

def sjoin(*items):
    "Join str() of all arguments."
    return "".join(map(str,items))


def submit(jobs, nodes=1, ppn=1, walltime=3, mem=None, vmem=None, name="myjob", \
           workdir=".", email=None, paths=(), setup="", \
           keep_environment=True, serial=None, nice=None, dryrun=False, \
           queue=None, parallel_environment=None, backend="torque"): 
    """Submit a sequence of commands to the pbs queue using qsub.
    
    Example use:
        from pjobs import submit, sjoin
        for n in (10,20,30):
            for h in (0.1, 0.01, 0.005):
                dt = 0.1*h
                jobs.append("pmpirun.openmpi mysimulator -n%d -h%f -dt%f" % (n, h, dt))
        submit(jobs, nodes=4, ppn=8, walltime=24*7)
    
    Arguments:
    @param jobs:
        A list of commands (strings) to pass to the queue for separate execution.
    @param nodes:
        Number of compute nodes to use.
        Default is 1 node.
    @param ppn:
        Number of processors per node to use.
        Default is 1 processor.
    @param walltime:
        Maximum amount of real-life time the job can use, in hours.
        Default is 3.
    @param mem:
        Optional maximum amount of physical memory the process will need.
        The syntax is mem='100mb', or mem='4gb', with allowed suffices 
        one of 'kb', 'mb' or 'gb' and only int values.
    @param vmem:
        Optional maximum amount of virtual memory the process will need.
        For syntax see mem.
    @param name:
        Basename of the job, just for the queue system.
        If a single string this will be combined with a
        counter to separate multiple jobs. A list or
        tuple of strings can be provided instead to
        define separate names for each job.
        Default is 'myjob'.
    @param email:
        Your email address if you want email when
        the jobs are started and stopped.
    @param workdir:
        The directory the jobs will be run in.
        Default is the place you commit the job from
    @param paths:
        A list of directories to add to the $PATH.
    @param setup:
        Optional additional job script contents for setting 
        up e.g. environment variables before each job.
    @param keep_environment:
        True if you want to copy all current environment
        variables to the job environment (PBS -V option).
        Default True.
    @param serial:
        True if you want to run the jobs as a serial sequence of processes
        instead of passing them to the qsub queue.
        Default is False, unless the environment variable PJOBS_SERIAL=1 
        is defined in which case the default is True.
    @param nice:
        Optional nice value (int) used for all jobs when running in serial.
        If provided, "nice -n %d" % nice is added before each job command.
    @param dryrun:
        True if you want to test this script but not send jobs to the queue.
    @param queue:
        Set a particular queue to be used when submitting
        (only available for the sge backend)
    @param parallel_environment:
        Set the parallel environment (PE) of the submitting backend 
        (only available for the sge backend)
    @param backend:
        Set the name of the actuall submitter.
        Available backends: 'torque' (default) and 'sge'
    """
    if isinstance(jobs, str):
        jobs = [jobs]
    assert all(isinstance(job, str) for job in jobs)
    assert isinstance(nodes, int)
    assert isinstance(ppn, int)
    assert isinstance(walltime, (int, float))
    if not check_mem_argument(mem):
        raise TypeError, "Wrong format of mem attribute, see docstring" 
    if not check_mem_argument(vmem):
        raise TypeError, "Wrong format of vmem attribute, see docstring"
    assert isinstance(name, (str, list, tuple))
    assert isinstance(workdir, str)
    assert email is None or isinstance(email, str)
    if serial is None:
        # Get environment variable to toggle serial behaviour
        serial = bool(os.environ.get("PJOBS_SERIAL", False))
    assert isinstance(serial, bool)

    # Check queue and PE
    queue = queue or ""
    parallel_environment = parallel_environment or ""
    assert isinstance(queue, str)
    assert isinstance(parallel_environment, str)

    # Check backend
    assert isinstance(backend, str)
    assert (backend in ["torque", "sge"])
    
    # Shared setup
    paths = "PATH=$PATH:" + ":".join(paths) if paths else ""
    
    # Modify options to make sense in serial
    if serial:
        if paths:
            print "TODO: Not using paths for serial job."
        if setup:
            print "TODO: Not using setup for serial job."
        if workdir == ".":
            workdir = os.getcwd()
        workdir = os.path.abspath(workdir)
        if nice is not None:
            jobs = ["nice -n %d %s" % (nice, job) for job in jobs]
    else:
        workdir = os.path.abspath(workdir)
    
    # Define range of names if only a single name is provided
    if isinstance(name, str):
        names = ["%s_%d" % (name, k) for k in range(len(jobs))]
    else:
        names = name

    # Enter workdir
    origdir = os.getcwd()
    try:
        if serial:
            os.chdir(workdir)
        # Queue sequence of jobs
        for job, jobname in zip(jobs, names):
            
            if serial:
                # Execute job in separate process
                cmd = job
                if dryrun:
                    print "NOT calling job command:", cmd
                else:
                    print "Calling job command:", cmd
                    # Pass process output to file f
                    #f = open("%s_stdout" % jobname, "w")
                    subprocess.call(cmd.split())
                    #f.close()
            else:
                job_script = globals()[backend].job_script(job, jobname, nodes, \
                                                           ppn, walltime, mem, \
                                                           vmem, workdir, email, paths,\
                                                           setup, keep_environment,\
                                                           queue, parallel_environment)
                
                # Write job script to file
                scriptfilename = "run_%s.sh" % jobname
                f = open(scriptfilename, "w")
                f.write(job_script)
                f.close()
                print "Wrote file ", scriptfilename
                
                # Queue the generated job script
                cmd = ["qsub"] + [scriptfilename]
                if dryrun:
                    print "NOT calling queue command:", " ".join(cmd)
                else:
                    print "Calling queue command:", " ".join(cmd)
                    subprocess.call(cmd)
                    os.unlink(scriptfilename)
    finally:
        os.chdir(origdir)

