# -*- coding: utf-8 -*-
"""
Created on Thu Oct 17 09:23:10 2024

@author: diana
"""
# -*- coding: utf-8 -*-

import glob
import os
import rlxnix as rlx
import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as sig
from scipy.integrate import quad


### FUNCTIONS ###
def binary_spikes(spike_times, duration, dt): 
    """ Converts the spike times to a binary representation. 
    Zeros when there is no spike, One when there is.     

    Parameters
    ----------
    spike_times : np.array
        The spike times.
    duration : float
        The trial duration.
    dt : float
        the temporal resolution.

    Returns
    -------
    binary : np.array
        The binary representation of the spike times.

    """
    binary = np.zeros(int(np.round(duration / dt))) #Vektor, der genauso lang ist wie die stim time
    spike_indices = np.asarray(np.round(spike_times / dt), dtype=int)
    binary[spike_indices] = 1
    return binary


def firing_rate(binary_spikes, box_width, dt=0.000025):
    """Calculate the firing rate from binary spike data.

    This function computes the firing rate using a boxcar (moving average) 
    filter of a specified width.

    Parameters
    ----------
    binary_spikes : np.array
        A binary array representing spike occurrences.
    box_width : float
        The width of the box filter in seconds.
    dt : float, optional
        The temporal resolution (time step) in seconds. Default is 0.000025 seconds.

    Returns
    -------
    rate : np.array
        An array representing the firing rate at each time step.
    """
    box = np.ones(int(box_width // dt))
    box /= np.sum(box) * dt #Normalization of box kernel to an integral of 1
    rate = np.convolve(binary_spikes, box, mode="same")
    return rate


def powerspectrum(rate, dt):
    """Compute the power spectrum of a given firing rate.

    This function calculates the power spectrum using the Welch method.

    Parameters
    ----------
    rate : np.array
        An array of firing rates.
    dt : float
        The temporal resolution (time step) in seconds.

    Returns
    -------
    frequency : np.array
        An array of frequencies corresponding to the power values.
    power : np.array
        An array of power spectral density values.
    """
    frequency, power = sig.welch(rate, fs=1/dt, nperseg=2**15, noverlap=2**14)
    return frequency, power


def prepare_harmonics(frequencies, categories, num_harmonics, colors):
    points_categories = {}
    for idx, (freq, category) in enumerate(zip(frequencies, categories)):
        points_categories[category] = [freq * (i + 1) for i in range(num_harmonics[idx])]

    points = [p for harmonics in points_categories.values() for p in harmonics]
    color_mapping = {category: colors[idx] for idx, category in enumerate(categories)}

    return points, color_mapping, points_categories


def plot_power_spectrum_with_integrals(frequency, power, points, delta, color_mapping, points_categories):
    """Create a figure of the power spectrum with integrals highlighted around specified points.

    This function creates a plot of the power spectrum and shades areas around
    specified harmonic points to indicate the calculated integrals.

    Parameters
    ----------
    frequency : np.array
        An array of frequencies corresponding to the power values.
    power : np.array
        An array of power spectral density values.
    points : list
        A list of harmonic frequencies to highlight.
    delta : float
        Half-width of the range for integration around each point.
    color_mapping : dict
        A mapping of point categories to colors.
    points_categories : dict
        A mapping of categories to lists of points.

    Returns
    -------
    fig : matplotlib.figure.Figure
        The created figure object.
    """
    fig, ax = plt.subplots()
    ax.plot(frequency, power)
    integrals = []

    for point in points:
        indices = (frequency >= point - delta) & (frequency <= point + delta)
        integral = np.trapz(power[indices], frequency[indices])
        integrals.append(integral)

        # Get color based on point category
        color = next((c for cat, c in color_mapping.items() if point in points_categories[cat]), 'gray')
        ax.axvspan(point - delta, point + delta, color=color, alpha=0.3, label=f'{point:.2f} Hz')
        print(f"Integral around {point:.2f} Hz: {integral:.5e}")

    ax.set_xlim([0, 1200])
    ax.set_xlabel('Frequency (Hz)')
    ax.set_ylabel('Power')
    ax.set_title('Power Spectrum with marked Integrals')
    ax.legend()

    return fig




### Data retrieval ###
datafolder = "../data" # Geht in der Hierarchie einen Ordern nach oben (..) und dann in den Ordner 'data'
example_file = os.path.join("..", "data", "2024-10-16-ad-invivo-1.nix")
dataset = rlx.Dataset(example_file)
sams = dataset.repro_runs("SAM") 
sam = sams[2]

## Daten für Funktionen
df = sam.metadata["RePro-Info"]["settings"]["deltaf"][0][0]
stim = sam.stimuli[1]
potential, time = stim.trace_data("V-1")
spikes, _ = stim.trace_data("Spikes-1")
duration = stim.duration
dt = stim.trace_info("V-1").sampling_interval


### Anwendung Functionen ### 
b = binary_spikes(spikes, duration, dt)
rate = firing_rate(b, box_width=0.05, dt=dt)
frequency, power = powerspectrum(b, dt)

## Important stuff 
eodf = stim.metadata[stim.name]["EODf"][0][0]
stimulus_frequency = eodf + df
AM = 50 # Hz
#print(f"EODf: {eodf}, Stimulus Frequency: {stimulus_frequency}, AM: {AM}")

frequencies = [AM, eodf, stimulus_frequency]
categories = ["AM", "EODf", "Stimulus frequency"]
num_harmonics = [4, 2, 2]
colors = ["green", "orange", "red"]
delta = 2.5


### Peaks im Powerspektrum finden ###
points, color_mapping, points_categories = prepare_harmonics(frequencies, categories, num_harmonics, colors)
fig = plot_power_spectrum_with_integrals(frequency, power, points, delta, color_mapping, points_categories)
plt.show()