# HG changeset patch
# User Paul Boddie <paul@boddie.org.uk>
# Date 1443825178 -7200
# Node ID 3a950efbc68f8f732a92692ce0166b06ed712d71
# Parent  b645174fcbb838a2b991fb34b07091de67c7a30e
Sum colour probabilities instead of counting requested colours.
Find the best alternative for each unallocated colour from the allocated ones.
Removed the preview image saving.

diff -r b645174fcbb8 -r 3a950efbc68f optimiser.py
--- a/optimiser.py	Fri Oct 02 13:52:42 2015 +0200
+++ b/optimiser.py	Sat Oct 03 00:32:58 2015 +0200
@@ -1,9 +1,10 @@
 #!/usr/bin/env python
 
 from random import random
-from os.path import extsep, splitext
+from os.path import splitext
 import EXIF
 import PIL.Image
+import itertools
 import math
 import sys
 
@@ -17,32 +18,30 @@
     r2, g2, b2 = rgb2
     return math.sqrt(pow(r1 - r2, 2) + pow(g1 - g2, 2) + pow(b1 - b2, 2))
 
-def factor(start, end, rgb):
-    r1, g1, b1 = start
-    r2, g2, b2 = end
-    gr, gg, gb = r2 - r1, g2 - g1, b2 - b1
-    r, g, b = rgb
-    pr, pg, pb = r - r1, g - g1, b - b1
-    dp = pr * gr + pg * gg + pb * gb
-    return float(dp) / pow(distance(start, end), 2)
+def distribution(rgb, values=None):
+    l = []
+    total = 0
+    for c in values or corners:
+        d = distance(rgb, c)
+        if d == 0:
+            return [(1, c)]
+        l.append((pow(d, -3), c))
+        total += pow(d, -3)
+    return [(d / total, c) for d, c in l]
 
-def nearest(rgb, values=None):
-    l = map(lambda c: (distance(rgb, c), c), values or corners)
-    l.sort()
+def pattern(rgb, values=None):
+    l = distribution(rgb, values or corners)
+    l.sort(reverse=True)
     return l
 
-def pattern(rgb, values=None):
-    l = nearest(rgb, values)
-    start, end = l[0][1], l[1][1]
-    f = factor(start, end, rgb)
-    return start, end, f
-
 def get_value(rgb, values=None):
-    rgb1, rgb2, f = pattern(rgb, values)
-    if random() < pow(f, 2):
-        return rgb2
-    else:
-        return rgb1
+    choose = random()
+    threshold = 0
+    for f, c in pattern(rgb, values):
+        threshold += f
+        if choose < threshold:
+            return c
+    return c
 
 def sign(x):
     return x >= 0 and 1 or -1
@@ -96,20 +95,15 @@
 
     input_filename, output_filename = sys.argv[1:3]
     basename, ext = splitext(output_filename)
-    preview_filename = extsep.join([basename + "_preview", ext])
 
-    preview = "-p" in sys.argv[3:]
     rotate = "-r" in sys.argv[3:]
     saturate = sys.argv[3:].count("-s")
     desaturate = sys.argv[3:].count("-d")
 
     x = EXIF.process_file(open(input_filename))
-    im = PIL.Image.open(input_filename)
+    im = PIL.Image.open(input_filename).convert("RGB")
     im = rotate_and_scale(im, width, height, rotate)
 
-    if preview:
-        im_preview = im.copy()
-
     width, height = im.size
 
     colours = []
@@ -125,24 +119,18 @@
                 rgb = saturate_rgb(rgb, saturate and math.pow(0.5, saturate) or math.pow(2, desaturate))
                 im.putpixel((x, y), rgb)
 
-            # Count the number of requested colours.
+            # Sum the colour probabilities.
 
-            value = get_value(rgb)
-            if not c.has_key(value):
-                c[value] = 1
-            else:
-                c[value] += 1
-
-            if preview:
-                im_preview.putpixel((x, y), value)
+            for f, value in distribution(rgb):
+                if not c.has_key(value):
+                    c[value] = f
+                else:
+                    c[value] += f
 
         c = [(n, value) for value, n in c.items()]
         c.sort(reverse=True)
         colours.append(c)
 
-    if preview:
-        im_preview.save(preview_filename)
-
     for y, c in enumerate(colours):
         most = [value for n, value in c[:4]]
         least = [value for n, value in c[4:]]
@@ -155,7 +143,9 @@
 
             value = get_value(rgb)
             if value in least:
-                value = get_value(rgb, most)
+                rgb = im.getpixel((x, y))
+                value = pattern(value, most)[0][1]
+
             im.putpixel((x, y), value)
 
     im.save(output_filename)