data_preprocessing.py

import os
import pickle
import numpy as np
import pandas as pd
from scipy.signal import resample
import neurokit2 as nk

# ---------------- CONFIG ----------------
DATA_DIR = "./dataset"   # Change this to your dataset path
EDA_TEMP_RATE = 4        # Hz, for EDA & TEMP
BVP_RATE = 64            # Hz, for BVP/HRV
LABEL_RATE = 700         # Hz, for the 'label' array in PKL file
WINDOW_SIZE_SEC = 60     # seconds
OVERLAP = 0.5
OUTPUT_FILE = "wesad_wrist_clean_hr.csv"

# ---------------- HELPERS ----------------
def normalize(sig):
    """Z-score normalization"""
    return (sig - np.mean(sig)) / np.std(sig)

def process_subject(file_path):
    """Load a single subject .pkl and extract wrist sensor data"""
    with open(file_path, 'rb') as f:
        # FIX: Use latin1 encoding
        data = pickle.load(f, encoding='latin1')

    signals = data['signal']['wrist']
    labels = data['label'] # THIS IS THE 700 HZ ARRAY

    # Extract raw signals
    eda = np.array(signals['EDA'])
    temp = np.array(signals['TEMP'])
    bvp = np.array(signals['BVP'])  # 64 Hz signal
    
    # We don't need to align labels here. We need to align the signals themselves.

    # Resample TEMP to match EDA rate (already both 4Hz, so len_target = len(eda))
    temp = resample(temp, len(eda))

    # Normalize
    eda = normalize(eda)
    temp = normalize(temp)
    bvp = normalize(bvp)  # for simplicity

    # RETURN THE FULL 700 HZ LABELS ARRAY
    return eda, temp, bvp, labels 


# ---------------- HR & HRV Extraction ----------------
def extract_hr_hrv_and_indices(bvp_signal, sampling_rate=BVP_RATE, label_rate=LABEL_RATE, window_size_sec=WINDOW_SIZE_SEC):
    """Extract HR and HRV (SDNN) from BVP signal in sliding windows.
    Computes HR/HRV and returns the features, along with the start indices 
    in the BVP array (64 Hz) and the corresponding label indices (700 Hz).
    """
    n_samples = len(bvp_signal)
    window_samples = window_size_sec * sampling_rate
    step_samples = int(window_samples * (1 - OVERLAP))

    hr_list, hrv_list = [], []
    bvp_start_indices = []

    # 1. Calculate the rate conversion factor from BVP rate (64Hz) to Label rate (700Hz)
    rate_conversion_factor = label_rate / sampling_rate

    for start_bvp in range(0, n_samples - window_samples + 1, step_samples):
        segment = bvp_signal[start_bvp:start_bvp+window_samples]
        hr, hrv = np.nan, np.nan
        
        # Calculate the corresponding start index in the 700 Hz label array
        label_index_start = int(start_bvp * rate_conversion_factor)

        try:
            # Clean & process BVP signal
            cleaned = nk.ppg_clean(segment, sampling_rate=sampling_rate)
            peaks = nk.ppg_findpeaks(cleaned, sampling_rate=sampling_rate)['PPG_Peaks']

            # Mean HR and HRV (SDNN) using the simplified method that successfully generated finite values
            if len(peaks) >= 2:
                rr_intervals = np.diff(peaks) / sampling_rate # in seconds
                hr = 60 / np.mean(rr_intervals)               # HR in bpm
                hrv = np.std(rr_intervals) * 1000            # SDNN in ms

        except Exception as e:
            pass  # Keep as NaN if error

        hr_list.append(hr)
        hrv_list.append(hrv)
        bvp_start_indices.append(start_bvp) # Save BVP index to look up EDA/TEMP later

    return hr_list, hrv_list, bvp_start_indices

# ---------------- WINDOW CREATION ----------------
def create_windows(eda, temp, hr, hrv, bvp_indices, labels):
    """Combine features per window and assign labels using the correct 700Hz alignment"""
    X, y = [], []
    
    # Calculate the scale factor from BVP_RATE (64 Hz) to EDA/TEMP_RATE (4 Hz)
    eda_scale_factor = BVP_RATE / EDA_TEMP_RATE
    
    # Calculate the duration of the window in 700 Hz samples
    label_window_duration = WINDOW_SIZE_SEC * LABEL_RATE

    for i in range(len(hr)):
        # 1. Get the starting index of the window in the 64 Hz BVP signal
        start_bvp = bvp_indices[i]
        
        # 2. Find the corresponding window in the 700 Hz label array
        # We need the label that represents the majority/center of the window.
        label_window_start = int(start_bvp * (LABEL_RATE / BVP_RATE))
        label_window_end = label_window_start + label_window_duration
        
        # Ensure we don't go past the end of the labels array
        label_segment = labels[label_window_start : min(label_window_end, len(labels))]
        
        # 3. Determine the majority label (mode)
        # Filter for valid labels (1, 2, 3, 4) and exclude 0, 5, 6, 7
        valid_labels = label_segment[np.isin(label_segment, [1, 2, 3, 4])]
        
        if len(valid_labels) == 0:
             # Skip transient/undefined segments
             continue 

        # Find the most frequent label in the valid segment
        mode_label = np.argmax(np.bincount(valid_labels.astype(int)))
        
        # 4. Check for valid HR/HRV features (if you still want to filter noise)
        is_valid_features = np.isfinite(hr[i]) and np.isfinite(hrv[i])
        
        if is_valid_features:
            # 5. Get the corresponding 4 Hz EDA/TEMP feature point
            # Since EDA/TEMP are also 4Hz, they are shorter. We use the BVP index
            # scaled down by the BVP_RATE/EDA_RATE ratio (64/4 = 16)
            eda_index = int(start_bvp / eda_scale_factor) 

            # Append the features
            X.append([eda[eda_index][0], temp[eda_index][0], hr[i], hrv[i]])
            y.append(mode_label)

    return X, y

# ---------------- MAIN ----------------
all_X, all_y = [], []

for subject in sorted(os.listdir(DATA_DIR)):
    if subject.startswith("S") and subject.endswith(".pkl"):
        pkl_path = os.path.join(DATA_DIR, subject)
        print(f"Processing {subject} ...")

        # Returns eda/temp (4Hz), bvp (64Hz), and labels (700Hz)
        eda, temp, bvp, labels = process_subject(pkl_path)

        # Extract HR & HRV and BVP starting indices
        hr_list, hrv_list, bvp_indices = extract_hr_hrv_and_indices(bvp)

        # Use the new create_windows function with proper label alignment
        X, y = create_windows(eda, temp, hr_list, hrv_list, bvp_indices, labels)
        
        all_X.extend(X)
        all_y.extend(y)

# ---------------- SAVE ----------------
df = pd.DataFrame(all_X, columns=['EDA','TEMP','HR','HRV'])
df['label'] = all_y
label_map = {1:'neutral',2:'stressed',3:'happy',4:'relaxed'}
df['emotion'] = df['label'].map(label_map)

df.to_csv(OUTPUT_FILE, index=False)
print(f"\n✅ Saved dataset with {len(df)} rows to {OUTPUT_FILE}")