import numpy as np
import os
from reg_define import *


class ZChipConfig(object):
    CHANNEL_OFFSET = 0x00600000
    SCALE_FACTOR = 2 ** 31

    def __init__(self, mk_instance, channel_id=0, **kwargs):
        self.mk = mk_instance
        self.channel_id = channel_id
        self.FolderName = kwargs.get('FolderName')
        self.config_file = kwargs.get('config_file')
        self._setup_environment()
        self._init_logger()

    def _setup_environment(self):
        if self.FolderName:
            os.makedirs(self.FolderName, exist_ok=True)
        if self.config_file and os.path.exists(self.config_file):
            os.remove(self.config_file)

    def _init_logger(self):
        import logging
        self.logger = logging.getLogger(f"ChipConfig.ch{self.channel_id}")
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)
            self.logger.setLevel(logging.INFO)

    def get_channel_addr(self, base_addr):
        if isinstance(base_addr, str):
            if base_addr.startswith('0x'):
                base_addr = int(base_addr, 16)
            else:
                base_addr = int(base_addr)

        if not hasattr(self, '_channel_offset'):
            self._channel_offset = self.CHANNEL_OFFSET * self.channel_id

        return hex(base_addr + self._channel_offset)

    def _configure_tc_mode(self):
        """配置TC模式 - 分组配置系数"""
        # self.logger.info(
        #     "Configuring TC mode with group-wise coefficient setup")

        # 写入TC_BYPASS
        self.write_register(
            self.get_register_addr('BYPASS'),
            self.REGISTER_VALUES['BYPASS_TC']
        )

        # 计算TC系数
        coefficients = self._calculate_tc_coefficients()

        # 分组配置TC系数并设置配置完成标志
        self._configure_tc_coefficients_by_groups(coefficients)

        # self.logger.info("TC mode configuration completed")

    def _configure_tc_coefficients_by_groups(self, coefficients):
        """分组配置TC系数，每组配置完成后设置对应的配置完成标志"""

        # 定义7个系数组，每组包含一个系数的4个寄存器(alpha_re, alpha_im, beta_re, beta_im)
        tc_register_groups = self._build_tc_register_groups(coefficients)

        for group_idx, group_registers in enumerate(tc_register_groups):
            # 配置当前组的所有寄存器
            # self.logger.info(f"配置系数组 {group_idx}")
            self.write_registers_batch(group_registers)

            # 设置当前组的配置完成标志
            self._set_coef_config_done_for_group(group_idx)

            # self.logger.info(f"系数组 {group_idx} 配置完成")

    def set_tc_coefficient_set(self, coeff_set_name):
        """设置要使用的TC系数组"""
        if coeff_set_name not in self.TC_COEFFICIENT_SETS:
            available = list(self.TC_COEFFICIENT_SETS.keys())
            raise ValueError(f"未知的系数组: {coeff_set_name}. 可用的系数组: {available}")

        self.tc_coeff_set = coeff_set_name
        # self.logger.info(
        #     f"切换到TC系数组: {coeff_set_name} - {self.TC_COEFFICIENT_SETS[coeff_set_name]['description']}")

    def _calculate_tc_coefficients(self, **kwargs):
        tc_coef_set = kwargs.pop('tc_coef_set')
        coef_set = TC_COEFFICIENT_SETS[tc_coef_set]

        amp_real = coef_set['amp_real']
        amp_imag = coef_set['amp_imag']
        time_real = coef_set['time_real']
        time_imag = coef_set['time_imag']

        sampling_rate = 3e9
        coef1, coef2 = [], []
        for ar, ai, tr, ti in zip(amp_real, amp_imag, time_real, time_imag):
            amp_coef = ar + 1j * ai
            time_coef = tr + 1j * ti
            coef1.append(amp_coef * np.exp(1e9 / (sampling_rate) / 2 / (1 - amp_coef) * time_coef) /
                         (1 - amp_coef))
            coef2.append(np.exp(1e9 / (sampling_rate) / (1 - amp_coef) * time_coef))
        # print('coef1_real = ')
        # for c in coef1:
        #     print(c.real)
        #     # print(hex(int(c.real*(2**31-1)) & 0xFFFFFFFF))
        # print('coef1_imag = ')
        # for c in coef1:
        #     # print(hex(int(c.imag*(2**31-1)) & 0xFFFFFFFF))
        #     print(c.imag)
        # print('coef2_real = ')
        # for c in coef2:
        #     # print(hex(int(c.real*(2**31-1)) & 0xFFFFFFFF))
        #     print(c.real)
        # print('coef2_imag = ')
        # for c in coef2:
        #     # print(hex(int(c.imag*(2**31-1)) & 0xFFFFFFFF))
        #     print(c.imag)
        return {
            'alpha_re': [int(c.real * self.SCALE_FACTOR) for c in coef1],
            'alpha_im': [int(c.imag * self.SCALE_FACTOR) for c in coef1],
            'beta_re': [int(c.real * self.SCALE_FACTOR) for c in coef2],
            'beta_im': [int(c.imag * self.SCALE_FACTOR) for c in coef2]
        }

    def _build_tc_register_groups(self, coefficients):
        """构建分组的TC寄存器配置"""
        tc_register_groups = []

        # 获取所有寄存器地址
        alpha_re_addrs = self.get_register_addr('TC_ALPHA_RE')
        alpha_im_addrs = self.get_register_addr('TC_ALPHA_IM')
        beta_re_addrs = self.get_register_addr('TC_BETA_RE')
        beta_im_addrs = self.get_register_addr('TC_BETA_IM')

        # 按组分配寄存器（假设有8个系数，但只使用前7组）
        num_groups = len(coefficients['alpha_re'])

        for i in range(num_groups):
            group_registers = [
                (alpha_re_addrs[i], coefficients['alpha_re'][i]),
                (alpha_im_addrs[i], coefficients['alpha_im'][i]),
                (beta_re_addrs[i], coefficients['beta_re'][i]),
                (beta_im_addrs[i], coefficients['beta_im'][i])
            ]
            tc_register_groups.append(group_registers)

        return tc_register_groups

    def _set_coef_config_done_for_group(self, group_idx):
        """为指定组设置配置完成标志"""
        if not (0 <= group_idx <= 7):  # 8组，索引0-7
            self.logger.warning(f"无效的组索引: {group_idx}")
            return

        # 计算当前组的配置完成标志值：2^group_idx
        config_done_value = 1 << group_idx
        config_done_hex = f'0x{config_done_value:02x}'

        # 读取当前配置完成寄存器的值，进行或运算以保留其他组的标志
        try:
            # 如果需要保留其他组的标志，这里需要先读取当前值
            # current_value = self._read_register(self.get_register_addr('COEF_CONFIG_DONE'))
            # new_value = current_value | config_done_value

            # 简化版本：直接写入当前组的标志
            self.write_register(
                self.get_register_addr('COEF_CONFIG_DONE'),
                config_done_hex
            )

            # self.logger.info(f"系数组 {group_idx} 配置完成标志已设置: {config_done_hex}")

        except Exception as e:
            self.logger.error(f"设置系数组 {group_idx} 配置完成标志失败: {e}")
            raise

    def write_registers_batch(self, register_values):
        """
        批量写入寄存器

        参数:
            register_values: [(address, value), ...] 或 {address: value, ...}
        """
        if isinstance(register_values, dict):
            register_values = register_values.items()

        for address, value in register_values:
            try:
                self.mk.rw_once('w', address, value, self.config_file)
            except Exception as e:
                self.logger.error(f"Failed to write register {address}: {e}")
                raise

    def write_register(self, address, value):
        """写入单个寄存器"""
        try:
            self.mk.rw_once('w', address, value, self.config_file)
        except Exception as e:
            self.logger.error(f"Failed to write register {address}: {e}")
            raise




    def _freq2hex(self, freq):
        fs = 750
        fcw = int(freq / fs / 4 * 2 ** 32)
        return int(fcw)

    def _deg2hex(self, deg):
        return int(deg / 360 * (2 ** 16 - 1))

    # def _system_reg_config(self, **kwargs):
    # 中断屏蔽寄存器
    # mk_instance.rw_once('w', '0x14', '0x10000000', config.config_file)

    def _general_reg_config(self, **kwargs):
        ctrl_registers = []
        # 自定义寄存器配置支持（元组/字典/列表三种格式）
        if 'custom_registers' in kwargs:
            custom_registers = kwargs.pop('custom_registers')
            first_item = custom_registers[0]
            # 元组格式：[(addr, val), ...] 批量写入
            if isinstance(first_item, tuple):
                self.write_registers_batch(custom_registers)
            # 字典格式：[{"addr": addr, "values": val}, ...] 单条写入
            elif isinstance(first_item, dict):
                for item in custom_registers:
                    addr = item['addr']
                    values = item['values']
                    self.write_register(addr, values)
            # 列表格式：[[addr, val], ...] 单条写入
            elif isinstance(first_item, list):
                for item in custom_registers:
                    addr = item[0]
                    values = item[1]
                    self.write_register(addr, values)

        # 芯片工作模式配置
        if 'chip_mode' in kwargs:
            # TODO: 部分模式可以合并，双频点与单音NCO
            chip_mode = kwargs.pop('chip_mode')
            if chip_mode == 'RAMP':
                self._ramp_config(**kwargs)
            elif chip_mode == 'AWG':
                self._awg_config(**kwargs)
    def _ramp_config(self, **kwargs):
        self.write_register(addr_base['CTRL0_BASE'] + reg_define['ctrl_reg']['MODDOTR'], 8)  # 切换到RAMP输出
        ramp_ctrl = kwargs.pop('ramp_ctrl')
        if ramp_ctrl == 'MCU':     # RAMP连到 mcu_regfile
            self.write_register(addr_base['CTRL0_BASE'] + reg_define['ctrl_reg']['SPI_RAMPENR'], 1 << 30)
            # 之后来用汇编操作RAMP四个参数寄存器
        elif ramp_ctrl == 'SPI':   # RAMP连到 ctrl_regfile
            ramp_spi_registers = []
            fixed_enable = kwargs.pop('fixed_enable')
            if fixed_enable:   # 打开固定值使能位，把值填进去就好
                fixed_value = kwargs.pop('fixed_value')
                ramp_spi_registers += [1 << 15 | fixed_value << 16]
                self.write_register(addr_base['CTRL0_BASE'] + reg_define['ctrl_reg']['SPI_RAMPFIXR'],   # 常数寄存器
                                    ramp_spi_registers)
            else:
                ramp_spi_registers += [0 << 15]   # 不是固定值模式，常数寄存器给0就好
                height = kwargs.pop('height', 0)
                length = kwargs.pop('step_time', 0)
                ramp_spi_registers += [height << 16]
                ramp_spi_registers += [length]
                self.write_register(addr_base['CTRL0_BASE'] + reg_define['ctrl_reg']['SPI_RAMPFIXR'],
                                    ramp_spi_registers)
            # 必须最后配使能
            self.write_register(addr_base['CTRL0_BASE'] + reg_define['ctrl_reg']['SPI_RAMPENR'], 1 << 31 | 0 << 30)




    def _awg_config(self, **kwargs):
            #数据选择寄存器300108配置开始
            # 模式与寄存器数值的映射字典 (Bit 2 和 Bit[1:0])
            mode_map = {
                'nco': 6,
                'nco_fm': 7,
                'env': 4,
                'mod': 5
            }
            # 1. 在内存中先计算 mode 对应的基础寄存器值
            mode = kwargs.pop('mode', None)
            moddotr_val = mode_map.get(mode, 0)  # 如果 mode 不在字典里，默认基础值为 0
            # 2. DSP（拖尾矫正）开关配置
            tail_en = kwargs.pop('tail_en', False)
            if tail_en:
                coefficients = self._calculate_tc_coefficients(**kwargs)
                # 拼接系数列表
                tc_coef_registers = (
                        coefficients['alpha_re'] +
                        coefficients['alpha_im'] +
                        coefficients['beta_re'] +
                        coefficients['beta_im']
                )
                self.write_register(addr_base['TCCO0_BASE'] + reg_define['tc_reg']['TCPARR0'], tc_coef_registers)
                tccdr_addr = addr_base['TCCO0_BASE'] + reg_define['tc_reg']['TCCDR']
                for i in range(8):
                    self.write_register(tccdr_addr, 1 << i)
                # 如果开启了 tail_en，在内存里直接给第 4 位 (Bit 3) 置 1
                moddotr_val &= ~4    #把[2]变0

            # 3. 所有逻辑判断完毕后，只触发一次物理写操作！
            moddotr_addr = addr_base['CTRL0_BASE'] + reg_define['ctrl_reg']['MODDOTR']
            self.write_register(moddotr_addr, moddotr_val)

            #调制使能寄存器300104配置开始
            amp_mod_enable = kwargs.pop('amp_mod_enable', False)
            freq_mod_enable = kwargs.pop('freq_mod_enable', False)
            bias_enable = kwargs.pop('bias_enable', False)
            # 低电平使能逻辑：True(开启) -> 0，False(关闭) -> 1
            mod_enable = (
                    (int(not bias_enable) << 0)  # Bit 0: Bias (低电平有效)
                    | (int(not freq_mod_enable) << 1)  # Bit 1: Freq (加括号，先取非再左移 1 位)
                    | (int(not amp_mod_enable) << 2)  # Bit 2: Amp  (加括号，先取非再左移 2 位)
            )
            self.write_register(addr_base['CTRL0_BASE'] + reg_define['ctrl_reg']['MODENR'], mod_enable)

            #mcu_regfile配置开始
            mcu_registers = []
            fcw_list = kwargs.pop('fcw', 0)
            mcu_reg_clr = kwargs.pop('mcu_reg_clr', 0)
            pcw_list = kwargs.pop('pcw', 0)
            fcw = int(fcw_list[0] / fs / 4 * 2 ** 32)
            for fcw in fcw_list:
                mcu_registers.append(int(fcw / fs / 4 * 2 ** 32))
            if mcu_reg_clr:
                mcu_registers += [1 << 31]
            else:
                mcu_registers += [0 << 31]
            for pcw in pcw_list:
                mcu_registers.append(int(pcw / 360 * (2 ** 16 - 1)) << 16)
            rz_pha = kwargs.pop('rz_pha', 0)
            mcu_registers.append(int(rz_pha / 360 * (2 ** 16 - 1)))
            #幅度
            ff_amp_list = kwargs.pop('ff_amp', 0)
            fm_amp_list = kwargs.pop('fm_amp', 0)
            for ff, fm in zip(ff_amp_list, fm_amp_list):
                mcu_registers.append(((ff & 0xFFFF) << 16) | (fm & 0xFFFF))
            #偏置
            bias_list = kwargs.pop('bias', 0)
            for bias in bias_list:
                mcu_registers.append(bias << 16)
            self.write_register(addr_base['DTCM0_BASE'] + reg_define['mcu_reg']['CWFR0'], mcu_registers)
            #FMER
            fm_en = kwargs.pop('fm_en', False)
            if fm_en:
                self.write_register(addr_base['DTCM0_BASE'] + reg_define['mcu_reg']['FMER'], 1 << 31)
            elif not fm_en:
                self.write_register(addr_base['DTCM0_BASE'] + reg_define['mcu_reg']['FMER'], 0 << 31)
            #mcu_regfile配置到此完成
            return None


def config_chip_reg(mk_instance, **kwargs):
    config = ZChipConfig(mk_instance, **kwargs)
    config._general_reg_config(**kwargs)
