#!/usr/bin/env python
from __future__ import print_function
import argparse
import json
import os
import pkg_resources
from distutils.dir_util import copy_tree
import shutil
import subprocess
import yaml
# from mako.template import Template
from pycon import pycon

try:
<<<<<<< HEAD
    from hyperband.hyperband import hyperband, get_random_hyperparameter_configuration
=======
    from numpy import argsort, ceil, log, random
    import pandas
    import hyperband
>>>>>>> NoDAG
except ImportError:
    pass

try:
    import progress.bar
    hasProgress = True
except ImportError:
    hasProgress = False
    pass

resource_package = 'pycon';
resource_path_submit = os.path.join('templates','process.mako')
submit_template_string = pkg_resources.resource_string(resource_package, resource_path_submit)

p = argparse.ArgumentParser()
# Required Positional Arguments
p.add_argument('stub')
# Flags
p.add_argument('-l','--local_data',action='store_true')
p.add_argument('-p','--posix_paths',action='store_true',help='Force the path-separater character to be / (rather than \ on Windows).')
# Options
p.add_argument('-o','--offset_index',type=int,nargs=1,default=0,help='Number of the first directory to affect/create. Intended to help add jobs to an existing set.')
p.add_argument('-r','--root_dir',type=str,nargs=1,default='.',help='Directory in which job directories will be constructed.')

args = p.parse_args()

STUBYAML = args.stub
OFFSET = args.offset_index
ROOTDIR = args.root_dir

FLAG_LOCALDATA = args.local_data
FLAG_POSIXPATHS = args.posix_paths
if FLAG_POSIXPATHS:
    import posixpath
    pathjoin = posixpath.join
    normpath = posixpath.normpath
else:
    pathjoin = os.path.join
    normpath = os.path.normpath

<<<<<<< HEAD
if SUBMITYAML:
    FLAG_SETUPDAG = True
    #dag_template = Template(dag_template_string)
    submit_template = Template(submit_template_string)
    with open(SUBMITYAML,'rb') as f:
        ProcessInfo = yaml.load(f, Loader=yaml.FullLoader)

    if not 'PRESCRIPT' in ProcessInfo:
        ProcessInfo['PRESCRIPT'] = ''

    if not 'POSTSCRIPT' in ProcessInfo:
        ProcessInfo['POSTSCRIPT'] = ''

else:
    FLAG_SETUPDAG = False

SPECIALFIELDS = ['EXPAND','URLS','COPY','CODE','DATA']
=======
SPECIALFIELDS = ['EXPAND','URLS','COPY']
>>>>>>> NoDAG
StubAsMaster = False

with open(STUBYAML,'rb') as f:
    stub = yaml.load(f, Loader=yaml.FullLoader)

try:
    EXPAND = pycon.utils.flatten(stub['EXPAND'])
except KeyError:
    print("Stub does not include an EXPAND list. Treating stub as master...")
    StubAsMaster = True

sharedir = os.path.join(ROOTDIR,'shared')
if not os.path.isdir(sharedir):
    os.makedirs(sharedir)
sharedir = pathjoin(ROOTDIR,'shared')

#############################################################
#                HYPERBAND initialization                   #
#############################################################
if 'HYPERBAND' in stub:
    budget = stub['HYPERBAND']['budget']
    eta = stub['HYPERBAND']['aggressiveness']
    hyperparams = stub['HYPERBAND']['hyperparameters']
    BRACKETS = hyperband.hyperband(budget, eta)
    for p in hyperparams:
        distribution = stub[p]['distribution']
        dist_args = stub[p]['args']
        stub[p] = [[
            hyperband.get_random_hyperparameter_configuration(distribution, dist_args)
            for i in range(s['n'][0]) ]
            for s in BRACKETS ]

    stub['BRACKETS'] = BRACKETS
    stub['EXPAND'].append(['BRACKETS'] + hyperparams)

    print("Writing stub_hb.yaml...")
    stub['SearchWithHyperband'] = True
    with open('stub_hb.yaml', 'w') as f:
        yaml.dump(dict((k,v) for k,v in stub.items() if k!='HYPERBAND'), f)
        #yaml.dump({k:v for k,v in stub.items() if k!='HYPERBAND'}, f)

else:
    if not 'SearchWithHyperband' in stub:
        stub['SearchWithHyperband'] = False

#############################################################
#         Copy files to be shared into ./shared             #
#############################################################
# After copying, the file path is updated so that the
# executable will reference the copied version of the file. If
# FLAG_LOCALDATA is True, then the path will be set to look
# for the file in the job's current directory. This is for jobs
# that run on HTCondor, since all files are transfered into the
# job directory (directories themselves are not transfered).
if 'COPY' in stub and not stub['COPY'] is None:
    COPY_shared = [field for field in stub['COPY'] if field not in EXPAND]
else:
    COPY_shared = []

if 'CODE' in stub and not stub['CODE'] is None:
    CODE_shared = [field for field in stub['CODE'] if field not in EXPAND]
else:
    CODE_shared = []

if 'DATA' in stub and not stub['DATA'] is None:
    DATA_shared = [field for field in stub['DATA'] if field not in EXPAND]
else:
    DATA_shared = []

for field in COPY_shared:
    if isinstance(stub[field],list):
        SourceList = stub[field]
        for iSource,source in enumerate(SourceList):
            target = os.path.join(sharedir,os.path.basename(source))
            if os.path.isdir(source):
                from distutils.dir_util import copy_tree
                copy_tree(source, target)
            else:
                shutil.copyfile(source, target)
            target = pathjoin(sharedir,os.path.basename(source))
            if FLAG_LOCALDATA:
                stub[field][iSource] = os.path.basename(target)
            else:
                stub[field][iSource] = normpath(pathjoin('..',sharedir,os.path.basename(target)))
    else:
        source = stub[field]
        target = os.path.join(sharedir,os.path.basename(source))
        shutil.copyfile(source, target)
        target = pathjoin(sharedir,os.path.basename(source))
        if FLAG_LOCALDATA:
            stub[field] = os.path.basename(target)
        else:
            stub[field] = normpath(pathjoin('..',sharedir,os.path.basename(target)))

for field in CODE_shared:
    if isinstance(stub[field],list):
        SourceList = stub[field]
        for iSource,source in enumerate(SourceList):
            target = os.path.join(sharedir,os.path.basename(source))
            if os.path.isdir(source):
                copy_tree(source, target)
            else:
                shutil.copyfile(source, target)
            target = pathjoin(sharedir,os.path.basename(source))
            if FLAG_LOCALDATA:
                stub[field][iSource] = os.path.basename(target)
            else:
                stub[field][iSource] = normpath(pathjoin('..',sharedir,os.path.basename(target)))
    else:
        source = stub[field]
        target = os.path.join(sharedir,os.path.basename(source))
        shutil.copyfile(source, target)
        target = pathjoin(sharedir,os.path.basename(source))
        if FLAG_LOCALDATA:
            stub[field] = os.path.basename(target)
        else:
            stub[field] = normpath(pathjoin('..',sharedir,os.path.basename(target)))

for field in DATA_shared:
    if isinstance(stub[field],list):
        SourceList = stub[field]
        for iSource,source in enumerate(SourceList):
            target = os.path.join(sharedir,os.path.basename(source))
            if os.path.isdir(source):
                copy_tree(source, target)
            else:
                shutil.copyfile(source, target)
            target = pathjoin(sharedir,os.path.basename(source))
            if FLAG_LOCALDATA:
                stub[field][iSource] = normpath(pathjoin('..',sharedir,os.path.basename(target)))
            else:
                stub[field][iSource] = os.path.basename(target)
    else:
        source = stub[field]
        target = os.path.join(sharedir,os.path.basename(source))
        shutil.copyfile(source, target)
        target = pathjoin(sharedir,os.path.basename(source))
        if FLAG_LOCALDATA:
            stub[field] = normpath(pathjoin('..',sharedir,os.path.basename(target)))
        else:
            stub[field] = os.path.basename(target)

if StubAsMaster:
    master = [stub]
else:
    master = pycon.expand_stub(stub)

#############################################################
#                  Write Queue Input File                   #
#############################################################
njobs = len(master)

if hasProgress:
    bar = progress.bar.Bar('', max=njobs)

print("Composing {njobs:d} individual jobs...".format(njobs=njobs))
width = pycon.utils.ndigits(njobs-1)
JobDirs = [os.path.join(ROOTDIR, "{job:0{w}d}".format(job=i,w=width)) for i in range(njobs)]
<<<<<<< HEAD
if 'CODE' in stub and not stub['CODE'] is None:
    CODE_uniq = [field for field in stub['CODE'] if field in EXPAND]
else:
    CODE_uniq = []

if 'DATA' in stub and not stub['DATA'] is None:
    DATA_uniq = [field for field in stub['DATA'] if field in EXPAND]
else:
    DATA_uniq = []

=======

if 'URLS' in stub:
    URLS = stub['URLS']
    with open('./queue_input.header','w') as f:
        f.write(','.join(['jobdir']+URLS) + '\n')

    with open('./queue_input.csv','w') as f:
        for i, (iJob, config) in enumerate(zip(JobDirs,master)):
            ForQueue = []
            for field in URLS:
                if isinstance(config[field], list):
                    config[field] = list(config[field])
                    SourceList = config[field]
                    ForQueue.append('"' + '","'.join(SourceList) + '"')
                    for iSource,source in enumerate(SourceList):
                        if FLAG_LOCALDATA:
                            config[field][iSource] = normpath(source)
                            # master[i][field][iSource] = normpath(source)
                        else:
                            config[field][iSource] = os.path.basename(source)
                            # master[i][field][iSource] = os.path.basename(source)
                else:
                    source = config[field]
                    ForQueue.append(source)
                    if FLAG_LOCALDATA:
                        # master[i][field] = normpath(source)
                        config[field] = normpath(source)
                    else:
                        # master[i][field] = os.path.basename(source)
                        config[field] = os.path.basename(source)

            if 'URLS' in config:
                del config['URLS']

            f.write(','.join([os.path.basename(iJob)] + ForQueue) +  '\n')

#############################################################
#              Setup individual jobs structure              #
#############################################################
>>>>>>> NoDAG
if 'COPY' in stub and not stub['COPY'] is None:
    COPY_uniq = [field for field in stub['COPY'] if field in EXPAND]
else:
    COPY_uniq = []

if 'URLS' in stub:
    URLS_uniq = [field for field in stub['URLS'] if field in EXPAND]
else:
    URLS_uniq = []

for iJob,config in zip(JobDirs,master):
    if not os.path.isdir(iJob):
        os.makedirs(iJob)
    #############################################################
    #                       Copy files                          #
    #############################################################
    for field in COPY_uniq:
        if isinstance(config[field], list):
            config[field] = list(config[field])
            SourceList = config[field]
            for iSource,source in enumerate(SourceList):
                target = os.path.join(iJob,os.path.basename(source))
                if os.path.isdir(source):
                    copy_tree(source, target)
                else:
                    shutil.copyfile(source, target)

                target = pathjoin(iJob,os.path.basename(source))
                if FLAG_LOCALDATA:
                    config[field][iSource] = os.path.basename(target)
                else:
                    config[field][iSource] = normpath(target)
        else:
            source = config[field]
            target = os.path.join(iJob,os.path.basename(source))
            shutil.copyfile(source, target)
            target = pathjoin(iJob,os.path.basename(source))
            if FLAG_LOCALDATA:
                config[field] = os.path.basename(target)
            else:
                config[field] = normpath(target)

    for field in CODE_uniq:
        if isinstance(config[field], list):
            config[field] = list(config[field])
            SourceList = config[field]
            for iSource,source in enumerate(SourceList):
                target = os.path.join(iJob,os.path.basename(source))
                if os.path.isdir(source):
                    copy_tree(source, target)
                else:
                    shutil.copyfile(source, target)

                target = pathjoin(iJob,os.path.basename(source))
                if FLAG_LOCALDATA:
                    config[field][iSource] = os.path.basename(target)
                else:
                    config[field][iSource] = normpath(target)
        else:
            source = config[field]
            target = os.path.join(iJob,os.path.basename(source))
            shutil.copyfile(source, target)
            target = pathjoin(iJob,os.path.basename(source))
            if FLAG_LOCALDATA:
                config[field] = os.path.basename(target)
            else:
                config[field] = normpath(target)

    for field in DATA_uniq:
        if isinstance(config[field], list):
            config[field] = list(config[field])
            SourceList = config[field]
            for iSource,source in enumerate(SourceList):
                target = os.path.join(iJob,os.path.basename(source))
                if os.path.isdir(source):
                    copy_tree(source, target)
                else:
                    shutil.copyfile(source, target)

                target = pathjoin(iJob,os.path.basename(source))
                if FLAG_LOCALDATA:
                    config[field][iSource] = normpath(target)
                else:
                    config[field][iSource] = os.path.basename(target)
        else:
            source = config[field]
            target = os.path.join(iJob,os.path.basename(source))
            shutil.copyfile(source, target)
            target = pathjoin(iJob,os.path.basename(source))
            if FLAG_LOCALDATA:
                config[field] = normpath(target)
            else:
                config[field] = os.path.basename(target)

    #############################################################
    #           Distribute params.json file to each job         #
    #############################################################
    paramfile = os.path.join(iJob,'params.json')
    if 'COPY' in config:
        del config['COPY']

    with open(paramfile, 'w') as f:
        json.dump(config, f, sort_keys = True, indent = 4)

    if hasProgress:
        bar.next()


if hasProgress:
    bar.finish()

with open('.njobs','w') as f:
    f.write(str(njobs)+'\n')

print("Done.")
