#! /usr/bin/env python

# This file is part of the SpikeCoding repository - MAVLab TU Delft
#
# MIT License
#
# Copyright (c) 2021 Julien Dupeyroux
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# @author Julien Dupeyroux, Stein Stroobants

import numpy as np
from scipy.stats import norm


def temporal_contrast(data, factor):
    # Based on algorithm provided in:
    #   Sengupta et al. (2017)
    #   Petro et al. (2020)
    diff = np.zeros(len(data)-1)
    spikes = np.zeros(len(data))
    for i in range(len(data)-1):
        diff[i] = data[i+1] - data[i]
    threshold = np.mean(diff) + factor * np.std(diff)
    diff = np.insert(diff, 0, diff[1])
    for i in range(len(data)):
        if diff[i] > threshold:
            spikes[i] = 1
        elif diff[i] < -threshold:
            spikes[i] = -1
    return spikes, threshold


def step_forward(data, threshold):
    # Based on algorithm provided in:
    #   Petro et al. (2020)
    startpoint = data[0]
    spikes = np.zeros(len(data))
    base = startpoint
    for i in range(1,len(data)):
        if data[i] > base + threshold:
            spikes[i] = 1
            base = base + threshold
        elif data[i] < base - threshold:
            spikes[i] = -1
            base = base - threshold
    return spikes, startpoint


def moving_window(data, threshold, window):
    # Based on algorithm provided in:
    #   Petro et al. (2020)
    startpoint = data[0]
    spikes = np.zeros(len(data))
    base = np.mean(data[0:window+1])
    for i in range(window+1):
        if data[i] > base + threshold:
            spikes[i] = 1
        elif data[i] < base - threshold:
            spikes[i] = -1
    for i in range(window+2, len(data)):
        base = np.mean(data[(i-window-1):(i-1)])
        if data[i] > base + threshold:
            spikes[i] = 1
        elif data[i] < base - threshold:
            spikes[i] = -1
    return spikes, startpoint


def hough_spike(data, fir):
    # Based on algorithm provided in:
    #   Schrauwen et al. (2003)
    spikes = np.zeros(len(data))
    shift = min(data)
    data = data - shift*np.ones(len(data))
    for i in range(len(data)):
        count = 0
        for j in range(len(fir)):
            if i+j < len(data):
                if data[i+j] >= fir[j]:
                    count = count + 1
        if count == len(fir):
            spikes[i] = 1
            for j in range(len(fir)):
                if i+j < len(data):
                    data[i+j] = data[i+j] - fir[j]
    return spikes, shift


def modified_hough_spike(data, fir, threshold):
    # Based on algorithm provided in:
    #   Schrauwen et al. (2003)
    spikes = np.zeros(len(data))
    shift = min(data)
    data = data - shift*np.ones(len(data))
    for i in range(len(data)):
        error = 0
        for j in range(len(fir)):
            if i+j < len(data):
                if data[i+j] < fir[j]:
                    error = error + fir[j] - data[i+j]
        if error <= threshold:
            spikes[i] = 1
            for j in range(len(fir)):
                if i+j < len(data):
                    data[i+j] = data[i+j] - fir[j]
    return spikes, shift


def ben_spike(data, fir, threshold):
    # Based on algorithm provided in:
    #   Petro et al. (2020)
    #   Sengupta et al. (2017)
    #   Schrauwen et al. (2003)
    spikes = np.zeros(len(data))
    shift = min(data)
    data = data - shift*np.ones(len(data))
    for i in range(len(data)-len(fir)+1):
        err1 = 0
        err2 = 0
        for j in range(len(fir)):
            err1 = err1 + abs(data[i+j] - fir[j])
            err2 = err2 + abs(data[i+j-1])
        if err1 <= err2*threshold:
            spikes[i] = 1
            for j in range(len(fir)):
                if i+j+1 < len(data):
                    data[i+j+1] = data[i+j+1] - fir[j]
    return spikes, shift


def grf_spike(data, m, min_input, max_input):
    # Adapted from algorithm provided in:
    #   Bohté et al. (2002)
    # Modifications: definition of sigma, removal of beta constant,
    #                and modified WTA process

    if np.isscalar(data):
        data = [data]

    spikes = np.zeros((len(data),m))
    neuron_outputs = np.zeros(m)

    for j in range(len(data)):
        for i in range(m):
            mu = min_input + (2*(i + 1)-3)/2*(max_input - min_input)/(m-2)
            sigma = (max_input - min_input)/(m-2)
            neuron_outputs[i] = norm.pdf(data[j], mu, sigma)

        spikes[j,np.argmax(neuron_outputs)] = 1
    return spikes

def one_hot_place_spike(data, m, min_input, max_input):
    # Simple population coding algorithm adapted from Stagsted et al. (2020) that represents inputs by a location. 
    # An input is assigned to the neuron that is closest to its value. 
    # Only one neuron fires at every timestep

    if np.isscalar(data):
        data = [data]

    spikes = np.zeros((len(data),m))

    for j in range(len(data)):
        size_change = 1/2*(max_input - min_input)/(m-2) # to make sure it has the same lower/upper bounds as the Bohte paper
        idx = int(np.round(((data[j] - (min_input - size_change)) / ((max_input + size_change) - (min_input - size_change))) * (m - 1)))
        spikes[j, idx] = 1
    
    return spikes


def grf_spike_with_internal_timesteps(data, min_input, max_input, neurons=10, timesteps=10, beta=1.5):
    """Create a series of spikes based on Gaussian Receptive Fields
    Adapted from algorithm provided in:
        Bohté et al. (2002)
    
    Keyword arguments:
    data -- 
    neurons -- numbers of neurons (default 10)
    timesteps -- number of timesteps (default 10)
    min_input -- minimal value
    max_input -- maximum value
    beta -- tuning parameter that determines the width of the receptive fields
    """

    if np.isscalar(data):
        data = [data]
        
    spikes = np.zeros((len(data), timesteps, neurons))
    responses = np.zeros(neurons)

    # Calculation of mu and sigma of the Gaussian receptive fields
    mu = min_input + (2*(np.arange(neurons)+1)-3)/2*(max_input - min_input)/(neurons-2)
    sigma = 1/beta*(max_input - min_input)/(neurons-2)
    max_prob = norm.pdf(mu[0], mu[0], sigma)

    for j in range(len(data)):
        for i in range(neurons):
            responses[i] = norm.pdf(data[j], mu[i], sigma)
            size_change = max_prob / (2 * timesteps)
            new = int(np.round(((responses[i] + size_change) / (max_prob + 2 * size_change) * (timesteps + 1)) + 0.0001)) # 0.0001 for roundoff errors...
            spiking_time = timesteps - new
            if spiking_time < timesteps - 1:
                spikes[j, spiking_time, i] = 1
    spikes = spikes.reshape([len(data) * timesteps, neurons])
    return spikes
:
