import numpy as np
import os
import copy
from FourChZreg_define import *

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

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



    def write_registers_batch(self, register_values):
        """批量写入寄存器"""
        if isinstance(register_values, dict):
            register_values = register_values.items()
        for address, value in register_values:
            self.mk.rw_once('w', address, value, self.config_file)


    def write_register(self, address, value):
        self.mk.rw_once('w', address, value, self.config_file)

    def _calculate_tc_coefficients(self, **kwargs):
        tc_coef_set = kwargs.get('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))

        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 _general_reg_config(self, **kwargs):
        # 深拷贝 kwargs 保护原始字典
        kwargs = copy.deepcopy(kwargs)
        # 自定义寄存器配置支持
        if 'custom_registers' in kwargs:
            custom_registers = kwargs.pop('custom_registers')
            first_item = custom_registers[0]
            if isinstance(first_item, tuple):
                self.write_registers_batch(custom_registers)
            elif isinstance(first_item, dict):
                for item in custom_registers:
                    self.write_register(item['addr'], item['values'])
            elif isinstance(first_item, list):
                for item in custom_registers:
                    self.write_register(item[0], item[1])

        # 芯片工作模式配置
        if 'chip_mode' in kwargs:
            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):
        # 统一使用相对基地址（动态适配当前 Channel）
        channel = kwargs.get('channel')
        ctrl_base = addr_base['CTRL0_BASE'] + channel * 0x600000

        self.write_register(ctrl_base + FourChZreg_define['ctrl_reg']['MODDOTR'], 8)  # 切换到RAMP输出
        ramp_ctrl = kwargs.get('ramp_ctrl')

        if ramp_ctrl == 'MCU':  # RAMP连到 mcu_regfile
            self.write_register(ctrl_base + FourChZreg_define['ctrl_reg']['SPI_RAMPENR'], 1 << 30)
        elif ramp_ctrl == 'SPI':  # RAMP连到 ctrl_regfile
            ramp_spi_registers = []
            fixed_enable = kwargs.get('fixed_enable', False)
            if fixed_enable:  # 固定值模式
                fixed_value = kwargs.get('fixed_value', 0)
                ramp_spi_registers += [1 << 15 | fixed_value << 16]
                # 写寄存器函数自带了通道偏移
                self.write_register(ctrl_base + FourChZreg_define['ctrl_reg']['SPI_RAMPFIXR'], ramp_spi_registers)
            else:  # 斜坡模式
                height = kwargs.get('height', 0)
                length = kwargs.get('step_time', 0)
                ramp_spi_registers += [0 << 15]
                ramp_spi_registers += [height << 16]
                ramp_spi_registers += [length]
                self.write_register(ctrl_base + FourChZreg_define['ctrl_reg']['SPI_RAMPFIXR'], ramp_spi_registers)

            # 配置使能
            self.write_register(ctrl_base + FourChZreg_define['ctrl_reg']['SPI_RAMPENR'], 1 << 31 | 0 << 30)

    def _awg_config(self, **kwargs):
        channel = kwargs.get('channel', 0)
        ctrl_base = addr_base['CTRL0_BASE'] + channel*0x600000
        tcco_base = addr_base['TCCO0_BASE'] + channel*0x600000
        dtcm_base = addr_base['DTCM0_BASE'] + channel*0x600000

        mode_map = {'nco': 6, 'nco_fm': 7, 'env': 4, 'mod': 5}
        mode = kwargs.get('mode', None)
        moddotr_val = mode_map.get(mode, 0)

        # 1. DSP (拖尾矫正) 配置
        tail_en = kwargs.get('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(tcco_base + FourChZreg_define['tc_reg']['TCPARR0'], tc_coef_registers)
            tccdr_addr = tcco_base + FourChZreg_define['tc_reg']['TCCDR']
            for i in range(8):
                self.write_register(tccdr_addr, 1 << i)
            moddotr_val &= ~4  # 把[2]清零

        # 2. 数据选择寄存器 MODDOTR 配置
        self.write_register(ctrl_base + FourChZreg_define['ctrl_reg']['MODDOTR'], moddotr_val)

        # 3. 调制使能寄存器 MODENR 配置
        amp_mod_enable = kwargs.get('amp_mod_enable', False)
        freq_mod_enable = kwargs.get('freq_mod_enable', False)
        bias_enable = kwargs.get('bias_enable', False)

        mod_enable = (
                (int(not bias_enable) << 0)
                | (int(not freq_mod_enable) << 1)
                | (int(not amp_mod_enable) << 2)
        )


        self.write_register(ctrl_base + FourChZreg_define['ctrl_reg']['MODENR'], mod_enable)

        # 4. mcu_regfile 配置
        mcu_registers = []
        fcw_list = kwargs.get('fcw', [0])
        mcu_reg_clr = kwargs.get('mcu_reg_clr', False)
        pcw_list = kwargs.get('pcw', [0])

        for fcw in fcw_list:
            mcu_registers.append(int(fcw / fs / 4 * 2 ** 32))

        mcu_registers.append(1 << 31 if mcu_reg_clr else 0 << 31)

        for pcw in pcw_list:
            mcu_registers.append(int(pcw / 360 * (2 ** 16 - 1)) << 16)

        rz_pha = kwargs.get('rz_pha', 0)
        mcu_registers.append(int(rz_pha / 360 * (2 ** 16 - 1)))

        ff_amp_list = kwargs.get('ff_amp', [0])
        fm_amp_list = kwargs.get('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.get('bias', [0])
        for bias in bias_list:
            mcu_registers.append(bias << 16)

        self.write_register(dtcm_base + FourChZreg_define['mcu_reg']['CWFR0'], mcu_registers)

        # 5. FMER 配置
        fm_en = kwargs.get('fm_en', False)
        self.write_register(dtcm_base + FourChZreg_define['mcu_reg']['FMER'], 1 << 31 if fm_en else 0)
        return None


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