# Create a script to run a random hyperparameter search.

import copy
import getpass
import os
import random
import numpy as np
import gflags
import sys

import datetime
now = datetime.datetime.now()
date = now.strftime("%m%d")

NYU_NON_PBS = False
NAME = date + "_20s"
SWEEP_RUNS = 100

LIN = "LIN"
EXP = "EXP"
BOOL = "BOOL"
CHOICE = "CHOICE"
SS_BASE = "SS_BASE"

FLAGS = gflags.FLAGS

# 240k, 540k, 990k
gflags.DEFINE_string("training_data_path", "spinn/data/listops/train_d20s.tsv", "")
gflags.DEFINE_string("eval_data_path", "spinn/data/listops/test_d20s.tsv", "")
gflags.DEFINE_string("log_path", "/misc/vlgscratch4/BowmanGroup/nikita/listops/rl_tree_baseline", "")
# /misc/vlgscratch4/BowmanGroup/nikita/listops
FLAGS(sys.argv)

# Instructions: Configure the variables in this block, then run
# the following on a machine with qsub access:
# python make_sweep.py > my_sweep.sh
# bash my_sweep.sh

# - #

# Non-tunable flags that must be passed in.

FIXED_PARAMETERS = {
    "data_type":     "listops",
    "model_type":      "RLSPINN",
    "training_data_path":    FLAGS.training_data_path,
    "eval_data_path":    FLAGS.eval_data_path,
    "log_path": FLAGS.log_path,
    "ckpt_path":  FLAGS.log_path,
    "word_embedding_dim":   "128",
    "model_dim":   "128",
    "seq_length": "100",
    "eval_seq_length":  "3000",
    "eval_interval_steps": "1000",
    "statistics_interval_steps": "100",
    "batch_size":  "64",
    "encode": "pass",
    "mlp_dim": "16",
    "num_mlp_layers": "2",
    "semantic_classifier_keep_rate": "1.0",
    "sample_interval_steps": "1000",
    "nocomposition_ln": "",
    "optimizer_type": "Adam",
    "early_stopping_steps_to_wait": "10000", 
    "gpu": "0",
    "rl_reward": "standard",
    "rl_baseline": "lbtree",
    "norl_wake_sleep": "",
    "use_internal_parser": "",
    "tracking_lstm_hidden_dim": "48",
    "transition_weight": "0",
    "write_eval_report": "",
    "cp_metric": "",
}

# Tunable parameters.

SWEEP_PARAMETERS = {
    "l2_lambda":          ("l2", EXP, 8e-7, 1e-3),
    "learning_rate": ("lr", EXP, 0.0001, 0.01),
    "learning_rate_decay_when_no_progress": ("ld", CHOICE, ['0.1', '0.25', '0.5', '1.0'], None),
    "model_dim": ("dim", CHOICE, ['48', '128'], None),
    "tracking_lstm_hidden_dim": ("tdim", CHOICE, ['12', '48'], None),
    "rl_weight":  ("rlwt", EXP, 0.5, 5.0),
    "rl_epsilon": ("eps", LIN, 0.1, 1.0),
    "rl_epsilon_decay": ("epsd", EXP, 1000, 1000000),
    "rl_confidence_penalty": ("rlconf", EXP, 0.00001, 0.01),
    "rl_confidence_interval": ("rlconfint", EXP, 10, 100),
    "rl_baseline": ("base", CHOICE, ["value", "pass", "greedy", "ema"], None),
    "rl_value_size": ("rlmlp", CHOICE, ["12", "24", "32", "48", "128", "256", "512", "1024"], None),
    "rl_value_lstm": ("rllstm", CHOICE, ["100", "200", "300"], None),
}

"""
# Narrow sweep
SWEEP_PARAMETERS = {
    "l2_lambda":          ("l2", EXP, 1e-5, 9e-4), # 0.00022 # Adam
    "learning_rate": ("lr", EXP, 0.0002, 0.0009), # 0.00041 # Adam
    "learning_rate_decay_when_no_progress": ("ld", CHOICE, ['1.0'], None),
    "model_dim": ("dim", CHOICE, ['128'], None),
    "tracking_lstm_hidden_dim": ("tdim", CHOICE, ['12', '48'], None),
    "rl_weight":  ("rlwt", EXP, 2.0, 4.0), # 2.97122920235 
    "rl_epsilon": ("eps", LIN, 0.48, 0.6), # 0.5097648659191173
    "rl_epsilon_decay": ("epsd", EXP, 550000, 750000), #614149
    "rl_confidence_penalty": ("rlconf", EXP, 0.00001, 0.0003),
    "rl_confidence_interval": ("rlconfint", LIN, 30, 42), # 34
    "rl_baseline": ("base", CHOICE, ["value"], None),
    "rl_value_size": ("rlmlp", CHOICE, ["48", "128"], None),
    "rl_value_lstm": ("rllstm", CHOICE, ["100", "200", "300"], None),
}
"""

sweep_name = "sweep_" + NAME + "_" + \
    FIXED_PARAMETERS["data_type"] + "_" + FIXED_PARAMETERS["model_type"]

# - #
print("# NAME: " + sweep_name)
print("# NUM RUNS: " + str(SWEEP_RUNS))
print("# SWEEP PARAMETERS: " + str(SWEEP_PARAMETERS))
print("# FIXED_PARAMETERS: " + str(FIXED_PARAMETERS))
print()

for run_id in range(SWEEP_RUNS):
    params = {}
    name = sweep_name + "_" + str(run_id)

    params.update(FIXED_PARAMETERS)
    # Any param appearing in both sets will be overwritten by the sweep value.

    for param in SWEEP_PARAMETERS:
        config = SWEEP_PARAMETERS[param]
        t = config[1]
        mn = config[2]
        mx = config[3]

        r = random.uniform(0, 1)
        if t == EXP:
            lmn = np.log(mn)
            lmx = np.log(mx)
            sample = np.exp(lmn + (lmx - lmn) * r)
        elif t == BOOL:
            sample = r > 0.5
        elif t==SS_BASE:
            lmn = np.log(mn)
            lmx = np.log(mx)
            sample = 1 - np.exp(lmn + (lmx - lmn) * r)
        elif t==CHOICE:
            sample = random.choice(mn)
        else:
            sample = mn + (mx - mn) * r

        if isinstance(mn, int):
            sample = int(round(sample, 0))
            val_disp = str(sample)
            params[param] = sample
        elif isinstance(mn, float):
            val_disp = "%.2g" % sample
            params[param] = sample
        elif t==BOOL:
            val_disp = str(int(sample))
            if not sample:
                params['no' + param] = ''
            else:
                params[param] = ''
        else:
            val_disp = sample
            params[param] = sample
        name += "-" + config[0] + val_disp

    # If tuning model dim, ensure the following for listops.
    params["word_embedding_dim"] = params["model_dim"]

    flags = ""
    for param in params:
        value = params[param]
        flags += " --" + param + " " + str(value)

    flags += " --experiment_name " + name
    
    if NYU_NON_PBS:
        print("cd spinn/python; python3 -m spinn.models.rl_classifier " + flags)
    else:
        print("SPINNMODEL=\"spinn.models.rl_classifier\" SPINN_FLAGS=\"" + flags + "\" bash ../scripts/sbatch_submit_cilvr_cpu.sh")
    print()
    