import numpy as np

def get_sweep_v_i_t_from_set(data_set, sweep_number):
    sweep_data = data_set.get_sweep(sweep_number)
    i = sweep_data["stimulus"] # in A
    v = sweep_data["response"] # in V
    i *= 1e12 # to pA
    v *= 1e3 # to mV
    sampling_rate = sweep_data["sampling_rate"] # in Hz
    t = np.arange(0, len(v)) * (1.0 / sampling_rate)
    return v, i, t

def get_sweeps_of_type(sweep_type, sweeps):
    sweeps = [ s for s in sweeps if s['ephys_stimulus']['description'].startswith( sweep_type )]
    sweep_numbers = [ s['sweep_number'] for s in sweeps ]
    statuses = [ s['workflow_state'] for s in sweeps ]
    
    return sweeps, sweep_numbers, statuses

def get_step_stim_characteristics(i, t):
    # Assumes that there is a test pulse followed by the stimulus step
    di = np.diff(i)
    up_idx = np.flatnonzero(di > 0)
    down_idx = np.flatnonzero(di < 0)

    # second step is the stimulus
    if len(up_idx) < 2 or len(down_idx) < 2:
        return (np.nan, np.nan, np.nan, np.nan, np.nan)

    if up_idx[1] < down_idx[1]: # positive step
        start_idx = up_idx[1] + 1 # shift by one to compensate for diff()
        end_idx = down_idx[1] + 1
    else: # negative step
        start_idx = down_idx[1] + 1
        end_idx = up_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)
