import numpy as np
from scipy.special import erf

def flattop(A, edge, length, fsn):
    '''
    生成平顶包络函数 (Python版本)
    
    Parameters:
    -----------
    A : float
        幅度
    edge : float
        边沿时间参数  
    length : int
        长度（采样点数）
    fsn : float
        采样频率
    
    Returns:
    --------
    numpy.ndarray
        平顶包络波形
    '''
    Ts = 0
    r_sigma = 0.21230
    
    T = length  / fsn
    t = np.arange(0, length + 2) / fsn
    
    mu = 0.5 * edge
    sigma = r_sigma * (edge - 1)
    p = T - 1 - edge
    
    x1 = (t - mu - Ts) / (np.sqrt(2) * sigma)
    x2 = (t - mu - p + Ts) / (np.sqrt(2) * sigma)
    
    f = A / 2 * (erf(x1) - erf(x2))
    # f_padded = np.pad(f, (1, 1), mode='constant', constant_values=0)

    return f
"""
# 使用示例
if __name__ == "__main__":
    # 测试参数
    A = 1.0        # 幅度
    edge = 0.1     # 边沿时间
    length = 1000  # 长度
    fsn = 10000    # 采样频率
    
    # 生成平顶包络
    envelope = flattop(A, edge, length, fsn)
    
    print(f"生成的包络长度: {len(envelope)}")
    print(f"最大值: {np.max(envelope):.6f}")
    print(f"最小值: {np.min(envelope):.6f}")
    
    # 可选：绘图显示
    try:
        import matplotlib.pyplot as plt
        
        plt.figure(figsize=(10, 6))
        time_axis = np.arange(len(envelope)) / fsn
        plt.plot(time_axis, envelope, 'b-', linewidth=2)
        plt.xlabel('时间 (s)')
        plt.ylabel('幅度')
        plt.title('平顶包络波形')
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()
        
    except ImportError:
        print("matplotlib 未安装，跳过绘图")

"""