import math
import random
import numpy as np

#this algorithm is a MINIMIZATION  specific implementation of a discrete method called Ant Colony Optimization.

#method/function space
def object_funct_eval(path, var_table):
    """
    Given a path and the array with all the discrete values per layer, the value of the obj. Function is returned as a result
    :param path: the path of a specific ant
    :param var_table: the array representing the discrete values of the layers
    :return: the value of the objective function for the specific path
    """
    length = len(path)
    my_vars = []
    for i in range(length):
        j = path[i]
        my_vars.append(var_table[i][j])
    #print(my_vars)
    #objective function - Ackley
    #my_value = -20*math.exp(-0.2*math.sqrt(0.5*((my_vars[0]**2 + (my_vars[1]**2))))) - math.exp(0.5*(math.cos(my_vars[0]*2*math.pi) + math.cos(my_vars[1]*2*math.pi))) +math.e + 20
    #Booth (1, 3)
    my_value = (my_vars[0] + 2*my_vars[1] -7)**2 + (2*my_vars[0] + my_vars[1] - 5)**2
    #my_value = my_vars[0]**2 - my_vars[0] - 11
    #beale (3, 0.5)
    #my_value = (1.5 - my_vars[0] + my_vars[0] * my_vars[1]) ** 2 + (2.25 - my_vars[0] + my_vars[0] * (my_vars[1] ** 2)) ** 2 + (2.625 - my_vars[0] + my_vars[0] * (my_vars[1] ** 3)) ** 2
    return my_value
def calc_discrete_values(upper, lower, r):
    """

    :param range: range of accuracy for the permissible discrete values
    :param upper: upper bounds
    :param lower: lower bounds
    :return: d the number of discrete values per variable
    """
    d = (upper[0] - lower[0]) /r + 1
    return d

def create_pheromone_table(n_var, disc):
    """
     My pheromone table acts like a neighbor-table for describing the arcs of my graph
     and the pheromones in the trails. Its cells contain the amount of pheromone in the arc
     connecting the row with the collumns. The pheromone value for all the arc is initialized as 1.
     IT MUST BE THE SAME FOR ALL THE ARCS INITIALLY. The last row remains table, it represents the food source.

     The node or layer  'zero' meaning the first row  is considered the home. (since we count beginning from zero)
    :param n_var: how many variables the object function has
    :param disc: number for permissible discrite values for each design variable

    :return: the initial form of the 2d-numpy array for the pheromone table

    """
    #numpy array N-dimensions . rows 1 for home + n_var, columns for the discrete values [0 ...d-1]
    ph_array = np.ones((n_var + 1, disc))
    #note : MIGHT need to change some nodes to zero
    return ph_array

def pheromone_trail_update(pher_table,bestants, worstants, fbest, fworst,bestpath):
    """

    :param pher_table: the pheromone array . *using a numpy array
    :param bestants: list of the best ants represented by their position in the colony
    :param worstants: the list of the worst ants
    :param fbest: the best obj.value for the current iteration
    :param fworst: the worst obj.value for the current iteration
    :param bestpath: the path the best ants have followed
    :return: the changes are directly made to the array (matrix) itself  thus i don't need to return a new array
    """
    #trail update, initialize my parameters
    scaling_par = 3
    decay_factor = 0.005
    #the added sum is only useful to the best path
    added_sum = 0
    b_len = len(bestants)
    w_len = len(worstants)

    added_sum = (scaling_par * b_len * fbest)/ (w_len * fworst)
    #print("added sum =", added_sum)
    length = len(pher_table) - 1
    #update the paths, check if an arc belongs to the best one
    for i in range(length):
        for j in range(len(pher_table[0])):
            if j == bestpath[i]:
                pher_table[i][j] += added_sum
            else:
                pher_table[i][j] = (1 - decay_factor) * pher_table[i][j]

def find_best_path(bestants, colony):
    """
    Given the ids of the best ants in the colony (at least one) ,the best path is returned.
    :param bestants: a list which contains at least one element, it's the ants with the best (min) value for the obj. function
    :param colony: the colony of the ants
    :return: a list of length the same as the number of variables of the obj.function. It's the best path in the iteration
    """
    my_ant = bestants[0]
    print("I m checking ant :", my_ant)
    path = colony[my_ant][0]
    return path

def create_layer_table(n_var, r, lower, upper):
    """
    Creates the design variables table. It represents the permitted values for each variable during the optimization.
    The permitted values get created according to a specific range given.
    :param n_var: number of variables
    :param lower, upper: lower and upper bounds
    :param r: we spread the values by adding r for each permitted discrete value (r:range)
    :return: 2-d array of n_var x p to represent the discrete values of each design variable
    """
    d = (upper[0] - lower[0]) / r + 1
    layer_array = np.zeros((n_var, d))
    #counter is used for future implementations when the limits for each variable change, for this first version the value
    #of d remains the same and the functions used have the same limit for each variable
    counter = 0
    iterator = lower[counter]
    for j in range(n_var):
        for i in range(int(d)):
            layer_array[j][i] = iterator
            iterator = iterator + r
        iterator = lower[counter]
    return layer_array


def create_ant(var_num):
    """

    :param var_num: number of variables needed for the obj.function
    :return: a combined list of 2 elements. Element 0 is the path as a list. Element 1 is the obj. value.
    """

    path = []
    #for initial values we believe the ants chose the 0-th discrete value. (count method 0 ...n-1)
    for i in range(var_num):
        path.append(0)
    #for initial value i choose a big one , but it depends on the nature of the obj. function
    value = 1000
    ant = []
    ant.append(path)
    ant.append(value)
    return ant

def create_colony(N, var_num):
    """
    creates the initial colony of ants, it is in a list form
    :param N: size of colony
    :param var_num: number of variables - it's needed to create ants individually before adding them to the colony
    :return:
    """
    colony = []
    for i in range(N):
        colony.append(create_ant(var_num))
    return colony

def prob_handling(var_table, var_tag, d, ph_table):
    """
    A roulette stohastic method of choosing which path to follow.
    :param var_table: the numpy array representing the discrete values
    :param var_tag: tag of the variable
    :param d: number of discrete values per variable
    :param ph_table: pheromone trail array
    :return: returns whick discrete value the ant k picked for a specific layer
    """
    #this is probability handling for ONE ant k regarding its choice for ONE variable
    prob_list=[]
    my_sum = 0
    #print("vartag",var_tag)
    for i in range(d):
        my_sum += ph_table[var_tag][i]

    #print("my sum",my_sum)
    for i in range(d):
        prob_list.append(ph_table[var_tag][i]/my_sum)
    #print("prob list", prob_list)
    #the way the ant picks a specific route is with a roulette selection on the prob.list
    roll = round(random.uniform(0, 1),3)
    #print("my roll:", roll)
    my_sum = 0
    my_pick= 0
    for i in range(d):
        my_sum += prob_list[i]
        if (roll<= my_sum or i == d-1):
            my_pick = i
            break
    #print("my pick is:",my_pick)
    #I return the ant's choice
    return my_pick
def count_ants(colony, value):
    ant_list = []
    for i in range(len(colony)):
        if colony[i][1] == value:
            ant_list.append(i)
    return ant_list

def check_convergence(antlist, N):
    flag = False
    if len(antlist) == int(N):
        flag = True
    return flag
def find_solution(bpath, var_table):
    """

    :param bpath: the best path
    :param var_table: table containing the discrete values
    :return: the discrete values representing given path
    """
    solution = []
    length = len(bpath)
    for i in range(length):
        j = bpath[i]
        solution.append(var_table[i][j])
    return solution

#main space
def mainloop():
    #create the parameters and run the algorithm via the functions
    N = 30
    myrange = 0.5
    l_bound = [-10, -10]
    u_bound = [10, 10]
    n_var = 2
    #d = number of discrete values for each variable
    d = int(calc_discrete_values(u_bound,l_bound,myrange))
    test_array_layer = create_layer_table(n_var, myrange, l_bound, u_bound)
    print("array layer:", test_array_layer)
    test_ph_array = create_pheromone_table(n_var, d)
    print("ph array:", test_ph_array)
    iteration = 1
    iteration_limit = 1000
    test_colony = create_colony(N,n_var)

    convergence_flag = False

    while convergence_flag == False and iteration <= iteration_limit:
        #for the colony each ant chooses a path going from Layer-to-Layer
        var_counter = 0
        while var_counter < n_var :
            #print("Choosing for variable:", var_counter)
            for ant in test_colony:
                #print("var counter",var_counter)
                ant[0][var_counter] = prob_handling(test_array_layer,var_counter,d,test_ph_array)
            var_counter += 1
        for ant in test_colony:
            ant[1] = object_funct_eval(ant[0], test_array_layer)

        #initialize best and worst values
        bestf = test_colony[0][1]
        worstf = test_colony[0][1]
        bestants = []
        worstants = []

        for ant in test_colony:
            if ant[1]<= bestf:
                bestf = ant[1]
        for ant in test_colony:
            if ant[1]>= worstf:
                worstf = ant[1]

        bestants = count_ants(test_colony,bestf)
        worstants = count_ants(test_colony, worstf)
        #print("my colony:", test_colony)
        print("best value and worst value ", bestf, ",", worstf)
        print("best ants", bestants)
        #print("worst ants", worstants)
        best_path = find_best_path(bestants,test_colony)
        print("best_path:", best_path)
        pheromone_trail_update(test_ph_array, bestants, worstants, bestf, worstf,best_path)
        print("new trails", test_ph_array)
        convergence_flag = check_convergence(bestants,N)
        #creating a new colony for the following iteration
        #test_colony = create_colony(N,n_var)
        iteration += 1

    solution = find_solution(best_path,test_array_layer)
    print("The min of the objective value is :", solution, "and the value is :", object_funct_eval(best_path, test_array_layer) )




if __name__ == '__main__':
    mainloop()
