#!/usr/bin/env python
###
# This script allows users to start a jupyter notebooks/labs session
# running on the UCLA Hoffman2 Cluster while displaying it on their
# local browser. 
#
# Users download this script on their local machine and run it on a local
# terminal via:
#
# python3 h2jupynb -u $H2-USERNAME [other options]
# 
# ***The script prompts the users twice for their Hoffman2 password***
#
#
# To see a full list of current available options:
#
# python3 h2jupynb --help
#
# See also:
#
# https://www.hoffman2.idre.ucla.edu/Using-H2/Connecting/Connecting.html#connecting-via-jupyter-notebook-lab
#
# For questions open a ticket on:
# 
# https://support.idre.ucla.edu/helpdesk/Tickets/New?categoryId=390006
# 
# or file an issue or discussion on:
#
# https://github.com/rdauria/jupyter-notebook
#
# Author: RD (hpc@ucla.edu) / Adapted from: https://github.com/pyHPC/ipynbhpc
#
###

from __future__ import print_function
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
import signal

random.seed(datetime.now().strftime("%S"))

script_version=3.1
username = ''
username = str(username)
timeinhours = 2
timeinhours = int(timeinhours)
memoryingb = 1
memoryingb = int(memoryingb) 
parenv = 'shared'
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 = 'python/3.9.6'
pythonver = str(pythonver)
usegpu = 'no'
gpu = ''
highp = 'no'
cudaver = '10.2'
exclusive = 'no'
arch = ''
mods = ''
lofmods = ''
modstoload = ''
cve = ''
pve = ''
#gputype = ''
gpumem = ''
sshlogs= 'no'

print("h2jupynb version ",script_version)

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)


def usage():
  print("Usage:")
  print("h2jupynb [-u <Hoffman2 user name>] [-v <python-version>] \n         [-t <time, integer number of hours>] [-m <memory, integer number of GB per core>] \n         [-e <parallel environment: 1 for shared, 2 for distributed>] [-s <number of slots>] \n         [-o <run on group owned nodes: yes/no>] [-x <run on an exclusively reserved node: yes/no>]  \n         [-a <CPU-type>] [-d <path to directory from which the jupyter NB/lab will start>] \n         [-g <run on a gpu node: yes/no>] [-c <gpu-card-type>] [-r <gpu-global-memory>] \n         [-l <cuda-version>]  [-p <port number>] [-j <conda virtual environment name>]\n         [-k <path/to/python/virtual-env>] \n         [-b <comma separated list of environmental modules>]\n         [-z <write ssh debug files?: yes/no>]\n")
  print("If no arguments are given to this script it is assumed that:\n")
  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 parallel environment is shared")
  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 3.9.6")
  print("\t the port on which the server is started is 8789")
  print("\t the starting directory on Hoffman2 is your $HOME")
  print("\t use GPU? default is no")
  print("\t GPU type default is RTX2080Ti (if -g yes)")
  print("\t GPU memory (if -g yes) not specified")
  print("\t CUDA version 10.2 (if -g yes)")
  print("\t not running on owned nodes")
  print("\t not running in exclusive mode")
  print("\t no specific CPU selected (see \"ARCH\" in output of \"qhost\")")
  print("\n\t Python versions currently available are: 2.7.18, 3.7.3, 3.9.6, 3.11.5,")
  print("\t anaconda3/2020.07, anaconda3 (2020.11), anaconda3/2021.11,")
  print("\t anaconda3/2022.05, anaconda3/2023.03, mamba/1.4.2, mamba/1.5.5\n")
  print("\t Architectures (CPUs) currently publicly available: see output of qhost")
  print("\t to request a specific architecture, e.g., intel-gold*: -a intel-gold\\\\*\n")
  print("\t Cuda versions currently available are: 9.2, 10.0, 10.2, 11.0, 11.7, 11.8, 12.3\n")
  print("\t GPU cards currently available are: P4, GTX1080Ti, RTX2080Ti, V100, A100, A6000, H100\n")
  print("\t GPU memory currently available are (in GB): 8 (P4), 12 (GTX1080Ti), 10 (RTX2080Ti),\n\t 32 (V100), 40 or 80 (A100), 48 (A6000), 96 (H100)\n")
  print("\t Environmental modules to be loaded as a comma (no space) list (e.g.:gsl/2.6,")
  print("\t or gsl,curl/7.70.0)\n")
  print("\t Conda virtual environment - assumes a version of anaconda3/mamba is requested")
  print("\t (e.g.:treemix)\n")
  print("\t Python virtual environment - assumes a version of python (not anaconda3/mamba)")
  print('\t is requested (e.g.: $HOME/PATH/TO/MYPYTHONENV)\n')
  print("\t **Should you specify non existing parameters default values will be assumed**\n")
  sys.exit(2)

    
def main(argv):
  global  username, timeinhours, memoryingb, numberofslots, pythonver, port, directory, usegpu, gpu, gpumem, cudaver, highp, exclusive, arch, parenv, mods, cve, pve, lofmods, modstoload, sshlogs
  try:
    opts, args = getopt.getopt(sys.argv[1:],"h:u:t:m:e:s:v:p:d:g:c:r:l:o:x:a:b:j:k:z:",["help","username=","timeinhours=","memoryingn=","parallel-enviromnment=","numberofslots=","pythonver=","port=","dir=","usegpu=","gpu_card=","gpu_mem","cudaver=","highp=","exclusive=","arch=","modules=","conda-virt-env=","python-venv-path=","ssh-debug-files="])
  except getopt.GetoptError:
    usage() 
  for opt, arg in opts:
    #print(opt, arg)
    if opt in ("-h", "--help"):
      usage()
      #print("h2jupynb [-u <Hoffman2 user name>] [-v <python-version>] \n         [-t <time, integer number of hours>] [-m <memory, integer number of GB per core>] \n         [-e <parallel environment: 1 for shared, 2 for distributed>] [-s <number of slots>] \n         [-o <run on group owned nodes: yes/no>] [-x <run on an exclusively reserved node: yes/no>]  \n         [-a <CPU-type>] [-d <path to directory from which the jupyter NB/lab will start>] \n         [-g <run on a gpu node: yes/no>] [-c <gpu-card-type>]  [-r <gpu-global-memory>] \n         [-l <cuda-version>] [-p <port number>]  [-j <conda virtual environment name>]\n         [-k <path/to/python/virtual-env>] \n         [-b <comma separated list of environmental modules>]\n         [-z <write ssh debug files?: yes/no>]\n")
      sys.exit()
    elif opt in ("-u", "--user"):
      username = arg
      username = str(username)
    elif opt in ("-w", "--what"):
      a_command = arg
      a_command = str(a_command)
      print("A_command= ",a_command)
    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 ("-e", "--parallel-environment"):
      parenv = arg
      try:
        parenv=int(parenv)
      except exceptions.KeyError:
        print("Setting the parallel environmen to shared...")
        parenv = 1
        if (parenv == 1):
          parenv = 'shared'
        elif (parenv == 2):
          parenv = 'dc\\*'
        else:
          print("Setting the parallel environmen to shared...")
          parenv = 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.18') and  (not pythonver == '3.7.3') and (not pythonver == '3.9.6') and (not pythonver == '3.11.5') and (not pythonver == 'anaconda3/2020.11') and (not pythonver == 'anaconda3/2020.07') and (not pythonver == 'anaconda3/2021.11') and (not pythonver == 'anaconda3/2022.05') and (not pythonver == 'anaconda3/2023.03') and (not pythonver == 'anaconda3') and (not pythonver == 'mamba/1.4.2') and (not pythonver == 'mamba/1.5.5') and (not pythonver == 'mamba'):
        print("version, ", pythonver,", of python not available")
        print("setting python version to 3.9.6")
        pythonver = 'python/3.9.6'
        if (pythonver == '3.7.3'):
          pythonver = "python/"+pythonver
        elif (pythonver == '2.7.18'):
          pythonver = "python/"+pythonver
        elif (pythonver == '3.9.6'):
          pythonver = "python/"+pythonver
        elif (pythonver == '3.11.5'):
          pythonver = "python/"+pythonver
        elif (pythonver == 'anaconda3/2020.07'):
          pythonver = "anaconda3/2020.07"
        elif (pythonver == 'anaconda3'):
          pythonver = "anaconda3"
        elif (pythonver == 'anaconda3/2020.11'):
          pythonver = "anaconda3/2020.11"
        elif (pythonver == 'anaconda3/2021.11'):
          pythonver = "anaconda3/2021.11"
        elif (pythonver == 'anaconda3/2022.05'):
          pythonver = "anaconda3/2022.05"
        elif (pythonver == 'anaconda3/2023.03'):
          pythonver = "anaconda3/2023.03"
        elif (pythonver == 'mamba/1.4.2'):
          pythonver = "mamba/1.4.2"
        elif (pythonver == 'mamba/1.5.5'):
          pythonver = "mamba/1.5.5"
        elif (pythonver == 'mamba'):
          pythonver = "mamba"
    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 ("-d", "--dir"):
      directory = arg
      directory = str(directory)
    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 == 'RTX2080Ti') and (not gpu == 'GTX1080Ti') and (not gpu == 'P4') and (not gpu == 'V100') and (not gpu == 'A100') and (not gpu == 'A6000') :
        print("version, ",gpu,", of GPU not available")
        print("setting GPU card to RTX2080Ti")
        gpu = 'RTX2080Ti'
        usegpu='yes'
    elif opt in ("-r", "--gpu_mem"):
      gpumem = arg
      gpumem = str(gpumem)
      if (gpumem == '8'):
        print("GPU global memory = ",gpumem," available for card P4")
        print("setting GPU card to P4")
        gpu = 'P4'
        usegpu='yes'
      elif (gpumem == '12'):
        print("GPU global memory = ",gpumem," available for card GTX1080Ti")
        print("setting GPU card to GTX1080Ti")
        gpu = 'GTX1080Ti'
        usegpu='yes'
      elif (gpumem == '10'):
        print("GPU global memory = ",gpumem," available for card RTX2080Ti")
        print("setting GPU card to RTX2080Ti")
        gpu = 'RTX2080Ti'
        usegpu='yes'
      elif (gpumem == '32'):
        print("GPU global memory = ",gpumem," available for card V100")
        print("setting GPU card to V100")
        gpu = 'V100'
        usegpu='yes'
      elif (gpumem == '40'):
        print("GPU global memory = ",gpumem," available for card A100")
        print("setting GPU card to A100")
        gpu = 'A100'
        usegpu='yes'
      elif (gpumem == '42'):
        print("GPU global memory = ",gpumem," available for card A6000")
        print("setting GPU card to A6000")
        gpu = 'A6000'
        usegpu='yes'
      elif ( gpumem == '80'):
        print("GPU global memory = ",gpumem," available for card A100")
        print("setting GPU card to A100")
        gpu = 'A100'
        usegpu='yes'
      else:
        print("GPU global memory = ",gpumem," not available")
        print("setting GPU global memory version to NONE")
        gpumem = ''
    elif opt in ("-l", "--cuda"):
      cudaver = arg
      cudaver = str(cudaver)
      if (not cudaver == '9.2') and (not cudaver == '10.0') and (not cudaver == '10.2') and (not cudaver == '11.0') and (not cudaver == '11.7') and (not cudaver == '11.8') and (not cudaver == '12.3') and (not cudaver == '12.4'):
        print("version, ",cudaver,", not available")
        print("setting cuda version to 10.2")
        cudaver = '10.2'
    elif opt in ("-o", "--highp"):
      highp = arg
      highp = str(highp)
      #print("HIGHP= ",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 ("-x", "--exclusive"):
      exclusive = arg
      #print("EXCL= ",exclusive)
      exclusive = str(exclusive)
      #print("EXCL= ",exclusive)
      if ( exclusive == 'NO') or (exclusive == 'No') or (exclusive == 'no') or (exclusive == 'nO') or (exclusive == 'N') or (exclusive == 'n'):
        print("Jupyter notebook or lab will not run in exclusive mode")
        exclusive = 'no'
      elif (exclusive == 'Yes') or (exclusive == 'yes') or (exclusive == 'YES') or (exclusive == 'Y') or (exclusive == 'y'):
        print("Jupyter notebook or lab will run in exclusive mode")
        exclusive='yes'
      else:
        print("EXCLUSIVE=",exclusive)
    elif opt in ("-a", "--arch"):
      arch = arg
      arch = str(arch)
      if ( not arch == 'amd-epyc-7642' ) and ( not arch == 'intel-E5-2650' ) and ( not arch == 'intel-E5-2650v2' ) and ( not arch == 'intel-E5-2650v4' ) and ( not arch == 'intel-E5-2670' ) and ( not arch == 'intel-E5-2670v2' ) and ( not arch == 'intel-E5-2670v3' ) and ( not arch == 'intel-E5-2697Av4' ) and ( not arch == 'intel-E5-4620v2' ) and ( not arch == 'intel-E7-8890v4' ) and ( not arch == 'intel-gold-5118' ) and ( not arch == 'intel-gold-5218' ) and ( not arch == 'intel-gold-6140' ) and ( not arch == 'intel-gold-6240' ) and ( not arch == 'lx-amd64' ) and (not arch == 'intel\\*') and (not arch == 'intel-gold\\*'):
        print("version, ",arch,", of CPU not available")
        print("no CPU-type will be requested")
        arch = ''
    elif opt in ("-b", "--modules"):
      mods = arg
      mods = str(mods)
      lofmods=mods.split(",")
      for l in lofmods:
        if ( modstoload == '' ):
          modstoload="module load " + l
        else:
          modstoload=modstoload + "; module load " + l
    elif opt in ("-j", "--conda-virt-env"):
      cve = arg
      cve = str(cve)
      print("cve=",cve)
    elif opt in ("-k", "--python-venv-path"):
      pve = arg
      pve = str(pve)
      print("pve=",pve)
    elif opt in ("-z", "--ssh-debug-files"):
      sshlogs = arg
      sshlogs = str(sshlogs)
      if ( sshlogs == 'NO') or (sshlogs == 'No') or (sshlogs == 'no') or (sshlogs == 'nO') or (sshlogs == 'N') or (sshlogs == 'n'): 
        print("SSHLOGS= ",sshlogs," SSH logfiles will not be written.")
        sshlogs = 'no'
      elif ( sshlogs == 'Yes') or ( sshlogs == 'yes') or ( sshlogs == 'YES') or ( sshlogs == 'Y') or ( sshlogs == 'y'):
        print("SSHLOGS= ",sshlogs," SSH logfiles will be written.")
        sshlogs='yes'
    else:
      usage()
          
  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)
  if (numberofslots > 1):
    print("The parallel environment is", parenv)
    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)
  if (gpumem != ''):
    print("Your Jupyter NB or lab will run on GPU card ",gpu," with ",gpumem,"GB of global GPU memory")
  if (gpumem == '') and (gpu != ''):
    print("Your Jupyter NB or lab will run on GPU card ",gpu)
  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)
  if (port != ''):
    print("Your Jupyter NB or lab will run on port ",port)
  if (exclusive == 'yes'):
    print("Your Jupyter NB or lab will reserve a node exclusively ")
  if (directory == '~'):
    print("The directory on Hoffman2 is $HOME")
  else:
    print("The directory on Hoffman2 is", directory)
  if (arch != ''):
    print("Your Jupyter NB or lab will run on ",arch," CPU")
  if (mods != ''):
    print("Your Jupyter NB or lab will load these environmental modules ",mods)
    #for l in lofmods:
    #  print("MODS=",l)
  if (cve != ''):
    print("Your Jupyter NB or lab will load this conda virtual environment ",cve)  
  if (pve != ''):
    print("Your Jupyter NB or lab will load this python virtual environment ",pve)
  if (sshlogs != ''):
    print("SSH log files will be written ",sshlogs)

  ## function main def ends here

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

print("usegpu= ",usegpu,"; highp= ",highp,"; exclusive= ",exclusive, "; gputype= ",gpu, "; gpu_mem= ",gpumem, "; arch= ",arch, "; time= ",timeinhours, "; mem= ",memoryingb, "; PE= ", parenv, "; NSLOTS= ",numberofslots,"sshlogs",sshlogs)

if ( arch != '' ):
  arch=',arch='+arch

if ( gpumem != '' ):
  gpumem=',gpumem='+gpumem

if (usegpu == 'no') and (highp == 'no') and (exclusive == 'no'):
#   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l i,h_rt=%d:00:00,h_data=%dG%s -pe dc\* %d")
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG%s -pe %s %d -now n")
elif (usegpu == 'no') and (highp == 'yes') and (exclusive == 'no'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,highp%s -pe %s %d -w e")
elif (usegpu == 'no') and (highp == 'no') and (exclusive == 'yes'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,exclusive%s -pe %s %d -w e")
elif (usegpu == 'no') and (highp == 'yes') and (exclusive == 'yes'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,exclusive,highp%s -pe %s %d -w e")
elif (usegpu == 'yes') and (highp == 'no') and (exclusive == 'no'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,gpu,%s%s%s -pe %s %d -w e")
elif (usegpu == 'yes') and (highp == 'yes') and (exclusive == 'no'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,gpu,%s%s,highp%s -pe %s %d -w e")
elif (usegpu == 'yes') and (highp == 'yes') and (exclusive == 'yes'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,gpu,%s%s,highp,exclusive%s -pe %s %d -w e")
elif (usegpu == 'yes') and (highp == 'no') and (exclusive == 'yes'):
   QSUB_TEMPLATE = os.environ.get("IPYNB_QSUB_TEMPLATE", "qrsh -N JUPYNB -l h_rt=%d:00:00,h_data=%dG,gpu,%s%s,exclusive%s -pe %s %d -w e")
MODULE_LOADCUDA_TEMP = "module load cuda/%s"
MODULE_LOAD_TEMP = "module load %s"
CONDA_ACTIVATE = "conda activate %s"
PYTHON_ACTIVATE= "source %s/bin/activate"
CD_TEMP = "cd %s"



STARTNBCMD = "jupyter notebook"

print('STARTNBCMD = ',STARTNBCMD)

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


if (sshlogs == 'no'):          
  pqsub=Popen(['ssh','-o','ServerAliveCountMax=5','-o','IPQoS=throughput','-o','ServerAliveInterval=30','-X','-Y','-t','-t','-4','%s@hoffman2.idre.ucla.edu' % username],stdin=PIPE,stdout=PIPE,stderr=PIPE,universal_newlines=True)
  pqsub.stdin.flush()
  pqsub.stdout.flush()
if (sshlogs == 'yes'):          
  pqsub=Popen(['ssh','-o','ServerAliveCountMax=5','-o','IPQoS=throughput','-o','ServerAliveInterval=30','-vvv','-E','myssh.log','-X','-Y','-t','-t','-4','%s@hoffman2.idre.ucla.edu' % username],stdin=PIPE,stdout=PIPE,stderr=PIPE,universal_newlines=True)
  pqsub.stdin.flush()
  pqsub.stdout.flush()



if (usegpu == 'no'):
   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,arch,parenv,numberofslots))+"\n")
elif (usegpu == 'yes'):
   pqsub.stdin.write((QSUB_TEMPLATE % (timeinhours,memoryingb,gpu,gpumem,arch,parenv,numberofslots))+"\n") 



if (usegpu == 'yes'):
  pqsub.stdin.write(MODULE_LOADCUDA_TEMP % (cudaver) +"\n")
if (mods != ''):
  pqsub.stdin.write(modstoload +"\n")
if (pythonver == 'python/3.11.5'):
  pqsub.stdin.write("module load gcc/7.5.0" +"\n")
pqsub.stdin.write(MODULE_LOAD_TEMP % (pythonver) +"\n")
if (cve != ''):
  pqsub.stdin.write(CONDA_ACTIVATE % (cve) +"\n")
if (pve != ''):
  pqsub.stdin.write(PYTHON_ACTIVATE % (pve) +"\n")
pqsub.stdin.write('module li\n')
pqsub.stdin.write(CD_TEMP % (directory) +"\n")
pqsub.stdin.write('echo HOSTNAME=`hostname`\n')
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("\n ############################################")
   print(" ******We seem to be on the wrong host!******")
   print(" Most likely the Jupyter session cannot start")
   print(" because the interactive session cannot be   ")
   print(" allocated. Please open a ticket at:         ")
   print(" https://support.idre.ucla.edu/helpdesk      ")
   print(" providing the exact command you used to     ")
   print(" start the session. Copying and pasting the  ")
   print(" output of your terminal from the issuing of ")
   print(" the h2jupynb command will expedite matters. ")
   print(" **NOTE: GPU requests will exit if a gpu node")
   print("         is not available when the request is")
   print("         made. Check GPU availability with:  ")
   print("         GPU_NODES_AT_A_GLANCE_WITH_USED_GPUs")
   print(" Exiting...                                  ")
   print(" #############################################\n")
   quit()

 
def gethttpcommand(line):
    global httpcommand
    if (line.find('http://%s:%s/'%(hostname,port))>0):
       httpcommand = line.split(' ')[-1].strip()
       return True


pqsub.stdin.write(STARTNBCMD + ' --no-browser --ip=0.0.0.0 --port=%s; echo exitcode: $?\n'%port)
pqsub.stdin.flush()


if (STARTNBCMD == 'jupyter notebook'):
  readwhile(pqsub.stdout, gethttpcommand)
else:
  readwhile(pqsub.stdout, lambda line: line.find('http://%s:%s/'%(hostname,port))>0)

print("httpcommand=",httpcommand)  
  
token = httpcommand.split("=")[-1]

newhttpcommand = 'http://localhost:%s/?token=%s'%(port,token)

print("httpcommand=",newhttpcommand)  
  

if (sshlogs == 'no'): 
  tunnel = ['ssh','-o','ServerAliveCountMax=5','-o','IPQoS=throughput','-o','ServerAliveInterval=30','-g', '-L', '%s:%s:%s'%(port,hostname,port), '-t', '%s@hoffman2.idre.ucla.edu' % username]
if (sshlogs == 'yes'):
  tunnel = ['ssh','-o','ServerAliveCountMax=5','-o','IPQoS=throughput','-o','ServerAliveInterval=30','-vvv','-E','mytunnel.log','-g', '-L', '%s:%s:%s'%(port,hostname,port), '-t', '%s@hoffman2.idre.ucla.edu' % username]

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",newhttpcommand])
   else:
     call(["cygstart",'http://localhost:%s'%(port)])
else:
   if (STARTNBCMD == 'jupyter notebook'):
     webbrowser.open(newhttpcommand)
   else:
     webbrowser.open('http://localhost:%s'%(port))

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


def handler(signum, frame):
    res = print("Detected Ctrl-c. Exiting...")
    pqsub.kill()
    ptunnel.kill()
    print("Succesfully cleaned up connections.")
    exit(1)
 
signal.signal(signal.SIGINT, handler)

#ptunnel.stdin.write('echo TESTTESTTESTTEST#############')
#ptunnel.stdin.flush()
#print(repr(ptunnel.stdout.readline()))
#readwhile(ptunnel.stdout, lambda line: line.startswith('TESTTESTTESTTEST#############')>0)

pver=sys.version_info[0]

if pver == 2:
 for i in xrange(timeinhours*3600,0,-10):
    sys.stdout.write(str(i)+' ')
    sys.stdout.flush()
    #ptunnel.stdin.write('echo T')
    #ptunnel.stdin.flush()
    #print(repr(ptunnel.stdout.readline()))
    time.sleep(10)
elif pver == 3:
 for i in range(timeinhours*3600,0,-10):
    sys.stdout.write(str(i)+' ')
    sys.stdout.flush()
    #ptunnel.stdin.write('echo T')
    #ptunnel.stdin.flush()
    #print(repr(ptunnel.stdout.readline()))
    time.sleep(10)

#time.sleep(timeinhours*3600)


pqsub.kill()
ptunnel.kill()

print("Elapsted runtime. Terminating...")
print("Succesfully cleaned up connections.")