#!/usr/bin/env python
from __future__ import print_function
"""In order to use the script you need to copy your SSH key to the target server
and also copy the server SSH public key (usually .ssh/id_rsa.pub) to .ssh/authorized_keys, 
so that the computing node can ssh passwordless to the login node"""

from subprocess import Popen, PIPE, call
import sys
import webbrowser
from getopt import getopt
import time
import os
import platform
import random
from datetime import datetime
import sys, getopt

random.seed(datetime.now())

username = ''
username = str(username)
timeinhours = 2
timeinhours = int(timeinhours)
memoryingb = 1
memoryingb = int(memoryingb) 
numberofslots = 1
numberofslots = int(numberofslots)
port = random.randrange(6222, 10000, 2)
port=int(port)
directory = ''
directory = str(directory)
opsys = platform.system()
opsys = str(opsys)
pythonver = '2.7.13'
pythonver = str(pythonver)
usegpu='no'
gpu = 'P4'
highp = 'no'
cudaver='9.1'
#qrsh_error='no'

if sys.version_info<(2,6,0):
  sys.stderr.write("You need python 2.6 or later to run this script\n")
  exit(1)


if opsys == 'Windows':
    try:
        ret=call(['where','ssh.exe'])
	#mytest = 'MINGW' in os.environ['MSYSTEM']
    except:
        print("It appears that Git for Windows is not installed")
        print("or that you are not running this script from its")
        print("Git BASH terminal. Please install either Git for")
        print("Windows, availabe at: ")
        print("https://git-for-windows.github.io/")
        print("and run this script from its git BASH.")
        sys.exit(2)

myssh = 'ssh'


def main(argv):
    global  username, timeinhours, memoryingb, numberofslots, pythonver, port, directory, usegpu, gpu, cudaver, highp
    try:
        opts, args = getopt.getopt(argv,"h:u:t:m:s:v:p:d:g:c:l:o:",["username=","timeinhours=","memoryingn=","numberofslots=","pythonver=","port=","dir=","usegpu=","gpu=","cudaver=","highp="])
    except getopt.GetoptError:
        print("Usage:")
        print("\n\t h2jupynb [-u <Hoffman2 user name>] [-t <time in hours>] [-m <memory in GB>] [-s <number of slots>] [-v <python-version>] [-p <port>] [-g <yes/no>] [-c <gpu-card-type>] [-l <cuda-version>] [-o <yes/no>] [-d <dir>] \n")
        print("If no arguments are given to this script it is assumed that:")
        print("\t your Hoffman2 user name is the same as on your client machine")
        print("\t the time duration for your session is of 2 hours")
        print("\t the memory per slot for your session is of 1GB")
        print("\t the number of slots for your session is of 1")
        print("\t the python version for your notebook is 2.7.13")
        print("\t the port on which the server is started is 8789")
        print("\t use GPU? default is no")
        print("\t GPU type default is P4 (if -g = yes)")
        print("\t CUDA version 9.1 (if -g = yes)")
        print("\t not running on owned nodes")
        print("\t the starting directory on Hoffman2 is your $HOME")
        print("\n\t python versions currently available are: 2.7.13, 3.6.1 and 3.7.0")
        print("\t python versions: 2.7.3 or 3.4 are available but deprecated.\n")
        print("\t cuda versions currently available are: 7.0, 7.5, 8.0 and 9.1\n")
        sys.exit(2)

    for opt, arg in opts:
        if opt == '-h':
            print("h2jupynb -u <Hoffman2 user name> -t <time in hours> [-p <port> [-d <dir>]")
            sys.exit()
        elif opt in ("-u", "--user"):
            username = arg
            username = str(username)
        elif opt in ("-t", "--time"):
            timeinhours = arg
            try:
                timeinhours=int(timeinhours)
            except exceptions.KeyError:
                print("Setting the time to 2 hours...")
                timeinhours = 2
        elif opt in ("-m", "--memory"):
            memoryingb = arg
            try:
                memoryingb=int(memoryingb)
            except exceptions.KeyError:
                print("Setting the memory to 1GB...")
                memoryingb = 1
        elif opt in ("-s", "--slots"):
            numberofslots = arg
            try:
                numberofslots=int(numberofslots)
            except exceptions.KeyError:
                print("Setting the number of slots to 1...")
                numberofslots = 1
        elif opt in ("-v", "--version"):
            pythonver = arg
            pythonver = str(pythonver)
            if (not pythonver == '2.7.3') and (not pythonver == '2.7.13') and (not pythonver == '3.4') and (not pythonver == '3.6.1') and (not pythonver == '3.7.0'):
                print("version, ", pythonver,", of python not available")
                print("setting python version to 2.7.13")
                pythonver = '2.7.13'
        elif opt in ("-p", "--port"):
            port = arg
            try:
                port=int(port)
            except exceptions.KeyError:
                print("Setting the port to 2 8789...")
                port = 8789
        elif opt in ("-g", "--usegpu"):
            usegpu = arg
            usegpu = str(usegpu)
            if ( usegpu == 'NO') or (usegpu == 'No') or (usegpu == 'no') or (usegpu == 'nO') or (usegpu == 'N') or (usegpu == 'n'):
                print("Jupyter notebook or lab will not run on a GPU node")
                usegpu = 'no'
            elif (usegpu == 'Yes') or (usegpu == 'yes') or (usegpu == 'YES') or (usegpu == 'Y') or (usegpu == 'y'):
                print("Jupyter notebook or lab will run on a GPU node")
                usegpu='yes'
            else:
                print("USEGPU=",usegpu)
        elif opt in ("-c", "--gpu_card"):
            gpu = arg
            gpu = str(gpu)
            if (not gpu == 'P4') and (not gpu == 'K40') and (not gpu == 'GTX1080Ti'):
                print("version, ",gpu,", of GPU not available")
                print("setting GPU version to P4")
                gpu = 'P4'
        elif opt in ("-l", "--cuda"):
            cudaver = arg
            cudaver = str(cudaver)
            if (not cudaver == '7.0') and (not cudaver == '7.5') and (not cudaver == '8.0') and (not cudaver == '9.1'):
                print("version, ",cudaver,", not available")
                print("setting cuda version to 9.1")
                gpu = '9.1'            
        elif opt in ("-o", "--highp"):
            highp = arg
            highp = str(highp)
            if ( highp == 'NO') or (highp == 'No') or (highp == 'no') or (highp == 'nO') or (highp == 'N') or (highp == 'n'): 
                print("HIGHP= ",highp," Jupyter notebook or lab will not run on owned node(s)")
                highp = 'no'
            elif ( highp == 'Yes') or ( highp == 'yes') or ( highp == 'YES') or ( highp == 'Y') or ( highp == 'y'):
                print("HIGHP= ",highp," Jupyter notebook or lab will run on owned node(s)")
                highp='yes'
        elif opt in ("-d", "--dir"):
            directory = (arg,None)
            directory = str(directory)
        else:
            print("Usage:")
            print("\n\t h2jupynb [-u <Hoffman2 user name>] [-t <time in hours>] [-m <memory in GB>] [-s <number of slots>] [-v <python-version>] [-g <yes/no>] [-c <gpu card>] [-o <yes/no>] [-p <port>] [-d <dir>] \n")
            print("If no arguments are given to this script it is assumed that:")
            print("\t your Hoffman2 user name is the same as on your client machine")
            print("\t the time duration for your session is of 2 hours")
            print("\t the memory per slot for your session is of 1GB")
            print("\t the number of slots for your session is of 1")
            print("\t the python version for your notebook is 2.7.13")
            print("\t use GPU? default is no")
            print("\t GPU type default is P4 (if -g = yes)")
            print("\t not running on owned nodes")
            print("\t the port on which the server is started is 8789")
            print("\t the starting directory on Hoffman2 is your $HOME")
            sys.exit(2)
            
    if not username:
        if opsys == 'Windows':
           username = str(os.environ["USERNAME"])
        else:
           username = str(os.environ["USER"])


    print("Your Hoffman2 user name is", username )
    print("The time in hours is", timeinhours )
    print("The memory in GB per slots is", memoryingb)
    print("The number of slots is", numberofslots)
    print("The version of python for the notebook is", pythonver)
    if (usegpu == 'no') and (highp == 'yes'):
       print("Your Jupyter NB or lab will run on owned resources")
    elif (usegpu == 'yes') and (highp == 'no'):
       print("Your Jupyter NB or lab will run on a GPU node")
       print("Your Jupyter NB or lab will load cuda version ",cudaver)
    elif (usegpu == 'yes') and (highp == 'yes'):
       print("Your Jupyter NB or lab will run on a GPU owned node")  
       print("Your Jupyter NB or lab will load cuda version ",cudaver)
    print("The port is" , port)
    if not directory:
        print("The directory on Hoffman2 is $HOME")
    else:
        print("The directory on Hoffman2 is", directory)

if __name__ == "__main__":
   main(sys.argv[1:])

if (usegpu == 'no') and (highp == 'no'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l i,h_rt=%d:00:00,h_data=%dg -pe dc\* %d")
elif (usegpu == 'no') and (highp == 'yes'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dg,highp -pe dc\* %d -w e")
elif (usegpu == 'yes') and (highp == 'no'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dg,gpu,%s -pe dc\* %d -w e")
elif (usegpu == 'yes') and (highp == 'yes'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dg,gpu,%s,highp -pe dc\* %d -w e")
MODULE_LOADCUDA_TEMP = "module load cuda/%s"
MODULE_LOAD_TEMP = "module load python/%s"
MODULE_LOAD_R = "module load R/3.5.1"

if (pythonver == '2.7.3'):
  STARTNBCMD = "ipython notebook"
else:
  STARTNBCMD = "jupyter notebook"


def readwhile(stream,func):
    while True:
        line = stream.readline()
        if line!='':
            print(line[:-1])
            if func(line): break
        else:
            raise Exception("Disconnected unexpectedly.")

pqsub=Popen(['ssh','-t','-t','-4','%s@hoffman2.idre.ucla.edu' % username],stdin=PIPE,stdout=PIPE,stderr=PIPE,universal_newlines=True)
pqsub.stdin.flush()
if (usegpu == 'no') and (highp == 'no'): 
   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,numberofslots))+"\n")
elif (usegpu == 'no') and (highp == 'yes'):
   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,numberofslots))+"\n")
elif (usegpu == 'yes') and (highp == 'no'):
   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,gpu,numberofslots))+"\n")
elif (usegpu == 'yes') and (highp == 'yes'):
   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,gpu,numberofslots))+"\n")
# check for error on qrsh:
#def getqrsherror(line):
#    global qrsh_error
#    if (line.find('error: no suitable queues')>0):
#       qrsh_error = 'yes'
#       return True
#
#readwhile(pqsub.stdout, getqrsherror)
##readwhile(pqsub.stdout, lambda line: line.find('error: no suitable queues')>0)
#if (qrsh_error == 'yes'):
#    print('The combination of resources requested is ivalid please try again...')
#    pqsub.kill()
#    #raise SystemExit
#else:
if (usegpu == 'yes'):
   pqsub.stdin.write(MODULE_LOADCUDA_TEMP % (cudaver) +"\n")
pqsub.stdin.write(MODULE_LOAD_TEMP % (pythonver) +"\n")
pqsub.stdin.write('echo HOSTNAME=`hostname`\n')
pqsub.stdin.write(MODULE_LOAD_R +"\n") #load R as well
pqsub.stdin.flush()

def gethostname(line):
    global hostname
    if line.startswith('HOSTNAME'):
        hostname = line.split('=')[1].strip()
        return True

readwhile(pqsub.stdout, gethostname)

if ('login' in hostname):
   print("We are on the wrong host. Exiting...")
   quit()

if directory: 
    pqsub.stdin.write('cd %s\n'%directory)
    pqsub.stdin.write('echo CD\n')
    readwhile(pqsub.stdout, lambda line: line.startswith('CD'))
 
def gethttpcommand(line):
    global httpcommand
    if (line.find('localhost')>0):
       httpcommand = line.split(' ')[-1]
       return True

pqsub.stdin.write(STARTNBCMD + ' --port=%s\n'%port)
pqsub.stdin.flush()

if (STARTNBCMD == 'jupyter notebook'):
  #readwhile(pqsub.stdout, lambda line: line.find('localhost')>0)
  #print("\n\tUse the token above to log into your notebook as needed\n")
  readwhile(pqsub.stdout, gethttpcommand)
  print("\n\t", httpcommand, "\n")
else:
  readwhile(pqsub.stdout, lambda line: line.find('NotebookApp')>0)

if myssh == 'ssh':
    tunnel = ['ssh','-4', '-t', '-Y', '%s@hoffman2.idre.ucla.edu' % username, '-L', '%s:localhost:%s'%(port,port), 'ssh', '-t', '-Y', hostname, '-L', '%s:localhost:%s'%(port,port)]
elif myssh == 'plink':
    tunnel = ['plink','-4', '-t', '-X', '%s@hoffman2.idre.ucla.edu' % username, '-L', '%s:localhost:%s'%(port,port), 'ssh', '-t', '-X', hostname, '-L', '%s:localhost:%s'%(port,port)]

print(' '.join(tunnel))

ptunnel = Popen(tunnel,stdout=PIPE,stdin=PIPE,universal_newlines=True)
ptunnel.stdin.write('echo TUNNEL\n')
ptunnel.stdin.flush()


readwhile(ptunnel.stdout,lambda line: line.startswith('TUNNEL'))

if ('CYGWIN' in opsys) or ('Cygwin' in opsys) or ('cygwin' in opsys):
   if (STARTNBCMD == 'jupyter notebook'):
     call(["cygstart",httpcommand])
   else:
     call(["cygstart",'http://localhost:%s'%(port)])
else:
   if (STARTNBCMD == 'jupyter notebook'):
     webbrowser.open(httpcommand)
   else:
     webbrowser.open('http://localhost:%s'%(port))

print("Succesfully opened notebook!")
print("Kill this process to end your notebook connection.")

time.sleep(timeinhours*3600)

pqsub.kill()
ptunnel.kill()


print("Succesfully cleaned up connections.")





