import yaml
from pathlib import Path

class ConfigLoader:
    def __init__(self, config_path="config/cell_config.yaml"):
        self.config_path = Path(__file__).parent / config_path
        with open(self.config_path, 'r') as f:
            self.config = yaml.safe_load(f)

    def get_cell_ids_for_subsegment(self, segment_id, subsegment_id):
        """Get cell IDs for a specific subsegment"""
        for segment in self.config['segments']:
            if segment['id'] == segment_id:
                for subsegment in segment['subsegments']:
                    if subsegment['id'] == subsegment_id:
                        return subsegment['cell_ids']  # Updated key name
        return []

    def get_segment_config(self, segment_id):
        """Get configuration for a specific segment"""
        for segment in self.config['segments']:
            if segment['id'] == segment_id:
                return segment
        return None