# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/05_lmu_dms_tuning.ipynb (unless otherwise specified).

__all__ = ['logger', 'generate_default_params', 'run_trial']

# Cell
import argparse
import logging
import pyopencl as cl
from sklearn import metrics
import numpy as np
import nengo
import nengo_ocl
import nengo.utils.least_squares_solvers as lss
import nni
from srnn_pfc.lmu import make_lmu_dms


logger = logging.getLogger('lmu_dms_tuning')


# Cell
def generate_default_params():
    return {"theta": 6.0, "q": 6, "n_neurons": 1000, "tau": 0.2,
            "max_rates": "uniform_low", "dales_law": True, "hetero_tau": True}


# Cell
def run_trial(params,
              ens_seed=1337,
              train_trials_seed=1337,
              test_trials_seed=1234,
              train_trials_per_cond=5,
              test_trials_per_cond=5,
              eval_strict=True,
              cl_platform_vendor='none'):

    cl_ctx = None
    for plat in cl.get_platforms():
        if plat.vendor.upper().startswith(cl_platform_vendor.upper()):
            cl_ctx = cl.Context(dev_type=cl.device_type.ALL,
                                properties=[(cl.context_properties.PLATFORM, plat)])
            break
    if cl_ctx is None:
        cl_ctx = cl.create_some_context(interactive=False)

    # Create model for training
    srate = 1000
    factory_kwargs = {
        'q': int(params['q']),
        'theta': params['theta'],
        'tau': params['tau'],
        'n_neurons': int(params['n_neurons']),
        'max_rates': params['max_rates'],
        'dales_law': params['dales_law'],
        'hetero_tau': params['hetero_tau'],
        'ssp_dim': 0
    }

    # Generate training data
    train_model, train_probes = make_lmu_dms(
        seed=ens_seed,
        out_transform=None,
        n_trials_per_cond=train_trials_per_cond,
        trial_seed=train_trials_seed,
        **factory_kwargs
    )
    n_train_trials = train_trials_per_cond * 8 * 2  # 16 conditions
    with nengo_ocl.Simulator(train_model, context=cl_ctx) as train_sim:
        train_sim.run(6 * n_train_trials)  # 6 seconds per trial

    # Solve for weights - Data are generally too large for nengo's filt so we slice
    #  them to only evaluate specific ranges
    if eval_strict:
        eval_ranges = [(0, 0.25), (0.75, 1.75), (2.75, 3.75), (4.75, 6.0)]
    else:
        eval_ranges = [(1.0, 1.5), (5.0, 6.0)]
    trial_tvec = np.arange(0, 6.0, 1/srate)
    b_eval = np.zeros(trial_tvec.shape, dtype=bool)
    for t_win in eval_ranges:
        b_eval = np.logical_or(b_eval, np.logical_and(t_win[0] <= trial_tvec,
                                                      trial_tvec < t_win[1]))
    b_eval = np.tile(b_eval, n_train_trials)

    filt = nengo.synapses.Lowpass(0.01)
    Y = filt.filt(train_sim.data[train_probes['ideal']][b_eval])
    A = filt.filt(train_sim.data[train_probes['ensemble']][b_eval])

    D, info = nengo.solvers.LstsqL2(solver=lss.LSMRScipy())(A, Y)
    Y = A = None

    # Create a new model with the learned weights
    test_model, test_probes = make_lmu_dms(
        seed=ens_seed,
        out_transform=D.T,
        n_trials_per_cond=test_trials_per_cond,
        trial_seed=test_trials_seed,
        **factory_kwargs)

    # Run test simulation - break after N trials to report cumulative accuracy
    def get_labels_from_sim(sim, n_trials):
        samps = 6 * srate * n_trials
        Y = sim.data[test_probes['ideal']][-samps:]
        A = sim.data[test_probes['output']][-samps:]
        b_slice = Y != 0
        label = Y[b_slice].reshape(n_trials, -1)[:, 0]
        score = np.mean(A[b_slice].reshape(n_trials, -1), axis=-1)
        return label, score

    total_test_trials = test_trials_per_cond * 8 * 2
    test_sim = nengo_ocl.Simulator(test_model, context=cl_ctx)

    test_ix = min(20, total_test_trials)
    test_sim.run(6 * test_ix)
    labels, scores = get_labels_from_sim(test_sim, test_ix)

    trial_step = 10
    while test_ix < total_test_trials:
        test_sim.run(6 * min(trial_step, total_test_trials - test_ix))
        label, score = get_labels_from_sim(test_sim, trial_step)
        labels = np.append(labels, label)
        scores = np.append(scores, score)
        fpr, tpr, thresholds = metrics.roc_curve(labels, scores)
        test_auc = metrics.auc(fpr, tpr)
        nni.report_intermediate_result(test_auc)
        test_ix += trial_step

    test_sim.close()
    logger.debug('test_sim.close() returned.')
    logger.debug('Final result is: %d', test_auc)
    nni.report_final_result(test_auc)
    logger.debug('Sent final result. Trial complete.')


# Cell
try:
    from nbdev.imports import IN_NOTEBOOK
except ModuleNotFoundError:
    IN_NOTEBOOK = False


if __name__ == "__main__" and not IN_NOTEBOOK:
    parser = argparse.ArgumentParser()
    parser.add_argument("--ens_seed", type=int, default=1337,
                        help="seed for random initialization of ensemble", required=False)
    parser.add_argument("--train_trials_seed", type=int, default=1337,
                        help="seed for random initialization of training trial conditions", required=False)
    parser.add_argument("--test_trials_seed", type=int, default=1234,
                        help="seed for random initialization of test trial conditions", required=False)
    parser.add_argument("--train_trials_per_cond", type=int, default=5,
                        help="training trials per task condition", required=False)
    parser.add_argument("--test_trials_per_cond", type=int, default=5,
                        help="testing trials per task condition", required=False)
    parser.add_argument("--eval_window_strict", type=bool, default=False,
                        help="True: expect 0 output for all segs except response; False: expect 0 for Cue only.",
                        required=False)
    parser.add_argument("--cl_platform_vendor", type=str, default='NVIDIA',
                        help="Use the first CL platform matching this vendor (case insensitive)",
                        required=False)
    args, unknown = parser.parse_known_args()

    try:
        params = {**generate_default_params(), **nni.get_next_parameter()}
        logger.debug(params)
        run_trial(params,
                  ens_seed=args.ens_seed,
                  train_trials_seed=args.train_trials_seed,
                  test_trials_seed=args.test_trials_seed,
                  train_trials_per_cond=args.train_trials_per_cond,
                  test_trials_per_cond=args.test_trials_per_cond,
                  eval_strict=args.eval_window_strict,
                  cl_platform_vendor=args.cl_platform_vendor
                 )
    except Exception as e:
        logger.exception(e)
        raise

# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/05_lmu_dms_tuning.ipynb (unless otherwise specified).

__all__ = ['logger', 'generate_default_params', 'run_trial']

# Cell
import argparse
import logging
import pyopencl as cl
from sklearn import metrics
import numpy as np
import nengo
import nengo_ocl
import nengo.utils.least_squares_solvers as lss
import nni
from srnn_pfc.lmu import make_lmu_dms


logger = logging.getLogger('lmu_dms_tuning')


# Cell
def generate_default_params():
    return {"theta": 6.0, "q": 6, "n_neurons": 1000, "tau": 0.2,
            "max_rates": "uniform_low", "dales_law": True, "hetero_tau": True}


# Cell
def run_trial(params,
              ens_seed=1337,
              train_trials_seed=1337,
              test_trials_seed=1234,
              train_trials_per_cond=5,
              test_trials_per_cond=5,
              eval_strict=True,
              cl_platform_vendor='none'):

    cl_ctx = None
    for plat in cl.get_platforms():
        if plat.vendor.upper().startswith(cl_platform_vendor.upper()):
            cl_ctx = cl.Context(dev_type=cl.device_type.ALL,
                                properties=[(cl.context_properties.PLATFORM, plat)])
            break
    if cl_ctx is None:
        cl_ctx = cl.create_some_context(interactive=False)

    # Create model for training
    srate = 1000
    factory_kwargs = {
        'q': int(params['q']),
        'theta': params['theta'],
        'tau': params['tau'],
        'n_neurons': int(params['n_neurons']),
        'max_rates': params['max_rates'],
        'dales_law': params['dales_law'],
        'hetero_tau': params['hetero_tau'],
        'ssp_dim': 0
    }

    # Generate training data
    train_model, train_probes = make_lmu_dms(
        seed=ens_seed,
        out_transform=None,
        n_trials_per_cond=train_trials_per_cond,
        trial_seed=train_trials_seed,
        **factory_kwargs
    )
    n_train_trials = train_trials_per_cond * 8 * 2  # 16 conditions
    with nengo_ocl.Simulator(train_model, context=cl_ctx) as train_sim:
        train_sim.run(6 * n_train_trials)  # 6 seconds per trial

    # Solve for weights - Data are generally too large for nengo's filt so we slice
    #  them to only evaluate specific ranges
    if eval_strict:
        eval_ranges = [(0, 0.25), (0.75, 1.75), (2.75, 3.75), (4.75, 6.0)]
    else:
        eval_ranges = [(1.0, 1.5), (5.0, 6.0)]
    trial_tvec = np.arange(0, 6.0, 1/srate)
    b_eval = np.zeros(trial_tvec.shape, dtype=bool)
    for t_win in eval_ranges:
        b_eval = np.logical_or(b_eval, np.logical_and(t_win[0] <= trial_tvec,
                                                      trial_tvec < t_win[1]))
    b_eval = np.tile(b_eval, n_train_trials)

    filt = nengo.synapses.Lowpass(0.01)
    Y = filt.filt(train_sim.data[train_probes['ideal']][b_eval])
    A = filt.filt(train_sim.data[train_probes['ensemble']][b_eval])

    D, info = nengo.solvers.LstsqL2(solver=lss.LSMRScipy())(A, Y)
    Y = A = None

    # Create a new model with the learned weights
    test_model, test_probes = make_lmu_dms(
        seed=ens_seed,
        out_transform=D.T,
        n_trials_per_cond=test_trials_per_cond,
        trial_seed=test_trials_seed,
        **factory_kwargs)

    # Run test simulation - break after N trials to report cumulative accuracy
    def get_labels_from_sim(sim, n_trials):
        samps = 6 * srate * n_trials
        Y = sim.data[test_probes['ideal']][-samps:]
        A = sim.data[test_probes['output']][-samps:]
        b_slice = Y != 0
        label = Y[b_slice].reshape(n_trials, -1)[:, 0]
        score = np.mean(A[b_slice].reshape(n_trials, -1), axis=-1)
        return label, score

    total_test_trials = test_trials_per_cond * 8 * 2
    test_sim = nengo_ocl.Simulator(test_model, context=cl_ctx)

    test_ix = min(20, total_test_trials)
    test_sim.run(6 * test_ix)
    labels, scores = get_labels_from_sim(test_sim, test_ix)

    trial_step = 10
    while test_ix < total_test_trials:
        test_sim.run(6 * min(trial_step, total_test_trials - test_ix))
        label, score = get_labels_from_sim(test_sim, trial_step)
        labels = np.append(labels, label)
        scores = np.append(scores, score)
        fpr, tpr, thresholds = metrics.roc_curve(labels, scores)
        test_auc = metrics.auc(fpr, tpr)
        nni.report_intermediate_result(test_auc)
        test_ix += trial_step

    test_sim.close()
    logger.debug('test_sim.close() returned.')
    logger.debug('Final result is: %d', test_auc)
    nni.report_final_result(test_auc)
    logger.debug('Sent final result. Trial complete.')


# Cell
try:
    from nbdev.imports import IN_NOTEBOOK
except ModuleNotFoundError:
    IN_NOTEBOOK = False


if __name__ == "__main__" and not IN_NOTEBOOK:
    parser = argparse.ArgumentParser()
    parser.add_argument("--ens_seed", type=int, default=1337,
                        help="seed for random initialization of ensemble", required=False)
    parser.add_argument("--train_trials_seed", type=int, default=1337,
                        help="seed for random initialization of training trial conditions", required=False)
    parser.add_argument("--test_trials_seed", type=int, default=1234,
                        help="seed for random initialization of test trial conditions", required=False)
    parser.add_argument("--train_trials_per_cond", type=int, default=5,
                        help="training trials per task condition", required=False)
    parser.add_argument("--test_trials_per_cond", type=int, default=5,
                        help="testing trials per task condition", required=False)
    parser.add_argument("--eval_window_strict", type=bool, default=False,
                        help="True: expect 0 output for all segs except response; False: expect 0 for Cue only.",
                        required=False)
    parser.add_argument("--cl_platform_vendor", type=str, default='NVIDIA',
                        help="Use the first CL platform matching this vendor (case insensitive)",
                        required=False)
    args, unknown = parser.parse_known_args()

    try:
        params = {**generate_default_params(), **nni.get_next_parameter()}
        logger.debug(params)
        run_trial(params,
                  ens_seed=args.ens_seed,
                  train_trials_seed=args.train_trials_seed,
                  test_trials_seed=args.test_trials_seed,
                  train_trials_per_cond=args.train_trials_per_cond,
                  test_trials_per_cond=args.test_trials_per_cond,
                  eval_strict=args.eval_window_strict,
                  cl_platform_vendor=args.cl_platform_vendor
                 )
    except Exception as e:
        logger.exception(e)
        raise

# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/05_lmu_dms_tuning.ipynb (unless otherwise specified). # Cell # Cell # Cell # Create model for training # Generate training data # 16 conditions # 6 seconds per trial # Solve for weights - Data are generally too large for nengo's filt so we slice #  them to only evaluate specific ranges # Create a new model with the learned weights # Run test simulation - break after N trials to report cumulative accuracy # Cell