import numpy

def spiketimes(t,v,v_th = 0.5):
    """Given voltage and time, returns array of spike times

    INPUT:

    t - example arange(0,1000,0.1)
    v - voltage, same shape as t

    """

    v_above_th = v>v_th
    idx = numpy.nonzero((v_above_th[:-1]==False)&(v_above_th[1:]==True))
    return t[idx[0]+1]
    
def f(TypeXY, I):
    """Return the firing rate of neuron TypeXY under input current I

    INPUT:

    TypeXY - Neuron model (with a Step function)
    I - input current in amp.

    """

    t,v,w,I = TypeXY.Step(I_amp=I,Step_tstart=100.0,Step_tend=1000.0,tend=1000.0)

    st = spiketimes(t,v)

    # no spikes or 1 spike
    if len(st)<2:
        return 0.0

    isi = st[1:]-st[:-1]

    f = 1000.0/numpy.mean(isi)

    return f