import pygame
import sys
import math
import time

# ============================================================
# 三棱镜光学模拟 - 完整修正版
# ============================================================

# ---------- 材质 ----------
MATERIALS = {
    "Crown Glass K9": {"A": 1.500, "B": 5000,  "color": (180, 220, 255), "desc": "冕牌玻璃"},
    "Flint Glass F2": {"A": 1.580, "B": 12000, "color": (200, 180, 255), "desc": "火石玻璃"},
    "Diamond":        {"A": 2.400, "B": 30000, "color": (180, 255, 255), "desc": "钻石"},
}

# ---------- pygame init ----------
pygame.init()
WIDTH, HEIGHT = 900, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("三棱镜光学模拟")
clock = pygame.time.Clock()
font_s = pygame.font.SysFont("simhei", 14)
font_m = pygame.font.SysFont("simhei", 16)

# ---------- 全局状态 ----------
material_name = "Crown Glass K9"
prism_angle = math.radians(60)
prism_rotation = 0.0
incident_angle = 0.0
light_pos = [120, HEIGHT // 2]
prism_center = [WIDTH // 2, HEIGHT // 2]

auto_rotate = False
paused = False
show_trails = True
show_guide = True
wave_toggles = [True] * 16
trails = {}

drag_mode = None
last_mouse_pos = None


# ============================================================
# 物理 / 数学
# ============================================================

def get_refractive_index(wl_nm, mat):
    m = MATERIALS[mat]
    return m["A"] + m["B"] / (wl_nm ** 2)


def wl_to_rgb(wl):
    gamma = 0.8
    if 380 <= wl < 440:
        r, g, b = -(wl - 440) / 60, 0.0, 1.0
    elif 440 <= wl < 490:
        r, g, b = 0.0, (wl - 440) / 50, 1.0
    elif 490 <= wl < 510:
        r, g, b = 0.0, 1.0, -(wl - 510) / 20
    elif 510 <= wl < 580:
        r, g, b = (wl - 510) / 70, 1.0, 0.0
    elif 580 <= wl < 645:
        r, g, b = 1.0, -(wl - 645) / 65, 0.0
    elif 645 <= wl <= 780:
        r, g, b = 1.0, 0.0, 0.0
    else:
        r, g, b = 0.0, 0.0, 0.0
    if 380 <= wl < 420:
        factor = 0.3 + 0.7 * (wl - 380) / 40
    elif 420 <= wl <= 700:
        factor = 1.0
    elif 700 < wl <= 780:
        factor = 0.3 + 0.7 * (780 - wl) / 80
    else:
        factor = 0.0

    def adj(c):
        return int(255 * (c * factor) ** gamma) if c > 0 else 0
    return (adj(r), adj(g), adj(b))


def get_prism_vertices(cx, cy, angle, rot):
    size = 120
    ha = angle / 2
    v = [(0, -size),
         (-size * math.cos(ha), size * math.sin(ha)),
         (size * math.cos(ha), size * math.sin(ha))]
    cr, sr = math.cos(rot), math.sin(rot)
    out = []
    for vx, vy in v:
        rx = vx * cr - vy * sr
        ry = vx * sr + vy * cr
        out.append((cx + rx, cy + ry))
    return out


def point_in_tri(pt, tri):
    def sgn(p1, p2, p3):
        return (p1[0]-p3[0])*(p2[1]-p3[1])-(p2[0]-p3[0])*(p1[1]-p3[1])
    d1, d2, d3 = sgn(pt, tri[0], tri[1]), sgn(pt, tri[1], tri[2]), sgn(pt, tri[2], tri[0])
    return not ((d1 < 0 or d2 < 0 or d3 < 0) and (d1 > 0 or d2 > 0 or d3 > 0))


def line_seg_intersect(a1, a2, b1, b2):
    x1,y1,x2,y2 = *a1,*a2
    x3,y3,x4,y4 = *b1,*b2
    denom = (y4-y3)*(x2-x1)-(x4-x3)*(y2-y1)
    if abs(denom) < 1e-10:
        return None
    ua = ((x4-x3)*(y1-y3)-(y4-y3)*(x1-x3))/denom
    ub = ((x2-x1)*(y1-y3)-(y2-y1)*(x1-x3))/denom
    if 0 <= ua <= 1 and 0 <= ub <= 1:
        return (x1+ua*(x2-x1), y1+ua*(y2-y1))
    return None


def get_line_normal(p1, p2, pt_on_line, inside_pt):
    """返回指向三角形内部的法线"""
    dx = p2[0]-p1[0]
    dy = p2[1]-p1[1]
    ln = math.sqrt(dx*dx+dy*dy)
    if ln == 0:
        return (0, -1)
    nx, ny = -dy/ln, dx/ln
    # inside_pt 一定是 (x,y) 元组
    dot = nx*(inside_pt[0]-pt_on_line[0]) + ny*(inside_pt[1]-pt_on_line[1])
    if dot < 0:
        nx, ny = -nx, -ny
    return (nx, ny)


def refract(inc_dir, normal, n1, n2):
    cos_i = -(inc_dir[0]*normal[0]+inc_dir[1]*normal[1])
    if cos_i < 0:
        normal = (-normal[0], -normal[1])
        cos_i = -cos_i
    sin_t = (n1/n2)*math.sqrt(max(0, 1-cos_i*cos_i))
    if sin_t > 1.0:
        return None
    cos_t = math.sqrt(1-sin_t*sin_t)
    rx = (n1/n2)*inc_dir[0] + (n1/n2*cos_i-cos_t)*normal[0]
    ry = (n1/n2)*inc_dir[1] + (n1/n2*cos_i-cos_t)*normal[1]
    rl = math.sqrt(rx*rx+ry*ry)
    return (rx/rl, ry/rl) if rl > 0 else None


# ============================================================
# 光线追踪
# ============================================================

def trace_ray(wl_nm, wave_i):
    path = []
    n_out, n_in = 1.0, get_refractive_index(wl_nm, material_name)
    verts = get_prism_vertices(prism_center[0], prism_center[1], prism_angle, prism_rotation)

    # 棱镜重心（传给 get_line_normal 用）
    cx = (verts[0][0]+verts[1][0]+verts[2][0])/3
    cy = (verts[0][1]+verts[1][1]+verts[2][1])/3
    centroid = (cx, cy)

    ang = math.radians(incident_angle)
    inc_dir = (math.cos(ang), math.sin(ang))
    start = (light_pos[0], light_pos[1])

    # 1. 入射 → 棱镜
    hit_pt = hit_edge = hit_nrm = None
    for e1,e2 in [(0,1),(1,2),(2,0)]:
        inter = line_seg_intersect(start, (start[0]+inc_dir[0]*2000, start[1]+inc_dir[1]*2000), verts[e1], verts[e2])
        if inter:
            hit_pt, hit_edge, hit_nrm = inter, (e1,e2), get_line_normal(verts[e1], verts[e2], inter, centroid)
            break
    if not hit_pt:
        return path
    path.extend([start, hit_pt])

    # 2. 折射进玻璃
    refr = refract(inc_dir, hit_nrm, n_out, n_in)
    if not refr:
        return path

    # 3. 在棱镜内 → 找出射面
    out_pt = out_edge = out_nrm = None
    for e1,e2 in [(0,1),(1,2),(2,0)]:
        if (e1,e2)==hit_edge or (e2,e1)==hit_edge:
            continue
        inter = line_seg_intersect(hit_pt, (hit_pt[0]+refr[0]*500, hit_pt[1]+refr[1]*500), verts[e1], verts[e2])
        if inter:
            out_pt, out_edge, out_nrm = inter, (e1,e2), get_line_normal(verts[e1], verts[e2], inter, centroid)
            break
    if not out_pt:
        return path
    path.append(out_pt)

    # 4. 折射出玻璃
    exit_dir = refract(refr, out_nrm, n_in, n_out)
    if not exit_dir:
        # 全内反射 → 找下一个面
        refl = (refr[0]-2*(refr[0]*out_nrm[0]+refr[1]*out_nrm[1])*out_nrm[0],
                refr[1]-2*(refr[0]*out_nrm[0]+refr[1]*out_nrm[1])*out_nrm[1])
        for e1,e2 in [(0,1),(1,2),(2,0)]:
            if (e1,e2)==out_edge or (e2,e1)==out_edge:
                continue
            inter2 = line_seg_intersect(out_pt, (out_pt[0]+refl[0]*500, out_pt[1]+refl[1]*500), verts[e1], verts[e2])
            if inter2:
                path.append(inter2)
                n2 = get_line_normal(verts[e1], verts[e2], inter2, centroid)
                ed = refract(refl, n2, n_in, n_out)
                if ed:
                    path.append((inter2[0]+ed[0]*400, inter2[1]+ed[1]*400))
                break
        return path

    path.append((out_pt[0]+exit_dir[0]*400, out_pt[1]+exit_dir[1]*400))
    return path


# ============================================================
# 绘制
# ============================================================

def draw_prism():
    v = get_prism_vertices(prism_center[0], prism_center[1], prism_angle, prism_rotation)
    c = MATERIALS[material_name]["color"]
    if len(v)==3:
        pygame.draw.polygon(screen, (*c,50), v)
        pygame.draw.polygon(screen, c, v, 2)
    for i,p in enumerate(v):
        pygame.draw.circle(screen, (255,255,255), (int(p[0]),int(p[1])), 3)
    return v


def draw_light_src():
    pygame.draw.circle(screen, (255,255,200), (int(light_pos[0]),int(light_pos[1])), 8)
    pygame.draw.circle(screen, (255,255,255), (int(light_pos[0]),int(light_pos[1])), 12, 1)


def draw_incident():
    ang = math.radians(incident_angle)
    ex, ey = light_pos[0]+math.cos(ang)*200, light_pos[1]+math.sin(ang)*200
    pygame.draw.line(screen, (255,255,255), light_pos, (ex,ey), 3)


def draw_beams():
    wls = [380+i*(400/15) for i in range(16)]
    for i,wl in enumerate(wls):
        if not wave_toggles[i]:
            continue
        col = wl_to_rgb(wl)
        path = trace_ray(wl, i)
        if len(path)<2:
            continue
        if show_trails:
            trails.setdefault((i,),[]).append(path[-1])
            if len(trails[(i,)])>60:
                trails[(i,)] = trails[(i,)][-60:]
        for j in range(len(path)-1):
            pygame.draw.line(screen, col, path[j], path[j+1], 2)
        if len(path)>=2:
            pygame.draw.circle(screen, col, (int(path[-1][0]),int(path[-1][1])), 3)


def draw_trails_f():
    if not show_trails:
        trails.clear()
        return
    for key,trail in list(trails.items()):
        if len(trail)<2:
            continue
        wi = key[0] if isinstance(key,tuple) else key
        col = wl_to_rgb(380+wi*(400/15))
        for i in range(1,len(trail)):
            a = int(100*i/len(trail))
            pygame.draw.line(screen, (*col,a), trail[i-1], trail[i], 2)


def draw_spectrum():
    by, bh = HEIGHT-25, 20
    bx, bw = 20, WIDTH-40
    for i in range(bw):
        pygame.draw.line(screen, wl_to_rgb(380+i/400*bw), (bx+i,by), (bx+i,by+bh))
    for wl in [380,450,500,550,600,650,700]:
        x = bx+int((wl-380)/400*bw)
        pygame.draw.line(screen,(255,255,255),(x,by-5),(x,by+bh+5),1)
    sw = bw/16
    for i in range(16):
        if not wave_toggles[i]:
            pygame.draw.line(screen,(0,0,0),(bx+int(i*sw),by),(bx+int(i*sw),by+bh),2)


def draw_hud():
    y = 10
    lines = [
        f"材质:{material_name}",
        f"顶角:{math.degrees(prism_angle):.0f}° 旋转:{math.degrees(prism_rotation):.0f}°",
        f"入射角:{incident_angle:.1f}°",
        f"n(红)={get_refractive_index(700,material_name):.4f} n(紫)={get_refractive_index(400,material_name):.4f}",
    ]
    if auto_rotate: lines.append("自动旋转:ON")
    if paused: lines.append("暂停")
    for l in lines:
        screen.blit(font_s.render(l,True,(220,220,220)), (10,y))
        y += 16


def draw_guide():
    if not show_guide:
        return
    lines = ["←→入射角 ↑↓顶角 1/2/3材质","滚轮旋转 拖拽棱镜/光源","T拖尾 M测量 A自动旋转","空格暂停 R重置 S截图 G指南","点击底部光谱条开关单色光"]
    py = HEIGHT-10-len(lines)*15
    for l in lines:
        screen.blit(font_s.render(l,True,(180,180,180)), (WIDTH-200, py))
        py += 15


# ============================================================
# main
# ============================================================

def main():
    global prism_angle, prism_rotation, incident_angle
    global light_pos, prism_center, auto_rotate, paused, show_trails, show_guide, wave_toggles, trails, drag_mode, last_mouse_pos

    while True:
        dt = clock.tick(60)/1000.0

        for e in pygame.event.get():
            if e.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif e.type == pygame.KEYDOWN:
                if e.key == pygame.K_ESCAPE: pygame.quit(); sys.exit()
                elif e.key == pygame.K_LEFT: incident_angle -= 2
                elif e.key == pygame.K_RIGHT: incident_angle += 2
                elif e.key == pygame.K_UP: prism_angle = min(math.radians(90), prism_angle+math.radians(2))
                elif e.key == pygame.K_DOWN: prism_angle = max(math.radians(20), prism_angle-math.radians(2))
                elif e.key == pygame.K_1: material_name="Crown Glass K9"; trails.clear()
                elif e.key == pygame.K_2: material_name="Flint Glass F2"; trails.clear()
                elif e.key == pygame.K_3: material_name="Diamond"; trails.clear()
                elif e.key == pygame.K_a: auto_rotate = not auto_rotate
                elif e.key == pygame.K_SPACE: paused = not paused
                elif e.key == pygame.K_t: show_trails = not show_trails; trails.clear()
                elif e.key == pygame.K_g: show_guide = not show_guide
                elif e.key == pygame.K_r:
                    prism_angle=math.radians(60); prism_rotation=0; incident_angle=0
                    light_pos=[120,HEIGHT//2]; prism_center=[WIDTH//2,HEIGHT//2]
                    auto_rotate=False; paused=False; trails.clear(); wave_toggles=[True]*16
                elif e.key == pygame.K_s:
                    pygame.image.save(screen, f"prism_{int(time.time())}.png")
            elif e.type == pygame.MOUSEBUTTONDOWN:
                mp = e.pos
                if e.button == 1:
                    # 光谱条点击
                    if HEIGHT-25 <= mp[1] <= HEIGHT-5:
                        bw = WIDTH-40; bx=20
                        if bx <= mp[0] <= bx+bw:
                            seg = max(0,min(15,int((mp[0]-bx)/bw*16)))
                            wave_toggles[seg] = not wave_toggles[seg]
                            continue
                    # 光源
                    if math.hypot(mp[0]-light_pos[0], mp[1]-light_pos[1]) < 20:
                        drag_mode="light"; last_mouse_pos=mp; continue
                    # 棱镜内部
                    v = get_prism_vertices(prism_center[0],prism_center[1],prism_angle,prism_rotation)
                    if point_in_tri(mp, v):
                        drag_mode="prism"; last_mouse_pos=mp
                    else:
                        drag_mode="rot"; last_mouse_pos=mp
                elif e.button == 4: prism_rotation -= math.radians(3)
                elif e.button == 5: prism_rotation += math.radians(3)
            elif e.type == pygame.MOUSEBUTTONUP:
                drag_mode=None; last_mouse_pos=None
            elif e.type == pygame.MOUSEMOTION:
                if drag_mode=="light":
                    light_pos[0]=max(20,min(WIDTH-20,e.pos[0]))
                    light_pos[1]=max(20,min(HEIGHT-50,e.pos[1]))
                    trails.clear()
                elif drag_mode=="prism":
                    prism_center[0]=max(100,min(WIDTH-100,prism_center[0]+e.pos[0]-last_mouse_pos[0]))
                    prism_center[1]=max(100,min(HEIGHT-80,prism_center[1]+e.pos[1]-last_mouse_pos[1]))
                    last_mouse_pos=e.pos; trails.clear()
                elif drag_mode=="rot":
                    prism_rotation += (e.pos[0]-last_mouse_pos[0])*0.005
                    last_mouse_pos=e.pos; trails.clear()

        if auto_rotate and not paused:
            prism_rotation += 0.3*dt

        screen.fill((10,10,15))
        # 网格
        for x in range(0,WIDTH,50): pygame.draw.line(screen,(25,25,30),(x,0),(x,HEIGHT))
        for y in range(0,HEIGHT,50): pygame.draw.line(screen,(25,25,30),(0,y),(WIDTH,y))

        draw_prism()
        draw_light_src()
        draw_incident()
        draw_beams()
        draw_trails_f()
        draw_spectrum()
        draw_hud()
        draw_guide()

        pygame.display.flip()


if __name__ == "__main__":
    main()