'''This script contains all functions for various plots that could be relevant 
for the presentation or protocol of the Grewe GP 2024'''

import os
import matplotlib.pyplot as plt
import numpy as np
import rlxnix as rlx
from useful_functions import power_spectrum
import sys

'''IMPORT DATA'''
datafolder = '../data' #./ wo ich gerade bin; ../ eine ebene höher; ../../ zwei ebenen höher
example_file = os.path.join('..', 'data', '2024-10-16-ac-invivo-1.nix')

'''EXTRACT DATA'''
dataset = rlx.Dataset(example_file)

# get sams
sams = dataset.repro_runs('SAM')
sam = sams[2]

# get stims
stimulus = sam.stimuli[-1]
stim_count = sam.stimulus_count

'''PLOTS'''
# create colormap
colors = plt.cm.prism(np.linspace(0, 1, stim_count))

# plot timeline of whole rec
dataset.plot_timeline()

# plot voltage over time for whole trace
def plot_vt_spikes(t, v, spike_t):
    fig = plt.figure(figsize=(5, 2.5))
    # alternative to ax = axs[0]
    ax = fig.add_subplot()
    # plot vt diagram
    ax.plot(t[t<0.1], v[t<0.1])
    # plot spikes into vt diagram, at max V
    ax.scatter(spike_t[spike_t<0.1], np.ones_like(spike_t[spike_t<0.1]) * np.max(v))
    plt.show()

# plot scatter plot for one sam with all 3 stims
def scatter_plot(colormap, stimuli_list, stimulus_count):
    fig = plt.figure()
    ax = fig.add_subplot()
        
    ax.eventplot(stimuli_list, colors=colormap)
    ax.set_xlabel('Spike Times [ms]')
    ax.set_ylabel('Loop #')
    ax.set_yticks(range(stimulus_count))
    ax.set_title('Spikes of SAM 3')
    plt.show()
    
# calculate power spectrum
freq, power = power_spectrum(stimulus)

# plot power spectrum
def power_spectrum_plot(f, p):
    # plot power spectrum
    fig = plt.figure()
    ax = fig.add_subplot()
    ax.plot(freq, power)
    ax.set_xlabel('Frequency [Hz]')
    ax.set_ylabel('Power [1/Hz]')
    ax.set_xlim(0, 1000)
    plt.show()
    
# DIANAS POWER SPECTRUM PLOT
functions_path = r"C:\Users\diana\OneDrive - UT Cloud\Master\GPs\GP1_Grewe\Projekt\gpgrewe2024\code"
sys.path.append(functions_path)
import useful_functions as u
import matplotlib.ticker as ticker

def float_formatter(x, _):
    """Format the y-axis values as floats with a specified precision."""
    return f'{x:.5f}'

def plot_highlighted_integrals(ax, frequency, power, points, color_mapping, points_categories, delta=2.5):
    """
    Highlight integrals on the existing axes of the power spectrum.

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        The axes on which to plot the highlighted integrals.
    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 check and highlight.
    delta : float
        Half-width of the range for integration around each point.
    color_mapping : dict
        A dictionary mapping each category to its color.
    points_categories : dict
        A mapping of categories to lists of points.

    Returns
    -------
    None
    """
    ax.plot(frequency, power)  # Plot power spectrum on the existing axes

    for point in points:
        # Calculate the integral and local mean
        integral, local_mean = u.calculate_integral_2(frequency, power, point)
        
        # Check if the point is valid
        valid = u.valid_integrals(integral, local_mean, point)
        
        if valid:  
            # Define color based on the category of the point
            point_category = next((cat for cat, pts in points_categories.items() if point in pts), "Unknown")
            color = next((c for cat, c in color_mapping.items() if point in points_categories[cat]), 'gray')
            
            # Shade the region around the point where the integral was calculated
            ax.axvspan(point - delta, point + delta, color=color, alpha=0.3, label=f'{point_category}')

            # Print out point, category, and color
            #ax.text(f"{point_category}: Integral: {integral:.5e}, Color: {color}")

    ax.set_xlim([0, 1200])
    ax.set_ylim([0, 6e-5])
    ax.set_xlabel('Frequency (Hz)')
    ax.set_ylabel('Power')
    ax.set_title('Power Spectrum with Highlighted Integrals')

    # Apply float formatting to the y-axis
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(float_formatter))