from pathlib import Path
from PIL import Image, ImageDraw, ImageFont

OUT = Path(r"D:\Project\spinal_stim\pl\spinal_stim_v2\tmp\fpga_doc2\submodule_diagrams")
OUT.mkdir(parents=True, exist_ok=True)

try:
    FONT = ImageFont.truetype(r"C:\Windows\Fonts\msyh.ttc", 24)
    SMALL = ImageFont.truetype(r"C:\Windows\Fonts\msyh.ttc", 19)
except Exception:
    FONT = SMALL = ImageFont.load_default()

def diagram(name, title, boxes, suffix):
    im = Image.new("RGB", (1400, 520), "white")
    d = ImageDraw.Draw(im)
    d.text((30, 20), title, fill="black", font=FONT)
    y = 180
    w, h, gap = 235, 100, 45
    total = len(boxes)*w + (len(boxes)-1)*gap
    x = (1400-total)//2
    for i, b in enumerate(boxes):
        d.rounded_rectangle((x,y,x+w,y+h), radius=8, outline="#1f4e79", width=3, fill="#d9eaf7")
        lines = b.split("\n")
        yy = y + 20
        for line in lines:
            bb = d.textbbox((0,0), line, font=SMALL)
            d.text((x+(w-(bb[2]-bb[0]))//2, yy), line, fill="black", font=SMALL)
            yy += 28
        if i < len(boxes)-1:
            x2 = x+w+gap
            d.line((x+w+5,y+h//2,x2-10,y+h//2), fill="#c00000", width=4)
            d.polygon([(x2-10,y+h//2),(x2-25,y+h//2-10),(x2-25,y+h//2+10)], fill="#c00000")
        x += w+gap
    p = OUT / f"{name}_{suffix}.png"
    im.save(p)

mods = {
 "eeg_top": ["BRAM config", "SPI init\nstate machine", "DRDY sample\n24-bit x 8", "AXIS record"],
 "emg_top": ["BRAM config", "ADS1299 reset\n/start", "DRDY sample\n8 channels", "AXIS EMG"],
 "upload_top": ["Input packet", "Header + ID", "Channel map\n/filter", "CRC + TLAST"],
 "eeg_emg_upload": ["EEG FIFO", "Select mode", "Drop head/tail\nmerge", "CRC + output"],
 "cur_mon_top": ["CNV trigger", "SPI capture\nAD7689", "Map/filter\nstatistics", "Threshold\nfault latch"],
 "rhs2116_top": ["Stim config", "DDS/timer", "SPI distribute", "Monitor/event"],
 "channel_switch": ["PS GPIO/SPI", "Decode select", "ADG1414\nshift", "Switch state"],
 "trigger_top": ["Input trigger", "Synchronize", "Detect edge\n/code", "Interrupt/stream"],
}
for name, boxes in mods.items():
    diagram(name, f"{name} 状态/处理流程图", boxes, "flow")
    diagram(name, f"{name} 数据流图", [boxes[0], boxes[1], boxes[2], boxes[3]], "data")
    diagram(name, f"{name} 时序处理图", ["IDLE", "CONFIG", "RUN", "END"], "state")