import os
import pandas as pd
import pyedflib
from pathlib import Path

def get_time_to_seconds(time_str):
    h, m, s = map(int, time_str.split(":"))
    return h * 3600 + m * 60 + s

base_event = ["Hypopnea", "Central apnea", "Obstructive apnea", "Mixed apnea"]

# 输入设置
sampID = 888
dir_path = r"E:\data_annotation\6SleepApnea_annotation_GUI_demo\data"

# 数据路径设置
PSG_Data_Path = Path(os.path.join(dir_path, "PSG"))
PSG_Label_Path = Path(os.path.join(dir_path, "PSG_label"))
BCG_Data_Path = Path(os.path.join(dir_path, "BCG"))
BCG_Label_Path = Path(os.path.join(dir_path, "BCG_label"))

# 读取PSG标签
df_PSG_label = pd.read_csv(PSG_Label_Path / (f"export" + str(sampID) + ".csv"), encoding="gbk")
df_PSG_label = df_PSG_label.loc[:, ~df_PSG_label.columns.str.contains('^Unnamed')]
df_PSG_label = df_PSG_label[df_PSG_label["Event type"].isin(base_event)]
df_PSG_label['Duration'] = df_PSG_label['Duration'].str.replace(r' \(.*?\)', '', regex=True)

# 读取EDF文件
edf_File = pyedflib.EdfReader(str(PSG_Data_Path / f"A{str(sampID).rjust(7, '0')}.edf"))

# 获取PSG记录开始时间
start_time = str(edf_File.getStartdatetime()).split(" ")[1]
start_time_abs = get_time_to_seconds(start_time)

# 计算起始时间秒数和终止时间秒数
df_PSG_label['Start'] = (df_PSG_label['Time'].apply(get_time_to_seconds) - start_time_abs).apply(lambda x: x + 24 * 3600 if x < 0 else x).astype(int)
df_PSG_label['End'] = df_PSG_label['Start'] + df_PSG_label['Duration'].astype(float).round(0).astype(int)

# 打印结果
print(df_PSG_label)

# 写入csv文件
df_PSG_label.to_csv(r"E:\data_annotation\6SleepApnea_annotation_GUI_demo\data\BCG_label\export" + str(sampID) + "_all.csv", index=False, encoding="gbk")

# 打印结果
print("sampID_" + str(sampID) + "写入csv成功")