import pygame
import math

# =========================
# 基本設定與物理參數
# =========================
WIDTH, HEIGHT = 1100, 760
FPS = 60
BALL_COUNT = 5
BALL_RADIUS = 30
PIVOT_Y = 150
STRING_LENGTH = 300
CENTER_X = WIDTH // 2
BALL_SPACING = BALL_RADIUS * 2.2
MAX_ANGLE = math.radians(40)
# 將速度提升 1.25 倍
SPEED_FACTOR = 1.25

# 顏色設定
WHITE = (255, 255, 255)
BLACK = (20, 20, 20)
GRAY = (180, 180, 180)
BLUE = (70, 130, 180)
HIGHLIGHT = (230, 80, 80) # 高能量顏色
BUTTON_COLOR = (220, 220, 220)

pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("牛頓擺物理教學 - 能量傳遞演示")
clock = pygame.time.Clock()

def get_font(size):
    font_name = pygame.font.match_font('microsoftjhenghei')
    return pygame.font.Font(font_name, size)

font_main = get_font(24)
font_title = get_font(32)

class Button:
    def __init__(self, x, y, w, h, text, action):
        self.rect = pygame.Rect(x, y, w, h)
        self.text = text
        self.action = action
    def draw(self, screen):
        pygame.draw.rect(screen, BUTTON_COLOR, self.rect, border_radius=5)
        pygame.draw.rect(screen, BLACK, self.rect, 2, border_radius=5)
        label = font_main.render(self.text, True, BLACK)
        screen.blit(label, label.get_rect(center=self.rect.center))
    def is_clicked(self, pos): return self.rect.collidepoint(pos)

buttons = [
    Button(250, 650, 120, 50, "1 顆球", 1),
    Button(400, 650, 120, 50, "2 顆球", 2),
    Button(550, 650, 120, 50, "3 顆球", 3),
    Button(700, 650, 120, 50, "重置", "reset")
]

moving_count = 1
phase = "left_out"
t = 0
status_text = "等待開始"
rest_positions = [CENTER_X - (BALL_COUNT - 1) * BALL_SPACING / 2 + i * BALL_SPACING for i in range(BALL_COUNT)]

def update_physics():
    global phase, t, status_text
    t += (1.0 / FPS) * SPEED_FACTOR
    
    if phase == "left_out":
        status_text = "拉起球體 (儲存重力勢能)"
        if t >= 0.6: phase, t = "left_to_center", 0
    elif phase == "left_to_center":
        status_text = "釋放球體 (轉換為動能)"
        if t >= 1.0: phase, t = "collision", 0
    elif phase == "collision":
        status_text = "碰撞點：動量轉移"
        if t >= 0.3: phase, t = "right_out", 0
    elif phase == "right_out":
        status_text = "另一端球體受力擺動"
        if t >= 1.0: phase, t = "right_to_center", 0
    elif phase == "right_to_center":
        status_text = "能量回傳"
        if t >= 1.0: phase, t = "left_out", 0

def get_angles():
    angles = [0.0] * BALL_COUNT
    if phase == "left_out":
        for i in range(moving_count): angles[i] = -MAX_ANGLE
    elif phase == "left_to_center":
        angle = -MAX_ANGLE * math.cos(t * math.pi / 2)
        for i in range(moving_count): angles[i] = angle
    elif phase == "right_out":
        angle = MAX_ANGLE * math.sin(t * math.pi / 2)
        for i in range(BALL_COUNT - moving_count, BALL_COUNT): angles[i] = angle
    elif phase == "right_to_center":
        angle = MAX_ANGLE * math.cos(t * math.pi / 2)
        for i in range(BALL_COUNT - moving_count, BALL_COUNT): angles[i] = angle
    return angles

running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT: running = False
        if event.type == pygame.MOUSEBUTTONDOWN:
            for btn in buttons:
                if btn.is_clicked(event.pos):
                    if btn.action != "reset": moving_count = btn.action
                    phase, t = "left_out", 0

    update_physics()
    angles = get_current_angles = get_angles()
    
    screen.fill(WHITE)
    
    # 顯示教學文字
    title = font_title.render("牛頓擺能量傳遞模擬", True, BLACK)
    status = font_main.render(f"目前狀態: {status_text}", True, BLUE)
    screen.blit(title, (50, 30))
    screen.blit(status, (50, 80))
    
    # 繪製支架
    pygame.draw.line(screen, BLACK, (rest_positions[0]-50, PIVOT_Y), (rest_positions[-1]+50, PIVOT_Y), 5)
    
    for i in range(BALL_COUNT):
        x = rest_positions[i] + STRING_LENGTH * math.sin(angles[i])
        y = PIVOT_Y + STRING_LENGTH * math.cos(angles[i])
        pygame.draw.line(screen, GRAY, (rest_positions[i], PIVOT_Y), (x, y), 2)
        
        # 繪製球體與能量條
        color = HIGHLIGHT if abs(angles[i]) > 0.05 else BLUE
        pygame.draw.circle(screen, color, (int(x), int(y)), BALL_RADIUS)
        
        # 能量條視覺化
        bar_w = 40
        bar_h = 8
        pygame.draw.rect(screen, BLACK, (x - bar_w/2, y + BALL_RADIUS + 10, bar_w, bar_h))
        energy_level = min(bar_w, abs(angles[i]) * 100)
        pygame.draw.rect(screen, color, (x - bar_w/2, y + BALL_RADIUS + 10, energy_level, bar_h))

    for btn in buttons: btn.draw(screen)
    
    pygame.display.flip()
    clock.tick(FPS)

pygame.quit()