import numpy as np
import matplotlib.pyplot as plt

t = np.loadtxt('f2.txt')


def random_result(dim):
    res = np.zeros(t.size, dtype=np.int8)
    y = np.random.randint(0, dim[0], dim[1])
    x = np.arange(0, dim[1], dtype=np.int8)
    res2 = res.reshape(dim)
    res2[y, x] = 1
    return res2


def wage(res):
    return np.max(np.sum(t * res, 1))


def mutate(res):
    new = res.copy()
    np.random.shuffle(new[:, np.random.randint(0, res.shape[1])])
    return new


def start_temperature(res, prob):
    sum_wage = 0.0
    res_wage = wage(res)
    ind = 0
    while ind < 10:
        mod_wage = wage(mutate(res))
        if mod_wage > res_wage:
            sum_wage += (res_wage - mod_wage)
            ind += 1
    return (sum_wage/10.)/np.log(prob)


def next_temperature(step, start):
    return start/(step + 1)


def simulated_annealing(plot, temp_prob):
    result = random_result(t.shape)
    result_wage = wage(result)

    max_iteration = 1000
    start_temp = start_temperature(result, temp_prob)

    iteration_array = np.arange(max_iteration)
    temperature_array = np.zeros(max_iteration)
    result_wage_array = np.zeros(max_iteration)

    for i in range(0, max_iteration):
        temperature = next_temperature(i, start_temp)
        new = mutate(result)
        new_wage = wage(new)
        delta = new_wage - result_wage
        if delta < 0:
            result = new
            result_wage = new_wage
        elif np.random.sample() < np.exp(-delta / temperature):
            result = new
            result_wage = new_wage

        temperature_array[i] = temperature
        result_wage_array[i] = result_wage

    if plot:
        plt.subplot(2, 1, 1)
        plt.plot(iteration_array, result_wage_array)
        plt.subplot(2, 1, 2)
        plt.plot(iteration_array, temperature_array)
        plt.show()

    return result_wage

if __name__ == "__main__":
    print simulated_annealing(True, 0.8)