import xml.etree.ElementTree as ET
ANNOTATION_MAP = {
    "Wake|0": 0,
    "Stage 1 sleep|1": 1,
    "Stage 2 sleep|2": 2,
    "Stage 3 sleep|3": 3,
    "Stage 4 sleep|4": 4,
    "REM sleep|5": 5,
    "Unscored|9": 9,
    "Movement|6": 6
}

SA_EVENTS = ['Central apnea', 'Hypopnea', 'Obstructive apnea']

def parse_sleep_annotations(annotation_path):
    """解析睡眠分期注释"""
    try:
        tree = ET.parse(annotation_path)
        root = tree.getroot()
        events = []
        for scored_event in root.findall('.//ScoredEvent'):
            event_type = scored_event.find('EventType').text
            if event_type != "Stages|Stages":
                continue
            description = scored_event.find('EventConcept').text
            start = float(scored_event.find('Start').text)
            duration = float(scored_event.find('Duration').text)
            if description not in ANNOTATION_MAP:
                continue
            events.append({
                'onset': start,
                'duration': duration,
                'description': description,
                'stage': ANNOTATION_MAP[description]
            })
        return events
    except Exception as e:
        return None


def extract_osa_events(annotation_path):
    """提取睡眠呼吸暂停事件"""
    try:
        tree = ET.parse(annotation_path)
        root = tree.getroot()
        events = []
        for scored_event in root.findall('.//ScoredEvent'):
            event_concept = scored_event.find('EventConcept').text
            event_type = event_concept.split('|')[0].strip()
            if event_type in SA_EVENTS:
                start = float(scored_event.find('Start').text)
                duration = float(scored_event.find('Duration').text)
                if duration >= 10:
                    events.append({
                        'start': start,
                        'duration': duration,
                        'end': start + duration,
                        'type': event_type
                    })
        return events
    except Exception as e:
        return []