import math
import os
from string import Template
import subprocess

from mycluster import get_data

def scheduler_type():
    return 'pbs'

def name():
    return 'pbs'

def queues():
    queue_list = []
    try:
        output = subprocess.check_output(['qstat', '-Q'])
        lines = output.splitlines()[2:]
        for queue in lines:
            queue_list.append(queue.split(' ')[0])
    except Exception,e:
        print "ERROR"
        print e
        pass
    return queue_list

def available_tasks(queue_id):
    free_tasks = 0
    max_tasks = 0
    return {'available': free_tasks, 'max tasks': max_tasks}

def tasks_per_node(queue_id):
    return 2

def min_tasks_per_node(queue_id):
    return 1

def node_config(queue_id):
    return {'max thread': 1, 'max memory': "Unknown"}

def create_submit(queue_id,**kwargs):

    queue_name = queue_id

    num_tasks = 1
    if 'num_tasks' in kwargs:
        num_tasks = kwargs['num_tasks']

    tpn = tasks_per_node(queue_id)
    queue_tpn = tpn
    if 'tasks_per_node' in kwargs:
        tpn = min(tpn,kwargs['tasks_per_node'])

    nc = node_config(queue_id)
    qc = available_tasks(queue_id)

    num_tasks = min(num_tasks,qc['max tasks'])

    num_threads_per_task = nc['max thread']
    if 'num_threads_per_task' in kwargs:
        num_threads_per_task = kwargs['num_threads_per_task']
    num_threads_per_task = min(num_threads_per_task,int(math.ceil(float(nc['max thread'])/float(tpn))))

    my_name = "myclusterjob"
    if 'my_name' in kwargs:
        my_name = kwargs['my_name']
    my_output = "myclusterjob.out"
    if 'my_output' in kwargs:
        my_output = kwargs['my_output']
    if 'my_script' not in kwargs:
        raise ValueError("Script name not set")
    my_script = kwargs['my_script']
    if 'mycluster-' in my_script:
        my_script = get_data(my_script)
    if 'user_email' not in kwargs:
        raise ValueError("Email not set")
    user_email = kwargs['user_email']

    project_name = 'default'
    if 'project_name' in kwargs:
        project_name = kwargs['project_name']

    wall_clock = '12:00:00'
    if 'wall_clock' in kwargs:
        if ':' not in str(kwargs['wall_clock']):
            wall_clock = str(kwargs['wall_clock'])+':00:00'
        else:
            wall_clock = str(kwargs['wall_clock'])

    num_nodes = int(math.ceil(float(num_tasks)/float(tpn)))

    if num_nodes == 0:
        raise ValueError("Must request 1 or more nodes")

    # For exclusive node use total number of slots required
    # is number of nodes x number of slots offer by queue
    num_queue_slots = num_nodes*queue_tpn
    if 'shared' in kwargs:
        if kwargs['shared'] and num_nodes == 1: # Assumes fill up rule
            num_queue_slots = num_nodes*max(tpn,min_tasks_per_node(queue_id))

    record_job = "True"
    if 'no_syscribe' in kwargs:
        record_job = ""

    if 'openmpi_args' not in kwargs:
        openmpi_args = "-bysocket -bind-to-socket"
    else:
        openmpi_args =  kwargs['openmpi_args']

    script=Template(r"""#!/bin/bash
#
# PBS job submission script generated by MyCluster
#
# Job name
#PBS -N $my_name
# Redirect output stream to this file.
#PBS -o $my_output.$$PBS_JOBID
# Join the error stream to the output stream.
#PBS -j oe
# Send status information to this email address.
#PBS -M $user_email
# Send me an e-mail when the job has finished.
#PBS -m abe
# Queue name
#PBS -q $queue_name
# Project name
#PBS -A $project_name
# Maximum wall clock
#PBS -l walltime=$wall_clock
# Request number of nodes
#PBS -l nodes=$num_nodes:ppn=$tpn


export MYCLUSTER_QUEUE=$queue_name
export MYCLUSTER_JOB_NAME=$my_name
export NUM_TASKS=$num_tasks
export TASKS_PER_NODE=$tpn
export THREADS_PER_TASK=$num_threads_per_task
export NUM_NODES=$num_nodes

# OpenMP configuration
export OMP_NUM_THREADS=$$THREADS_PER_TASK
export OMP_PROC_BIND=true
export OMP_PLACES=sockets

# OpenMPI
export OMPI_CMD="mpiexec -n $$NUM_TASKS -npernode $$TASKS_PER_NODE $openmpi_args"

# MVAPICH2
export MV2_CPU_BINDING_LEVEL=SOCKET
export MV2_CPU_BINDING_POLICY=scatter
export MVAPICH_CMD="mpiexec -n $$NUM_TASKS -ppn $$TASKS_PER_NODE -bind-to-socket"

# Intel MPI
# The following variables define a sensible pinning strategy for Intel MPI tasks -
# this should be suitable for both pure MPI and hybrid MPI/OpenMP jobs:
export I_MPI_PIN_DOMAIN=omp:compact # Domains are $$OMP_NUM_THREADS cores in size
export I_MPI_PIN_ORDER=scatter # Adjacent domains have minimal sharing of caches/sockets
#export I_MPI_FABRICS=shm:ofa
export IMPI_CMD="mpiexec -n $$NUM_TASKS -ppn $$TASKS_PER_NODE"

cd $$PBS_O_WORKDIR

# Summarise environment
echo -e "JobID: $$PBS_JOBID\n======"
echo "Time: `date`"
echo "Running on master node: `hostname`"
echo "Current directory: `pwd`"

if [ "$$PBS_NODEFILE" ]; then
        #! Create a machine file:
        cat $$PBS_NODEFILE | awk '{print $$1, " slots=" $$2}' > machine.file.$$PBS_JOBID
        echo -e "\nNodes allocated:\n================"
        echo `cat machine.file.$$PBS_JOBID | sed -e 's/\..*$$//g'`
fi

echo -e "\nnumtasks=$num_tasks, numnodes=$num_nodes, tasks_per_node=$tpn (OMP_NUM_THREADS=$$OMP_NUM_THREADS)"

echo -e "\nExecuting command:\n==================\n$my_script\n"

# Run user script
. $my_script

# Report on completion
echo -e "\nJob Complete:\n==================\n"
if [ $record_job ]; then
    echo -e "\nRecording hardware setup\n==================\n"
    mycluster --sysscribe $$PBS_JOBID
    if [ "$$MYCLUSTER_APP_NAME" ]; then
        mycluster --jobid $$PBS_JOBID --appname=$$MYCLUSTER_APP_NAME
    fi
    if [ "$$MYCLUSTER_APP_DATA" ]; then
        mycluster --jobid $$PBS_JOBID --appdata=$$MYCLUSTER_APP_DATA
    fi
fi
echo -e "Complete========\n"
""")
    script_str = script.substitute({'my_name':my_name,
                                   'my_script':my_script,
                                   'my_output':my_output,
                                   'user_email':user_email,
                                   'queue_name':queue_name,
                                   'num_queue_slots':num_queue_slots,
                                   'num_tasks':num_tasks,
                                   'tpn':tpn,
                                   'num_threads_per_task':num_threads_per_task,
                                   'num_queue_slots':num_queue_slots,
                                   'num_nodes':num_nodes,
                                   'project_name': project_name,
                                   'wall_clock' : wall_clock,
                                   'openmpi_args': openmpi_args,
                                   'record_job' : record_job,
                                   })

    return script_str



def submit(script_name, immediate, depends=None):
    job_id = None
    with os.popen('qsub' + script_name) as f:
        job_id = 0
        try:
            job_id = f.readline.strip()
        except:
            print 'Failed to launch job'
            print f.readline()
            pass
    return job_id

def delete(job_id):
    with os.popen('qdel ' + job_id) as f:
        pass

def status():
    status_dict = {}
    with os.popen('qstat') as f:
        pass
    return status_dict

def job_stats(job_id):
    stats_dict = {}

    return stats_dict

def running_stats(job_id):
    stats_dict = {}

    return stats_dict

