#going to need to take preprocessed dictionaries and model configuration and create 
#a preprocessed model configuration and preprocessed_config file
#
#something will have to tell it what parameters to take out of the preprocessed dict
import os
import sys
import logging
import time
import numpy as np
import argparse
import allensdk.core.json_utilities as ju
import find_sweeps as fs

class ModelConfigurationException( Exception ): pass

DEFAULT_NEURON_PARAMETERS = {
    "type": "GLIF", 
    "dt": 5e-05, 
    "El": 0,
    "asc_tau_array": [ 1, 1 ],
    "asc_amp_array": [ 0, 0 ],
    "init_AScurrents": [ 0.0, 0.0 ],
    "init_threshold": 0.02, 
    "init_voltage": 0.0,
    "extrapolation_method_name": "endpoints",
    "dt_multiplier": 1
    }

DEFAULT_OPTIMIZER_PARAMETERS = { 
    "xtol": 1e-05, 
    "ftol": 1e-05, 
    "sigma_outer": 0.3, 
    "sigma_inner": 0.01, 
    "inner_iterations": 3,  
    "outer_iterations": 3, 
    "internal_iterations": 10000000,
    "iteration_info": [],
    "param_fit_names": [],
    "cut": 0,
    "bessel": { 'N': 4, 'freq': 10000 }
    }
    

def specify_parameter_groups(dictionary, dict_specifer, neuron_type):
    '''Specifies which values from the preprocessor will be used in the model configuration. 
    This is helpful if the preprocessor calculates many different values. 
        
    Parameters
    ----------
    dictionary: dict
        dictionary from preprocessor
    dict_specifier: string
        The following are available model levels 
        'LIF' (GLIF1)
        'LIF_R' (GLIF2)
        'LIF_ASC' (GLIF3)
        'LIF_R_ASC'  (GLIF_4) 
        'LIF_R_ASC_AT'  (GLIF_5)
    neuron_type: string
        'simple_neuron' is the only available option however here would be a good place for 
        the user to implement their own configurations.
    
    Returns
    -------
    output_dict: dict
        dictionary containing model configuration
    '''

    output_dict={'El_reference':dictionary['El']['El_noise']['measured']['mean'],
             'El':0.,
             'dt':dictionary['dt_used_for_preprocessor_calculations'], 
             'spike_cut_length':dictionary['spike_cutting']['NOdeltaV']['cut_length'],
             'spike_cutting_intercept':dictionary['spike_cutting']['NOdeltaV']['intercept'],
             'spike_cutting_slope':dictionary['spike_cutting']['NOdeltaV']['slope'],
             'asc_amp_array':dictionary['asc']['amp'],
             'asc_tau_array':(1./np.array(dictionary['asc']['k'])).tolist(),
             'th_inf': dictionary['th_inf']['via_Vmeasure']['from_zero'],
             'deltaV': None,
             'threshold_adaptation': {'a_spike_component_of_threshold': dictionary['threshold_adaptation']['a_spike_component_of_threshold'],
                                      'b_spike_component_of_threshold':dictionary['threshold_adaptation']['b_spike_component_of_threshold'],
                                      'a_voltage_component_of_threshold':dictionary['threshold_adaptation']['a_voltage_comp_of_thr_from_fitab'],
                                      'b_voltage_component_of_threshold': dictionary['threshold_adaptation']['b_voltage_comp_of_thr_from_fitab']},
             'MLIN': dictionary['MLIN'],
             'spike_inds': {
                'noise1': [ ],
                'noise2': [ ]
              }
             } 
    
    # specify specific values different for different levels.  Although there is only one neuron 
    # type here, this would be a good place to add other user defined neuron types
    if neuron_type=='simple_neuron':
        output_dict['C']=dictionary['capacitance']['C_test_list']['mean']
        if dict_specifer in ['LIF', 'LIF_R']:
            output_dict['R_input']=dictionary['resistance']['R_test_list']['mean']        
        elif dict_specifer in ['LIF_ASC', 'LIF_R_ASC', 'LIF_R_ASC_AT']:
            output_dict['R_input']=dictionary['resistance']['R_fit_ASC_and_R']['mean'] 
 
    for k,v in dictionary['sweep_properties']['noise1'].iteritems():
        output_dict['spike_inds']['noise1'].append( v['spike_ind'] )
        output_dict['spike_inds']['noise2'].append( v['spike_ind'] )
                
    return output_dict


def validate_method_requirements(method_config_name, has_mss):
    '''Confirm that the neuron has the specific sweeps required for the specified configuration
        
    Parameters
    ----------
    method_config_name: string
        Specifies the model level. Options are: 
            'LIF' (GLIF1)
            'LIF_R' (GLIF2)
            'LIF_ASC' (GLIF3)
            'LIF_R_ASC'  (GLIF_4) 
            'LIF_R_ASC_AT'  (GLIF_5)
    has_mss: boolean
        Specifies if the neuron has a multi short square sweep (for fitting spike component of threshold).
    '''
    if not has_mss:
        valid_configs = ['LIF', 'LIF_ASC']  
    else:
        valid_configs = ['LIF', 'LIF_ASC','LIF_R', 'LIF_R_ASC', 'LIF_R_ASC_AT']

    if method_config_name not in valid_configs:
        raise ModelConfigurationException("Model type %s cannot be configured due to missing data (mss: %s)" % ( method_config_name, str(has_mss)))

def update_neuron_method(method_type, arg_method_name, neuron_config):
    #TODO: documentation
    neuron_config[method_type] = { 'name': arg_method_name, 'params': None }


def configure_model(method_config, preprocessor_values):
    '''Configures the model from the specified method configuration and preprocessor values.
    
    Parameters
    ----------
    method_config:  dictionary
        contains values needed to configure the methods for the specified level within the dictionary
    preprocessor_values: dictionary
        dictionary from preprocessor
    '''
    
    preprocessor_values = specify_parameter_groups(preprocessor_values, method_config['name'], 'simple_neuron')

    neuron_config = {}
    neuron_config.update(DEFAULT_NEURON_PARAMETERS)
    optimizer_config = {}
    optimizer_config.update(DEFAULT_OPTIMIZER_PARAMETERS)

    #a) select values want to use out of the preprocessor_values via specifying parameter_gropus
    #b) look what levels are available via the levels available in the preprocessor_values.
    
    # Skip trace if subthreshold noise has a spike in it.
    noise1_ind =  [ n1i for n1i in preprocessor_values['spike_inds']['noise1'] if n1i is not None ]
    noise1_ind = np.concatenate(noise1_ind)
    if np.any(noise1_ind * preprocessor_values['dt'] < 8.0):
        raise ModelConfigurationException("Subthreshold region of noise1 stimulus contains spikes.")

    # check if there is a short square triple
    if preprocessor_values['threshold_adaptation']['b_spike_component_of_threshold'] and preprocessor_values['threshold_adaptation']['a_spike_component_of_threshold']:
        has_mss=True
    else:
        has_mss=False
        
    # make sure that the requested method config meets minimum requirements
    validate_method_requirements(method_config['name'], has_mss)

    update_neuron_method('AScurrent_dynamics_method', method_config['AScurrent_dynamics_method'], neuron_config)
    update_neuron_method('voltage_dynamics_method', method_config['voltage_dynamics_method'], neuron_config)
    update_neuron_method('threshold_dynamics_method', method_config['threshold_dynamics_method'], neuron_config)
    update_neuron_method('AScurrent_reset_method', method_config['AScurrent_reset_method'], neuron_config)
    update_neuron_method('voltage_reset_method', method_config['voltage_reset_method'], neuron_config)
    update_neuron_method('threshold_reset_method', method_config['threshold_reset_method'], neuron_config)

    neuron_config['El_reference'] = preprocessor_values['El_reference']
    neuron_config['C'] = preprocessor_values['C']           
    neuron_config['El'] = preprocessor_values['El']  
    neuron_config['spike_cut_length'] = preprocessor_values['spike_cut_length']
    neuron_config['asc_amp_array'] = preprocessor_values['asc_amp_array']
    neuron_config['asc_tau_array'] = preprocessor_values['asc_tau_array']
    neuron_config['R_input'] = preprocessor_values['R_input']
    neuron_config['th_inf'] = preprocessor_values['th_inf']

    optimizer_config['error_function'] = method_config['error_function']
    optimizer_config['param_fit_names'] = method_config['param_fit_names']

    #b) choose the sets want from the preprocessor_values 
    configure_method_parameters(neuron_config,
                                optimizer_config,
                                preprocessor_values['spike_cutting_slope'],
                                preprocessor_values['spike_cutting_intercept'],
                                preprocessor_values['threshold_adaptation']['a_spike_component_of_threshold'],
                                preprocessor_values['threshold_adaptation']['b_spike_component_of_threshold'],
                                preprocessor_values['threshold_adaptation']['a_voltage_component_of_threshold'],
                                preprocessor_values['threshold_adaptation']['b_voltage_component_of_threshold'],
                                preprocessor_values['MLIN']['var_of_section'],
                                preprocessor_values['MLIN']['sv_for_expsymm'],
                                preprocessor_values['MLIN']['tau_from_AC'])

    return {
        'neuron': neuron_config,
        'optimizer': optimizer_config
        }

def configure_method_parameters(neuron_config,
                                optimizer_config,
                                v_reset_slope,
                                v_reset_intercept,
                                a_spike_component_of_threshold,
                                b_spike_component_of_threshold,
                                a_voltage_component_of_threshold,
                                b_voltage_component_of_threshold,
                                var_of_section,
                                sv_for_expsymm,
                                tau_from_AC):
    '''Configures the methods used to run the models
    
    Parameters
    ----------
    neuron_config: dict
        contains neuron parameters
    optimizer_config: dict
        contains parameters for optimizaton
    v_reset_slope: float
        slope of the line in voltage reset
    v_reset_intercept: float
        intercept of the line in voltage reset
    a_spike_component_of_threshold: float or None
        amplitude of spike component of the threshold
    b_spike_component_of_threshold: float or None
        time course of spike component of the threshold
    a_voltage_component_of_threshold: float or None
        a parameter in voltage component of threshold
    b_voltage_component_of_threshold: float or None
        b parameter in voltage component of threshold
    var_of_section: float
        variance in noise of highest amplitude subthreshold long square pulse
    sv_for_expsymm: float
        parameter in MLIN optimization
    tau_from_AC: float
        time course of exponential fit to the autocorrelation
    '''
    
    # configure voltage reset rules
    method_config = neuron_config['voltage_reset_method']
    if method_config.get('params', None) is None:
        if method_config['name'] == 'zero':
            method_config['params'] = {}
        elif method_config['name'] == 'v_before':                        
            method_config['params'] = {
                'a': v_reset_slope,
                'b': v_reset_intercept
                }
    
        elif method_config['name'] == 'i_v_before':
            method_config['params'] = {
                'a': 1,
                'b': 2,
                'c': 3
                }
            raise ModelConfigurationException('i_v_before of voltage reset method is not yet implemented')
        elif method_config['name'] == 'fixed':
            raise ModelConfigurationException('cannot use fixed voltage reset method in preprocessor')
        else:
            method_config['params'] = {}
    
    # configure threshold reset rules  
    method_config = neuron_config['threshold_reset_method']        
    if method_config.get('params', None) is None:
        coeff_th_inf = neuron_config.get('coeffs', {}).get('th_inf',1.0)
        adjusted_th_inf = neuron_config['th_inf'] * coeff_th_inf
            
        if method_config['name'] == 'max_v_th':
            raise ModelConfigurationException('max_v_th threshold reset rule is not currently in use')
            
        elif method_config['name'] == 'th_before':
            raise ModelConfigurationException('th_before is not currently in use')
            
        elif method_config['name'] == 'inf':
            method_config['params'] = {}
            neuron_config['init_threshold'] = adjusted_th_inf
                
        elif method_config['name'] == 'three_components':
            method_config['params'] = { 'a_spike': a_spike_component_of_threshold,
                                       'b_spike': b_spike_component_of_threshold }
            neuron_config['init_threshold'] = adjusted_th_inf
                
        elif method_config['name'] == 'fixed':
            raise ModelConfigurationException("cannot use fixed threshold reset method in preprocessor")
        else:  
            raise ModelConfigurationException("unknown threshold reset method: ", method_config['name'])
    
    # configure voltage dynamics rules
    
    method_config = neuron_config['voltage_dynamics_method']
    if method_config.get('params', None) is None:
        if method_config['name'] == 'quadratic_i_of_v':
            raise ModelConfigurationException('quadraticIofV of voltage_dynamics_method preprocessing is not yet implemented')
        elif method_config['name'] == 'linear_forward_euler':
            method_config['params'] = {}
        elif method_config['name'] == 'linear_exact':
            method_config['params'] = {}

        else:
            raise ModelConfigurationException("unknown voltage dynamics method: ", method_config['name'])
    
    # configure threshold dynamics rules    
    method_config = neuron_config['threshold_dynamics_method']
    
    if method_config.get('params', None) is None:
        if method_config['name'] == 'three_components_forward':
            method_config['params'] = {
                    'a_spike': a_spike_component_of_threshold,
                    'b_spike': b_spike_component_of_threshold,
                    'a_voltage': a_voltage_component_of_threshold,
                    'b_voltage': b_voltage_component_of_threshold
                    }
            
        elif method_config['name'] == 'three_components_exact':
            method_config['params'] = {
                    'a_spike': a_spike_component_of_threshold,
                    'b_spike': b_spike_component_of_threshold,
                    'a_voltage': a_voltage_component_of_threshold,
                    'b_voltage': b_voltage_component_of_threshold
                    }

        elif method_config['name'] == 'spike_component':
            method_config['params'] = {
                    'a_spike': a_spike_component_of_threshold,
                    'b_spike': b_spike_component_of_threshold,
                    'a_voltage': 0,
                    'b_voltage': 0
                    }
    
        elif method_config['name'] == 'inf':
            method_config['params'] = {}
    
        else:
            raise ModelConfigurationException("unknown threshold dynamics method: ", method_config['name'])
    
    # configure ascurrent dynamics rules
    method_config = neuron_config['AScurrent_dynamics_method']
    if method_config.get('params', None) is None:
        # TODO: rename 'vector' to something more specific
        if method_config['name'] == 'vector': 
            method_config['params'] = {
                'vector': [1, 2, 3]
                }
            raise ModelConfigurationException('vector of AScurrent_dynamics_method is not yet implemented')  
        elif method_config['name'] == 'none':
            method_config['params'] = {}
        elif method_config['name'] == 'exp':
            method_config['params'] = {}
        else:
            raise ModelConfigurationException("unknown AScurrent dynamics method: ", method_config['name'])
    
    # configure ascurrent reset rule 
    # this is down here because it depends on numbers computed for the AScurrent_dynamics_method
    method_config = neuron_config['AScurrent_reset_method']
    if method_config.get('params', None) is None:
        if method_config['name'] == 'sum':
            method_config['params'] = {
                'r': np.ones(len(neuron_config['asc_tau_array']))
                }
        elif method_config['name'] == 'none':
            method_config['params'] = {}
        else:
            raise ModelConfigurationException("unknown AScurrent reset method: ", method_config['name'])
    
    # configure parameters for MLIN optimization
    if optimizer_config['error_function']=='MLIN':
        optimizer_config['error_function_data'] = {
            'subthreshold_long_square_voltage_variance': var_of_section,
            'sv_for_expsymm': sv_for_expsymm,
            'tau_from_AC': tau_from_AC
            }

    # validation

    # make sure that the initial ascurrents have the correct size
    if len(neuron_config['init_AScurrents']) != len(neuron_config['asc_tau_array']):
        raise ModelConfigurationException("init_AScurrents have incorrect length.")
    
    spike_cut_length = neuron_config.get('spike_cut_length', None)
    
    if spike_cut_length is None:
        raise ModelConfigurationException("Spike cut length must be set, but it is not.")
    
    if spike_cut_length < 0:
        raise ModelConfigurationException("Spike cut length must be non-negative.")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('preprocessor_values_path', help='path to preprocessor values json')
    parser.add_argument('method_config_path', help='path to method configuration json')
    parser.add_argument('output_path', help='path to store final model configuration')

    args = parser.parse_args()

    preprocessor_values = ju.read(args.preprocessor_values_path)
    method_config = ju.read(args.method_config_path)
    out_config = configure_model(method_config, preprocessor_values)

    ju.write(args.output_path, out_config)


if __name__ == "__main__": main()
