import math
import string
import random
import time
from PIL import Image
from PIL import ImageDraw
from operator import attrgetter

POSITION_MULTIPLIER = 1
RADIUS_MULTIPLIER = 1
two_pass = False

class Graph:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.mse = []
        self.max = 755**2
        self.image = Image.new("RGB", (self.width,self.height), (255,255,255))

    def draw(self):
        self.image = Image.new("RGB", (self.width,self.height),(255,255,255))
        self.max = max(self.mse) + 20
        draw = ImageDraw.Draw(self.image)
        for i in range(len(self.mse)):
            if i > 1:
                x1 = int(float(self.width)/(len(self.mse)-1)*(i))
                x2 = int(float(self.width)/(len(self.mse)-1)*(i-1))
                y1 = self.height - (int(float(self.height) / self.max * self.mse[i]))
                y2 = self.height - (int(float(self.height) / self.max * self.mse[i-1]))
                draw.line([(x1,y1),(x2,y2)],(0,0,0))

        self.image.save("/tmp/graph.png")
        return


class Stroke:
    def __init__(self, x, y, r, c):
        self.x = x
        self.y = y
        self.r = r
        self.c = c

    def soft_mutate(self, magnitude):
        if random.randint(0,1) == 1:
            self.x += int(random.uniform( -magnitude, magnitude) * POSITION_MULTIPLIER)
            self.y += int(random.uniform( -magnitude, magnitude) * POSITION_MULTIPLIER)
            if self.x < 0:
                self.x = 0
            if self.x > art.width - 1:
                self.x = art.width -1
            if self.y < 0:
                self.y = 0
            if self.y > art.height - 1:
                self.y = art.height - 1
            self.c = art.getpixel((self.x, self.y))
        if random.randint(0,1) == 1:
            self.r += int(random.uniform( -magnitude, magnitude) * RADIUS_MULTIPLIER)
            if self.r < MIN_RADIUS:
                self.r = MIN_RADIUS
            if self.r > MAX_RADIUS:
                self.r = MAX_RADIUS
        return


class Painting:
    def __init__(self, strokes, image):
        self.strokes = strokes
        self.image = image
        self.error = 90000

    def mean_square_error( self ):
        im1_pix = self.image.load() # get a pixelAccess object to read pixels faster
        im2_pix = art.load()

        errorTotal = 0 # total error of the image - adding an int is faster than making a list
        for y in range(1,art.height-1,1):
            for x in range(1,art.width-1,1):
                ar,ag,ab = im1_pix[x,y] # unpack the pixel RGB
                br,bg,bb = im2_pix[x,y]
                errorTotal += (( ar - br ) ** 2)
                errorTotal += (( ag - bg ) ** 2)
                errorTotal += (( ab - bb ) ** 2)
        mean_error = errorTotal/ ( art.width * art.height )
        return mean_error

    def duplicate( self ):

        strokes = []
        for stroke in self.strokes:
            strokes.append(Stroke(stroke.x, stroke.y, stroke.r, stroke.c))
        return Painting(strokes, Image.new("RGB", self.image.size, (255,255,255)))

    def create_image(self):
        image = Image.new("RGB", self.image.size, (255,255,255))
        draw = ImageDraw.Draw(image)
        for stroke in self.strokes:
            draw_circle( draw, stroke.x, stroke.y, stroke.r, stroke.c)
        self.image = image
        return image


def draw_circle(draw, x, y, r, c):

    draw.ellipse([x-r,y-r,x+r,y+r],c,c)
    return

def scale_painting( painting, factor ):
    image = Image.new("RGB", (painting.image.width*factor,painting.image.height*factor),(255,255,255))
    strokes = []
    for stroke in painting.strokes:
        strokes.append(Stroke(stroke.x*factor,stroke.y*factor,stroke.r*factor,stroke.c))
    painting = Painting(strokes,image)
    painting.create_image()
    painting.image.save("./scaled_painting.png")
    return painting

def paint( source, num_of_strokes, num_of_tries ):
    global art, MIN_RADIUS, MAX_RADIUS
    graph = Graph(600,400)
    art = Image.open(source)
    art.convert("RGB")
    MIN_RADIUS = int(float(2 )/695*art.width)
    if MIN_RADIUS < 1:
        MIN_RADIUS = 1
    MAX_RADIUS = int(float(66)/695*art.width)
    layers = []
    strokes = []
    prev_error = 999999
    painting = Painting(strokes, Image.new( "RGB", art.size, (255, 255,255)))
    old_image = painting.image
    for stroke in range(num_of_strokes):
        print("STROKE #" + str(stroke))
        error = painting.mean_square_error()
        tmp_error = error + 1
        while tmp_error >= error:
            tmp_painting = painting.duplicate()
            x = random.randint(0,art.width - 1)
            y = random.randint(0,art.height- 1)
            r = random.randint(MIN_RADIUS,MAX_RADIUS)
            c = art.getpixel((x,y))
            tmp_painting.strokes.append(Stroke(x,y,r,c))
            tmp_painting.create_image()
            tmp_error = tmp_painting.mean_square_error()
        painting.strokes.append(Stroke(x,y,r,c))
        prev_error = painting.mean_square_error()
        loop_num = 0
        while loop_num < num_of_tries:
            old_painting = painting.duplicate()
            painting.strokes[stroke].soft_mutate(5)
            painting.create_image()
            error = painting.mean_square_error()
            graph.mse.append(error)
            graph.draw()
            print("MSE: " + str(error))
            if error < prev_error:
                loop_num = 0
                painting.create_image()
                prev_error = error
                painting.image.save("./painting.png")
            else:
                painting = old_painting
                loop_num += 1


painting = paint('./src.jpg', 500, 50)

scale_painting(painting, 10)