firing_plot.py

'''
@Author: Abhinav

Function im using to plot raw_data for the first plot Gaussian smoothing for RA and Savitzky-Golay smoothing for SA.
Then second plot is for the median or (Avg) data.

'''
import sys
sys.path.append('dependencies')
from models.model_constants import MC_GROUPS
from scipy.ndimage import gaussian_filter1d
from scipy.signal import savgol_filter
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
from lmfit import minimize, fit_report, Parameters
from scipy.stats import ttest_ind, ttest_rel
import scipy.stats as stats
from models.spike_utils import get_mod_spike, stress_to_group_current
from models.model_utils import strain_to_stress
import argparse
import models.Izhikevich_Interactive as iz_module

# Mod spike function located in dependencies/aim2_population_model_sptaial_aff_parallel
'''
def get_mod_spike(lmpars, groups, time, stress, g =.4, h = 1):
    # print(f"DEBUG(get_mod_spike): tau1={lmpars['tau1'].value}")
    params = lmpars_to_params(lmpars)
    mod_spike_time, mod_fr_inst = stress_to_fr_inst(time, stress,
                                                    groups,g=g,h=h,**params)
    return (mod_spike_time, mod_fr_inst)
'''
# Global figure layout parameters
FIGURE_PARAMS = {
    'figsize': (14, 6),
    'left': 0.098,
    'right': 0.95,
    'top': 0.937,
    'bottom': 0.097
}

# Global color scheme
COLOR_MAP = {
    3.61: '#3E26A8',
    4.08: '#475BF9',
    4.31: '#2797EB',
    4.56: '#12BEB9',
}


def style_axes(ax):
    """Applying consistent styling to axes"""
    # Remove top and right spines
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    # Set to light gray
    ax.spines['left'].set_color('#666666')
    ax.spines['bottom'].set_color('#666666')
    # Remove grid lines
    ax.grid(False)
    # Ensure 0 is included in y-axis ticks
    yticks = ax.get_yticks()
    if 0 not in yticks:
        yticks = np.append(yticks, 0)
        ax.set_yticks(np.sort(yticks))
    return ax




def plot_raw_data(afferent_type="SA",
                  save_path=None, 
                  param_arr=None, 
                  data_path=None, 
                  E=20,
                  model_type="LIF"
                  ):
                  
    """
    Plot firing rates calculated from raw stress data for each von Frey tip size.
    Uses data from aggregated_data/{vf}_raw_agg_stress.csv
    Creates three plots: raw data, Gaussian smoothed, and Savitzky-Golay smoothed.

    """

    np.random.seed(1)
    vf_tip_sizes = [3.61, 4.08, 4.31, 4.56]
    
    # Creating directory
    #os.makedirs('Figure3/F3plots', exist_ok=True)
    
    # Create single figure
    fig = plt.figure(figsize=FIGURE_PARAMS['figsize'])
    ax = fig.add_subplot(111)
    ax = style_axes(ax)
    
    # Create another figure for average data
    fig_avg = plt.figure(figsize=FIGURE_PARAMS['figsize'])
    ax_avg = fig_avg.add_subplot(111)
    ax_avg = style_axes(ax_avg)
    
    # Set up parameters based on afferent type
    lmpars = Parameters()
    if afferent_type == "SA":
        lmpars.add('tau1', value=8, vary=False) #tauRI(ms)
        lmpars.add('tau2', value=200, vary=False) #tauSI(ms)
        lmpars.add('tau3', value=1744.6, vary=False)#tauUSI(ms)
        lmpars.add('tau4', value=0.87, vary=False)# kPeak
        lmpars.add('k1', value=0.2442, vary=False, min=0) #a constant
        lmpars.add('k2', value=0.0792, vary=False, min=0) #b constant
        lmpars.add('k3', value=0.0231, vary=False, min=0) #c constant
        lmpars.add('k4', value=0.13, vary=False, min=0)# Ksteady
        g, h = 0.4, 1.0  # SA afferent parameters
    else:  # RA
        lmpars.add('tau1', value = 2.5, vary=False) #tauRI(ms)
        lmpars.add('tau2', value = 1, vary=False) #tauSI(ms)
        lmpars.add('tau3', value=1, vary=False)#tauUSI(ms)
        lmpars.add('tau4', value= 1, vary=False) #Kpeak
        lmpars.add('k1', value= 21 , vary=False, min=0) #a constant
        lmpars.add('k2', value=0, vary=False, min=0) #b constant
        lmpars.add('k3', value=0, vary=False, min=0) #c constant
        lmpars.add('k4', value=0, vary=False, min=0) #Kstead
        g, h = 0.4, 1.0  # RA afferent parameters
    
    mean_peak_iff = {
        3.61: [],
        4.08: [],
        4.31: [],
        4.56: []
    }

    mean_static_hold = {
        3.61: [],
        4.08: [],
        4.31: [],
        4.56: []
    }
    
    # Dictionary to store individual IFF values for 4.08 (avg only)
    vf_4_08_iff_data = {
        'time': [],
        'avg_iff': []
    }

    
    for vf in vf_tip_sizes:
            # Option 1. Read old stress data
            # Read stress data from aggregated_data directory

            # stress_data = pd.read_csv(f"aggregated_data/{vf}_raw_agg_stress_new.csv")
            # #stress_data = pd.read_csv(f"aggregated_data/{vf}_raw_agg_stress.csv")
            # time = stress_data['Time (ms)'].to_numpy()


            # stress1 = stress_data['Stress 1 (kPa)'].to_numpy()
            # stress2 = stress_data['Stress 2 (kPa)'].to_numpy()
            # stress3 = stress_data['Stress 3 (kPa)'].to_numpy()
            


            # Option 2. Read new Strain data
            
            """
            up = f"data/P2_Current/Compressive stress/{vf}/Up/Up_Agg_Radial_Strain_VF_{vf}.csv"
            avg = f"data/P2_Current/Compressive stress/{vf}/Avg/Avg_Agg_Radial_Strain_VF_{vf}.csv"
            low = f"data/P2_Current/Compressive stress/{vf}/Low/Low_Agg_Radial_Strain_VF_{vf}.csv"

            time = pd.read_csv(up).iloc[:,0].to_numpy()

            # Testing with Youngs Modulus of 100 kPa
            upper_stress = strain_to_stress(up, E= E).iloc[:,1].to_numpy()
            avg_stress = strain_to_stress(avg, E= E).iloc[:,1].to_numpy()
            lower_stress = strain_to_stress(low, E= E).iloc[:,1].to_numpy()
            stress1 = avg_stress
            stress2 = upper_stress
            stress3 = lower_stress
            """

            # Option 3. Read new Compressive stress data
            
            avg_stress = pd.read_csv(f"data/P2_Current/Compressive stress/Realistic/{vf}/Avg/Avg_Agg_Radial_Strain_VF_{vf}.csv")
            upper_stress = pd.read_csv(f"data/P2_Current/Compressive stress/Realistic/{vf}/Up/Up_Agg_Radial_Strain_VF_{vf}.csv")
            lower_stress = pd.read_csv(f"data/P2_Current/Compressive stress/Realistic/{vf}/Low/Low_Agg_Radial_Strain_VF_{vf}.csv")
            time = avg_stress.iloc[:,0].to_numpy()
            stress1 = avg_stress.iloc[:,1].to_numpy()
            stress2 = upper_stress.iloc[:,1].to_numpy()
            stress3 = lower_stress.iloc[:,1].to_numpy()
            print(f"Stress 1: {stress1.shape}")
            print(f"Stress 2: {stress2.shape}")
            print(f"Stress 3: {stress3.shape}")
            print(f"Time: {time.shape}")
            
            groups = MC_GROUPS
            print(f"Stress 1: {stress1.shape}")
            print(f"Stress 2: {stress2.shape}")
            print(f"Stress 3: {stress3.shape}")
            print(f"Time: {time.shape}")

            print(f"Model type: {model_type}")
            print(f"Param arr: {param_arr}")
            avg_spike_time, avg_fr = get_mod_spike(lmpars, groups, time, stress1, g=g, h=h, param_arr=param_arr, model_type=model_type)
            upper_spike_time, upper_fr = get_mod_spike(lmpars, groups, time, stress2, g=g, h=h, param_arr=param_arr, model_type=model_type)
            lower_spike_time, lower_fr = get_mod_spike(lmpars, groups, time, stress3, g=g, h=h, param_arr=param_arr, model_type=model_type)

            # DEBUG: Print results from get_mod_spike
            try: 
                print(f"DEBUG: Results for vf={vf}:")
                print(f"DEBUG:   avg_fr shape = {avg_fr.shape}")
                print(f"DEBUG:   avg_fr min = {np.min(avg_fr):.6f}, max = {np.max(avg_fr):.6f}")
                print(f"DEBUG:   avg_fr unique values = {np.unique(avg_fr)}")
                print(f"DEBUG:   avg_spike_time shape = {avg_spike_time.shape}")
                print(f"DEBUG:   avg_spike_time range = [{np.min(avg_spike_time):.2f}, {np.max(avg_spike_time):.2f}]")
                print("-" * 30)
            except ValueError as e:
                print(f"Error for vf={vf}: {e}")

            print(f"Avg Spike Times: {avg_spike_time}")

            # Apply smoothing based on afferent type (for other purposes, not for raw data plotting)
            if afferent_type == "RA":
                # Add Gaussian smoothing for RA
                sigma = 2  # You can adjust this value to control smoothing strength
                avg_fr_smooth = gaussian_filter1d(avg_fr, sigma)
                upper_fr_smooth = gaussian_filter1d(upper_fr, sigma)
                lower_fr_smooth = gaussian_filter1d(lower_fr, sigma)
            else:  #SA
                # Savitzky-Golay smoothing for SA 
                # Find the minimum length among all arrays to ensure compatibility
                min_length = min(len(avg_fr), len(upper_fr), len(lower_fr))
                window_length = min(11, min_length)  # Must be odd and not larger than data size
                if window_length >= 3:  # Need at least 3 points for Savitzky-Golay
                    polyorder = min(2, window_length - 1)  # polyorder must be < window_length
                    avg_fr_smooth = savgol_filter(avg_fr, window_length, polyorder)
                    upper_fr_smooth = savgol_filter(upper_fr, window_length, polyorder)
                    lower_fr_smooth = savgol_filter(lower_fr, window_length, polyorder)
                else:
                    # If not enough data points, use original data
                    avg_fr_smooth = avg_fr
                    upper_fr_smooth = upper_fr
                    lower_fr_smooth = lower_fr
            # Key statistics for largest tip size
            if vf == 4.56 : 
                first_spike_time = avg_spike_time[0]
                print(f"First spike time: {first_spike_time}")
            if vf == 4.56 and afferent_type == "SA":
                peak_fr_idx = np.argmax(avg_fr)
                peak_fr_time = avg_spike_time[peak_fr_idx]
                peak_fr_value = avg_fr[peak_fr_idx] * 1e3

                first_second_mask = (avg_spike_time >= 0) & (avg_spike_time <= 1000)
                if np.any(first_second_mask):
                    first_peak_idx = np.argmax(avg_fr[first_second_mask])
                    first_peak_time = avg_spike_time[first_second_mask][first_peak_idx]
                else:
                    first_peak_time = 0
                last_second_mask = (avg_spike_time >= 4000) & (avg_spike_time <= 5000)
                if np.any(last_second_mask):
                    last_peak_idx = np.argmax(avg_fr[last_second_mask])
                    last_peak_time = avg_spike_time[last_second_mask][last_peak_idx]
                else:
                    last_peak_time = 0



                #Calculate Ramp on 
                first_spike_time = avg_spike_time[0]
                last_spike_time = avg_spike_time[-1]
                ramp_on = first_peak_time-first_spike_time
                ramp_off = last_spike_time-last_peak_time

                static_hold = last_peak_time - first_peak_time

                
                print(f"Peak firing rate: {peak_fr_value:.2f} Hz")
                print(f"Time of peak firing rate: {peak_fr_time:.2f} ms")
                print(f"Static hold time: {static_hold:.2f} ms")
                print(f"Ramp on time: {ramp_on:.2f} ms")
                print(f"Ramp off time: {ramp_off:.2f} ms")
            
            
            # Combining all spike times
            all_spike_times = np.sort(np.unique(np.concatenate([avg_spike_time, upper_spike_time, lower_spike_time])))
            
            #interpolating firing rates
            if len(avg_fr) > 1:

                # Interpolate raw firing rates (not smoothed) for raw data plot
                avg_fr_interp = np.interp(all_spike_times, avg_spike_time, avg_fr)
                upper_fr_interp = np.interp(all_spike_times, upper_spike_time, upper_fr)
                lower_fr_interp = np.interp(all_spike_times, lower_spike_time, lower_fr)


            else:
                avg_fr_interp = np.zeros_like(all_spike_times)
                upper_fr_interp = np.zeros_like(all_spike_times)
                lower_fr_interp = np.zeros_like(all_spike_times)
                avg_fr_smooth_interp = np.zeros_like(all_spike_times)
                upper_fr_smooth_interp = np.zeros_like(all_spike_times)
                lower_fr_smooth_interp = np.zeros_like(all_spike_times)
            
            # Plot raw data with confidence intervals (using unsmoothed data)
            all_s = np.vstack([avg_fr_interp, upper_fr_interp, lower_fr_interp])
            mean_s = np.mean(all_s, axis=0)
            stderr_s = np.std(all_s, axis=0, ddof=1)/np.sqrt(all_s.shape[0])
            # tval = stats.t.ppf(0.975,df=all_s.shape[0]-1) #two tailed 95% CI
            # CI95_s = tval*stderr_s
            up_s = mean_s + stderr_s
            low_s = mean_s - stderr_s

            ax.plot(all_spike_times, mean_s * 1e3, color=COLOR_MAP[vf])
            # ax.plot(all_spike_times, avg_fr_smooth_interp * 1e3, color=COLOR_MAP[vf])
            # ax.plot(all_spike_times, lower_fr_smooth_interp * 1e3, color='red')
            # ax.plot(all_spike_times, upper_fr_smooth_interp * 1e3, color='blue')
            ax.fill_between(all_spike_times, low_s * 1e3, up_s * 1e3, 
                          color=COLOR_MAP[vf], alpha=0.2)
            
            # Plot average firing rate (using unsmoothed data)
            ax_avg.plot(all_spike_times, mean_s * 1e3, color=COLOR_MAP[vf], label=f"{vf}mm")
            #ax_avg.plot(all_spike_times, avg_fr_smooth_interp * 1e3, color=COLOR_MAP[vf], label=f"{vf}mm")

            #adding mean_peak_iff and mean_static_hold to the dict
            mean_peak_iff[vf].append(np.max(mean_s)*1e3)
            #finding index closest to 2750ms in all_spike_times
            target_time = 2750
            closest_idx = np.abs(all_spike_times - target_time).argmin()
            mean_static_hold[vf].append(mean_s[closest_idx]*1e3)
            
            # Save individual IFF data for 4.08 (avg only)
            if vf == 4.08:
                vf_4_08_iff_data['time'] = all_spike_times
                vf_4_08_iff_data['avg_iff'] = avg_fr_interp * 1e3

            print(f"Stress 1: {np.max(stress1)}")
            print(f"Stress 2: {np.max(stress2)}")   
            print(f"Stress 3: {np.max(stress3)}")  
            



    # Statistical testing between adjacent tip sizes
    print("\nStatistical Analysis Results:")
    print("=" * 50)
    
    for i in range(len(vf_tip_sizes) - 1):
        vf_current = vf_tip_sizes[i]
        vf_next = vf_tip_sizes[i + 1]
        
        print(f"\nComparing {vf_current}mm vs {vf_next}mm:")
        print("-" * 30)
        
        # Peak IFF comparison
        current_peak = np.array(mean_peak_iff[vf_current])
        next_peak = np.array(mean_peak_iff[vf_next])
        
        # Static hold comparison
        current_static = np.array(mean_static_hold[vf_current])
        next_static = np.array(mean_static_hold[vf_next])
        
        # Perform paired t-test for peak IFF
        t_stat_peak, p_value_peak = stats.ttest_rel(current_peak, next_peak)
        print(f"Peak IFF:")
        print(f"  T-statistic: {t_stat_peak:.4f}")
        print(f"  P-value: {p_value_peak:.4f}")
        print(f"  Significant difference: {p_value_peak < 0.05}")
        
        # Perform paired t-test for static hold
        t_stat_static, p_value_static = stats.ttest_rel(current_static, next_static)
        print(f"Static Hold:")
        print(f"  T-statistic: {t_stat_static:.4f}")
        print(f"  P-value: {p_value_static:.4f}")
        print(f"  Significant difference: {p_value_static < 0.05}")
        
        # Print mean values for context
        print(f"\nMean values:")
        print(f"  Peak IFF - {vf_current}mm: {np.mean(current_peak):.2f} Hz")
        print(f"  Peak IFF - {vf_next}mm: {np.mean(next_peak):.2f} Hz")
        print(f"  Static Hold - {vf_current}mm: {np.mean(current_static):.2f} Hz")
        print(f"  Static Hold - {vf_next}mm: {np.mean(next_static):.2f} Hz")

    ax.set_ylabel('IFF (Hz)')
    ax.set_xlabel('Time (ms)')
    ax.set_xlim(left=0) 
    ax.set_ylim(bottom=0, top=300) 

    ax_avg.set_ylabel('IFF (Hz)')
    ax_avg.set_xlabel('Time (ms)')
    ax_avg.set_xlim(left=0)
    ax_avg.set_ylim(bottom=0)
    ax_avg.set_ylim(0, 300)

    # Save figures
    fig.tight_layout()
    if save_path is not None:
        fig.savefig(f'{save_path}/raw_firing_rates_{afferent_type}.svg', format='svg', bbox_inches='tight')
    else:
        fig.savefig(f'Figure3/F3plots/raw_firing_rates_{afferent_type}.svg', format='svg', bbox_inches='tight')
    
    fig_avg.tight_layout()
    if save_path is not None:
        fig_avg.savefig(f'{save_path}/avg_firing_rates_{afferent_type}.svg', format='svg', bbox_inches='tight')
    else:
        fig_avg.savefig(f'Figure3/F3plots/avg_firing_rates_{afferent_type}.svg', format='svg', bbox_inches='tight')
    
    
    # Save 4.08 IFF data to CSV
    if len(vf_4_08_iff_data['time']) > 0:  # Check if data was collected
        vf_4_08_df = pd.DataFrame(vf_4_08_iff_data)
        if save_path is not None:
            csv_filename = f'{save_path}/vf_4_08_individual_iff_{afferent_type}.csv'
        else:
            csv_filename = f'Figure3/F3plots/vf_4_08_individual_iff_{afferent_type}.csv'
        
        # Create directory if it doesn't exist
        os.makedirs(os.path.dirname(csv_filename), exist_ok=True)
        vf_4_08_df.to_csv(csv_filename, index=False)
        print(f"Saved 4.08 average IFF data to: {csv_filename}")
        print(f"Data shape: {vf_4_08_df.shape}")
        print(f"Columns: {list(vf_4_08_df.columns)}")

    plt.show()
    plt.close(fig)
    plt.close(fig_avg)



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Plot single unit model results.')
    parser.add_argument('--afferent_type', choices=['SA', 'RA'], default='SA', help='Type of afferent (SA or RA)')
    parser.add_argument('--save_path', default=None,
                       help='Path to save the figure')
    parser.add_argument('--param_arr', type=float, nargs='+',
                       help='Space-separated list of values for the parameter array')
    parser.add_argument('--data_path', default=None,
                       help='Path to the data file')
    parser.add_argument('--E', type=float, default=100,
                       help='Youngs Modulus')
    parser.add_argument('--model_type', choices=['LIF', 'Izhikevich'], default='Izhikevich',
                       help='Model type (LIF or Izhikevich)')
    args = parser.parse_args()
    plot_raw_data(args.afferent_type, args.save_path, args.param_arr, args.data_path, args.E, args.model_type)