import numpy as np
import matplotlib.pyplot as plt

class MandelbrotExplorer:
    def __init__(self, width=800, height=800, max_iter=100):
        self.width = width
        self.height = height
        self.max_iter = max_iter
        
        # 初始范围：复平面上的坐标
        self.x_min, self.x_max = -2.0, 0.5
        self.y_min, self.y_max = -1.25, 1.25
        
        self.fig, self.ax = plt.subplots(figsize=(10, 10))
        self.fig.canvas.manager.set_window_title('无限分形浏览器：点击左键放大，右键缩小')
        
        # 绑定鼠标点击事件
        self.cid = self.fig.canvas.mpl_connect('button_press_event', self.on_click)
        
        self.render()
        plt.show()

    def mandelbrot(self, x_min, x_max, y_min, y_max, w, h, max_iter):
        """核心数学计算：计算曼德博集合矩阵"""
        x = np.linspace(x_min, x_max, w)
        y = np.linspace(y_min, y_max, h)
        X, Y = np.meshgrid(x, y)
        C = X + 1j * Y
        Z = np.zeros_like(C)
        img = np.zeros(C.shape, dtype=float)
        mask = np.full(C.shape, True, dtype=bool)

        for i in range(max_iter):
            Z[mask] = Z[mask] * Z[mask] + C[mask]
            # 找出已经发散的点
            diverged = np.abs(Z) > 2
            # 记录发散时的迭代次数（用于着色）
            img[mask & diverged] = i
            # 更新掩码，只计算未发散的点
            mask[diverged] = False
            
        return img

    def render(self):
        """渲染并绘制图像"""
        print(f"正在计算范围: X[{self.x_min:.10f}, {self.x_max:.10f}]...")
        img = self.mandelbrot(self.x_min, self.x_max, self.y_min, self.y_max, 
                             self.width, self.height, self.max_iter)
        
        self.ax.clear()
        # 使用 'magma' 渐变色，看起来非常深邃复杂
        self.ax.imshow(img, extent=[self.x_min, self.x_max, self.y_min, self.y_max], 
                       cmap='magma', origin='lower')
        self.ax.set_axis_off()
        self.fig.canvas.draw()

    def on_click(self, event):
        """处理鼠标点击缩放"""
        if event.xdata is None or event.ydata is None:
            return
        
        # 计算当前跨度
        x_range = self.x_max - self.x_min
        y_range = self.y_max - self.y_min
        
        if event.button == 1:  # 左键放大
            zoom_factor = 0.2
            self.max_iter += 20 # 放大时增加迭代次数以保持细节
        elif event.button == 3:  # 右键缩小
            zoom_factor = 5.0
            self.max_iter = max(50, self.max_iter - 20)
        else:
            return

        # 以点击点为中心重新计算范围
        self.x_min = event.xdata - (event.xdata - self.x_min) * zoom_factor
        self.x_max = event.xdata + (self.x_max - event.xdata) * zoom_factor
        self.y_min = event.ydata - (event.ydata - self.y_min) * zoom_factor
        self.y_max = event.ydata + (self.y_max - event.ydata) * zoom_factor
        
        self.render()

if __name__ == "__main__":
    MandelbrotExplorer()