import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import matplotlib.colors as mcolors

class FractalExplorer:
    def __init__(self):
        # --- 画布设置 ---
        self.width, self.height = 800, 600
        self.fig, self.ax = plt.subplots(figsize=(10, 8))
        self.fig.canvas.manager.set_window_title('Python 分形图形浏览器 - 滚轮缩放/拖拽平移')
        
        # --- 曼德博集合参数 ---
        # 初始视图范围 (xmin, xmax, ymin, ymax)
        self.x_min, self.x_max = -2.5, 1.0
        self.y_min, self.y_max = -1.25, 1.25
        
        # 最大迭代次数 (决定细节程度)
        self.max_iter = 100
        
        # --- 初始化界面 ---
        self.ax.set_title("曼德博集合 (Mandelbrot Set)", fontsize=16)
        self.ax.set_xlabel("实部 (Re)")
        self.ax.set_ylabel("虚部 (Im)")
        
        # 绑定鼠标事件
        self.fig.canvas.mpl_connect('scroll_event', self.on_scroll)
        self.fig.canvas.mpl_connect('button_press_event', self.on_click)
        
        # 绘制初始图形
        self.draw_fractal()
        
        # 显示操作说明
        self.show_help()

    def compute_mandelbrot(self):
        """
        核心算法：使用 NumPy 向量化计算曼德博集合
        这比使用 for 循环快得多
        """
        # 1. 生成网格坐标
        # 在 x 和 y 轴上生成线性空间
        x = np.linspace(self.x_min, self.x_max, self.width)
        y = np.linspace(self.y_min, self.y_max, self.height)
        # 创建网格 (X, Y 都是 800x600 的矩阵)
        X, Y = np.meshgrid(x, y)
        
        # 2. 初始化复数 c = x + iy
        C = X + 1j * Y
        
        # 3. 初始化 z = 0
        Z = np.zeros_like(C)
        
        # 4. 迭代计算
        # 我们需要记录每个点是在第几次迭代“逃逸”（模长 > 2）的
        # 如果达到最大迭代次数仍未逃逸，则认为在集合内（黑色）
        
        # 用于存储每个点的迭代次数
        iteration_counts = np.zeros(C.shape, dtype=int)
        # 用于标记哪些点已经逃逸（不再计算）
        mask = np.zeros(C.shape, dtype=bool)
        
        for i in range(self.max_iter):
            # 核心公式：z = z^2 + c
            Z = Z**2 + C
            
            # 检查逃逸条件：|z| > 2
            escaped = np.abs(Z) > 2.0
            
            # 找出本次新逃逸的点（之前没逃逸，现在逃逸了）
            new_escape = escaped & (~mask)
            
            # 记录这些点的迭代次数
            iteration_counts[new_escape] = i
            
            # 更新掩码
            mask |= escaped
            
            # 如果所有点都逃逸了，提前结束
            if mask.all():
                break
        
        return iteration_counts

    def draw_fractal(self):
        """绘制图形"""
        # 清除旧图
        self.ax.clear()
        
        # 计算数据
        data = self.compute_mandelbrot()
        
        # 使用特定的颜色映射 (cmap) 来美化边缘
        # 'hot', 'magma', 'inferno', 'twilight' 都是很好的选择
        cmap = plt.cm.get_cmap('magma')
        
        # 显示图像
        # extent 定义了图像在坐标轴上的范围
        self.ax.imshow(data, cmap=cmap, extent=[self.x_min, self.x_max, self.y_min, self.y_max], origin='lower')
        
        # 设置标题显示当前视图范围
        self.ax.set_title(f"曼德博集合 - 缩放级别: {self.get_zoom_level()}x")
        
        # 刷新画布
        self.fig.canvas.draw_idle()

    def get_zoom_level(self):
        # 简单的缩放级别计算
        return int(1.0 / (self.x_max - self.x_min) * 10) / 10

    def on_scroll(self, event):
        """处理鼠标滚轮事件 (缩放)"""
        if event.inaxes != self.ax:
            return

        # 缩放系数
        base_scale = 1.1
        if event.button == 'up':
            scale_factor = 1 / base_scale
        elif event.button == 'down':
            scale_factor = base_scale
        else:
            return

        # 计算鼠标位置相对于坐标轴的比例
        xdata = event.xdata
        ydata = event.ydata
        
        if xdata is None or ydata is None:
            return

        # 计算新的范围
        # 1. 计算当前宽高
        dx = self.x_max - self.x_min
        dy = self.y_max - self.y_min
        
        # 2. 计算鼠标位置在 0-1 之间的比例
        mx = (xdata - self.x_min) / dx
        my = (ydata - self.y_min) / dy
        
        # 3. 缩小范围
        new_dx = dx * scale_factor
        new_dy = dy * scale_factor
        
        # 4. 更新边界，保持鼠标位置不变
        self.x_min = xdata - mx * new_dx
        self.x_max = xdata + (1 - mx) * new_dx
        self.y_min = ydata - my * new_dy
        self.y_max = ydata + (1 - my) * new_dy
        
        # 随着缩放加深，增加迭代次数以看清细节
        if scale_factor < 1: # 放大
            self.max_iter = min(1000, self.max_iter + 10)
        
        self.draw_fractal()

    def on_click(self, event):
        """处理鼠标点击 (简单的平移或重置)"""
        if event.button == 3: # 右键重置
            self.x_min, self.x_max = -2.5, 1.0
            self.y_min, self.y_max = -1.25, 1.25
            self.max_iter = 100
            self.draw_fractal()
        
        # 左键拖拽平移逻辑较为复杂，这里简化为点击重绘
        # Matplotlib 自带了导航工具栏，通常不需要手写拖拽逻辑

    def show_help(self):
        print("="*40)
        print("🎨 Python 分形图形浏览器已启动")
        print("="*40)
        print("操作指南:")
        print("1. 鼠标滚轮向上: 放大 (探索细节)")
        print("2. 鼠标滚轮向下: 缩小 (查看全貌)")
        print("3. 鼠标右键点击: 重置视图")
        print("4. 关闭窗口: 退出程序")
        print("="*40)
        print("正在渲染初始图形...")

if __name__ == "__main__":
    explorer = FractalExplorer()
    plt.show()