from loguru import logger
from qlsdk.core.entity import ImpedancePacket, RscPacket
from qlsdk.persist import RscEDFHandler
from qlsdk.rsc.device.base import QLBaseDevice, intersection_positions
from qlsdk.rsc.device.device_factory import DeviceFactory
from qlsdk.rsc.interface import IDevice
from qlsdk.rsc.device.base import QLBaseDevice

class C16R(QLBaseDevice):
    
    device_type = 0x339  # C16R设备类型标识符 0x3903
    
    def __init__(self, socket):
        super().__init__(socket)
    
        # 存储通道反向映射位置值
        self._reverse_ch_pos  = None    
        
        # 存储通道反向映射位置值
        self._impedance_ch_pos  = None    
        self._impedance_channels_origin = None
        
        self.channel_name_mapping = {
            "FP1": 1, 
            "FP2": 2, 
            "F7": 3, 
            "F8": 4, 
            "P3": 5, 
            "P4": 6, 
            "O1": 7, 
            "O2": 8, 
            "CZ": 9, 
            "A1": 10, 
            "A2": 11, 
        }
        
        self.channel_mapping = {
            "1": 59, 
            "2": 60, 
            "3": 53, 
            "4": 50, 
            "5": 63, 
            "6": 24, 
            "7": 51, 
            "8": 64,
            "9": 49, 
            "10": 14, 
            "11": 10, 
        }
        
        self.channel_display_mapping = {
            59: 1, 
            60: 2, 
            53: 3, 
            50: 4, 
            63: 5, 
            24: 6, 
            51: 7, 
            64: 8,
            49: 9, 
            14: 10, 
            10: 11, 
        }
        
    @classmethod
    def from_parent(cls, parent:IDevice) -> IDevice:
        rlt = cls(parent.socket)
        rlt.device_id = parent.device_id
        rlt._device_no = parent.device_no
        return rlt        
        
    def init_edf_handler(self):
        self._edf_handler = RscEDFHandler(self.sample_rate,  self.sample_range * 1000 , - self.sample_range * 1000, self.resolution)  
        self._edf_handler.set_device_type(self.device_type)
        self._edf_handler.set_device_no(self.device_no)
        self._edf_handler.set_storage_path(self._storage_path)
        self._edf_handler.set_file_prefix(self._file_prefix if self._file_prefix else 'C8R')
        logger.debug(f"EDF Handler initialized")
        # 设置采集参数
    def set_acq_param(self, channels:list, sample_rate: int = 500, sample_range:int = 188):
        # 通道合法性校验
        self._channel_validate(channels)
        # 保存原始通道参数
        self._acq_param["original_channels"] = channels
        
        # 名称转换为数字通道  
        channels = [self.channel_name_mapping.get(str(i).upper(), i) for i in channels]
        
        # 根据映射关系做通道转换-没有映射的默认到第一个通道
        # 先设置不存在的通道为-1，再把-1替换为第一个通道，避免第一个通道也不合法的情况
        channels = [self.channel_mapping.get(str(i), -1) for i in channels]
        channels = [i if i != -1 else channels[0] for i in channels]
                
        # 更新采集参数
        self._acq_param["channels"] = channels
        self._acq_param["sample_rate"] = sample_rate
        self._acq_param["sample_range"] = sample_range
        self._acq_channels = channels
        self._sample_rate = sample_rate
        self._sample_range = sample_range
        
        logger.debug(f"C16RS: set_acq_param: {self._acq_param}")   
        
        # 参数改变后，重置通道位置映射
        self._reverse_ch_pos  = None  
    
    @property
    def acq_channels(self):
        if self._acq_channels is None:
            # 初始化通道参数为1-11
            self.set_acq_param([i for i in range(1, 12)])
        return self._acq_channels  
    
    def set_impedance_channels(self, channels:list):
        # 通道合法性校验
        self._channel_validate(channels)
        # 保存原始通道参数
        self._impedance_channels_origin = channels
        
        # 名称转换为数字通道  
        channels = [self.channel_name_mapping.get(str(i).upper(), i) for i in channels]
        
        # 根据映射关系做通道转换-没有映射的默认到第一个通道
        # 先设置不存在的通道为-1，再把-1替换为第一个通道，避免第一个通道也不合法的情况
        channels = [self.channel_mapping.get(str(i), -1) for i in channels]
        channels = [i if i != -1 else channels[0] for i in channels]
        self._impedance_channels = channels
        
        # 初始化位置映射
        self._impedance_ch_pos  = None  
        
    def get_impedance_channels(self): 
        if self._impedance_channels is None or len(self._impedance_channels) == 0:
            self.set_impedance_channels([i for i in range(1, 12)])
        return self._impedance_channels   
    
    def _channel_validate(self, channels: list):
        # 名称转换为数字通道  
        temp_channels = [self.channel_name_mapping.get(str(i).upper(), i) for i in channels]
        
        # 根据映射关系做通道转换-没有映射的默认到第一个通道
        # 先设置不存在的通道为-1，再把-1替换为第一个通道，避免第一个通道也不合法的情况
        temp_channels = [self.channel_mapping.get(str(i), -1) for i in temp_channels]
        for i in range(len(temp_channels)):
            if temp_channels[i] == -1:
                raise ValueError(f"通道参数有误，通道应为1-11的数字或对应名称，如：1、2或FP1、FP2等")
    
    def _impedance_wrapper(self, body: bytes):        
        packet = ImpedancePacket().transfer(body)   
        impedance_channels = self.get_impedance_channels() 
        if impedance_channels is not None and len(impedance_channels) > 0:
            # 只保留设置的阻抗通道
            channel_pos = intersection_positions(packet.channels, impedance_channels)
            packet.impedance = [packet.impedance[i] for i in channel_pos]
            packet.channels = [packet.channels[i] for i in channel_pos]
            
        # 升级为类变量，减少计算
        if self._impedance_ch_pos is None:
            self._impedance_ch_pos = map_indices(impedance_channels, packet.channels)
            
        packet.channels = self._impedance_channels_origin
        packet.impedance = [packet.impedance[i] for i in self._impedance_ch_pos]
            
        return packet
    
    # 信号数据转换(默认不处理)
    def _signal_wrapper(self, body: bytes):
        if body is None:
            return
        data = RscPacket().transfer(body)
        # 根据映射关系做通道转换-（注意数据和通道的一致性）
        # data.channels = [self.channel_display_mapping.get(i, i) for i in data.channels]
        
        # 升级为类变量，减少计算
        if self._reverse_ch_pos is None:
            self._reverse_ch_pos = map_indices(self._acq_param["channels"], data.channels)
        
        # 更新通道（数据）顺序和输入一致
        data.channels = self._acq_param["original_channels"]
        data.eeg = [data.eeg[i] for i in self._reverse_ch_pos]
        
        return data
        
    def start_stimulation(self):
        raise NotImplementedError("Stimulation function is not supported")
        
def map_indices(A, B):
    """
    
    参数:
        A: 源数组（无重复值）
        B: 目标数组（无重复值）
    
    返回:
        C: 与A长度相同的数组，元素为A中对应值在B中的索引（不存在则为-1）
    """
    # 创建B的值到索引的映射字典（O(n)操作）
    b_map = {value: idx for idx, value in enumerate(B)}
    
    # 遍历A，获取每个元素在B中的位置（O(m)操作）
    return [b_map.get(a, -1) for a in A]

DeviceFactory.register(C16R.device_type, C16R) 