import os
import numpy as np
import pandas as pd
from scipy import signal
import mne
import torch

class EEGDataAugmentation:
    def __init__(self, noise_factor=0.02, shift_limit=0.05):
        """
        EEG数据增强类
        
        参数:
        - noise_factor: 高斯噪声的强度
        - shift_limit: 时间平移的最大比例
        """
        self.noise_factor = noise_factor
        self.shift_limit = shift_limit
    
    def add_noise(self, data):
        """添加高斯噪声"""
        noise = np.random.normal(0, self.noise_factor, data.shape)
        return data + noise
    
    def time_shift(self, data):
        """时间平移"""
        shift = int(data.shape[-1] * self.shift_limit)
        if shift > 0:
            direction = np.random.choice([-1, 1])
            shift_value = np.random.randint(1, shift)
            shifted = np.roll(data, direction*shift_value, axis=-1)
            return shifted
        return data
    
    def apply_augmentation(self, data, augment_prob=0.5):
        """应用数据增强
        
        参数:
        - data: shape (trials, channels, time_points)
        - augment_prob: 应用每种增强的概率
        """
        augmented_data = data.copy()
        
        # 随机应用高斯噪声
        if np.random.random() < augment_prob:
            augmented_data = self.add_noise(augmented_data)
        
        # 随机应用时间平移
        if np.random.random() < augment_prob:
            augmented_data = self.time_shift(augmented_data)
        
        return augmented_data

def process_eeg_data(data, sfreq=250, augment=False):
    """
    对EEG数据进行全面的预处理
    
    参数:
    - data: shape (trials, channels, time_points)
    - sfreq: 采样频率，默认250Hz
    - augment: 是否启用数据增强
    
    返回:
    - processed_data: 预处理后的数据
    """
    processed_data = np.zeros_like(data)
    augmenter = EEGDataAugmentation() if augment else None
    
    for trial in range(data.shape[0]):
        # 1. 创建MNE Raw对象
        ch_names = [f'EEG{i+1}' for i in range(data.shape[1])]
        ch_types = ['eeg'] * data.shape[1]
        info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
        raw = mne.io.RawArray(data[trial], info)
        
        # 2. 信号滤波
        # 带通滤波 4-40Hz
        raw.filter(l_freq=4, h_freq=40, method='iir')
        
        try:
            # 使用scipy的iirnotch滤波器替代MNE的notch_filter
            nyq = sfreq / 2
            freq = 50 / nyq
            Q = 30
            b, a = signal.iirnotch(freq, Q)
            
            # 应用陷波滤波器
            processed = raw.get_data()
            for ch in range(processed.shape[0]):
                processed[ch] = signal.filtfilt(b, a, processed[ch])
            
            # 数据标准化 (每个通道独立)
            for ch in range(processed.shape[0]):
                ch_data = processed[ch]
                ch_mean = np.mean(ch_data)
                ch_std = np.std(ch_data)
                processed[ch] = (ch_data - ch_mean) / (ch_std + 1e-10)
            
            # 应用数据增强（如果启用）
            if augment and augmenter is not None:
                processed = augmenter.apply_augmentation(processed)
            
            processed_data[trial] = processed
            
        except Exception as e:
            print(f"Error in trial {trial}: {e}")
            continue

    return processed_data

def get_subject_id(filename):
    """从文件名中提取被试ID"""
    # 文件名格式为 "MI-EEG-A01T.csv" 或 "MI-EEG-A01E.csv"
    # 提取"01"部分并转换为整数
    subject_num = filename.split('-')[-1][1:3]  # 提取"01"
    return int(subject_num)

def load_and_process_data():
    # 设置路径
    train_data_folder = '/root/autodl-tmp/train'
    train_label_folder = '/root/autodl-tmp/train_label'
    test_data_folder = '/root/autodl-tmp/test'
    test_label_folder = '/root/autodl-tmp/test_label'
    save_root = '/root/autodl-tmp/datas'
    os.makedirs(save_root, exist_ok=True)

    # 处理训练数据
    train_data_list = []
    train_labels_list = []
    train_pid_list = []
    
    train_files = sorted([f for f in os.listdir(train_data_folder) if f.endswith('T.csv')])
    for file in train_files:
        subject_id = get_subject_id(file)
        print(f"Processing training file: {file}, Subject ID: {subject_id}")
        
        # 读取数据和标签
        data = pd.read_csv(os.path.join(train_data_folder, file))
        labels = pd.read_csv(os.path.join(train_label_folder, f"Etiquetas{file.split('-')[-1]}"))
        labels = labels.values.flatten() - 1  # 将1-4转换为0-3
        
        # 确保标签在正确范围内
        if not np.all((labels >= 0) & (labels <= 3)):
            print(f"Warning: Invalid label values found in {file}")
            labels = np.clip(labels, 0, 3)
        
        print(f"Label values range: {np.min(labels)} to {np.max(labels)}")
        
        # 打印数据维度信息
        print(f"Data shape: {data.values.shape}")
        print(f"Number of trials (labels): {len(labels)}")
        
        # 重塑数据为(trials, channels=22, time_points=1000)格式
        data_array = data.values  # shape: (287, 22000)
        n_trials = len(labels)
        n_channels = 22
        time_points = 1000  # 每个通道的时间点数
        
        print(f"Reshaping to: ({n_trials}, {n_channels}, {time_points})")
        
        try:
            # 重塑数据
            reshaped_data = np.zeros((n_trials, n_channels, time_points))
            for trial in range(n_trials):
                for channel in range(n_channels):
                    start_idx = channel * time_points
                    end_idx = (channel + 1) * time_points
                    reshaped_data[trial, channel, :] = data_array[trial, start_idx:end_idx]
            
            print(f"Reshaped data shape: {reshaped_data.shape}")
            
            # 数据预处理
            processed_data = process_eeg_data(reshaped_data, augment=True)
            
            # 添加基线校正
            baseline_period = slice(0, 100)  # 使用前100个时间点作为基线
            baseline_mean = np.mean(processed_data[..., baseline_period], axis=-1, keepdims=True)
            processed_data = processed_data - baseline_mean
            
            train_data_list.append(processed_data)
            train_labels_list.append(labels)
            train_pid_list.extend([subject_id] * len(labels))
            
        except Exception as e:
            print(f"Error processing data: {e}")
            continue

    # 处理测试数据
    test_data_list = []
    test_labels_list = []
    test_pid_list = []
    
    test_files = sorted([f for f in os.listdir(test_data_folder) if f.endswith('E.csv')])
    for file in test_files:
        subject_id = get_subject_id(file)
        print(f"Processing test file: {file}, Subject ID: {subject_id}")
        
        data = pd.read_csv(os.path.join(test_data_folder, file))
        labels = pd.read_csv(os.path.join(test_label_folder, f"Etiquetas{file.split('-')[-1]}"))
        labels = labels.values.flatten() - 1  # 将1-4转换为0-3
        
        # 确保标签在正确范围内
        if not np.all((labels >= 0) & (labels <= 3)):
            print(f"Warning: Invalid label values found in {file}")
            labels = np.clip(labels, 0, 3)
        
        data_array = data.values
        n_trials = len(labels)
        n_channels = 22
        time_points = 1000
        
        try:
            # 使用相同的重塑逻辑
            reshaped_data = np.zeros((n_trials, n_channels, time_points))
            for trial in range(n_trials):
                for channel in range(n_channels):
                    start_idx = channel * time_points
                    end_idx = (channel + 1) * time_points
                    reshaped_data[trial, channel, :] = data_array[trial, start_idx:end_idx]
            
            processed_data = process_eeg_data(reshaped_data)
            
            # 添加基线校正
            baseline_period = slice(0, 100)  # 使用前100个时间点作为基线
            baseline_mean = np.mean(processed_data[..., baseline_period], axis=-1, keepdims=True)
            processed_data = processed_data - baseline_mean
            
            test_data_list.append(processed_data)
            test_labels_list.append(labels)
            test_pid_list.extend([subject_id] * len(labels))
            
        except Exception as e:
            print(f"Error processing data: {e}")
            continue

    # 合并数据
    print("\nMerging training data...")
    X_train = np.vstack(train_data_list) if train_data_list else np.array([])
    y_train = np.concatenate(train_labels_list) if train_labels_list else np.array([])
    pid_train = np.array(train_pid_list)

    print("\nMerging test data...")
    X_test = np.vstack(test_data_list) if test_data_list else np.array([])
    y_test = np.concatenate(test_labels_list) if test_labels_list else np.array([])
    pid_test = np.array(test_pid_list)

    if len(X_train) == 0 or len(X_test) == 0:
        print("Error: No data was successfully processed!")
        return

    # 最终的数据标准化（在所有trials上）
    print("\nNormalizing data...")
    X = np.concatenate((X_train, X_test))
    print("X shape after concatenation:", X.shape)
    
    # 对每个通道分别进行标准化
    for ch in range(X.shape[1]):
        ch_mean = np.mean(X[:, ch, :])
        ch_std = np.std(X[:, ch, :])
        X[:, ch, :] = (X[:, ch, :] - ch_mean) / ch_std
    
    # 合并标签和ID
    y = np.concatenate((y_train, y_test))
    pid = np.concatenate((pid_train, pid_test))

    print("\nData shape before saving:")
    print("X shape:", X.shape)
    print("y shape:", y.shape)
    print("pid shape:", pid.shape)

    # 保存处理后的数据
    print("\nSaving data...")
    np.save(os.path.join(save_root, 'x.npy'), X)#将 X 数组保存为 x.npy 文件。
    np.save(os.path.join(save_root, 'y.npy'), y)
    np.save(os.path.join(save_root, 'pid.npy'), pid)

    print("\nData saved to:", save_root)

if __name__ == "__main__":
    load_and_process_data()