# Allen Institute Software License - This software license is the 2-clause BSD
# license plus a third clause that prohibits redistribution for commercial
# purposes without further permission.
#
# Copyright 2017. Allen Institute. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Redistributions for commercial purposes are not permitted without the
# Allen Institute's written permission.
# For purposes of this license, commercial purposes is the incorporation of the
# Allen Institute's software into anything for which you will charge fees or
# other compensation. Contact terms@alleninstitute.org for commercial licensing
# opportunities.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
import numpy as np
import logging
import six
from . import ephys_extractor as efex
from . import ephys_features as ft

HERO_MIN_AMP_OFFSET = 39.0
HERO_MAX_AMP_OFFSET = 61.0

SHORT_SQUARE_TYPES = ["Short Square",
                      "Short Square - Triple",
                      "Short Square - Hold -60mv",
                      "Short Square - Hold -70mv",
                      "Short Square - Hold -80mv"]

SHORT_SQUARE_THRESH_FRAC_FLOOR = 0.1

MEAN_FEATURES = [ "upstroke_downstroke_ratio", "peak_v", "peak_t", "trough_v", "trough_t",
                  "fast_trough_v", "fast_trough_t", "slow_trough_v", "slow_trough_t",
                  "threshold_v", "threshold_i", "threshold_t", "peak_v", "peak_t" ]


def extract_sweep_features(data_set, sweeps_by_type):
    # extract sweep-level features
    sweep_features = {}

    for stimulus_type, sweep_numbers in six.iteritems(sweeps_by_type):
        logging.debug("%s:%s" % (stimulus_type, ','.join(map(str, sweep_numbers))))

        if stimulus_type == "Short Square - Triple":
            tmp_ext = efex.extractor_for_nwb_sweeps(data_set, sweep_numbers)
            t_set = [s.t for s in tmp_ext.sweeps()]
            v_set = [s.v for s in tmp_ext.sweeps()]

            # IT-14530
            # triple-sweeps to use different window
            win_start = efex.SHORT_SQUARE_TRIPLE_WINDOW_START
            win_end = efex.SHORT_SQUARE_TRIPLE_WINDOW_END
            cutoff, thresh_frac = ft.estimate_adjusted_detection_parameters(
                                    v_set, t_set, win_start, win_end)
            thresh_frac = max(SHORT_SQUARE_THRESH_FRAC_FLOOR, thresh_frac)

            fex = efex.extractor_for_nwb_sweeps(data_set, sweep_numbers,
                                    dv_cutoff=cutoff, thresh_frac=thresh_frac)
        elif stimulus_type in SHORT_SQUARE_TYPES:
            tmp_ext = efex.extractor_for_nwb_sweeps(data_set, sweep_numbers)
            t_set = [s.t for s in tmp_ext.sweeps()]
            v_set = [s.v for s in tmp_ext.sweeps()]

            win_start = efex.SHORT_SQUARES_WINDOW_START
            win_end = efex.SHORT_SQUARES_WINDOW_END
            cutoff, thresh_frac = ft.estimate_adjusted_detection_parameters(
                                     v_set, t_set, win_start, win_end)
            thresh_frac = max(SHORT_SQUARE_THRESH_FRAC_FLOOR, thresh_frac)

            fex = efex.extractor_for_nwb_sweeps(data_set, sweep_numbers,
                                                dv_cutoff=cutoff, thresh_frac=thresh_frac)
        else:
            fex = efex.extractor_for_nwb_sweeps(data_set, sweep_numbers)

        fex.process_spikes()

        sweep_features.update({ f.id:f.as_dict() for f in fex.sweeps() })

    return sweep_features

# if subthreshold minimum amplitude is known (e.g., for human cells) then
#   specify it. otherwise the default value will be used
def extract_cell_features(data_set,
                          ramp_sweep_numbers,
                          short_square_sweep_numbers,
                          long_square_sweep_numbers,
                          subthresh_min_amp = None):

    if subthresh_min_amp is None:
        fex = efex.cell_extractor_for_nwb(data_set,
                                          ramp_sweep_numbers,
                                          short_square_sweep_numbers,
                                          long_square_sweep_numbers)
    else:
        fex = efex.cell_extractor_for_nwb(data_set,
                                          ramp_sweep_numbers,
                                          short_square_sweep_numbers,
                                          long_square_sweep_numbers,
                                          subthresh_min_amp)

    fex.process()

    cell_features = fex.as_dict()

    # find hero sweep
    rheo_amp = cell_features['long_squares']['rheobase_i']
    hero_min, hero_max = rheo_amp + HERO_MIN_AMP_OFFSET, rheo_amp + HERO_MAX_AMP_OFFSET
    hero_amp = float("inf")
    hero_sweep = None
    for sweep in fex.long_squares_features("spiking").sweeps():
        nspikes = len(sweep.spikes())
        amp = sweep.sweep_feature("stim_amp")

        if nspikes > 0 and amp > hero_min and amp < hero_max and amp < hero_amp:
            hero_amp = amp
            hero_sweep = sweep

    if hero_sweep:
        adapt = hero_sweep.sweep_feature("adapt")
        latency = hero_sweep.sweep_feature("latency")
        mean_isi = hero_sweep.sweep_feature("mean_isi")
    else:
        raise ft.FeatureError("Could not find hero sweep.")

    # find the mean features of the first spike for the ramps and short squares
    ramps_ms0 = mean_features_spike_zero(fex.ramps_features().sweeps())
    ss_ms0 = mean_features_spike_zero(fex.short_squares_features().sweeps())

    # compute baseline from all long square sweeps
    v_baseline = np.mean(fex.long_squares_features().sweep_features('v_baseline'))

    cell_features['long_squares']['v_baseline'] = v_baseline
    cell_features['long_squares']['hero_sweep'] = hero_sweep.as_dict() if hero_sweep else None
    cell_features["ramps"]["mean_spike_0"] = ramps_ms0
    cell_features["short_squares"]["mean_spike_0"] = ss_ms0

    return cell_features

def mean_features_spike_zero(sweeps):
    """ Compute mean feature values for the first spike in list of extractors """

    output = {}
    for mf in MEAN_FEATURES:
        mfd = [ sweep.spikes()[0][mf] for sweep in sweeps if sweep.sweep_feature("avg_rate") > 0 ]
        output[mf] = np.mean(mfd)
    return output

def get_stim_characteristics(i, t, no_test_pulse=False):
    '''
    Identify the start time, duration, amplitude, start index, and
    end index of a general stimulus.
    This assumes that there is a test pulse followed by the stimulus square.
    '''

    di = np.diff(i)
    diff_idx = np.flatnonzero(di != 0)

    if len(diff_idx) == 0:
        return (None, None, 0.0, None, None)

    # skip the first up/down
    idx = 0 if no_test_pulse else 1

    # shift by one to compensate for diff()
    start_idx = diff_idx[idx] + 1
    end_idx = diff_idx[-1] + 1

    stim_start = float(t[start_idx])
    stim_dur = float(t[end_idx] - t[start_idx])
    stim_amp = float(i[start_idx])

    return (stim_start, stim_dur, stim_amp, start_idx, end_idx)

def get_ramp_stim_characteristics(i, t):
    ''' Identify the start time and start index of a ramp sweep. '''

    # Assumes that there is a test pulse followed by the stimulus ramp
    di = np.diff(i)
    up_idx = np.flatnonzero(di > 0)

    start_idx = up_idx[1] + 1 # shift by one to compensate for diff()
    return (t[start_idx], start_idx)

def get_square_stim_characteristics(i, t, no_test_pulse=False):
    '''
    Identify the start time, duration, amplitude, start index, and
    end index of a square stimulus.
    This assumes that there is a test pulse followed by the stimulus square.
    '''

    di = np.diff(i)
    up_idx = np.flatnonzero(di > 0)
    down_idx = np.flatnonzero(di < 0)

    idx = 0 if no_test_pulse else 1

    # second square is the stimulus
    if up_idx[idx] < down_idx[idx]: # positive square
        start_idx = up_idx[idx] + 1 # shift by one to compensate for diff()
        end_idx = down_idx[idx] + 1
    else: # negative square
        start_idx = down_idx[idx] + 1
        end_idx = up_idx[idx] + 1

    stim_start = float(t[start_idx])
    stim_dur = float(t[end_idx] - t[start_idx])
    stim_amp = float(i[start_idx])

    return (stim_start, stim_dur, stim_amp, start_idx, end_idx)
