import pygame
import pygame.midi

# ===================== 初始化 =====================
pygame.init()
pygame.midi.init()

# 打开系统默认MIDI输出设备
try:
    midi_out = pygame.midi.Output(0)
    midi_out.set_instrument(0)  # 0=原声大钢琴音色
except:
    print("未检测到MIDI设备，程序退出")
    pygame.midi.quit()
    pygame.quit()
    exit()

# 窗口设置
WIDTH, HEIGHT = 960, 420
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pygame 电子钢琴 | 键盘+鼠标弹奏")
clock = pygame.time.Clock()
FPS = 60

# 颜色
WHITE_KEY = (255, 255, 255)
WHITE_PRESS = (210, 230, 255)
BLACK_KEY = (20, 20, 20)
BLACK_PRESS = (80, 80, 100)
BG_COLOR = (45, 45, 55)
TEXT_COLOR = (255, 255, 255)

# 字体兜底兼容
try:
    font = pygame.font.Font("simhei.ttf", 16)
except:
    font = pygame.font.SysFont("Microsoft YaHei", 16)

# ===================== 音符与按键映射 =====================
# MIDI标准音符编号 C4=60 中央Do
WHITE_NOTES = [60, 62, 64, 65, 67, 69, 71, 72, 74, 76, 77, 79, 81, 83]
BLACK_NOTES = [61, 63, 66, 68, 70, 73, 75, 78, 80, 82]

# 电脑键盘 → MIDI音符
keyboard_map = {
    pygame.K_a: 60, pygame.K_s: 62, pygame.K_d: 64, pygame.K_f: 65,
    pygame.K_g: 67, pygame.K_h: 69, pygame.K_j: 71, pygame.K_k: 72,
    pygame.K_l: 74, pygame.K_SEMICOLON: 76, pygame.K_QUOTE: 77,
    pygame.K_z: 79, pygame.K_x: 81, pygame.K_c: 83,
    pygame.K_w: 61, pygame.K_e: 63, pygame.K_t: 66,
    pygame.K_y: 68, pygame.K_u: 70, pygame.K_o: 73,
    pygame.K_p: 75, pygame.K_LEFTBRACKET: 78, pygame.K_RIGHTBRACKET: 80
}

# 正在发声的音符集合（防止重复触发）
playing_notes = set()

# ===================== 琴键布局参数 =====================
white_w = 60
white_h = 320
black_w = 36
black_h = 200
start_x = 30

# 白键坐标列表
white_keys = []
for i, note in enumerate(WHITE_NOTES):
    x = start_x + i * white_w
    white_keys.append({"note": note, "rect": pygame.Rect(x, 60, white_w, white_h), "press": False})

# 黑键偏移规则：对应白键1、2、4、5、6、8、9、11、12、13上方
black_offset_pos = [0, 1, 3, 4, 5, 7, 8, 10, 11, 12]
black_keys = []
for idx, pos in enumerate(black_offset_pos):
    note = BLACK_NOTES[idx]
    x = start_x + pos * white_w + white_w - black_w // 2
    black_keys.append({"note": note, "rect": pygame.Rect(x, 60, black_w, black_h), "press": False})

# ===================== 发声工具函数 =====================
def note_on(note_num):
    if note_num not in playing_notes:
        midi_out.note_on(note_num, 110)  # 音量0~127
        playing_notes.add(note_num)

def note_off(note_num):
    if note_num in playing_notes:
        midi_out.note_off(note_num)
        playing_notes.remove(note_num)

# 关闭所有音符
def all_note_stop():
    for n in list(playing_notes):
        note_off(n)

# ===================== 主循环 =====================
running = True
while running:
    mx, my = pygame.mouse.get_pos()
    screen.fill(BG_COLOR)

    # 事件处理
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            all_note_stop()
            running = False

        # 键盘按下
        if event.type == pygame.KEYDOWN:
            if event.key in keyboard_map:
                n = keyboard_map[event.key]
                note_on(n)
                # 标记对应按键按下
                for wk in white_keys:
                    if wk["note"] == n:
                        wk["press"] = True
                for bk in black_keys:
                    if bk["note"] == n:
                        bk["press"] = True

        # 键盘松开
        if event.type == pygame.KEYUP:
            if event.key in keyboard_map:
                n = keyboard_map[event.key]
                note_off(n)
                for wk in white_keys:
                    if wk["note"] == n:
                        wk["press"] = False
                for bk in black_keys:
                    if bk["note"] == n:
                        bk["press"] = False

        # 鼠标按下弹奏
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            # 先检测黑键（层级在上）
            hit = False
            for bk in black_keys:
                if bk["rect"].collidepoint(mx, my):
                    note_on(bk["note"])
                    bk["press"] = True
                    hit = True
                    break
            if not hit:
                for wk in white_keys:
                    if wk["rect"].collidepoint(mx, my):
                        note_on(wk["note"])
                        wk["press"] = True

        # 鼠标松开停音
        if event.type == pygame.MOUSEBUTTONUP and event.button == 1:
            for wk in white_keys:
                if wk["press"]:
                    note_off(wk["note"])
                    wk["press"] = False
            for bk in black_keys:
                if bk["press"]:
                    note_off(bk["note"])
                    bk["press"] = False

    # 绘制白键
    for wk in white_keys:
        color = WHITE_PRESS if wk["press"] else WHITE_KEY
        pygame.draw.rect(screen, color, wk["rect"])
        pygame.draw.rect(screen, (0,0,0), wk["rect"], 2)

    # 绘制黑键（后绘制，覆盖在白键上方）
    for bk in black_keys:
        color = BLACK_PRESS if bk["press"] else BLACK_KEY
        pygame.draw.rect(screen, color, bk["rect"])

    # 顶部提示文字
    tip1 = font.render("A~L Z X C 弹奏白键 | W E T Y U O P [ ] 弹奏黑键 | 鼠标点击琴键演奏", True, TEXT_COLOR)
    tip2 = font.render("关闭窗口退出程序", True, TEXT_COLOR)
    screen.blit(tip1, (20, 15))
    screen.blit(tip2, (20, 38))

    pygame.display.update()
    clock.tick(FPS)

# 收尾释放资源
all_note_stop()
midi_out.close()
pygame.midi.quit()
pygame.quit()