import tkinter as tk
from tkinter import ttk

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['axes.unicode_minus'] = False


screen_distance = 1  # Distance from slits to screen (in meters)

# Create a grid of points on the screen
screen_width = 1e-2  # Width of the screen (in meters)
num_points = 10000
screen_positions = np.linspace(-screen_width / 2, screen_width / 2, num_points)
ymax = 1e-2


def wavelength_to_rgb(wavelength, gamma=0.8):
    '''This converts a given wavelength of light to an
    approximate RGB color value. The wavelength must be given
    in nanometers in the range from 380 nm through 750 nm
    (789 THz through 400 THz).
    Based on code by Dan Bruton
    http://www.physics.sfasu.edu/astro/color/spectra.html
    '''

    wavelength = float(wavelength)
    if wavelength >= 380 and wavelength <= 440:
        attenuation = 0.3 + 0.7 * (wavelength - 380) / (440 - 380)
        R = ((-(wavelength - 440) / (440 - 380)) * attenuation) ** gamma
        G = 0.0
        B = (1.0 * attenuation) ** gamma
    elif wavelength >= 440 and wavelength <= 490:
        R = 0.0
        G = ((wavelength - 440) / (490 - 440)) ** gamma
        B = 1.0
    elif wavelength >= 490 and wavelength <= 510:
        R = 0.0
        G = 1.0
        B = (-(wavelength - 510) / (510 - 490)) ** gamma
    elif wavelength >= 510 and wavelength <= 580:
        R = ((wavelength - 510) / (580 - 510)) ** gamma
        G = 1.0
        B = 0.0
    elif wavelength >= 580 and wavelength <= 645:
        R = 1.0
        G = (-(wavelength - 645) / (645 - 580)) ** gamma
        B = 0.0
    elif wavelength >= 645 and wavelength <= 750:
        attenuation = 0.3 + 0.7 * (750 - wavelength) / (750 - 645)
        R = (1.0 * attenuation) ** gamma
        G = 0.0
        B = 0.0
    else:
        R = 0.0
        G = 0.0
        B = 0.0
    return (R, G, B)


# Function to calculate interference pattern
def calculate_interference_pattern(d, theta, wavelength):
    # Calculate the intensity at each point on the screen
    x_values = screen_positions
    y_values = np.linspace(0, ymax, 50)

    X, Y = np.meshgrid(x_values, y_values)

    path_difference = (d * np.sin(theta)) + screen_positions * d / screen_distance
    intensity = (1 + np.cos(2 * np.pi * path_difference / wavelength)) ** 2

    return screen_positions, intensity


# Function to update the plot
def update_plot(value):
    wavelength = int(wl_entry.get()) * 1e-9
    d = float(d_entry.get()) * 1e-3
    theta = float(theta_entry.get())
    theta = np.deg2rad(theta)
    x, I = calculate_interference_pattern(d, np.deg2rad(theta), wavelength)
    ax.clear()
    # ax.plot(x, I, '.')
    I = np.tile(I, (10, 1))
    image = ax.imshow(I, extent=(-screen_width / 2 * 1e3, screen_width / 2 * 1e3, -ymax, ymax), aspect='auto',
                      origin='lower', cmap='coolwarm')
    ax.set_xlabel("观察位置(单位:mm)", fontsize=20)
    ax.set_ylabel("光强", fontsize=20)
    ax.set_yticks([])

    d0 = d / 0.02 * 0.8 * 2
    dd = 0.05
    x1 = [-0.5, -d0 - dd]
    y1 = [0, 0]

    x2 = [-d0, d0]
    x3 = [d0 + dd, 0.5]

    # lightx1 = [-d0 - dd / 2, -d0 - dd / 2 + 0.5 * np.tan(theta)]
    # lighty1 = [0, 0.5]
    # lightx2 = [d0 + dd / 2, d0 + dd / 2 + 0.5 * np.tan(theta)]
    ax2.clear()
    ax2.set_title('实验条件示意图(狭缝均为无限窄)', fontsize=20)
    ax2.axis('off')
    ax2.plot(x1, y1, c='black')
    ax2.plot(x2, y1, c='black')
    ax2.plot(x3, y1, c='black')
    rgb_color = wavelength_to_rgb(int(wavelength * 1e9))
    ax2.annotate('', xy=(-d0 - dd / 2, 0), xytext=(-d0 - dd / 2 + 0.5 * np.tan(theta), 0.5),
                 arrowprops=dict(facecolor=rgb_color, edgecolor=rgb_color, alpha=0.7, shrink=0.05, width=0.1,
                                 headwidth=5, headlength=5))
    ax2.annotate('', xy=(d0 + dd / 2, 0), xytext=(d0 + dd / 2 + 0.5 * np.tan(theta), 0.5),
                 arrowprops=dict(facecolor=rgb_color, edgecolor=rgb_color, alpha=0.7, shrink=0.05, width=0.1,
                                 headwidth=5, headlength=5))
    ax2.annotate('', xy=(0, 0), xytext=(0 + 0.5 * np.tan(theta), 0.5),
                 arrowprops=dict(facecolor=rgb_color, edgecolor=rgb_color, alpha=0.7, shrink=0.05, width=0.1,
                                 headwidth=5, headlength=5))
    ax2.annotate('', xy=(d0 / 2 + dd / 4, 0), xytext=(d0 / 2 + dd / 4 + 0.5 * np.tan(theta), 0.5),
                 arrowprops=dict(facecolor=rgb_color, edgecolor=rgb_color, alpha=0.7, shrink=0.05, width=0.1,
                                 headwidth=5, headlength=5))
    ax2.annotate('', xy=(-d0 / 2 - dd / 4, 0), xytext=(-d0 / 2 - dd / 4 + 0.5 * np.tan(theta), 0.5),
                 arrowprops=dict(facecolor=rgb_color, edgecolor=rgb_color, alpha=0.7, shrink=0.05, width=0.1,
                                 headwidth=5, headlength=5))
    ax2.set_xticks([])
    ax2.set_ylim(-0.2, 0.6)

    canvas.draw()


# Create the main window
root = tk.Tk()
root.title("双缝干涉演示实验")
root.geometry('1200x800')

custom_font = tk.font.Font(family="SimHei", size=20)
root.option_add("*Font", custom_font)
# Create and pack a frame for the controls
control_frame = ttk.Frame(root)
control_frame.pack(padx=10, pady=10)

credit = tk.Label(root, text="by 夏轩哲 复旦大学 物理系22级")
credit.pack()
distance = tk.Label(root, text='屏与狭缝垂直距离:1m')
distance.pack()
# Create and pack labels and entry widgets for parameters
# ttk.Label(control_frame, text="Distance between Slits (meters):").grid(row=0, column=0)
# d_entry = ttk.Entry(control_frame)
# d_entry.grid(row=0, column=1)
# d_entry.insert(0, "0.001")  # Default value
wl_entry = tk.Scale(root, from_=400, to=750, resolution=10, length=400, sliderlength=40, label='波长(400nm-750nm)',
                    orient=tk.HORIZONTAL, command=update_plot)
wl_entry.pack(pady=10)
wl_entry.set(value=600)

d_entry = tk.Scale(root, from_=1e-1, to=5, resolution=1e-2, length=400, sliderlength=40, label='狭缝间距d(0.1mm-5mm)',
                   orient=tk.HORIZONTAL, command=update_plot)
d_entry.pack(pady=10)
d_entry.set(value=1)

theta_entry = tk.Scale(root, from_=-30, to=30, resolution=1, length=400, sliderlength=40, label='入射角(-30°)-(+30°)',
                       orient=tk.HORIZONTAL, command=update_plot)
theta_entry.pack(pady=10)
theta_entry.set(value=0)

# ttk.Label(control_frame, text="Incident Angle (degrees):").grid(row=1, column=0)
# theta_entry = ttk.Entry(control_frame)
# theta_entry.grid(row=1, column=1)
# theta_entry.insert(0, "0")  # Default value

# Create and pack a button to update the plot
# update_button = ttk.Button(control_frame, text="Update Plot", command=update_plot)
# update_button.grid(row=2, columnspan=2)
# Create a Matplotlib figure and plot

# Create a tk button
# tell me the best figure size
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(212)
ax2 = fig.add_subplot(211)
canvas = FigureCanvasTkAgg(fig, master=root)
canvas_widget = canvas.get_tk_widget()
canvas_widget.pack()

# Initialize the plot
update_plot(0)


def on_closing():
    # 在关闭窗口时触发的函数
    root.quit()
    root.destroy()


root.protocol("WM_DELETE_WINDOW", on_closing)
# Start the Tkinter main loop
root.mainloop()