#!/usr/bin/env python3
"""
DA4008 波形生成与transaction帧封装工具
支持生成正弦波、三角波、方波、Flattop波形，并自动生成对应的控制命令。
将命令和波形数据封装为两个transaction帧，输出到指定文件。
同时生成一个LVDS数据帧（包含帧头、地址/长度、波形数据、CRC），用于LVDS_DRIVER发送。

LVDS帧格式（每行一个32位十六进制数）：
    [帧头]       0xBCBCBCBC
    [地址/长度]   (lvds_addr << 16) | num_words
    [波形数据0]   ...
    [波形数据N]
    [CRC32]      从地址/长度字到最后一个数据字的CRC校验

帧格式（64位头部 + 数据）：
    W/R=0, ARD_FLAG=0, CID=0, EXADDR=0
    cmd帧 ADDR=0x100000, wave帧 ADDR=0x200000
    长度字段根据实际数据字节数计算
输出文件每行一个32位十六进制数，顺序为：
    [cmd帧高32位]
    [cmd帧低32位]
    cmd命令数据（每行一个32位命令字）
    [wave帧高32位]
    [wave帧低32位]
    wave波形数据（每行一个32位字，512-bit块打包后的结果）
"""

import argparse
import math
import sys
import os

# 常数
FS = 40e9                     # 采样率 40 GHz
MAX_ADDR = 8192               # SRAM深度（字）
CMD_WAVE = 0x00000000         # 波形命令标志位（bit31=0）
CMD_HOLD = 0x80000000         # 保持命令标志位（bit31=1）

def parse_args():
    parser = argparse.ArgumentParser(description='为DA4008生成波形与命令并封装为transaction帧，同时生成LVDS数据帧')
    parser.add_argument('--type', required=True, choices=['sine', 'triangle', 'pulse', 'flattop'],
                        help='波形类型: sine, triangle, pulse, flattop')
    parser.add_argument('--start_addr', type=int, required=True,
                        help='波形起始地址 (0-8191)')
    parser.add_argument('--length', type=int,
                        help='波形样本点数 (对sine/triangle/pulse有效，必须为正整数)')
    parser.add_argument('--cycles', type=int, default=1,
                        help='循环次数 (0-31, 0表示无限循环, 默认1)')
    # 正弦波参数
    parser.add_argument('--freq', type=float, help='正弦波频率 (Hz)')
    # 三角波参数
    parser.add_argument('--step', type=int, help='三角波步进 (1-255)')
    parser.add_argument('--duration', type=int, help='三角波每个台阶的持续时间 (时钟周期数)')
    # 方波参数
    parser.add_argument('--duty', type=float, default=0.5,
                        help='方波占空比 (0.0-1.0, 默认0.5)')
    # Flattop参数
    parser.add_argument('--rise_samples', type=int,
                        help='Flattop上升沿样本数')
    parser.add_argument('--hold_cycles', type=int,
                        help='Flattop平台保持周期数')
    parser.add_argument('--fall_samples', type=int,
                        help='Flattop下降沿样本数')
    parser.add_argument('--hold_value', type=int, default=255,
                        help='Flattop平台值 (0-255, 默认255)')
    # 输出文件
    parser.add_argument('--output', required=True,
                        help='输出transaction文件路径')
    # LVDS参数
    parser.add_argument('--lvds_addr', type=lambda x: int(x, 0), required=True,
                    help='LVDS帧中的地址 (支持十六进制，如0x1234)')
    parser.add_argument('--lvds_output', type=str,
                        help='LVDS数据帧输出文件路径 (默认: 将--output的扩展名改为_lvds.txt)')
    return parser.parse_args()

# ================== 波形生成函数 ==================
def generate_sine(num_samples, freq):
    """生成正弦波样本 (8-bit无符号)"""
    if freq <= 0 or freq > FS/2:
        raise ValueError(f"频率必须介于0和{FS/2:.1f} Hz之间")
    samples_per_cycle = FS / freq
    if samples_per_cycle < 1:
        raise ValueError("频率太高，每个周期样本数小于1")
    samples = []
    for i in range(num_samples):
        phase = 2 * math.pi * i / samples_per_cycle
        val = int(128 + 127 * math.sin(phase))
        samples.append(val)
    return samples

def generate_triangle(num_samples, step, duration):
    """生成三角波样本 (8-bit无符号)"""
    if step < 1 or step > 255:
        raise ValueError("步进必须介于1和255之间")
    if duration < 1:
        raise ValueError("持续时间必须为正整数")
    samples = []
    direction = 1
    current = 0
    cnt = 0
    while len(samples) < num_samples:
        samples.append(current)
        cnt += 1
        if cnt >= duration:
            cnt = 0
            next_val = current + direction * step
            if next_val > 255:
                next_val = 255
                direction = -1
            elif next_val < 0:
                next_val = 0
                direction = 1
            current = next_val
    return samples

def generate_pulse(num_samples, freq, duty):
    """生成方波样本 (8-bit, 高=255, 低=0)"""
    if freq <= 0 or freq > FS/2:
        raise ValueError(f"频率必须介于0和{FS/2:.1f} Hz之间")
    if duty < 0 or duty > 1:
        raise ValueError("占空比必须介于0和1之间")
    period = FS / freq
    high_samples = int(round(period * duty))
    low_samples = int(round(period * (1 - duty)))
    if high_samples + low_samples == 0:
        raise ValueError("频率和占空比导致零样本周期")
    samples = []
    while len(samples) < num_samples:
        for _ in range(high_samples):
            samples.append(255)
            if len(samples) >= num_samples:
                break
        for _ in range(low_samples):
            samples.append(0)
            if len(samples) >= num_samples:
                break
    return samples

def generate_flattop(rise_samples, hold_cycles, fall_samples, hold_value=255):
    """生成Flattop波形上升沿和下降沿样本（平台不存储）
    返回 (rise_samples_list, fall_samples_list)
    """
    if rise_samples < 1 or fall_samples < 1:
        raise ValueError("上升沿和下降沿样本数必须至少为1")
    if hold_cycles < 0:
        raise ValueError("保持周期数不能为负")
    # 上升沿：线性从0到hold_value
    rise = []
    for i in range(rise_samples):
        val = int(hold_value * i / (rise_samples - 1))
        rise.append(val)
    # 下降沿：线性从hold_value到0
    fall = []
    for i in range(fall_samples):
        val = int(hold_value * (fall_samples - 1 - i) / (fall_samples - 1))
        fall.append(val)
    return rise, fall

def pack_samples_to_blocks(samples):
    """将样本列表按每64个一组打包成512-bit字列表
    返回列表，每个元素为16个32-bit字的列表（每个32-bit字以整数表示）
    """
    blocks = []
    if len(samples) == 0:
        return blocks
    # 补零到64的倍数
    if len(samples) % 64 != 0:
        samples += [0] * (64 - len(samples) % 64)
    for i in range(0, len(samples), 64):
        block_samples = samples[i:i+64]
        words = []
        for j in range(0, 64, 4):
            word = (block_samples[j+3] << 24) | (block_samples[j+2] << 16) | \
                   (block_samples[j+1] << 8) | block_samples[j]
            words.append(word)
        blocks.append(words)
    return blocks

def generate_wave_command(start_addr, length_words, cycles):
    """生成波形命令字 (32-bit)"""
    if start_addr < 0 or start_addr >= MAX_ADDR:
        raise ValueError(f"起始地址必须介于0和{MAX_ADDR-1}之间")
    if length_words < 1 or length_words > MAX_ADDR:
        raise ValueError(f"长度必须介于1和{MAX_ADDR}之间")
    if start_addr + length_words > MAX_ADDR:
        raise ValueError("地址范围超出SRAM深度")
    if cycles < 0 or cycles > 31:
        raise ValueError("循环次数必须介于0和31之间")
    cmd = CMD_WAVE | (cycles << 26) | (start_addr << 13) | length_words
    return cmd

def generate_hold_command(hold_cycles):
    """生成保持命令字"""
    if hold_cycles < 0 or hold_cycles > 0x7FFFFFFF:
        raise ValueError("保持周期数超出范围 (0-2^31-1)")
    return CMD_HOLD | hold_cycles

# ================== CRC32 计算 ==================
def crc32_next(crc, data):
    """按位计算CRC32，多项式0x04C11DB7，MSB优先"""
    poly = 0x04C11DB7
    new_crc = crc
    for i in range(32):
        # 取CRC的最高位
        bit = (new_crc >> 31) & 1
        # 取数据的最高位（MSB first）
        data_bit = (data >> (31 - i)) & 1
        b = bit ^ data_bit
        new_crc = (new_crc << 1) & 0xFFFFFFFF
        if b:
            new_crc ^= poly
    return new_crc

def crc32_compute(words):
    """计算一组32位字的CRC32（不包括初始异或，直接返回）"""
    crc = 0xFFFFFFFF
    for w in words:
        crc = crc32_next(crc, w)
    return crc

# ================== transaction帧生成 ==================
def generate_transaction_header(addr, length_bytes):
    """
    生成64位头部，拆分为两个32位字
    addr: 25位地址
    length_bytes: 20位长度（字节数）
    返回 (high32, low32)
    """
    # W/R=0, ARD_FLAG=0, CID=0, EXADDR=0
    header = (addr << 32) | length_bytes
    high = (header >> 32) & 0xFFFFFFFF
    low = header & 0xFFFFFFFF
    return high, low

def write_transaction_file(filename, cmd_words, wave_words):
    """
    将cmd命令列表和wave数据列表写入transaction文件
    cmd_words: 整数列表，每个元素是一个32位命令字
    wave_words: 整数列表，每个元素是一个32位波形数据字
    """
    cmd_len_bytes = len(cmd_words) * 4
    wave_len_bytes = len(wave_words) * 4

    high_cmd, low_cmd = generate_transaction_header(0x100000, cmd_len_bytes)
    high_wave, low_wave = generate_transaction_header(0x200000, wave_len_bytes)

    with open(filename, 'w') as f:
        # cmd帧头部
        f.write(f"{high_cmd:08x}\n")
        f.write(f"{low_cmd:08x}\n")
        # cmd数据
        for w in cmd_words:
            f.write(f"{w:08x}\n")
        # wave帧头部
        f.write(f"{high_wave:08x}\n")
        f.write(f"{low_wave:08x}\n")
        # wave数据
        for w in wave_words:
            f.write(f"{w:08x}\n")

    print(f"已生成transaction文件: {filename}")
    print(f"  cmd帧: {len(cmd_words)} 字, 长度 {cmd_len_bytes} 字节")
    print(f"  wave帧: {len(wave_words)} 字, 长度 {wave_len_bytes} 字节")

# ================== LVDS帧生成 ==================
def write_lvds_file(filename, lvds_addr, data_words):
    """
    将波形数据打包成LVDS帧写入文件
    帧格式: [帧头] [地址/长度] [数据字...] [CRC32]
    所有字均为32位，每行一个十六进制数
    """
    if len(data_words) == 0:
        print("警告：波形数据为空，不生成LVDS帧")
        return

    header = 0xBCBCBCBC
    addr_len = (lvds_addr << 16) | len(data_words)  # 低16位为字数
    crc = crc32_compute([addr_len] + data_words)

    with open(filename, 'w') as f:
        f.write(f"{header:08x}\n")
        f.write(f"{addr_len:08x}\n")
        for w in data_words:
            f.write(f"{w:08x}\n")
        f.write(f"{crc:08x}\n")

    print(f"已生成LVDS数据帧文件: {filename}")
    print(f"  LVDS地址: 0x{lvds_addr:04x}, 数据字数: {len(data_words)}")

# ================== 主流程 ==================
def main():
    args = parse_args()

    # 确定LVDS输出文件名
    if args.lvds_output is None:
        base, ext = os.path.splitext(args.output)
        args.lvds_output = base + "_lvds" + ext

    if args.type in ['sine', 'triangle', 'pulse']:
        if args.length is None:
            sys.exit(f"波形类型 {args.type} 需要指定 --length")

        # 生成样本
        if args.type == 'sine':
            if args.freq is None:
                sys.exit("正弦波需要指定 --freq")
            samples = generate_sine(args.length, args.freq)
        elif args.type == 'triangle':
            if args.step is None or args.duration is None:
                sys.exit("三角波需要指定 --step 和 --duration")
            samples = generate_triangle(args.length, args.step, args.duration)
        else:  # pulse
            if args.freq is None:
                sys.exit("方波需要指定 --freq")
            samples = generate_pulse(args.length, args.freq, args.duty)

        # 打包波形数据
        blocks = pack_samples_to_blocks(samples)
        wave_words = [word for block in blocks for word in block]  # 展平为32位字列表
        num_words = len(blocks)  # 512-bit块数

        # 检查地址范围
        if args.start_addr + num_words > MAX_ADDR:
            sys.exit(f"地址 {args.start_addr} + 块数 {num_words} 超出SRAM深度 {MAX_ADDR}")

        # 生成波形命令
        cmd_wave = generate_wave_command(args.start_addr, num_words, args.cycles)
        cmd_words = [cmd_wave]

        # 如果需要hold命令，添加
        if args.hold_cycles is not None:
            cmd_hold = generate_hold_command(args.hold_cycles)
            cmd_words.append(cmd_hold)

        # 写入transaction文件
        write_transaction_file(args.output, cmd_words, wave_words)

        # 写入LVDS文件
        write_lvds_file(args.lvds_output, args.lvds_addr, wave_words)

    elif args.type == 'flattop':
        if args.rise_samples is None or args.hold_cycles is None or args.fall_samples is None:
            sys.exit("Flattop波形需要指定 --rise_samples, --hold_cycles, --fall_samples")

        # 生成上升沿和下降沿样本
        rise_samples, fall_samples = generate_flattop(args.rise_samples, args.hold_cycles,
                                                       args.fall_samples, args.hold_value)

        # 打包上升沿
        rise_blocks = pack_samples_to_blocks(rise_samples)
        rise_words = [word for block in rise_blocks for word in block]
        num_rise_words = len(rise_blocks)

        # 打包下降沿
        fall_blocks = pack_samples_to_blocks(fall_samples)
        fall_words = [word for block in fall_blocks for word in block]
        num_fall_words = len(fall_blocks)

        # 检查地址范围
        if args.start_addr + num_rise_words > MAX_ADDR:
            sys.exit(f"上升沿地址 {args.start_addr} + 块数 {num_rise_words} 超出SRAM深度 {MAX_ADDR}")
        fall_start_addr = args.start_addr + num_rise_words
        if fall_start_addr + num_fall_words > MAX_ADDR:
            sys.exit(f"下降沿起始地址 {fall_start_addr} + 块数 {num_fall_words} 超出SRAM深度 {MAX_ADDR}")

        # 生成命令
        cmd_rise = generate_wave_command(args.start_addr, num_rise_words, args.cycles)
        cmd_hold = generate_hold_command(args.hold_cycles)
        cmd_fall = generate_wave_command(fall_start_addr, num_fall_words, args.cycles)
        cmd_words = [cmd_rise, cmd_hold, cmd_fall]

        # 波形数据 = 上升沿 + 下降沿
        wave_words = rise_words + fall_words

        # 写入transaction文件
        write_transaction_file(args.output, cmd_words, wave_words)

        # 写入LVDS文件
        write_lvds_file(args.lvds_output, args.lvds_addr, wave_words)

if __name__ == "__main__":
    main()