import glob
from collections import defaultdict
import cv2
import numpy as np
import random

char2num = {
    '0': 0, '1': 1, '2': 2, '3': 3, '4': 4, '5': 5, '6': 6, '7': 7, '8': 8,
    '9': 9, 'A': 10, 'B': 12, 'C': 13, 'D': 14, 'E': 15, 'F': 16, 'G': 17,
    'H': 18, 'I': 19, 'J': 20, 'K': 21, 'L': 23, 'M': 24, 'N': 25, 'O': 26,
    'P': 27, 'Q': 28, 'R': 29, 'S': 30, 'T': 31, 'U': 32, 'V': 34, 'W': 35,
    'X': 36, 'Y': 37, 'Z': 38,
}
resize_backgound = 512

number_height = 96
number_width = 512


def compute_check_digit(number):
    total = sum(char2num[c] * 2**x for x, c in enumerate(number))
    return (total-int(total/11.0)*11) % 10

def get_number():
    for cat in CATEGORY_INDENTIFIER:
        for first in range(0, 10):
            for second in range(0, 10):
                for third in range(0, 10):
                    for fourth in range(0, 10):
                        for fifth in range(0 ,10):
                            for sixth in range(0, 10):
                                yield (cat, first, second, fourth, fifth, sixth)

def get_random_number():
    cat = random.randint(0, len(CATEGORY_INDENTIFIER)-1)
    number = [ str(random.randint(0, 9)) for _ in range(6)]
    number.insert(0, CATEGORY_INDENTIFIER[cat])
    check_digit = compute_check_digit(''.join(number))
    number.append(str(check_digit))

    return ''.join(number)

CATEGORY_INDENTIFIER = ['EGHU', 'EGSU', 'EISU', 'EMCU', 'HMCU', 'PCIU', 'DRYU', 'EITU', 'WHLU', 'TCNU', 'IMTU', 'KKFU']

images = glob.glob('preprocess3/**/*.png')
backgrounds = glob.glob('background/*.png')
mapping = defaultdict(list)

# build label mapping
for im in images:
    label = im.split('/')[-2]
    mapping[label].append(im)



number = get_random_number()
print number[0]

im1 = cv2.imread(mapping[number[0]][random.randint(0, len(mapping[number[0]])-1)], 0)

blank_image = np.zeros((96,512), np.uint8)

start_row = 20
blank_image[start_row:im1.shape[0] + start_row,0:im1.shape[1]] = im1

height, width = blank_image.shape

for i in range(1, len(number)):
    im2 = cv2.imread(mapping[number[i]][random.randint(0, len(mapping[number[i]])-1)], 0)
    for j in range(width-1, -1, -1):
        if np.sum(blank_image[:,j])>0:
            col = j
            break
    #  print row, col
    padding = 0
    row_padding = 0
    if i < 4:
        padding = 4
    elif i == 4:
        padding = 20
    elif i == len(number) - 1:
        padding = 21
        # last number
        # put the bounding box around it
        im2 = np.pad(im2, ((1,1),(1,1)), mode='constant', constant_values=255)
        row_padding = -1

    elif i > 4 and i < len(number) - 1:
        padding = 5

    blank_image[start_row + row_padding:start_row + row_padding + im2.shape[0], col + padding: col + padding + im2.shape[1]] = im2

# get the bounding box of whole number
while np.sum(blank_image[0]) == 0:
    blank_image = blank_image[1:]

while np.sum(blank_image[:,0]) == 0:
    blank_image = np.delete(blank_image,0,1)

while np.sum(blank_image[-1]) == 0:
    blank_image = blank_image[:-1]

while np.sum(blank_image[:,-1]) == 0:
    blank_image = np.delete(blank_image,-1,1)

number_rows, number_cols = blank_image.shape

# put it to background image

# get background
background = cv2.imread(backgrounds[random.randint(0, len(backgrounds) - 1)])
b_height, b_width, _ = background.shape
if b_height < number_height or b_width < number_width:
    if b_height < number_height:
        background = cv2.resize(background, (b_width, resize_backgound))
    elif b_width < number_width:
        background = cv2.resize(background, (resize_backgound, b_height))

b_height, b_width, _ = background.shape
print background.shape
print blank_image.shape

startX = random.randint(0, b_width - number_cols - 1)
startY = random.randint(0, b_height - number_rows - 1)

for c in range(0, 3):
    for u in range(0, number_rows):
        for uu in range(0, number_cols):
            if blank_image[u][uu] > 60:
                background[startY + u][startX + uu] = blank_image[u][uu]






cv2.imshow('im1', blank_image)
cv2.imwrite('bb.png', background)
cv2.waitKey(0)
cv2.destroyAllWindows()