import numpy as np
import os
import logging
import matplotlib.pyplot as plt
import sys
# sys.path.append('D:/Work/EnvData')
# sys.path.append('D:/Work/EnvData/acz')
# sys.path.append('D:/Work/EnvData/accz')

from typing import Any, Dict, List, Optional, Union, Callable
from enum import Enum
from env_gen.flattop import flattop
from env_gen.acz import aczwave
from env_gen.accz_gen import accz_wave
from matplotlib import gridspec
from reg_define import *
class EnvelopeGenerator:

    _axes_list: List[tuple] = []
    
    PLOT_STYLES = [
        'ggplot', 'bmh', 'fivethirtyeight', 'Solarize_Light2',
        'fast', 'tableau-colorblind10', 'seaborn-poster', 'seaborn-bright'
    ]
    PLOT_COLORS = ['C0', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7']
    
    SINGLE_PLOT_CONFIG = {
        'figsize': (8, 6),
        'dpi': 220,
        'style': 'seaborn-v0_8-poster',
        'linewidth': 3,
        'fontsize': {'xlabel': 22, 'ylabel': 22, 'title': 20, 'legend': 18, 'tick': 18}
    }
    
    MULTI_PLOT_CONFIG = {
        'dpi': 200,
        'fontsize': {'xlabel': 18, 'ylabel': 18, 'title': 16, 'tick': 16},
        'grid_alpha': 0.3
    }

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

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


    # def _show_envelope_subplot(self, all_data):
    #     n = len(all_data)
    #     if n == 1:
    #         EnvelopeGenerator._show_single_plot(all_data[0])
    #     else:
    #         EnvelopeGenerator._show_multi_plots(all_data)
    
    # def _show_single_plot(self, plot_data):
    #     t, y, title, output_mode = plot_data
    #     config = EnvelopeGenerator.SINGLE_PLOT_CONFIG
        
    #     plt.figure(figsize=config['figsize'], dpi=config['dpi'])
    #     plt.style.use(config['style'])
        
    #     plt.plot(t, y, label=f'{output_mode} Mode', 
    #             linewidth=config['linewidth'], color=EnvelopeGenerator.PLOT_COLORS[0])
    #     plt.xlabel('Time (ns)', fontsize=config['fontsize']['xlabel'])
    #     plt.ylabel('Amplitude', fontsize=config['fontsize']['ylabel'])
    #     plt.title(title, fontsize=config['fontsize']['title'])
    #     plt.grid(True, alpha=0.6)
    #     plt.legend(fontsize=config['fontsize']['legend'])
    #     plt.tick_params(axis='both', labelsize=config['fontsize']['tick'])
    #     plt.tight_layout(pad=2.0)
    #     plt.show()
    
    # def _show_multi_plots(self, all_data):
    #     n = len(all_data)
    #     cols = int(np.ceil(np.sqrt(n)))
    #     rows = int(np.ceil(n / cols))
        
    #     config = EnvelopeGenerator.MULTI_PLOT_CONFIG
    #     fig = plt.figure(figsize=(5 * cols, 4 * rows), dpi=config['dpi'])
    #     gs = gridspec.GridSpec(rows, cols)
        
    #     for i, (t, y, title, output_mode) in enumerate(all_data):
    #         row = i // cols
    #         col = i % cols
            
    #         # 为每个子图分配不同风格和颜色
    #         style_idx = i % len(EnvelopeGenerator.PLOT_STYLES)
    #         color_idx = i % len(EnvelopeGenerator.PLOT_COLORS)
            
    #         with plt.style.context(EnvelopeGenerator.PLOT_STYLES[2]):
    #             ax = fig.add_subplot(gs[row, col])
    #             ax.plot(t, y, color=EnvelopeGenerator.PLOT_COLORS[color_idx])
    #             ax.set_xlabel('Time (ns)', fontsize=config['fontsize']['xlabel'])
    #             ax.set_ylabel('Amplitude', fontsize=config['fontsize']['ylabel'])
    #             ax.set_title(title, fontsize=config['fontsize']['title'])
    #             ax.grid(True, alpha=config['grid_alpha'])
    #             ax.tick_params(axis='both', labelsize=config['fontsize']['tick'])
    #     plt.tight_layout(pad=3.0)
    #     plt.show()
    
    def _generate_rect_envelope(self, **kwargs): 
        amp = kwargs.pop('amp')
        wave_time = kwargs.pop('wave_time')
        rect_wave = [amp]*wave_time
        return rect_wave

    def _generate_rect_hold_envelope(self, **kwargs): 
        amp = kwargs.pop('amp')
        rect_rising_edge = [amp]*4
        rect_falling_edge = [0]*4
        return rect_rising_edge, rect_falling_edge
    
    def _generate_flattop_envelope(self, **kwargs):
        amp = kwargs.pop('amp')
        edge_time = kwargs.pop('edge_time')
        wave_time = kwargs.pop('wave_time')
        flattop_wave = flattop(float(amp), float(edge_time), float(wave_time), 1.0)
        return flattop_wave
    
    def _generate_flattop_hold_envelope(self, **kwargs):
        amp = kwargs.pop('amp')
        edge_time = kwargs.pop('edge_time')
        wave_time = kwargs.pop('wave_time')
        flattop_wave = flattop(float(amp), float(edge_time), float(wave_time), 1.0)
        hold_value = flattop_wave.max()
        hold_idx = np.where(flattop_wave == hold_value)[0]
        rising_edge_end_idx = hold_idx[0]
        falling_edge_start_idx = hold_idx[-1]
        flattop_rising_edge = flattop_wave[:rising_edge_end_idx+1]
        flattop_falling_edge = flattop_wave[falling_edge_start_idx:-2]
        return flattop_rising_edge, flattop_falling_edge
    
    def _generate_acz_envelope(self, **kwargs):
        amp = kwargs.pop('amp')
        wave_time = kwargs.pop('wave_time')
        acz_wave = aczwave(amp, wave_time, 0.0, 0.0, 0.0, 0.864, 0.05, -0.18, 0.04)
        acz_wave_real = [val.real for val in acz_wave]
        return acz_wave_real
    
    def _generate_accz_envelope(self, **kwargs):
        wave_time = kwargs.pop('wave_time')
        env_accz = accz_wave(T=wave_time, A=0.8, plot=False)
        return env_accz
    
    def _cosine_envelope(self, **kwargs):
        alpha = kwargs.pop('alpha', 1)
        phi = kwargs.pop('phi', np.pi)
        amp = kwargs.pop('amp')
        wave_time = kwargs.pop('wave_time')
        t_norm = np.arange(wave_time) / wave_time
        cosine_wave = amp * (1 + alpha * np.cos(2 * np.pi * t_norm + phi)) / 2
        return cosine_wave
    
    def _generate_envelope_data(self, envelope_type, **kwargs):
        if envelope_type == 'rect':
            env_data = self._generate_rect_envelope(**kwargs)
        elif envelope_type == 'rect_hold':
            env_data = self._generate_rect_hold_envelope(**kwargs)
        elif envelope_type == 'flattop':
            env_data = self._generate_flattop_envelope(**kwargs)
        elif envelope_type == 'flattop_hold':
            env_data = self._generate_flattop_hold_envelope(**kwargs)
        elif envelope_type == 'acz':
            env_data = self._generate_acz_envelope(**kwargs)
        elif envelope_type == 'accz':
            env_data = self._generate_accz_envelope(**kwargs)
        elif envelope_type == 'cosine':
            env_data = self._cosine_envelope(**kwargs)
        elif envelope_type == 'file_read_direct':
            env_data = kwargs.get('external_envelope_data', [])
        elif envelope_type == 'file_read_txt':
            file_path = kwargs.get('file_path')
            if file_path is None:
                raise ValueError("Missing 'file_path'")

            txt_data = np.loadtxt(file_path)
            env_data = np.asarray(txt_data, dtype=float).reshape(-1).tolist()
        return env_data

    def _next_env_idx(self, idx_num, env_idx_mem, envelope_length):
        if idx_num == 0:
            return envelope_length
        else:
            last_env_idx = env_idx_mem[-1]
            env_base_addr = (last_env_idx >> 16) + ((last_env_idx & 0xFFFF) << 1)             
        return env_base_addr<<16 | envelope_length

    def _env_data_pack(self, float_data_array):
            
        data_int = np.round(float_data_array).astype(int)
        data_int[data_int < 0] += 65536
        hex_pairs = []
        for data0, data1 in zip(data_int[::2], data_int[1::2]):
            hex_pairs.append((data1 << 16) | data0)
            
        return hex_pairs

    def _generate_envelope_batch(self, **kwargs):
        envelope_configs = kwargs.pop('envelope_configs')
        env_data_mem = []
        env_idx_mem = []
        idx_num = 0
        for envelope_config in envelope_configs:
            envelope_type = envelope_config.pop('envelope_type')
            env_data = self._generate_envelope_data(envelope_type, **envelope_config)
            if isinstance(env_data, tuple):
                retun_param_count = len(env_data)
            else:
                retun_param_count = 1
            if retun_param_count == 1:
                envelope = env_data
                envelope_arr = np.asarray(envelope, dtype=float).reshape(-1)
                if envelope_arr.size % 4 != 0:
                    raise ValueError("Envelope length must be multiple of 4") 
                env_data_mem += envelope_arr.astype(int).tolist()
                envelope_length = int(envelope_arr.size)
                current_env_idx = self._next_env_idx(idx_num, env_idx_mem, envelope_length)
                env_idx_mem.append(current_env_idx)
                idx_num += 1 
            elif retun_param_count == 2:
                rising_edge, falling_edge = env_data            
                rising_edge_arr = np.asarray(rising_edge, dtype=float).reshape(-1)
                falling_edge_arr = np.asarray(falling_edge, dtype=float).reshape(-1)
                if rising_edge_arr.size % 4 != 0:
                    raise ValueError("Envelope length must be multiple of 4") 
                env_data_mem += rising_edge_arr.astype(int).tolist()
                rising_edge_length = int(rising_edge_arr.size)
                rising_edge_idx = self._next_env_idx(idx_num, env_idx_mem, rising_edge_length)
                env_idx_mem.append(rising_edge_idx)
                idx_num += 1 
                env_data_mem += falling_edge_arr.astype(int).tolist()
                falling_edge_length = int(falling_edge_arr.size)
                falling_edge_idx = self._next_env_idx(idx_num, env_idx_mem, falling_edge_length)
                env_idx_mem.append(falling_edge_idx)
                idx_num += 1
            
        env2mem_format = self._env_data_pack(env_data_mem)
        self.write_register(addr_base['ENVI0_BASE'], env_idx_mem)
        self.write_register(addr_base['ENVM0_BASE'], env2mem_format)
            # t = np.arange(len(processed_data)) / 3e9 * 1e9
            # EnvelopeGenerator._axes_list.append((t, processed_data, f'Processed Envelope - {output_mode} Mode', output_mode))
    # def _im_show(self, **kwargs):
        # EnvelopeGenerator._axes_list = []
        
        # if EnvelopeGenerator._axes_list:
        #     EnvelopeGenerator.show_envelope_subplot(EnvelopeGenerator._axes_list)
        #     EnvelopeGenerator._axes_list = []
def env_config(mk_instance, **kwargs):
    env_gen = EnvelopeGenerator(mk_instance, **kwargs) 
    env_gen._generate_envelope_batch(**kwargs)
        