import numpy as np
import threading
import queue
import time

INPUT_FILE = "data/medium.in"
POPULATION = 50
MUTATION_AMOUNT = 2000
ITERATIONS = 30
THREAD_COUNT = 20

data = [line for line in open(INPUT_FILE)]
params = list(map(int, data[0].split(" ")))
data = [[0 if x=="T" else 1 for x in line] for line in data[1:]]
data = np.array(data)[:, :-1]
clusters = np.arange(params[0]*params[1]).reshape((1, params[0], params[1]))
clusters = np.repeat(clusters, POPULATION, axis=0)+1
print(params)
print(data)
print(clusters[0])

values = {}
first = True

def get_fitness(vals, clean=False):
    fit = 0
    for key, val in vals.items():
        if key == 0:
            continue
        size = sum(val)
        if size <= params[3] and min(val) >= params[2]:
            fit += size
        if clean:
            continue
        size_diff = params[3]-size
        if size_diff < 0:
            fit += 1-size_diff**2
        elif size_diff > 0:
            fit += np.exp(-abs(size-params[3]))
    return fit


def get_left_bound(clust, y, x):
    val = clust[y, x]
    while x > 0 and val == clust[y, x-1]:
        x -= 1
    return (y, x)


def get_right_bound(clust, y, x):
    val = clust[y, x]
    while x+1 < clust.shape[1] and val == clust[y, x+1]:
        x += 1
    return (y, x)


def get_top_bound(clust, y, x):
    val = clust[y, x]
    while y > 0 and val == clust[y-1, x]:
        y -= 1
    return (y, x)


def get_bottom_bound(clust, y, x):
    val = clust[y, x]
    while y+1 < clust.shape[0] and val == clust[y+1, x]:
        y += 1
    return (y, x)


def set_area(clust, y1, x1, y2, x2, value):
    for y in range(y1, y2+1):
        for x in range(x1, x2+1):
            clust[y, x] = value
    return clust


def mutation(clust):
    for _ in range(np.random.random_integers(MUTATION_AMOUNT)):
        y = np.random.random_integers(params[0])-1
        x = np.random.random_integers(params[1])-1
        z = np.random.random()
        if z < 0.2:
            if y > 0:  # expand to top
                yn = y-1
                _, inner_left = get_left_bound(clust, y, x)
                _, outer_left = get_left_bound(clust, yn, inner_left)
                _, inner_right = get_right_bound(clust, y, x)
                _, outer_right = get_right_bound(clust, yn, inner_right)
                clust = set_area(clust, yn, outer_left, yn, outer_right, 0)
                clust = set_area(clust, yn, inner_left, yn, inner_right, clust[y, x])
        elif z < 0.4:
            if x > 0:  # expand to left
                xn = x-1
                inner_top, _ = get_top_bound(clust, y, x)
                outer_top, _ = get_top_bound(clust, inner_top, xn)
                inner_bot, _ = get_bottom_bound(clust, y, x)
                outer_bot, _ = get_bottom_bound(clust, inner_bot, xn)
                clust = set_area(clust, outer_top, xn, outer_bot, xn, 0)
                clust = set_area(clust, inner_top, xn, inner_bot, xn, clust[y, x])
        elif z < 0.6:
            if y < params[0]-1:  # expand to bottom
                yn = y+1
                _, inner_left = get_left_bound(clust, y, x)
                _, outer_left = get_left_bound(clust, yn, inner_left)
                _, inner_right = get_right_bound(clust, y, x)
                _, outer_right = get_right_bound(clust, yn, inner_right)
                clust = set_area(clust, yn, outer_left, yn, outer_right, 0)
                clust = set_area(clust, yn, inner_left, yn, inner_right, clust[y, x])
        elif z < 0.8:
            if x < params[1]-1:  # expand to right
                xn = x+1
                inner_top, _ = get_top_bound(clust, y, x)
                outer_top, _ = get_top_bound(clust, inner_top, xn)
                inner_bot, _ = get_bottom_bound(clust, y, x)
                outer_bot, _ = get_bottom_bound(clust, inner_bot, xn)
                clust = set_area(clust, outer_top, xn, outer_bot, xn, 0)
                clust = set_area(clust, inner_top, xn, inner_bot, xn, clust[y, x])
        else:
            pass#clust[y, x] = np.amax(clust)+1
    return clust






def myfunc(a, b):
    global first, values
    if first:
        first = False
        return
    if a not in values:
        values[a] = [0, 0]
    values[a][b] += 1
vfunc = np.vectorize(myfunc)

# mutation
for i in range(POPULATION):
    if i % 20 == 0:
        print("mutation", i)
    clusters[i] = mutation(clusters[i])

for iteration in range(ITERATIONS):
    # calc fitness
    fitnesses = np.zeros((POPULATION, ))
    for i, cluster in enumerate(clusters):
        if i % 20 == 0:
            print("fitness", i, iteration)
        values = {}
        first = True
        vfunc(cluster, data)
        fitnesses[i] = get_fitness(values)
    # select
    z_exp = [np.exp(i) for i in fitnesses]
    sum_z_exp = sum(z_exp)
    softmax = [i / sum_z_exp for i in z_exp]
    idx = np.random.choice(POPULATION, POPULATION, p=softmax)
    clusters = clusters[idx, :, :]

    # print best
    max_idx = np.argmax(fitnesses)
    print(clusters[max_idx])
    print(iteration, max(fitnesses))

    # mutation
    for i in range(POPULATION):
        clusters[i] = mutation(clusters[i])

fitnesses = np.zeros((POPULATION, ))
for i, cluster in enumerate(clusters):
    values = {}
    first = True
    vfunc(cluster, data)
    fitnesses[i] = get_fitness(values, clean=True)
max_idx = np.argmax(fitnesses)
print(clusters[max_idx])
print(max(fitnesses))