from Tkinter import *
import Image
import ImageTk
import ImageDraw


class Squirt:

    def __init__(self,filename=None):
        self.window = Tk()
        if filename:
            self.image = Image.open(filename).convert('RGB')
        else:
            self.image = Image.new('RGB',(500,500))
        self.draw = ImageDraw.Draw(self.image)
        self.W, self.H = self.image.size
        self.draw.line((self.W/2,0,self.W/2,self.H), fill=(64,64,64))
        self.draw.line((0,self.H/2,self.W,self.H/2), fill=(64,64,64))
        self.bitmap = ImageTk.PhotoImage(self.image)
        self.A = float(self.H)/float(self.W)
        self.canvas = Canvas(width=self.W, height=self.H)
        self.canvas.bind('<Button-1>', self.pen_down)
        self.canvas.bind('<ButtonRelease-1>', self.pen_up)
        self.canvas.bind('<Button-2>', self.disk_on)
        self.canvas.bind('<ButtonRelease-2>', self.disk_off)
        self.canvas.bind('<Motion>', self.motion)
        self.canvas.create_image(0,0, anchor=NW, image=self.bitmap)
        self.canvas.pack()
        self.drawing = 0
        self.disk = 0
        self.ellipse=None

    def complex_to_pixels(self, z):
        return (int(self.W*(.5 + .25*z.real)), int(self.H*(.5+.25*z.imag/self.A)))
                                                   

    def pixels_to_complex(self, m, n):
        return complex(  4.0*float(m)/float(self.W) - 2 ,
                         2*self.A - 4.0*self.A*float(n)/float(self.H) )

    def sqrt(self):
        newimage = Image.new('RGB',(self.W, self.H))
        for m in range(self.W):
            print m
            for n in range(self.H):
                z = self.pixels_to_complex(m,n)
                z = z*z
                xy = self.complex_to_pixels(z)
                try:
                    pixel = self.image.getpixel(xy)
                    newimage.putpixel((m,n), pixel)
                except IndexError:
                    pass
        self.image = newimage
        self.draw = ImageDraw.Draw(self.image)
        self.bitmap = ImageTk.PhotoImage(self.image)
        self.canvas.create_image(0,0, anchor=NW, image=self.bitmap) 

    def pen_down(self, event):
        self.drawing = 1
        self.x, self.y = event.x, event.y

    def pen_up(self, event):
        self.drawing = 0
        self.bitmap = ImageTk.PhotoImage(self.image)
        self.canvas.create_image(0,0, anchor=NW, image=self.bitmap) 

    def disk_on(self, event):
        self.disk = 1
        self.x, self.y = event.x, event.y

    def disk_off(self, event):
        self.disk = 0
        radius = max(abs(event.y - self.y), abs(event.x - self.x))/2
        self.draw.ellipse( (self.x - radius, self.y - radius,
                            self.x + radius, self.y + radius), fill=(255,0,0) )
        self.bitmap = ImageTk.PhotoImage(self.image)
        self.canvas.create_image(0,0, anchor=NW, image=self.bitmap) 
        
    def motion(self, event):
        if self.drawing == 1:
            self.canvas.create_line((self.x, self.y, event.x, event.y),
                                    fill='red')
            self.draw.line((self.x, self.y, event.x, event.y),
                           fill=(255,0,0))
            self.x, self.y = event.x, event.y
        if self.disk == 1:
            if self.ellipse:
                self.canvas.delete(self.ellipse)
            radius = max(abs(event.y - self.y), abs(event.x - self.x))/2
            self.ellipse = self.canvas.create_oval(
                self.x - radius, self.y - radius,
                self.x + radius, self.y + radius, fill = 'red')