import os, sys
import argparse

def mkdir(path):
    try:
        os.mkdir(path)
    except OSError:
        print ("Creation of the directory %s failed" % path)
    else:
        print ("Successfully created the directory %s " % path)


def start_experiment(exp):
    # clean up, repackage code
    gpu = 'v100'
    cmd = 'rm code.zip; zip code.zip **.py **.sh **.yml **/*.py'
    os.system(cmd)

    exp_files_path = 'experiment-files'
    exp_name = '{}_model-{}_embeddings-{}_annotations-{}_mode-{}_humantrain-{}_humantest-{}_synthtrain-{}_synthtest-{}_reshape-{}_batchsize-{}_itersperepoch'.format(exp['model'], exp['embedding_type'], exp['annotations'], exp['mode'], exp['max_train_human'], exp['max_test_human'], exp['max_train_synthetic'], exp['max_test_synthetic'], exp['reshape_type'], exp['batch_size'], exp['iters_per_epoch'])

    mkdir(exp_files_path)
    mkdir('experiment-files/' + exp_name)
    mkdir('experiment-files/' + exp_name + "/results/")

    outpath = '{}-results'.format(exp_name)
    pyfile = '{}.py'.format(exp_name)
    ymlfile = '{}.yml'.format(exp_name)
    outtxt = '{}.txt'.format(exp_name)

    exp_files_path = exp_files_path + "/" + exp_name

    # create python file
    py = '''\n
import os  \n
exp1 = 'python2 reinforcement.py --model {} --embedding_type {} --annotations {} --mode {} --max_train_human {} --max_test_human {} --max_train_synthetic {} --max_test_synthetic {} --bert_reshape_type {} --batch_size {} --iters_per_epoch {} --epochs 200'\n
os.system(exp1)\n'''.format(exp['model'], exp['embedding_type'], exp['annotations'], exp['mode'], exp['max_train_human'], exp['max_test_human'], exp['max_train_synthetic'], exp['max_test_synthetic'], exp['reshape_type'], exp['batch_size'], exp['iters_per_epoch'])

    with open(pyfile, 'w') as f:
        f.write(py)


    sh = '''\n
# >>> conda initialize >>>\n
__conda_setup="$('/opt/conda/bin/conda' 'shell.bash' 'hook' 2> /dev/null)"\n
if [ $? -eq 0 ]; then\n
    eval "$__conda_setup"\n
else\n
    if [ -f "/opt/conda/etc/profile.d/conda.sh" ]; then\n
        . "/opt/conda/etc/profile.d/conda.sh"\n
    else\n
        export PATH="/opt/conda/bin:$PATH"\n
    fi\n
fi\n
unset __conda_setup\n
# <<< conda initialize <<<\n

set -ex\n
exec 1>training-log.txt \n
exec 2>training-log-err.txt \n

mkdir -p ~/.conda_envs/spr\n
cd ~/.conda_envs\n
cp "${{DATA_DIR}}/spr.tar.gz" .\n
chown gpuuser: spr.tar.gz\n
chmod 600 spr.tar.gz\n
tar xfzo spr.tar.gz -C spr\n
conda activate ~/.conda_envs/spr\n


# cd /job/model-code/
rm -rf "${{DATA_DIR}}/{}"
mkdir -p "${{DATA_DIR}}/{}"
cp -R /job/model-code "${{DATA_DIR}}/{}/code"
cd "${{DATA_DIR}}/{}/code"
# cp -R /job/model-code "${{RESULT_DIR}}"/code
# cd "${{RESULT_DIR}}"/code

export PYTHONUNBUFFERED=1
set +e
bash -lc "$*"
result=$?
echo "${{result}}" > exit_code
exit "${{result}}"
'''.format(exp_name, exp_name, exp_name, exp_name)

    with open('run.sh', 'w') as f:
        f.write(sh)

    # create yml file
    yml ='''\n
model_definition:\n
    framework: \n
        name: pytorch \n
        version: '0.4' \n
        runtimes: \n
            name: python \n
            version: '2.7' \n
    name: {} \n
    description: This is a sample training-run \n
    execution:\n
        command: bash run.sh python {} \n
        compute_configuration: \n
            name: {} \n
training_data_reference: \n
    name: training-data-reference_name \n
    connection: \n
        endpoint_url: https://s3.us-south.objectstorage.softlayer.net \n
        access_key_id: 8ffab3dfa4fa44158659c33117043b58 \n
        secret_access_key: 876e94a6c02ee334a6a61c126e16e7898c260969d6b59e88 \n
    source:\n
        bucket: qq-data\n
    type: s3\n
training_results_reference:\n
    name: training-results-reference_name\n
    connection:\n
        endpoint_url: https://s3.us-south.objectstorage.softlayer.net \n
        access_key_id: 8ffab3dfa4fa44158659c33117043b58 \n
        secret_access_key: 876e94a6c02ee334a6a61c126e16e7898c260969d6b59e88 \n
    target: \n
        bucket: qq-results \n
    type: s3 \n
'''.format(exp_name, pyfile, gpu)

    with open(ymlfile, 'w') as f:
        f.write(yml)

    cmd = 'zip code.zip {} {} {}'.format(pyfile, ymlfile, 'run.sh')
    os.system(cmd)

    # start exp
    cmd = 'bx ml train code.zip {}'.format(ymlfile)
    os.system(cmd)

    cmd = 'mv {}.yml {}; mv {}.py {}; mv code.zip {}; mv run.sh {}'.format(exp_name, exp_files_path, exp_name, exp_files_path, exp_files_path, exp_files_path)
    print('moving cmd: {}'.format(cmd))
    os.system(cmd)

    sys.exit(0)





def main():
    datasets = {}

    # dataset specific (num examples per category)
    datasets['local-small'] = {'num_train_human': 10, 'num_test_human': 2, 'num_train_synthetic': 0, 'num_test_synthetic': 0}
    # datasets['local'] = {'num_train_human': 1600, 'num_test_human': 400, 'num_train_synthetic': 0, 'num_test_synthetic': 0}

    # datasets['global-small'] = {'num_train_human': 10, 'num_test_human': 2, 'num_train_synthetic': 0, 'num_test_synthetic': 0}
    # datasets['global-small'] = {'num_train_human': 1100, 'num_test_human': 200, 'num_train_synthetic': 0, 'num_test_synthetic': 0}


    # model specific (embeddings allowed per model, reshape layers type)
    model_configs = {}
    model_configs['bert-full'] = {'embedding_type':
                                                    [
                                                    'bert',
                                                    'one-hot',
                                                    'random',
                                                    'bert-fixed',
                                                    'bert-word',
                                                    'bert-word-fixed'
                                                    ],
                                  'reshape_type':
                                                    [
                                                    'linear',
                                                    'MLP'
                                                    ]}


    # model_configs['gpt2'] = model_configs['bert-full']
    # model_configs['gpt2-large'] = model_configs['bert-full']
    model_configs['lstm'] = {'embedding_type': 'lstm'}

    # agnostic
    batch_size = [1, 4, 8, 32]
    iters_per_epoch = ['default', '500', '1000']

    # TODO + options for learning rate schedule, + support for no-gradient models

    exp_params = {}
    i = 0
    for batch in batch_size:
        if i > 0:
            sys.exit(0)
        exp_params['batch_size'] = batch
        for iters in iters_per_epoch:
            exp_params['iters_per_epoch'] = iters
            for mode, info in datasets.items():
                if 'local' in mode:
                    exp_params['mode'] = 'local'
                elif 'global' in mode:
                    exp_params['mode'] = 'global'

                exp_params['max_train_human'] = int(info['num_train_human'])
                exp_params['max_test_human'] = int(info['num_test_human'])
                exp_params['max_train_synthetic'] = int(info['num_train_synthetic'])
                exp_params['max_test_synthetic'] = int(info['num_test_synthetic'])

                if info['num_train_synthetic'] == 0 and info['num_test_synthetic']:
                    exp_params['annotations'] = 'human'
                else:
                    exp_params['annotations'] = 'both'

                for model, config in model_configs.items():
                    exp_params['model'] = model
                    for embed_type in config['embedding_type']:
                        exp_params['embedding_type'] = embed_type
                        if 'reshape_type' in config:
                            for reshape_type in config['reshape_type']:
                                exp_params['reshape_type'] = reshape_type
                                start_experiment(exp_params)
                        else:
                            start_experiment(exp_params)
                        i += 1

if __name__ == '__main__':
    main()