import logging
import sys
import numpy as np
import h5py

from qc_support import *
from lab_notebook_reader import *

from allensdk.internal.core.lims_pipeline_module import PipelineModule
from allensdk.core.nwb_data_set import NwbDataSet


# manual keys are values that can be passed in through input.json.
# these values are used if the particular value cannot be computed.
# a better name might be 'DEFAULT_VALUE_KEYS'
MANUAL_KEYS = ['manual_seal_gohm', 'manual_initial_access_resistance_mohm', 'manual_initial_input_mohm' ]

# names of blocks used in output.json
# for sweep-specific data:
JSON_BLOCK_SWEEP_DATA = "sweep_data"
# for data that applies to the entire experiment:
JSON_BLOCK_EXPERIMENT_DATA = "experiment_data"

########################################################################
# bootstrapping code
# this module doesn't know anything about what's in the supplied NWB 
#   file and simply assumes that it's an IVSCC file. it must find and
#   fetch data as appropriate
# processing requires being able to pull out sweeps of specific types.
#   create an index of those types here, and provide accessor functions
#   for the indexed data


# local globals used to avoid having to pass parameters to functions that
#   can be several calls deep before they're needed
# consider refactoring into QC class to avoid this approach
sweep_stim_map = None
stim_sweep_map = None
sweep_list = None
nwb_file_name = None

# reads the NWB file and generates a mapping between sweep name and 
#   stimulus code, and vice versa
def build_sweep_stim_map():
    global sweep_stim_map, stim_sweep_map, nwb_file_name, sweep_list
    try:
        nwb_file = h5py.File(nwb_file_name, "r")
    except:
        raise Exception ("Unable to open input NWB file '%s'" % str(nwb_file_name))
    print("Opened '%s'" % str(nwb_file_name))
    sweep_stim_map = {}
    stim_sweep_map = {}
    sweep_list = []
    acq = nwb_file["acquisition/timeseries"]
    for sweep in acq:
        # if string storage is variable length, data appears to be stored
        #   or retrieved as an array of strings, so we need to take the
        #   first element. this happens with Igor-generated files.
        # if the string is stored with fixed length, data appears to be
        #   stored as a string, so we must take the entire value
        stim = acq[sweep]["aibs_stimulus_description"].value[0]
        if len(stim) == 1:
            stim = acq[sweep]["aibs_stimulus_description"].value
        stim_sweep_map[stim] = sweep
        #print "%s (%s) : %s (%s)" % (sweep, type(sweep), stim, type(stim))
        sweep_stim_map[sweep] = stim
        sweep_list.append(sweep)
    nwb_file.close()

# fetches stimulus code for a given sweep name, or None if no stimulus
#   was found for the specified sweep
def get_sweep_name_by_stimulus_code(stim_name):
    """ Returns the first sweep name that uses the specified stimulus
        type. 'First' does not mean lowest sweep number, only the first
        one found using a [random] dictionary search.

        Input: stimulus name (string)

        Output: sweep name (string), or None if no sweep found for this stim
    """
    global sweep_stim_map
    for k,v in stim_sweep_map.iteritems():
        if k.startswith(stim_name):
            return v
    return None

        
# returns True if stimulus name for specified sweep indicates the sweep
#   is a ramp and False otherwise
def sweep_is_ramp(sweep_name):
    """ Input: sweep name (string)

        Output: boolean (True if sweep is ramp, False otherwise)
    """
    global sweep_stim_map
    return sweep_stim_map[sweep_name].startswith('C1RP')


# old code based on using NwbDataSet objects. provide a way to
#   create them in order to leverage old code as much as possible
def get_sweep_data(sweep_name):
    """ Input: sweep name (string)
        
        Output: NwbDataSet object
    """
    global nwb_file_name
    try:
        num = int(sweep_name.split('_')[-1])
    except:
        print("Unable to parse sweep number from '%s'" % str(sweep_name))
        raise
    return NwbDataSet(nwb_file_name).get_sweep(num)


# functions to lookup a sweep having the desired stimulus code
# NOTE: if multiple instance exist then only one instance is returned
def get_blowout_sweep():
    """ Returns NwbDataSet for the blowout sweep, or None if it's absent
    """
    sweep_name = get_sweep_name_by_stimulus_code('EXTPBLWOUT')
    if sweep_name is None:
        return None
    return get_sweep_data(sweep_name)

def get_bath_sweep():
    """ Returns NwbDataSet for the bath sweep, or None if it's absent
    """
    sweep_name = get_sweep_name_by_stimulus_code('EXTPINBATH')
    if sweep_name is None:
        return None
    return get_sweep_data(sweep_name)

def get_seal_sweep():
    """ Returns NwbDataSet for the seal sweep, or None if it's absent
    """
    sweep_name = get_sweep_name_by_stimulus_code('EXTPCllATT')
    if sweep_name is None:
        return None
    return get_sweep_data(sweep_name)

def get_breakin_sweep():
    """ Returns NwbDataSet for the breakin sweep, or None if it's absent
    """
    sweep_name = get_sweep_name_by_stimulus_code('EXTPBREAKN')
    if sweep_name is None:
        return None
    return get_sweep_data(sweep_name)


########################################################################


########################################################################
# QC-relevant feature extraction code

# cell-level values (for ephys_roi_results) 
def cell_level_features(jin, jout, sweep_tag_list, manual_values):
    """
    """
    output_data = {}
    jout[JSON_BLOCK_EXPERIMENT_DATA] = output_data
    # measure blowout voltage
    try:
        blowout_data = get_blowout_sweep()
        blowout = measure_blowout(blowout_data['response'], 
                                  blowout_data['index_range'][0])
        output_data['blowout_mv'] = blowout
    except:
        msg = "Blowout is not available"
        sweep_tag_list.append(msg)
        logging.warning(msg)
        output_data['blowout_mv'] = None


    # measure "electrode 0"
    try:
        bath_data = get_bath_sweep()
        e0 = measure_electrode_0(bath_data['response'], 
                                 bath_data['sampling_rate'])
        output_data['electrode_0_pa'] = e0
    except:
        msg = "Electrode 0 is not available"
        sweep_tag_list.append(msg)
        logging.warning(msg)
        output_data['electrode_0_pa'] = None


    # measure clamp seal
    try:
        seal_data = get_seal_sweep()
        seal = measure_seal(seal_data['stimulus'], 
                            seal_data['response'], 
                            seal_data['sampling_rate'])
        # error may arise in computing seal, which falls through to
        #   exception handler. if seal computation didn't fail but
        #   computation generated invalid value, trigger same 
        #   exception handler with different error
        if seal is None or not np.isfinite(seal):
            raise Exception("Could not compute seal")
    except:
        # seal is not available, for whatever reason. log error
        msg = "Seal is not available"
        sweep_tag_list.append(msg)
        logging.warning(msg)
        # look for manual seal value and use it if it's available
        seal = manual_values.get('manual_seal_gohm', None)
        if seal is not None:
            logging.info("using manual seal value: %f" % seal)
            sweep_tag_list.append("Seal set using manual value")
    output_data["seal_gohm"] = seal


    # measure input and series resistance
    # this requires two steps -- finding the breakin sweep, and then 
    #   analyzing it
    # if the value is unavailable then check to see if it was set manually
    breakin_data = None
    try:
        breakin_data = get_breakin_sweep()
    except:
        logging.warning("Error reading breakin sweep.")
        sweep_tag_list.append("Breakin sweep not found")

    ir = None   # input resistance
    sr = None   # series resistance
    if breakin_data is not None:
        ###########################
        # input resistance
        try:
            ir = measure_input_resistance(breakin_data['stimulus'], 
                                          breakin_data['response'], 
                                          breakin_data['sampling_rate'])
        except:
            logging.warning("Error reading input resistance.")
        # apply manual value if it's available
        if ir is None:
            sweep_tag_list.append("Input resistance is not available")
            ir = manual_values.get('manual_initial_input_mohm', None)
            if ir is not None:
                msg = "Using manual value for input resistance"
                logging.info(msg)
                sweep_tag_list.append(msg);
        ###########################
        # initial access resistance
        try:
            sr = measure_initial_access_resistance(breakin_data['stimulus'], 
                                               breakin_data['response'], 
                                               breakin_data['sampling_rate'])
        except:
            logging.warning("Error reading initial access resistance.")
        # apply manual value if it's available
        if sr is None:
            sweep_tag_list.append("Initial access resistance is not available")
            sr = manual_values.get('manual_initial_access_resistance_mohm', None)
            if sr is not None:
                msg = "Using manual initial access resistance"
                logging.info(msg)
                sweep_tag_list.append(msg)
    #
    output_data['input_resistance_mohm'] = ir
    output_data["initial_access_resistance_mohm"] = sr

    sr_ratio = None # input access resistance ratio
    if ir is not None and sr is not None:
        try:
            sr_ratio = sr / ir
        except:
            pass    # let sr_ratio stay as None
    output_data['input_access_resistance_ratio'] = sr_ratio


##############################
def sweep_level_features(jin, jout, sweep_tag_list):
    """
    """
    global sweep_list
    # pull out features from each sweep (for ephys_sweeps)
    cnt = 0
    jout[JSON_BLOCK_SWEEP_DATA] = {}
    for sweep_name in sweep_list:
        # pull data streams from file
        sweep_num = int(sweep_name.split('_')[-1])
        try:
            sweep_data = NwbDataSet(nwb_file_name).get_sweep(sweep_num)
        except:
            logging.warning("Error reading sweep %d" % sweep_num)
            continue
        sweep = {}
        jout[JSON_BLOCK_SWEEP_DATA][sweep_name] = sweep

        # don't process voltage clamp sweeps
        if sweep_data["stimulus_unit"] == "Volts":
            continue    # voltage-clamp

        volts = sweep_data['response']
        current = sweep_data['stimulus']
        hz = sweep_data['sampling_rate']
        idx_start, idx_stop = sweep_data['index_range']

        # measure Vm and noise before stimulus
        idx0, idx1 = get_first_vm_noise_epoch(idx_start, current, hz)
        _, rms0 = measure_vm(1e3 * volts[idx0:idx1])

        sweep["pre_noise_rms_mv"] = float(rms0)

        # measure Vm and noise at end of recording
        # only do so if acquisition not truncated 
        # do not check for ramps, because they do not have enough time to recover
        mean1 = None
        sweep_not_truncated = ( idx_stop == len(current) - 1 )
        if sweep_not_truncated and not sweep_is_ramp(sweep_name):
            idx0, idx1 = get_last_vm_epoch(idx_stop, current, hz)
            mean1, _ = measure_vm(1e3 * volts[idx0:idx1])
            idx0, idx1 = get_last_vm_noise_epoch(idx_stop, current, hz)
            _, rms1 = measure_vm(1e3 * volts[idx0:idx1])
            sweep["post_vm_mv"] = float(mean1)
            sweep["post_noise_rms_mv"] = float(rms1)

        # measure Vm and noise over extended interval, to check stability
        stim_start = find_stim_start(idx_start, current)
        sweep['stimulus_start_time'] = stim_start / sweep_data['sampling_rate']

        idx0, idx1 = get_stability_vm_epoch(idx_start, stim_start, hz)
        mean2, rms2 = measure_vm(1000 * volts[idx0:idx1])

        slow_noise = float(rms2)
        sweep["slow_vm_mv"] = float(mean2)
        sweep["slow_noise_rms_mv"] = float(rms2)

        # for now (mid-feb 15), make vm_mv the same for pre and slow
        mean0 = mean2
        sweep["pre_vm_mv"] = float(mean0)
        if mean1 is not None:
            delta = abs(mean0 - mean1)
            sweep["vm_delta_mv"] = float(delta)
        else:
            # Use None as 'nan' still breaks the ruby strategies
            sweep["vm_delta_mv"] = None

        # compute stimulus duration, amplitude, interal
        stim_amp, stim_dur = find_stim_amplitude_and_duration(idx_start, current, hz)
        stim_int = find_stim_interval(idx_start, current, hz)

        sweep['stimulus_amplitude'] = stim_amp * 1e12
        sweep['stimulus_duration'] = stim_dur
        sweep['stimulus_interval'] = stim_int

        tag_list = []
        for i in range(len(sweep_tag_list)):
            tag = {}
            tag["name"] = sweep_tag_list[i]
            tag_list.append(tag)
        sweep["ephys_sweep_tags"] = tag_list


# create a summary table of sweeps and stimuli
def summarize_sweeps(jin, jout):
    global nwb_file_name
    # build stimulus name map
    stim_type_name_map = {}
    for group_name, raw_names in jin["ephys_raw_stimulus_names"].iteritems():
        for n in raw_names:
            stim_type_name_map[n] = group_name

    h5_file_name = jin.get("input_h5", None)
    notebook = create_lab_notebook_reader(nwb_file_name, h5_file_name)
    borg = h5py.File(nwb_file_name, 'r')

    # two json blocks to store data in
    exp_data = jout[JSON_BLOCK_EXPERIMENT_DATA]
    swp_data = jout[JSON_BLOCK_SWEEP_DATA]
    #jout["sweep_summary"] = output_data

#    # verify input file generated by Igor
#    generated_by = borg["general/generated_by"].value
#    igor = False
#    for row in generated_by:
#        if row[0] == "Program" and row[1].startswith('Igor'):
#            igor = True
#            break
#    if not igor:
#        print("Error -- File not recognized as Igor-generated NWB file")
#        return -1

    # validated nwb files can have different types of string storage
    # problem seems to be related to h5py and if string is stored as
    #   fixed- or variable-width. assume that string is more than one
    #   character and try to auto-correct for this issue
    session_date = borg["session_start_time"].value
    if len(session_date) == 1:
        session_date = session_date[0]
    exp_data['recording_date'] = session_date

    # get sampling rate
    # use same output strategy as h5-nwb converter
    # pick the sampling rate from the first iclamp sweep
    # TODO: figure this out for multipatch
    sampling_rate = None
    for sweep_name in borg["acquisition/timeseries"]:
        sweep_ts = borg["acquisition/timeseries"][sweep_name]
        ancestry = sweep_ts.attrs["ancestry"]
        if "CurrentClamp" in ancestry[-1]:
            if sampling_rate is None:
                sampling_rate = sweep_ts["starting_time"].attrs["rate"]
                break
    if sampling_rate is None:
        raise Exception("Unable to determine sampling rate from current clamp sweep.")
    exp_data['sampling_rate'] = sampling_rate
#    sweep_data = []
#    output_data["sweep_summary"] = sweep_data

    # read sweep-specific data
    for sweep_name in borg["acquisition/timeseries"]:
        # get h5 timeseries object, and the sweep number
        sweep_ts = borg["acquisition/timeseries"][sweep_name]
        sweep_num = int(sweep_name.split('_')[-1])
        #sweep_num = int(sweep_name[:-4].split('_')[-1]) # for reading igor nwb
        # fetch stim name from lab notebook
        stim_name = notebook.get_value("Stim Wave Name", sweep_num, "")
        if len(stim_name) == 0:
            raise Exception("Could not read stimulus wave name from lab notebook for sweep %d" % sweep_num)

        # stim units are based on timeseries type
        ancestry = sweep_ts.attrs["ancestry"]
        if "CurrentClamp" in ancestry[-1]:
            stim_units = 'pA'
        elif "VoltageClamp" in ancestry[-1]:
            stim_units = 'mV'
        else:
            # it's probably OK to skip this sweep and put a 'continue' 
            #   here instead of an exception, but wait until there's
            #   an actual error and investigate the data before doing so
            raise Exception("Unable to determine clamp mode in " + sweep_name)

        # stim name stored in database as, eg, C2SSTRIPLE150429
        # stim name in igor nwb stored as C2SSTRIPLE150429_DA_0
        # -> need to strip last 5 chars off to make match for lookup
        stim_type_name = stim_type_name_map.get(stim_name[:-5], None)
        if stim_type_name is None:
            raise Exception("Could not find stimulus raw name (\"%s\") for sweep %d." % (stim_name, sweep_num))

        # voltage-clamp sweeps shouldn't have a record yet -- make one
        if sweep_name not in swp_data:  
            swp_data[sweep_name] = {}
        info = swp_data[sweep_name]

        # sweep number
        info["sweep_number"] = sweep_num
        # bridge balance
        bridge_balance = notebook.get_value("Bridge Bal Value", sweep_num, None)
        # IT-14677
        # if bridge_balance is None, that's OK. do NOT change it to NaN

        info["bridge_balance_mohm"] = bridge_balance
        # stimulus units
        info["stimulus_units"] = stim_units
        # leak_pa (bias current)
        bias_current = notebook.get_value("I-Clamp Holding Level", sweep_num, None)
        # IT-14677
        # if bias_current is None, that's OK. do NOT change it to NaN

        info["leak_pa"] = bias_current
        #
        # ephys stim info
        scale_factor = notebook.get_value("Scale Factor", sweep_num, None)
        if scale_factor is None:
            raise Exception("Unable to read scale factor for " + sweep_name)
        # PBS-229 change stim name by appending set_sweep_count
        cnt = notebook.get_value("Set Sweep Count", sweep_num, 0)
        stim_name_ext = stim_name.split('_')[0] + "[%d]" % int(cnt)
        info["ephys_stimulus"] = {
            #'description': stim_name,
            'description': stim_name_ext,
            'amplitude': scale_factor,
            'ephys_stimulus_type': { 'name': stim_type_name }
        }
    #
    borg.close()


########################################################################
########################################################################


def main(jin):
    # to avoid passing arguments amongs the many functions and procedures,
    #   set a global value 'nwb_file_name' that each function can read
    global nwb_file_name
    nwb_file_name = jin["input_nwb"]

    # initialize index of stimuli and sweeps
    build_sweep_stim_map()

    # TODO Document manual keys, and what they're for
    manual_values = {}
    for k in MANUAL_KEYS:
        if k in jin:
            manual_values[k] = jin[k]

    # dictionary for json output
    jout = {}

    # list of messages (tags) that log information about this sweep set
    #  (eg, 'Seal not available')
    sweep_tag_list = []

    cell_level_features(jin, jout, sweep_tag_list, manual_values)
    # sweep level data. first pull out QC-relevant metrics, then store
    #   stimulus info with that data
    sweep_level_features(jin, jout, sweep_tag_list)
    summarize_sweeps(jin, jout)

    return jout



if __name__ == "__main__":
    # read module input. PipelineModule object automatically parses the 
    #   command line to pull out input.json and output.json file names
    module = PipelineModule()
    jin = module.input_data()   # loads input.json
    jout = main(jin)
    module.write_output_data(jout)  # writes output.json
