"""
三棱镜色散动画 + 光屏 (Prism Dispersion on Screen)
依赖: numpy, matplotlib
运行: python prism_dispersion_screen.py
"""

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# ---------- 中文字体 ----------
matplotlib.rcParams['font.sans-serif'] = [
    'Microsoft YaHei', 'SimHei', 'PingFang SC', 'Hiragino Sans GB',
    'Noto Sans CJK SC', 'Source Han Sans SC', 'WenQuanYi Micro Hei',
    'Arial Unicode MS', 'DejaVu Sans'
]
matplotlib.rcParams['axes.unicode_minus'] = False


# ============================================================
#  1. 波长 -> RGB 颜色
# ============================================================
def wavelength_to_rgb(wl):
    wl = np.clip(wl, 380, 780)
    if wl < 440:
        r, g, b = -(wl - 440) / 60.0, 0.0, 1.0
    elif wl < 490:
        r, g, b = 0.0, (wl - 440) / 50.0, 1.0
    elif wl < 510:
        r, g, b = 0.0, 1.0, -(wl - 510) / 20.0
    elif wl < 580:
        r, g, b = (wl - 510) / 70.0, 1.0, 0.0
    elif wl < 645:
        r, g, b = 1.0, -(wl - 645) / 65.0, 0.0
    else:
        r, g, b = 1.0, 0.0, 0.0
    if wl > 700:
        f = 0.35 + 0.65 * (780 - wl) / 80.0
    elif wl < 420:
        f = 0.45 + 0.55 * (wl - 380) / 40.0
    else:
        f = 1.0
    return (r * f, g * f, b * f)


# ============================================================
#  2. 玻璃色散模型 (Cauchy 公式, 数值放大以便肉眼可见)
# ============================================================
def n_glass(wl_nm):
    lam = wl_nm / 1000.0          # μm
    return 1.50 + 0.030 / lam ** 2


# ============================================================
#  3. 棱镜几何
# ============================================================
APEX = np.array([0.0, 1.0])
BL   = np.array([-0.6, -0.6])
BR   = np.array([0.6, -0.6])
CEN  = (APEX + BL + BR) / 3.0


def face_normal(p1, p2):
    """返回指向棱镜外侧的单位法向量"""
    e = p2 - p1
    n = np.array([e[1], -e[0]], dtype=float)
    n /= np.linalg.norm(n)
    if np.dot(n, 0.5 * (p1 + p2) - CEN) < 0:
        n = -n
    return n


N_LEFT  = face_normal(BL, APEX)
N_RIGHT = face_normal(APEX, BR)


# ============================================================
#  4. 光屏位置
# ============================================================
SCREEN_X    = 3.15
SCREEN_YMIN = -2.45
SCREEN_YMAX = 1.55


# ============================================================
#  5. 光学工具函数
# ============================================================
def refract(d, nrm, n1, n2):
    """斯涅尔定律矢量形式，返回折射方向；全反射时返回 None"""
    d = d / np.linalg.norm(d)
    nrm = nrm / np.linalg.norm(nrm)
    cosi = -np.dot(d, nrm)
    if cosi < 0:
        nrm, cosi = -nrm, -cosi
    eta = n1 / n2
    k = 1.0 - eta * eta * (1.0 - cosi * cosi)
    if k < 0:
        return None                      # 全反射
    return eta * d + (eta * cosi - np.sqrt(k)) * nrm


def ray_seg(origin, d, p1, p2):
    """射线与线段求交，返回 (t, s)"""
    e = p2 - p1
    M = np.array([[d[0], -e[0]], [d[1], -e[1]]])
    b = p1 - origin
    if abs(np.linalg.det(M)) < 1e-12:
        return None, None
    return np.linalg.solve(M, b)


def stop_at_screen(P, d):
    """光线停在光屏上（或画布后备边界）"""
    if d[0] > 1e-9:
        t = (SCREEN_X - P[0]) / d[0]
        Q = P + t * d
        if SCREEN_YMIN <= Q[1] <= SCREEN_YMAX:
            return Q
    # 后备：裁到屏幕底边或画布底边
    us = []
    if d[0] > 1e-9:
        us.append((SCREEN_X - P[0]) / d[0])
    if d[1] < -1e-9:
        us.append((SCREEN_YMIN - P[1]) / d[1])
    if not us:
        us.append(1.0)
    return P + min(us) * d


def path_upto(pts, s):
    """沿折线取前 s 长度的部分"""
    if s <= 0:
        return pts[:1]
    segs = np.diff(pts, axis=0)
    lens = np.linalg.norm(segs, axis=1)
    out = [pts[0]]
    acc = 0.0
    for i, L in enumerate(lens):
        if acc + L <= s:
            out.append(pts[i + 1])
            acc += L
        else:
            out.append(pts[i] + (s - acc) / L * segs[i])
            break
    return np.array(out)


# ============================================================
#  6. 光线追踪
# ============================================================
X_START, Y0 = -3.35, 0.0
WLS = np.linspace(400, 700, 9)

rays = []
for wl in WLS:
    n  = n_glass(wl)
    d0 = np.array([1.0, 0.0])
    P0 = np.array([X_START, Y0])

    t1, _ = ray_seg(P0, d0, BL, APEX)
    P1 = P0 + t1 * d0                        # 入射点

    d1 = refract(d0, N_LEFT, 1.0, n)         # 进入玻璃
    t2, _ = ray_seg(P1, d1, APEX, BR)
    P2 = P1 + t2 * d1                        # 出射点

    d2 = refract(d1, -N_RIGHT, n, 1.0)       # 射回空气
    P3 = stop_at_screen(P2, d2)              # 落在光屏上

    pts = np.array([P0, P1, P2, P3])
    L = np.sum(np.linalg.norm(np.diff(pts, axis=0), axis=1))
    rays.append((wl, wavelength_to_rgb(wl), pts, L))

L1    = np.linalg.norm(rays[0][2][1] - rays[0][2][0])
L_MAX = max(r[3] for r in rays)
WHITE = rays[0][2][:2]


# ============================================================
#  7. 画布
# ============================================================
fig, ax = plt.subplots(figsize=(11.5, 6.9), dpi=110)
fig.patch.set_facecolor('#07080d')
ax.set_facecolor('#07080d')
ax.set_xlim(-3.5, 3.9)
ax.set_ylim(-2.6, 1.75)
ax.set_aspect('equal')
ax.axis('off')
fig.subplots_adjust(left=0.005, right=0.995, top=0.995, bottom=0.005)

# --- 光屏 ---
ax.add_patch(plt.Rectangle(
    (SCREEN_X, SCREEN_YMIN), 0.07, SCREEN_YMAX - SCREEN_YMIN,
    facecolor=(0.80, 0.86, 0.96, 0.10),
    edgecolor=(0.72, 0.84, 1.00, 0.60),
    lw=1.3, zorder=1))

# --- 棱镜 ---
ax.add_patch(plt.Polygon(
    [APEX, BR, BL], closed=True,
    facecolor=(0.50, 0.70, 1.00, 0.11),
    edgecolor=(0.75, 0.88, 1.00, 0.95),
    lw=2.2, joinstyle='round', zorder=2))

# --- 标题与标签 ---
ax.text(-3.4, 1.62, "三棱镜色散", color='#ffffff',
        fontsize=21, fontweight='bold', ha='left', va='top')
ax.text(-3.4, 1.30, "Prism Dispersion   n(λ) = 1.50 + 0.030/λ²",
        color='#7f93aa', fontsize=10.5, ha='left', va='top')
ax.text(SCREEN_X + 0.16, SCREEN_YMIN + 0.12, "光屏",
        color='#8fa3bd', fontsize=11, ha='left', va='bottom')

# --- 光源 ---
ax.plot([X_START], [Y0], 'o', ms=13, color='white', alpha=0.12, zorder=5)
ax.plot([X_START], [Y0], 'o', ms=4,  color='white', zorder=6)

# --- 光线对象 ---
white_glow, = ax.plot([], [], color='white', lw=10, alpha=0.09,
                      solid_capstyle='round', zorder=3)
white_main, = ax.plot([], [], color='white', lw=2.6, alpha=1.0,
                      solid_capstyle='round', zorder=4)

line_objs = []
for wl, col, pts, L in rays:
    g, = ax.plot([], [], color=col, lw=7.5, alpha=0.10,
                 solid_capstyle='round', zorder=3)
    m, = ax.plot([], [], color=col, lw=2.0, alpha=1.0,
                 solid_capstyle='round', zorder=4)
    line_objs.append((g, m))

# --- 光屏上的光斑 ---
spot_objs = []
for wl, col, pts, L in rays:
    ys = pts[-1][1]
    g, = ax.plot([SCREEN_X, SCREEN_X], [ys - 0.055, ys + 0.055],
                 color=col, lw=15, alpha=0.0,
                 solid_capstyle='round', zorder=4)
    c, = ax.plot([SCREEN_X, SCREEN_X], [ys - 0.055, ys + 0.055],
                 color=col, lw=5.5, alpha=0.0,
                 solid_capstyle='round', zorder=5)
    spot_objs.append((g, c, L))


# ============================================================
#  8. 动画
# ============================================================
T_GROW = 170
T_HOLD = 80
FRAMES = T_GROW + T_HOLD


def update(frame):
    p = min(1.0, frame / max(1, T_GROW - 1))
    s = p * L_MAX

    # 白光段
    wp = path_upto(WHITE, min(s, L1))
    white_main.set_data(wp[:, 0], wp[:, 1])
    white_glow.set_data(wp[:, 0], wp[:, 1])

    # 色散段
    for (g, m), (wl, col, pts, L) in zip(line_objs, rays):
        if s <= L1:
            m.set_data([], [])
            g.set_data([], [])
        else:
            cp = path_upto(pts[1:], s - L1)
            m.set_data(cp[:, 0], cp[:, 1])
            g.set_data(cp[:, 0], cp[:, 1])

    # 光屏光斑：光线到达时渐显
    for g, c, L in spot_objs:
        frac = np.clip((s - (L - 0.10)) / 0.10, 0.0, 1.0)
        g.set_alpha(0.22 * frac)
        c.set_alpha(1.00 * frac)

    artists = [white_main, white_glow]
    artists += [o for pair in line_objs for o in pair]
    artists += [o for pair in spot_objs for o in pair[:2]]
    return artists


anim = FuncAnimation(fig, update, frames=FRAMES,
                     interval=33, blit=False, repeat=True, repeat_delay=900)

plt.show()
# anim.save('prism_dispersion_screen.gif', writer='pillow', fps=30)  # 需要保存时取消注释