import tkinter as tk
from tkinter import ttk, messagebox
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import matplotlib.cm as cm

# 主窗口类
class ComplexGraphApp:
    def __init__(self, root):
        self.root = root
        self.root.title("复杂图形绘制工具 - 曼德博分形")
        self.root.geometry("1000x700")
        
        # 初始化参数
        self.max_iter = tk.IntVar(value=100)  # 最大迭代次数
        self.zoom = tk.DoubleVar(value=1.0)   # 缩放比例
        self.center_x = tk.DoubleVar(value=-0.5)  # 中心X坐标
        self.center_y = tk.DoubleVar(value=0.0)   # 中心Y坐标
        
        # 创建界面布局
        self._create_widgets()
        # 初始绘制图形
        self.update_graph()

    def _create_widgets(self):
        # 1. 控制面板（左侧）
        control_frame = ttk.Frame(self.root, padding="10")
        control_frame.pack(side=tk.LEFT, fill=tk.Y)
        
        # 参数设置区域
        ttk.Label(control_frame, text="参数设置", font=("Arial", 12, "bold")).pack(pady=10)
        
        # 最大迭代次数
        ttk.Label(control_frame, text="最大迭代次数:").pack(anchor=tk.W)
        iter_spinbox = ttk.Spinbox(
            control_frame, from_=50, to=500, textvariable=self.max_iter, 
            width=10, command=self.update_graph
        )
        iter_spinbox.pack(pady=5)
        
        # 缩放比例
        ttk.Label(control_frame, text="缩放比例:").pack(anchor=tk.W)
        zoom_spinbox = ttk.Spinbox(
            control_frame, from_=0.1, to=5.0, increment=0.1, textvariable=self.zoom,
            width=10, command=self.update_graph
        )
        zoom_spinbox.pack(pady=5)
        
        # 重置按钮
        reset_btn = ttk.Button(
            control_frame, text="重置参数", command=self.reset_parameters
        )
        reset_btn.pack(pady=20)
        
        # 2. 绘图区域（右侧）
        self.fig, self.ax = plt.subplots(figsize=(8, 6), dpi=100)
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.root)
        self.canvas.get_tk_widget().pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)

    # 修复后的曼德博集合计算函数
    def mandelbrot(self, xmin, xmax, ymin, ymax, width, height, max_iter):
        # 生成复平面坐标
        x = np.linspace(xmin, xmax, width)
        y = np.linspace(ymin, ymax, height)
        X, Y = np.meshgrid(x, y)
        c = X + 1j * Y  # 形状 (height, width) = (600, 800)
        
        # 初始化数组（保持二维形状，避免广播错误）
        z = np.zeros(c.shape, dtype=np.complex128)  # 和c同形状
        div_time = np.zeros(c.shape, dtype=int)      # 迭代次数记录
        mask = np.ones(c.shape, dtype=bool)          # 未发散的点标记
        
        # 迭代计算
        for i in range(max_iter):
            # 只计算未发散的点，避免不必要的运算
            z[mask] = z[mask] **2 + c[mask]
            # 判断是否发散（模长>2）
            diverged = np.abs(z) > 2
            # 记录发散时的迭代次数
            div_time[diverged & mask] = i
            # 更新未发散的点标记
            mask[diverged] = False
            # 如果所有点都发散了，提前退出
            if not np.any(mask):
                break
        
        return div_time

    # 更新图形
    def update_graph(self):
        # 清空当前轴
        self.ax.clear()
        
        # 获取参数
        max_iter = self.max_iter.get()
        zoom = self.zoom.get()
        center_x = self.center_x.get()
        center_y = self.center_y.get()
        
        # 计算绘图范围
        scale = 2.0 / zoom
        xmin = center_x - scale
        xmax = center_x + scale
        ymin = center_y - scale
        ymax = center_y + scale
        
        # 生成曼德博集合数据
        width, height = 800, 600
        mandelbrot_set = self.mandelbrot(xmin, xmax, ymin, ymax, width, height, max_iter)
        
        # 绘制图形
        img = self.ax.imshow(
            mandelbrot_set, extent=(xmin, xmax, ymin, ymax),
            cmap=cm.plasma, interpolation="bilinear", origin="lower"  # 新增origin参数，修正坐标方向
        )
        self.fig.colorbar(img, ax=self.ax, label="迭代次数")
        self.ax.set_title(f"曼德博分形 (迭代次数: {max_iter}, 缩放: {zoom}x)")
        self.ax.set_xlabel("Re(c)")
        self.ax.set_ylabel("Im(c)")
        
        # 更新画布
        self.canvas.draw()

    # 重置参数
    def reset_parameters(self):
        self.max_iter.set(100)
        self.zoom.set(1.0)
        self.center_x.set(-0.5)
        self.center_y.set(0.0)
        self.update_graph()
        messagebox.showinfo("提示", "参数已重置为默认值！")

# 程序入口
if __name__ == "__main__":
    root = tk.Tk()
    app = ComplexGraphApp(root)
    root.mainloop()