import argparse, sys, logging

import allensdk.core.json_utilities as ju
import allensdk.internal.model.glif.find_sweeps as fs

from allensdk.internal.model.glif.glif_optimizer_neuron import GlifOptimizerNeuron
from allensdk.internal.model.glif.glif_experiment import GlifExperiment
from allensdk.internal.model.glif.glif_optimizer import GlifOptimizer

from allensdk.internal.model.data_access import load_sweeps
from allensdk.internal.model.glif.find_spikes import find_spikes_list
import allensdk.core.json_utilities as ju
import allensdk.internal.model.glif.preprocess_neuron as pn

def get_optimize_sweep_numbers(sweep_index): 
    #TODO: why is this here--why are sweep indicies being fed to a find_noise_sweeps sweeps and specifying
    #noise?--shouldn't the sweeps already be provided?   
    return fs.find_noise_sweeps(sweep_index)['noise1']

def optimize_neuron(model_config, sweep_index, nwb_file, save_callback=None):
    '''Optimizes a neuron.  
    1. Loads optimizer and neuron configuration data.
    2. Loads the voltage trace sweeps that will be optimized
    3. Configures the experiment and optimizer
    4. Runs the optimizer
    5. TODO: where is data saved
    
    Parameters
    ----------
    model_config : dictionary
        contains values of neuron and optimizer parameters
    sweep_index : list of integers
        indices (as labeled in the data configuration file) of sweeps that will be optimized
    save_callback : module
        saves output
    '''
    # define the neuron and optimizer dictionaries from the model configuration
    neuron_config = model_config['neuron']
    optimizer_config = model_config['optimizer']

    # load the neuron with along with the methods needed for optimization
    neuron = GlifOptimizerNeuron.from_dict(neuron_config)

    # TODO: not sure what this is doing
    optimize_sweeps = get_optimize_sweep_numbers(sweep_index)
    
    # load the sweeps to be optimized
    optimize_data = load_sweeps(nwb_file, optimize_sweeps, neuron.dt, 
                                optimizer_config["cut"], optimizer_config["bessel"])

    # needed to offset all voltages by El_reference
    El_reference = neuron_config['El_reference']

    # get indicies of spikes and voltage at those spikes
    spike_ind, spike_v = find_spikes_list(optimize_data['voltage'], neuron_config['dt'])

    # get times of spikes 
    grid_spike_times = [ si*neuron_config['dt'] for si in spike_ind ]
    
    # convert voltage at spikes into reference frame of El
    grid_spike_voltages_in_ref_to_zero = [ sv - El_reference for sv in spike_v ]

    # convert voltage into reference frame of El 
    resp_list = [ d - El_reference for d in optimize_data['voltage'] ]  

    # configure experiment
    experiment = GlifExperiment(neuron = neuron, 
                                dt = neuron.dt,
                                stim_list = optimize_data['current'],
                                resp_list = resp_list,
                                spike_time_steps = spike_ind, 
                                grid_spike_times = grid_spike_times,
                                grid_spike_voltages = grid_spike_voltages_in_ref_to_zero,
                                param_fit_names = optimizer_config['param_fit_names'])

    # configure optimizer
    optimizer = GlifOptimizer(experiment = experiment, 
                              dt = neuron.dt,              
                              outer_iterations = optimizer_config['outer_iterations'],
                              inner_iterations = optimizer_config['inner_iterations'],
                              sigma_inner = optimizer_config['sigma_inner'],
                              sigma_outer = optimizer_config['sigma_outer'],
                              param_fit_names = optimizer_config['param_fit_names'],
                              stim = optimize_data['current'],
                              error_function_data = optimizer_config['error_function_data'],
                              xtol = optimizer_config['xtol'],
                              ftol = optimizer_config['ftol'],
                              internal_iterations = optimizer_config['internal_iterations'],
                              init_params = optimizer_config.get('init_params', None),
                              bessel = optimizer_config['bessel'])

    def save(optimizer, outer, inner):
        logging.info('finished outer: %d inner: %d' % (outer, inner))
        if save_callback:
            save_callback(optimizer, outer, inner)
    
    # run the optimizer
    best_param, begin_param = optimizer.run_many(save) 

    # over write the the initial experiment parameters with the best found parameters
    # TODO: but why do this since it is not being returned
    experiment.set_neuron_parameters(best_param)
    
    return optimizer, best_param, begin_param
        
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('model_config_file')
    parser.add_argument('sweeps_file')
    parser.add_argument('output_file')
    parser.add_argument("--dt", default=pn.DEFAULT_DT)
    parser.add_argument("--bessel", default=pn.DEFAULT_BESSEL)
    parser.add_argument("--cut", default=pn.DEFAULT_CUT)

    args = parser.parse_args()

    model_config = ju.read(args.model_config_file)
    sweep_list = ju.read(args.sweeps_file)

    sweep_index = { s['sweep_number']:s for s in sweep_list }

    try:
        neuron, best_param, begin_param = optimize_neuron(model_config, sweep_index, dt, cut, bessel)
        ju.write(args.output_file, neuron.to_dict())
    except Exception, e:
        logging.error(e.message)

if __name__ == "__main__": main()
