#!/usr/bin/env python
# coding: utf-8

import tkinter as tk
from tkinter.filedialog import asksaveasfile
from tkinter import simpledialog
from PIL import Image, ImageTk
import math
import cmath

SHOW_FUNCTIONS = False

class PolynomialFractalApp:
    # Create and manage the tkinter interface.
    def __init__(self, redraw_on_resize=False):
        self.redraw_on_resize = redraw_on_resize
        self.first_time = True

        # Make the main interface.
        self.window = tk.Tk()
        self.window.title('polynomial_fractal')
        self.window.protocol('WM_DELETE_WINDOW', self.kill_callback)
        self.window.geometry('400x400')

        # Set windowing bounds.
        self.set_window_bounds()

        # Build the rest of the UI.
        self.build_ui()

        # Initially do not ignore changes to Z0 and Zn+1.
        self.ignore_changes = False

        # Display the window.
        self.window.focus_force()
        self.window.mainloop()

    def build_ui(self):
        self.build_menus()
        self.build_parameters()
        self.build_canvas()

        # Redraw on F5.
        self.window.bind('<F5>', lambda event: self.set_scale('redraw'))

        # Save on Ctrl+S.
        self.window.bind('<Control-s>', lambda event: self.save_as())

        # Start with the Mandelbrot set.
        self.preset('Mandelbrot Set', False)

    def build_menus(self):
        '''Build the menus.'''
        # Menubar
        menubar = tk.Menu(self.window)
        self.window.config(menu=menubar)
        self.window.option_add('*tearOff', False) # Prevent tear off in menus.

        # File menu
        file_menu = tk.Menu(menubar)
        file_menu.add_command(label='Save As...', accelerator="Ctrl+S", command=self.save_as)
        file_menu.add_separator()
        file_menu.add_command(label='Exit', command=self.kill_callback)
        menubar.add_cascade(label='File', menu=file_menu)

        # Scale menu
        scale_menu = tk.Menu(menubar)
        scale_menu.add_command(label='Redraw', accelerator="F5", command=lambda: self.set_scale('redraw'))
        scale_menu.add_separator()
        scale_menu.add_command(label='x2', command=lambda: self.set_scale('2'))
        scale_menu.add_command(label='x4', command=lambda: self.set_scale('4'))
        scale_menu.add_command(label='x8', command=lambda: self.set_scale('8'))
        scale_menu.add_command(label='x16', command=lambda: self.set_scale('16'))
        scale_menu.add_command(label='Full Scale', command=lambda: self.set_scale('full scale'))
        scale_menu.add_separator()
        scale_menu.add_command(label='Set Bounds...', command=self.set_bounds)
        menubar.add_cascade(label='Scale', menu=scale_menu)

        # Parameters menu
        parameters_menu = tk.Menu(menubar)
        self.use_smooth_colors = tk.IntVar(value=1)
        self.repeat_colors = tk.IntVar(value=1)
        parameters_menu.add_checkbutton(label="Use Smooth Colors", variable=self.use_smooth_colors, command=lambda: self.set_scale('redraw'))
        parameters_menu.add_checkbutton(label="Repeat Colors", variable=self.repeat_colors, command=lambda: self.set_scale('redraw'))
        menubar.add_cascade(label='Parameters', menu=parameters_menu)

        # Presets menu
        self.preset_selection_var = tk.IntVar(value=-1)
        presets_menu = tk.Menu(menubar)
        menubar.add_cascade(label='Presets', menu=presets_menu)
        i= 0
        presets_menu.add_radiobutton(
            label="Mandelbrot Set",
            variable=self.preset_selection_var, value=i,
            command=lambda: self.preset('Mandelbrot Set'))
        i += 1

        julia_menu = tk.Menu(presets_menu)
        presets_menu.add_cascade(menu=julia_menu, label='Julia Sets')
        names = ['Basic Julia Set'] + [f'Julia Set {i}' for i in range(1, 13)]
        for name in names:
            julia_menu.add_radiobutton(
                label=name,
                variable=self.preset_selection_var, value=i,
                command=lambda n=name: self.preset(n))
            i += 1

        juliaish_menu = tk.Menu(presets_menu)
        presets_menu.add_cascade(menu=juliaish_menu, label='Juliaish Sets')
        names = [f'Juliaish Set {i}' for i in range(1, 5)]
        for name in names:
            juliaish_menu.add_radiobutton(
                label=name,
                variable=self.preset_selection_var, value=i,
                command=lambda n=name: self.preset(n))
            i += 1

        other_menu = tk.Menu(presets_menu)
        presets_menu.add_cascade(menu=other_menu, label='Other')
        names = [f'Other {i}' for i in range(1, 6)]
        for name in names:
            other_menu.add_radiobutton(
                label=name,
                variable=self.preset_selection_var, value=i,
                command=lambda n=name: self.preset(n))
            i += 1

    def build_parameters(self):
        '''Build the basic parameter widgets.'''
        frame = tk.Frame(self.window)
        frame.pack(side=tk.TOP, padx=5, pady=5, anchor=tk.NW,
                   expand=False, fill=tk.X)

        self.max_steps_var = tk.StringVar(value='64')
        entry = tk.Entry(frame, width=4, textvariable=self.max_steps_var,
                         justify=tk.RIGHT)
        entry.pack(side=tk.RIGHT)
        label = tk.Label(frame, text='Max Steps:', width=8, anchor=tk.W)
        label.pack(side=tk.RIGHT, anchor=tk.W, padx=(20,0))

        label = tk.Label(frame, text='z0:', width=5, anchor=tk.W)
        label.pack(side=tk.LEFT, anchor=tk.W)
        self.z0_var = tk.StringVar()
        self.z0_var.trace_add('write', self.text_changed)
        entry = tk.Entry(frame, textvariable=self.z0_var)
        entry.pack(side=tk.LEFT, expand=True, fill=tk.X)

        frame = tk.Frame(self.window)
        frame.pack(side=tk.TOP, padx=5, anchor=tk.NW,
                   expand=False, fill=tk.X)
        label = tk.Label(frame, text='zn+1:', width=5, anchor=tk.W)
        label.pack(side=tk.LEFT, anchor=tk.NW)
        self.zn_plus_1_var = tk.StringVar()
        self.zn_plus_1_var.trace_add('write', self.text_changed)
        entry = tk.Entry(frame, textvariable=self.zn_plus_1_var)
        entry.pack(side=tk.LEFT, expand=True, fill=tk.X)

        self.max_mag_var = tk.StringVar(value='2')
        entry = tk.Entry(frame, width=4, textvariable=self.max_mag_var,
                         justify=tk.RIGHT)
        entry.pack(side=tk.RIGHT)
        label = tk.Label(frame, text='Max Mag:', width=8, anchor=tk.W)
        label.pack(side=tk.RIGHT, anchor=tk.W, padx=(20,0))

    def build_canvas(self):
        '''Make the canvas.'''
        # Drawing canvas
        self.canvas = tk.Canvas(self.window, bg='black', borderwidth=1,
                                highlightthickness=0, cursor='cross')
        self.canvas.pack(side=tk.TOP, padx=5, pady=5,
                         expand=True, fill=tk.BOTH)

        # Initially we are not selecting an area.
        self.start_pos = None
        self.end_pos = None

        # Track Button-1 down, Button-1 motion, and Button-1 up.
        self.canvas.bind("<Button-1>", self.mouse_down)
        self.canvas.bind("<B1-Motion>", self.mouse_move)
        self.canvas.bind("<ButtonRelease-1>", self.mouse_up)

        # Watch for resize.
        self.canvas.bind("<Configure>", self.configured)

    def text_changed(self, *_):
        '''The user changed Z0 or Zn+1. Clear the preset menu checkboxes.'''
        if self.ignore_changes:
            return
        self.preset_selection_var.set(-1)

    def preset(self, name, redraw=True):
        '''Load equations for a preset curve.'''
        # print(f'preset({name})')
        self.ignore_changes = True

        if name == 'Mandelbrot Set':
            self.z0_var.set('0 + 0j')
            self.zn_plus_1_var.set('zn ** 2 + x + y * 1j')

        elif name == 'Basic Julia Set':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + 1j')
        elif name == 'Julia Set 1':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (-0.4 + 0.6j)')
        elif name == 'Julia Set 2':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (0.285 + 0j)')
        elif name == 'Julia Set 3':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (0.285 + 0.01j)')
        elif name == 'Julia Set 4':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (0.45 + 0.1428j)')
        elif name == 'Julia Set 5':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (-0.70176 - 0.3842j)')
        elif name == 'Julia Set 6':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (-0.835 - 0.2321j)')
        elif name == 'Julia Set 7':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (-0.8 + 0.156j)')
        elif name == 'Julia Set 8':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (-0.7269 + 0.1889j)')
        elif name == 'Julia Set 9':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (0.0 - 0.8j)')
        elif name == 'Julia Set 10':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (-0.74543 + 0.11301j)')
        elif name == 'Julia Set 11':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (-0.75 + 0.11j)')
        elif name == 'Julia Set 12':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (-0.1 + 0.651j)')

        elif name == 'Juliaish Set 1':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('(1 - zn ** 3 / 6) / (zn - zn ** 2 / 2) ** 2 + 0.15')
        elif name == 'Juliaish Set 2':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('1 - zn ** 2 + zn ** 5 / (2 + 4 * zn)')
        elif name == 'Juliaish Set 3':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('(1 - zn ** 2 + zn ** 5) / (2 + 4 * zn)')
        elif name == 'Juliaish Set 4':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('1 - zn ** 2 + zn ** 5 / (2 + 4 * zn) + 0.2')

        elif name == 'Other 1':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 3 - 2 * zn ** 2 + zn')
        elif name == 'Other 2':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 4 - 3 * zn ** 2')
        elif name == 'Other 3':#@@@
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 1.5 + x + y * 1j')
        elif name == 'Other 4':
            self.z0_var.set('y + x * 1j')
            self.zn_plus_1_var.set('zn ** 2.5 + y + x * 1j')
        elif name == 'Other 5':
            self.z0_var.set('x + y * 1j')
            self.zn_plus_1_var.set('zn ** 2 + (x - y * 1j) / (x**2 + y**2)**0.5')

        # Redraw.
        if redraw:
            self.set_scale('redraw')

        self.ignore_changes = False

    def set_window_bounds(self):
        '''Set the world coordinate bounds.'''
        # full_wXXX gives the full scale coordinates for the window.
        # desired_wXXX gives the desired world window coordinates.
        # wXXX gives the current world window coordinates.
        # vXXX gives the current viewport coordinates.
        self.full_wxmin = -2.2
        self.full_wxmax = 1.0
        self.full_wymin = -1.2
        self.full_wymax = 1.2
        self.min_width = 1e-10
        self.min_height = 1e-10
        self.desired_wxmin = self.full_wxmin
        self.desired_wxmax = self.full_wxmax
        self.desired_wymin = self.full_wymin
        self.desired_wymax = self.full_wymax

    def save_as(self):
        file = asksaveasfile(initialfile = 'Mandelbrot.png',
            defaultextension='.png', filetypes=[('Image Files', '*.png *.jpg, *.bmp'), ('All Files', '*.*')])
        if file is None: return
        self.image.save(file.name)

    def set_bounds(self):
        '''Let the user set the bounds.'''
        old_bounds = f'{self.wxmin}, {self.wxmax}, {self.wymin}, {self.wymax}'
        new_bounds = simpledialog.askstring("Bounds", "Enter bounds:\t\t\t\t\t\t\t\t", initialvalue=old_bounds)
        if new_bounds is None: return
        new_bounds = new_bounds.replace(',', ' ')
        self.desired_wxmin = float(new_bounds.split()[0])
        self.desired_wxmax = float(new_bounds.split()[1])
        self.desired_wymin = float(new_bounds.split()[2])
        self.desired_wymax = float(new_bounds.split()[3])
        self.set_scale('redraw')

    def set_scale(self, scale_factor):
        '''The user has invoked a Scale menu item. Set the new world bounds and then redraw.'''
        # Get the desired world width and height.
        if scale_factor == 'redraw':
            # Just redraw. Leave the bounds unchanged.
            pass
        elif scale_factor == 'full scale':
            # Reset to full scale.
            self.desired_wxmin = self.full_wxmin
            self.desired_wxmax = self.full_wxmax
            self.desired_wymin = self.full_wymin
            self.desired_wymax = self.full_wymax
        else:
            # Scale by a number.
            scale = int(scale_factor)
            wid = (self.desired_wxmax - self.desired_wxmin) * scale
            hgt = (self.desired_wymax - self.desired_wymin) * scale
            cx = (self.desired_wxmax + self.desired_wxmin) / 2
            cy = (self.desired_wymax + self.desired_wymin) / 2
            self.desired_wxmin = cx - wid / 2
            self.desired_wxmax = self.desired_wxmin + wid
            self.desired_wymin = cy - hgt / 2
            self.desired_wymax = self.desired_wymin + hgt

        # Do not bound results for non-Mandelbrot sets.
        # # If the width/height is larger than that of
        # # full scale, use the full scale width/height.
        # if self.desired_wxmax  - self.desired_wxmin > self.full_wxmax - self.full_wxmin:
        #     self.desired_wxmax = self.full_wxmax
        #     self.desired_wxmin = self.full_wxmin
        # if self.desired_wymax  - self.desired_wymin > self.full_wymax - self.full_wymin:
        #     self.desired_wymax = self.full_wymax
        #     self.desired_wymin = self.full_wymin

        # Get the desired window bounds.
        wxmin = self.desired_wxmin
        wxmax = self.desired_wxmax
        wymin = self.desired_wymin
        wymax = self.desired_wymax

        # Get the width and height.
        w_wid = (wxmax - wxmin)
        w_hgt = (wymax - wymin)

        # Adjust to preserve aspect ratio.
        w_aspect = w_wid / w_hgt
        v_wid = self.vxmax - self.vxmin
        v_hgt = self.vymax - self.vymin
        v_aspect = v_wid / v_hgt
        if w_aspect > v_aspect:
            # World window bounds are too short and wide. Make taller.
            w_hgt = w_wid / v_aspect
        else:
            # World window bounds are too tall and thin. Make wider.
            w_wid = w_hgt * v_aspect

        # Set the new world window bounds.
        cx = (wxmin + wxmax) / 2
        cy = (wymin + wymax) / 2
        self.wxmin = cx - w_wid / 2
        self.wxmax = self.wxmin + w_wid
        self.wymin = cy - w_hgt / 2
        self.wymax = self.wymin + w_hgt

        # Print the bounds so we can use them later if desired.
        print(f'{self.wxmin}, {self.wxmax}, {self.wymin}, {self.wymax}')

        # Redraw.
        self.draw()

    def mouse_down(self, event):
        self.start_pos = (event.x, event.y)
        self.end_pos = (event.x, event.y)

    def mouse_move(self, event):
        # Ignore extraneous events.
        if self.start_pos is None: return

        # Delete any previous selection rectangle.
        self.canvas.delete('selection_rectangle')

        # Create a new selection rectangle.
        self.end_pos = (event.x, event.y)
        self.canvas.create_rectangle(self.start_pos, self.end_pos,
            outline='red', tag='selection_rectangle')
        self.canvas.create_rectangle(
            self.start_pos, self.end_pos, dash=(1, 5),
            outline='yellow', tag='selection_rectangle')

    def mouse_up(self, event):
        # Ignore extraneous events.
        if self.start_pos is None: return
        self.canvas.delete('selection_rectangle')

        # Get the selected bounds.
        xmin = min(self.start_pos[0], self.end_pos[0])
        xmax = max(self.start_pos[0], self.end_pos[0])
        ymin = min(self.start_pos[1], self.end_pos[1])
        ymax = max(self.start_pos[1], self.end_pos[1])
        self.start_pos = None
        self.end_pos = None
        
        # Do nothing if the selected area is empty.
        if xmin == xmax or ymin == ymax: return

        # Call area_selected to redraw.
        self.area_selected(xmin, xmax, ymin, ymax)

    def area_selected(self, xmin, xmax, ymin, ymax):
        # Convert into world coordinates.
        wxmin, wymin = self.d_to_w(xmin, ymin)
        wxmax, wymax = self.d_to_w(xmax, ymax)

        # Make sure the width and height aren't too small.
        if wxmax - wxmin < self.min_width:
            cx = (wxmax + wxmin) / 2
            wxmin = cx - self.min_width / 2
            wxmax = wxmin + self.min_width
        if wymax - wymin < self.min_height:
            cy = (wymax + wymin) / 2
            wymin = cy - self.min_height / 2
            wymax = wymin + self.min_height

        # Set the desired world coordinates.
        self.desired_wxmin = wxmin
        self.desired_wxmax = wxmax
        self.desired_wymin = wymin
        self.desired_wymax = wymax

        # Redraw.
        self.set_scale('redraw')

    def configured(self, event):
        '''The window may have been resized. Reset the viewport coordinates and redraw.'''
        # Do nothing unless the window is in a normal or zoomed state.
        if self.window.state() != 'normal' and self.window.state() != 'zoomed': return

        self.vxmin = 0
        self.vxmax = self.canvas.winfo_width() - 1
        self.vymin = 0
        self.vymax = self.canvas.winfo_height() - 1

        # Redraw.
        if self.first_time or self.redraw_on_resize:
            self.set_scale('redraw')
            self.first_time = False

    def kill_callback(self):
        self.window.destroy()

    def draw(self):
        '''Draw an image.'''
        # Get basic fractal parameters.
        max_iterations = int(self.max_steps_var.get())
        escape_radius = int(self.max_mag_var.get())

        # Set scale factors for drawing.
        self.x_scale = (self.vxmax - self.vxmin) / (self.wxmax - self.wxmin)
        self.y_scale = (self.vymax - self.vymin) / (self.wymax - self.wymin)

        # Make an image to fit the viewport.
        mode = 'RGB'
        image_wid = self.vxmax - self.vxmin
        image_hgt = self.vymax - self.vymin
        bg_color = (0, 0, 0)
        self.image = Image.new(mode, (image_wid, image_hgt), bg_color)

        # Access the image's pixels.
        self.pixels = self.image.load()

        # Get the viewport's bounds in world coordinates.
        w_vxmin, w_vymin = self.d_to_w(0, 0)
        w_vxmax, w_vymax = self.d_to_w(image_wid - 1, image_hgt - 1)

        # Get the Z0 and Zn+1 functions.
        function_def = \
f'''def func_z0(x, y):
    try:
        return {self.z0_var.get()}
    except:
        return 1000'''
        exec(function_def, globals())
        if SHOW_FUNCTIONS:
            print(function_def, '\n')

        function_def = \
f'''def func_zn_plus_1(zn, x, y):
    try:
        return {self.zn_plus_1_var.get()}
    except:
        return 1000'''
        exec(function_def, globals())
        if SHOW_FUNCTIONS:
            print(function_def, '\n')

        # Work until |Zn| >= escape_radius
        log_escape = math.log(escape_radius)

        # Define the colors.
        black      = (  0,   0,   0)
        red        = (255,   0,   0)
        orange     = (255, 128,   0)
        yellow     = (255, 255,   0)
        dark_green = (  0, 192,   0)
        cyan       = (  0, 255, 255)
        blue       = (  0,   0, 255)
        colors = ( black, red, orange, yellow, dark_green, cyan, blue )

        # Loop over the pixels.
        for pix_x in range(image_wid):
            for pix_y in range(image_hgt):                
                # Calculate this point's color.
                x, y = self.d_to_w(pix_x, pix_y)
                z0 = func_z0(x, y) # analysis:ignore 

                z = z0
                iteration = 0
                while ((iteration < max_iterations) and (abs(z) < escape_radius)):
                    # Calculate Z(clr).
                    z = func_zn_plus_1(z, x, y) # analysis:ignore
                    iteration += 1

                # Set the pixel's color.
                if (iteration >= max_iterations):
                    color = black
                else:
                    if self.use_smooth_colors.get():
                        # Reduce the error in mu.
                        for i in range(3):
                            z = func_zn_plus_1(z, x, y) # analysis:ignore
                            iteration += 1
                        try:
                            # cmath.log provides a more interesting picture
                            # because it is defined for negative numbers
                            # while math.log is not.
                            mu = abs(iteration + 1 - 
                                     cmath.log(cmath.log(abs(z))) / log_escape)
                        except:
                            mu = 0

                        if not self.repeat_colors.get():
                            mu = mu / max_iterations * len(colors)
                        color = self.get_color(colors, mu)
                    else:
                        color = colors[iteration % len(colors)]

                self.pixels[pix_x, pix_y] = color

        # Convert the PIL image into an ImageTk.PhotoImage.
        self.photo_image = ImageTk.PhotoImage(self.image)

        # Display the result.
        self.canvas.delete('all')
        self.display_image = self.canvas.create_image(0, 0, anchor=tk.NW, image=self.photo_image)

    def get_color(self, colors, mu):
        '''Get a color for this pixel.'''
        clr1 = int(mu)
        t2 = mu - clr1
        t1 = 1 - t2
        clr1 = clr1 % len(colors)
        clr2 = (clr1 + 1) % len(colors)

        r = int(colors[clr1][0] * t1 + colors[clr2][0] * t2)
        g = int(colors[clr1][1] * t1 + colors[clr2][1] * t2)
        b = int(colors[clr1][2] * t1 + colors[clr2][2] * t2)
        return (r, g, b)

    def w_to_d(self, x, y):
        '''Convert a point from world to device coordinates.'''
        vx = self.vxmin + (x - self.wxmin) * self.x_scale
        vy = self.vymin + (y - self.wymin) * self.y_scale
        return vx, vy

    def d_to_w(self, x, y):
        '''Convert a point from device to world coordinates.'''
        vx = self.wxmin + (x - self.vxmin) / self.x_scale
        vy = self.wymin + (y - self.vymin) / self.y_scale
        return vx, vy


# In[18]:


# Do not allow auto-redraw on resize.
PolynomialFractalApp(False)
print('Done')


# In[ ]:




