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

# In[17]:


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

# Layout constants.
MARGIN = 5
WINDOW_WID = 400
WINDOW_HGT = 400

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

        # Make the main interface.
        self.window = tk.Tk()
        self.window.title('julia_set')
        self.window.protocol('WM_DELETE_WINDOW', self.kill_callback)
        self.window.geometry(f'{WINDOW_WID}x{WINDOW_HGT}')

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

        # Set windowing bounds.
        self.set_window_bounds()

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

    def build_ui(self):
        # Menubar
        menubar = tk.Menu(self.window)
        self.window.config(menu=menubar)

        # File menu
        file_menu = tk.Menu(menubar, tearoff=False)
        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, tearoff=False)
        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, tearoff=False)
        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'))
        parameters_menu.add_command(label="Max Iterations", command=self.set_max_iterations)
        menubar.add_cascade(label='Parameters', menu=parameters_menu)

        # C.
        frame = tk.Frame(self.window)
        frame.pack(side=tk.TOP, padx=5, pady=(3,0), anchor=tk.NW)
        label = tk.Label(frame, text='C:')
        label.pack(side=tk.LEFT)
        self.c_re_var = tk.StringVar(value='-0.5251993')
        entry = tk.Entry(frame, justify=tk.RIGHT, width=12, textvariable=self.c_re_var)
        entry.pack(side=tk.LEFT, padx=5)
        label = tk.Label(frame, text='+')
        label.pack(side=tk.LEFT)
        self.c_im_var = tk.StringVar(value='-0.5251993')
        entry = tk.Entry(frame, justify=tk.RIGHT, width=12, textvariable=self.c_im_var)
        entry.pack(side=tk.LEFT, padx=5)
        label = tk.Label(frame, text='j')
        label.pack(side=tk.LEFT)

        # Drawing canvas
        self.canvas = tk.Canvas(self.window, bg='black', borderwidth=1, highlightthickness=0, cursor='cross')
        self.canvas.pack(side=tk.TOP, padx=MARGIN, pady=MARGIN, 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)

        # 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())

    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_max_iterations(self):
        '''Let the user set max_iterations.'''
        old_max = f'{self.max_iterations}'
        new_max = simpledialog.askinteger("Max Iterations", "Max Iterations:\t\t\t\t\t\t\t\t", initialvalue=old_max)
        if new_max is None: return
        self.max_iterations = new_max
        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

        # We don't restrict this to the full scale because
        # we don't know how big the interesting area may be.

        # 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.'''
        # 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)

        # Work until |Zn| >= escape_radius
        escape_radius = 2.0
        log_escape = math.log(escape_radius)
        try:
            c_re = float(self.c_re_var.get())
        except:
            print(f'Error parsing "{self.c_re_var.get()}"')
            return
        try:
            c_im = float(self.c_im_var.get())
        except:
            print(f'Error parsing "{self.c_im_var.get()}"')
            return
        c = c_re + c_im * 1j

        # 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 x in range(image_wid):
            for y in range(image_hgt):
                # Calculate this point's color.
                re, im = self.d_to_w(x, y)
                z = re + im * 1j
                iteration = 0
                while ((iteration < self.max_iterations) and (abs(z) < escape_radius)):
                    # Calculate Z(clr).
                    z = z * z + c
                    iteration += 1

                # Set the pixel's color.
                if (iteration >= self.max_iterations):
                    color = black
                else:
                    if self.use_smooth_colors.get():
                        # Reduce the error in mu.
                        for i in range(3):
                            z = z * z + c
                            iteration += 1
                        mu = iteration + 1 - math.log(math.log(abs(z))) / log_escape

                        if not self.repeat_colors.get():
                            mu = mu / self.max_iterations * len(colors)
                        color = self.get_color(colors, mu)
                    else:
                        color = colors[iteration % len(colors)]
                self.pixels[x, 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.
JuliaSetApp(False)
print('Done')


# In[ ]:




