#!/usr/bin/env python3

from PIL import Image
import numpy as np
from random import randint, random
import sys

im_memo = {}
im_s_memo = {}
def load_char(n, size):
    if (n,size) in im_s_memo:
        im = im_s_memo[(n,size)]
    else:
        if n in im_memo:
            im = im_memo[n]
        else:
            im = Image.open("chars/%s.png" % n)
            im_memo[n] = im
        orig_h = im.size[1]
        r = size / orig_h
        new_sizes = (max(int(r*s),1) for s in im.size)
        im = im.resize(new_sizes)
        im_s_memo[(n,size)] = im

    char = np.asarray(im, dtype=np.uint8)
    if char.shape[2] is not 4:
        print("CHAR ERROR: chars must be RGBA right now")
        sys.exit(0)
    return char

def numberify(in_data, chars, size_range, disjoint=False, use_color=False, num_chars=3000, max_fails=1000):
    output = np.zeros(in_data.shape, dtype=np.uint8)
    alpha = in_data[:,:,3] > 100

    chars_added = 0
    consecutive_failures = 0
    while consecutive_failures < max_fails and chars_added < num_chars:
        n = chars[randint(0,len(chars)-1)]
        size = randint(*size_range)
        char = load_char(n, size)

        y = randint(0, output.shape[0] - char.shape[0])
        x = randint(0, output.shape[1] - char.shape[1])

        output_region = output[y:y+char.shape[0],x:x+char.shape[1]]

        if disjoint and np.sum(output_region[:,:,3]) > 0:
            consecutive_failures += 1
            continue

        if not alpha[y+char.shape[0]//2,x+char.shape[1]//2]:
            consecutive_failures += 1
            continue


        char_mask = (char[:,:,3] > 0)
        if use_color:
            color = in_data[y+char.shape[0]//2, x+char.shape[1]//2]
            char_map = np.multiply.outer(char_mask, color)
        else:
            char_map = char
        
        output_region[char_mask] = char_map[char_mask]
        output_region = np.clip(output_region, 0, 255)

        chars_added += 1
        consecutive_failures = 0
    
    print("gave up after",chars_added,"chars placed")

    return output

if __name__ == "__main__":
    from optparse import OptionParser

    parser = OptionParser()
    parser.add_option("-i", "--input", dest="input")
    parser.add_option("-o", "--output", dest="output")
    parser.add_option("-c", "--chars", dest="chars", default="0,1,2,3,4,5,6,7,8,9", help="comma-separated list of characters to use")
    parser.add_option("-d", "--disjoint", action="store_true", dest="disjoint", default=False, help="whether should be disjoint or not")
    parser.add_option("-n", "--num_chars", dest="num_chars", default=100, help="number of characters to add")
    parser.add_option("-m", "--max_fails", dest="max_fails", default=100, help="number consecutive failed attemps to add a char before quitting")
    parser.add_option("-s", "--size_range", dest="size_range", default="10:50", help="colon-separated range for character size")
    parser.add_option("-u", "--use_color", action="store_true", dest="use_color", default=False, help="color chars with original")

    (options, args) = parser.parse_args()

    chars = options.chars.split(",")
    size_range = list(map(float,options.size_range.split(":")))

    input_name = options.input or args[0]
    im = Image.open(input_name)
    print("Charifying %s: %s %s %s" % (input_name, im.format, im.size, im.mode))
    print("   chars:    ", " ".join(chars))
    print("   size:     ", " : ".join(map(lambda x:str(int(x))+"px",size_range)))
    print("   disjoint: ", options.disjoint)
    print("   output:   ", options.output)

    in_data = np.asarray(im, dtype=np.uint8)

    if in_data.shape[2] is not 4:
        print("INPUT ERROR: input must be RGBA right now")
        sys.exit(0)

    output = numberify(in_data, chars, size_range, options.disjoint, options.use_color, int(options.num_chars), int(options.max_fails))

    output = np.uint8(output)
    out_im = Image.fromarray(output)
    if options.output:
        out_im.save(options.output, im.format)
    else:
        out_im.show()
