#AUTOGENERATED! DO NOT EDIT! File to edit: dev/om2.ipynb (unless otherwise specified).

__all__ = ['extract_tag', 'regex_tag', 'contains_tag', 'is_nbx', 'is_nbx_cell', 'is_magic_or_shell', 'regex_magic',
           'strip', 'parse_xarg_expr', 'regex_xarg', 'Bunch', 'load_nb', 'parse_src_with_parse_dict', 'parse_none',
           'parse_nbx', 'parse_xarg', 'parse_xuse', 'consume_line_below', 'parse_nbx_cell_with_parse_dict',
           'PARSE_DICT', 'concat', 'unzip', 'negate', 'is_constarg', 'get_item', 'get_items', 'not_constarg',
           'parse_nb_with_parse_dict', 'get_arrays', 'init_job', 'cont_job', 'chain_jobs', 'add_if_necessary',
           'create_script', 'create_om_files', 'create_folders', 'create_run_script', 'create_job_script', 'check_nb',
           'create_experiment_script', 'tpath', 'INSTRUCTIONS']

#Cell
#default_exp om2
import re
regex_tag = re.compile(r"^\s*#([a-zA-Z_]+).*$")

def extract_tag(line):
    """Returns the name of a tag (#name), if it
    occurs at the beginning of the line, or None."""
    m = regex_tag.match(line)
    if m is not None: return m.group(1)
    else: return None

#Cell
def contains_tag(name):
    return lambda line: extract_tag(line) == name

is_nbx = contains_tag("nbx")

#Cell
def is_nbx_cell(cell):
    if cell['cell_type'] != 'code': return False
    if not cell['source']: return False
    line0 = cell['source'][0]
    return is_nbx(line0)

#Cell
regex_magic =  re.compile(r"^\s*%{1,2}|^\s*!")

def is_magic_or_shell(line):
    """Checks if line contains a jupyter
    magic function or shell command"""
    m = regex_magic.match(line)
    return m is not None

#Cell
regex_xarg = re.compile(r"""
^
([^=]+)
=
([^;]+)
;?
(.*)
$""", re.VERBOSE)

def strip(s):
    return s.strip()

def parse_xarg_expr(line):
    """Parses the line below an `xarg` tag, e.g.
        ('x', '0', '[1,2,3,4]') = parse_xarg_expr("x = 0; [1,2,3,4]")
    """
    m = regex_xarg.match(line)
    name, val, sweep = map(strip, m.groups())
    return name, val, sweep

#Cell
import json
from argparse import Namespace

class Bunch(object):
    def __init__(self, adict={}):
        self.__dict__.update(adict)

    def __repr__(self):
        return str(self.__dict__.keys())

def load_nb(fname):
    nbdict = json.load(open(fname,'r',encoding="utf-8"))
    nb = Bunch(nbdict)
    nb.name = fname
    return nb

#Cell
def parse_src_with_parse_dict(a, src, parse_dict):
    if len(src) == 0: return a, []

    tag = extract_tag(src[0])
    if tag is None or \
       tag not in parse_dict:
            a, rest = parse_none(a, src)
    else:
            a, rest = parse_dict[tag](a, src)

    return parse_src_with_parse_dict(a, rest, parse_dict)


def parse_none(a, src):
    if not is_magic_or_shell(src[0]):
        a['none'].append(src[0])
    rest = src[1:]
    return a, rest

def parse_nbx(a, src):
    a["nbx"].append(src[0])
    rest = src[1:]
    return a, rest

def parse_xarg(a, src):
    a["xarg"].append(src[1])
    rest = src[2:]
    return a, rest

def parse_xuse(a, src):
    a["xuse"].append(src[1])
    rest = src[2:]
    return a, rest

def consume_line_below(tag, basket=None):
    if basket is None: basket = tag

    def parse(a, src):
        a[basket].append(src[1])
        rest = src[2:]
        return a, rest

    return parse


PARSE_DICT = {
    'xarg': consume_line_below('xarg', basket=None)
}

def parse_nbx_cell_with_parse_dict(cell, parse_dict=PARSE_DICT):
    a = dict([(t,[]) for t in parse_dict.keys()])
    a['none'] = []
    a, _ = parse_src_with_parse_dict(a, cell['source'], parse_dict)

    return a


#Cell
from functools import reduce

def concat(list1, list2):
    return list1 + list2

def unzip(zipped):
    return zip(*zipped)

def negate(func):
    return lambda x: not func(x)

def is_constarg(a):
    return len(a[2]) == 0

not_constarg = negate(is_constarg)

def get_item(i):
    return lambda x: x[i]

def get_items(*I):
    return lambda x: tuple([x[i] for i in I])

#Cell
def parse_nb_with_parse_dict(nb, parse_dict=PARSE_DICT):
    nbx_cells = filter(is_nbx_cell, nb.cells)

    keys = parse_dict
    A = dict([(k,[]) for k in parse_dict.keys()])
    A['func_body'] = []
    for cell in nbx_cells:
        a = parse_nbx_cell_with_parse_dict(cell, parse_dict)

        for k in parse_dict.keys():
            A[k].extend(a[k])
        A['func_body'].extend(a['none'])

    A['xarg'] = [parse_xarg_expr(line) for line in A['xarg']]
    A['args'] = list(map(get_items(0,1), A['xarg']))
    A['const_args'] = list(map(get_items(0,1), filter(is_constarg, A['xarg'])))
    A['sweep_args'] = list(map(get_items(0,2), filter(not_constarg, A['xarg'])))
    A['name'] = nb.name


    return A


#Cell
def get_arrays(num, m=1000):
    if num < m: return [[1,num]]

    arrays = []
    for i in range(num//m): arrays.append([i*m+1, (i+1)*m])
    last = arrays[-1][1]
    if last < num: arrays.append([last+1, num])

    return arrays

#Cell
def init_job(start, end, step):
    return f"job_0=`sbatch --array={start}-{end}%{step} job.sh | awk '{{ print $4 }}'`"
def cont_job(j, start, end, step):
    return f"job_{j}=`sbatch --array={start}-{end}%{step} --dependency=afterok:$job_{j-1} job.sh | awk '{{ print $4 }}'`"

def chain_jobs(arrays, step):
    s = ""
    for i, arr in enumerate(arrays):
        if i ==0: s += init_job(arr[0], arr[1], step)
        else: s += cont_job(i, arr[0], arr[1], step)
        s += "\n"
    return s

#Cell
from pathlib import PurePosixPath as Path
import pkg_resources
import importlib
from .templ import *
import os

def add_if_necessary(d, k, v):
    if k not in d:
        d[k] = v


def create_script(tpath, tname, fname, vars):
    print(f"Creating... {fname} \n\tfrom {tname}")
    create_file_from_template(tpath/tname, fname, vars)


tpath = Path(pkg_resources.resource_filename(__name__, "templates/"))


def create_om_files(target_dir, lang, num_jobs, simg, job_header, arr_size=1000, step=20):
    """
    Creates a bundle folder and all the scripts
    needed to run an experiment script on OM...

    Example usage:

    >> create_om_files(
            target_dir = "./_EXAMPLE_BUNDLE",
            lang = "py",
            num_jobs = 10,
            simg = "pytorch.simg",
            job_header = {
                "--time": "01:20:00",
                "--partition": "fiete",
                "--mem": "32gb",
                "--cpus-per-task": 4,
                "--mail-user": "me@somewhere.com"})

    >> create_experiment_script(
            nbname = "my_notebook.ipynb",
            target_dir = "./_EXAMPLE_BUNDLE",
            lang = "py")

    """
    print(f"Creating om ... files...\n\tfrom {tpath}")

    create_folders(target_dir, lang)
    create_run_script(target_dir, num_jobs, arr_size, step)
    create_job_script(target_dir, lang, simg, job_header)

    print(render_template_from_string(INSTRUCTIONS, {"path": target_dir}))


INSTRUCTIONS = """
** Instructions: **
    Copy to remote, run, and pull the results:
    - `!scp -r {{path}} $om:$omx`
    - `!ssh $om sbatch -D $omx/{{path}} $omx/{{path}}/run.sh`
    - `!scp -r $om:$omx/{{path}}/results/* ./results`

    For this to work you have to set a few environment variables...

"""


def create_folders(path, lang):
    path=Path(path)


    for p in [path, path/'io', path/'results']:
        if not os.path.exists(p): os.makedirs(p)


    if os.path.exists('./src'):
        if not os.path.exists(path/'src'):
            os.makedirs(path/'src')
        os.system(f"cp -r src/* {path/'src'}")

    if os.path.exists('./data'):
        if not os.path.exists(path/'data'):
            os.makedirs(path/'data')
        os.system(f"cp -r data/* {path/'data'}")

    if lang==".py":
        open(path/'__init__.py', 'a').close()


def create_run_script(target_dir, num_jobs, arr_size, step):
    assert arr_size <= 1000, "Maximum number of queued jobs on OM is 1000"
    fname = Path(target_dir)/'run.sh'
    with open(fname, "w", newline="\n") as f:
        f.write("#!/bin/sh\n\n")
        f.write("#SBATCH --out=io/runner_out__%A\n")
        f.write("#SBATCH --error=io/runner_err__%A\n\n")
        f.write(chain_jobs(get_arrays(num_jobs, arr_size), step))


def create_job_script(target_dir, lang, simg, job_header):

    simg        = Path(os.environ['omsimg'])/simg
    nbx_folder  = Path(os.environ['omx'])
    results_dir = Path("./results")

    add_if_necessary(job_header, "--out", "io/out_%a")
    add_if_necessary(job_header, "--error", "io/err_%a")
    add_if_necessary(job_header, "--mail-type", "END")
    add_if_necessary(job_header, "--exclude", "node030,node016,node015")

    tname = f"job_{lang}.tpl"
    fname = Path(target_dir)/'job.sh'
    create_script(tpath, tname, fname, {
        'job_header': job_header.items(),
        'nbx_folder': nbx_folder,
        'simg': simg,
        'results_dir': results_dir
    })


def check_nb(pnb):
    keys = list(map(get_item(0), pnb['args']))
    if "task_id" not in keys: raise KeyError("You didn't specify `task_id`!!")
    if "results_dir" not in keys: raise KeyError("You didn't specify `results_dir`!!")


def create_experiment_script(nbname, target_dir=".", lang="py"):
    print("** Creating Experiment script and folder **")
    nb = load_nb(nbname)
    nb = parse_nb_with_parse_dict(nb, parse_dict=PARSE_DICT)
    check_nb(nb)

    path=Path(target_dir)

    if not os.path.exists(path):
        os.makedirs(path)

    tname = f"experiment_{lang}.tpl"
    fname = path/f"experiment.{lang}"
    create_script(tpath, tname, fname, nb)


    if lang == "py":
        open(path/'__init__.py', 'a').close()
        exp = ".".join((path/'experiment').parts)
        m =  importlib.import_module(exp)
        num_params = len(m.sweep_params)

    print(f"Number of params: {num_params}")

    return {"num_jobs": num_params, "target_dir": target_dir, "lang": lang}



