import time
import networkx as nx
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import json
import os  # Import the os module here
from statistics import mean, stdev

class Enviroment():

    def __init__(self, file_name, init_pheromone, min_pheromone):
        self.min_pheromone = min_pheromone
        #Read Data
        self.data = self.__readData(file_name)
        self.num_jobs = len(self.data)
        self.num_machines = len(self.data[0])
        #Create Graph 
        self.G, self.node_names = self.__buildGraph(init_pheromone)


    def __readData(self, file_name):
        jobs = []
        file_path = os.path.join("C:\\Users\\HP\\Desktop\\ACO & PSO\\Ant-Colony-Optimization-main\\Ant-Colony-Optimization-main\\test_instances\\la41", file_name)
        with open(file_path, 'r') as file:
            for line in file:
                machines = {}
                this_machine = None
                for j, value in enumerate(line.split()):
                    if j % 2 != 0:
                        machines.update({this_machine: value})
                    else:
                        this_machine = value
                jobs.append(machines)
        return jobs


    def __buildGraph(self, init_pheromone):
        """
        Uses the networkx library to build a Direct, Weighted, and Full Connected Graph using the matrix supplied by the __readData() method.

        Working with edges will be much easier using this library.
        A virtual initial node (-1,-1) that links to every other node will be added.

        yields: Graph and names of every node.

        """
        num_nodes = self.num_jobs*self.num_machines + 1 #Every machineXjobs pair in addition to the first virtual node

        #Make a linked, directed graph with zero weight on each edge:
        unnamed_graph = nx.complete_graph(num_nodes,  nx.DiGraph())

        #Giving the nodes appropriate names:
        node_names = [(-1,-1)] #Add the first virtual node. 
        for job in range(self.num_jobs):
            for machine in range(self.num_machines):
                node_names.append((job, machine))
        mapping = {i:nodename for i,nodename in enumerate(node_names)} 
        G = nx.relabel_nodes(unnamed_graph,mapping)
        
        #Edges pointing to the initial virtual node should be removed so that they only extend from the initial to other nodes and not the other way around:
 
        for job in range(self.num_jobs):
            for machine in range(self.num_machines):
                G.remove_edge((job, machine), (-1,-1))

        #Pheromone and desirability updates:
        for from_job in range(self.num_jobs):
            for from_machine in range(self.num_machines):
                
                #Add pheromone and desirability edge properties for edges originating from virtual nodes.
                execution_time = self.data[from_job][str(from_machine)]
                desirability = 1/int(execution_time)
                G[(-1,-1)][(from_job,from_machine)]['pheromone'] = init_pheromone
                G[(-1,-1)][(from_job,from_machine)]['desirability'] = desirability

                for to_job in range(self.num_jobs):
                    for to_machine in range(self.num_machines):
                     # For edges from the operational node, add edge attributes (pheromone and desirability).
                        if from_job == to_job and from_machine == to_machine:
                            pass
                        else:
                            execution_time = self.data[to_job][str(to_machine)]  # To_machine conversion to string
                            desirability = 1/int(execution_time)
                            G[(from_job, from_machine)][(to_job, to_machine)]['pheromone'] = init_pheromone
                            G[(from_job, from_machine)][(to_job, to_machine)]['desirability'] = desirability


        node_names.remove((-1,-1)) #Delete the first virtual node from the names
        return G, node_names


    def getGraph(self):
        """
        Returns Graph with updated
        pheromones and node names
        """
        return self.G
    
    def getNodeNames(self):
        """
        Returns the node names
        of the Graph
        """
        return self.node_names

    def getTimeOfExecutions(self):
        """
        Returns the data readed
        with execution time of
        each node.
        """
        return self.data

    def getEdges(self):
        return [edge for edge in self.G.edges]

    
    def updatePheromone(
        self,
        evaporation_rate,
        cycle_edge_contribution): 
        """
        Multiplies the evaporation rate to the old value to simulate the pheromone evaporation at each edge. 

        pheromone concentrations.
        It replicates the contribution of the ants' pheromone trails when the sum of the inverse of the path's time that traveled over that edge is added.

        """
        for edge in self.G.edges:
            from_node = edge[0]
            to_node = edge[1]
            old_pheromone = self.G[from_node][to_node]['pheromone']
            new_pheromone = cycle_edge_contribution[edge] + (evaporation_rate * old_pheromone)
            if new_pheromone > self.min_pheromone:
                self.G[from_node][to_node]['pheromone'] = new_pheromone
            else:
                self.G[from_node][to_node]['pheromone'] = self.min_pheromone    


    def calculateMakespanTime(self, path):
        """
        determines the best time (makespan) for the work shop scheduling problem's entrance path.


        comes back: makespam time

        """
        #Start the scheduler.
        machine_task_moments = [] #Display each task at the appropriate time on the selected machine.
        machine_moments = [] #indicates the instant in which the machine is
        for i in range(self.num_machines):
            machine_task_moments.append([])
            machine_moments.append(0)

        for edge in path:
            this_job = edge[1][0]
            this_machine = edge[1][1]
            this_task_time = int(self.data[this_job][str(this_machine)])
            moment = machine_moments[this_machine]
            moment_for_task_not_found = True
            #Check the time at which the task can start.
            while moment_for_task_not_found:
                foud_other_machine_with_same_task = False
                for other_machine in range(self.num_machines):
                    if other_machine == this_machine:
                        pass
                    else:
                        try:
                            if this_job == machine_task_moments[other_machine][moment]:
                                foud_other_machine_with_same_task = True
                                break
                        except:
                            pass
                
                if foud_other_machine_with_same_task == False:
                    #Stops the loop
                    moment_for_task_not_found = False
                    #fill the job time for that machine
                    for i in range(this_task_time):
                        machine_task_moments[this_machine].append(this_job)
                    machine_moments[this_machine] = this_task_time + moment + 1
                    
                else:
                    #Make the machine wait till another machine has finished the task.
                    machine_task_moments[this_machine].append('-')
                    moment+=1
                    
       

        machine_execution_lengths = []
        for i in range(self.num_machines):
            machine_execution_lengths.append(len(machine_task_moments[i]))
        return max(machine_execution_lengths)

            


    def printGraph(self):
        """
       makes a picture of the constructed graph

        outcomes:

            png image: 'code/graph.png'
        """
        options = {
            'node_color': 'blue',
            'node_size': 2000,
            'width': 2.5,
            'arrowstyle': '-|>',
            'arrowsize': 20,
        }
        matplotlib.use('Agg')
        
        #Producing edge labels
        edge_labels = {}
        for from_node, to_node, edge in self.G.edges(data=True):
            desirability = round(edge['desirability'], 4)
            pheromone = round(edge['pheromone'], 3)
            weight = "(" + str(desirability) + " - " +  str(pheromone)+")"
            edge_labels.update({(from_node, to_node) : weight})

        #Producing a figure    
        fig = plt.figure()        
        ax = fig.add_subplot(111)
        pos = nx.spring_layout(self.G)
        nx.draw_networkx_edge_labels(self.G,pos,edge_labels=edge_labels, font_size=7)
        nx.draw_networkx(self.G, pos, arrows=True, ax=ax, **options)
        fig.savefig('graph.png')

class Ant():
    def __init__(self, Graph, node_names, ALPHA, BETA, seed, extended_seed):
        self.seed = seed + extended_seed
        self.ALPHA = ALPHA
        self.BETA = BETA
        self.G = Graph 
        self.not_visited = node_names.copy()
        self.ant_path = []

    def walk(self):
        np.random.seed(self.seed)
        current_node = (-1, -1)
        while self.not_visited:
            if len(self.not_visited) == 1:
                next_node = self.not_visited[0]
            else:
                next_node = self.__chooseNextNode(current_node)
            self.ant_path.append((current_node, next_node))
            current_node = next_node
            self.not_visited.remove(next_node)
        return self.ant_path

    def __chooseNextNode(self, current_node):
        node_probabilities = self.__calculateNodeProbabilityChoices(current_node)
        nodes = list(node_probabilities.keys()) 
        normalized_probabilities = self.__normalizeProbabilities(node_probabilities.values())
        indexs = [i for i in range(len(normalized_probabilities))]
        next_node_index = np.random.choice(indexs, p=normalized_probabilities)
        return nodes[next_node_index]
        
    def __normalizeProbabilities(self, node_probabilities):
        round_probabilities = [round(prob, 6) for prob in node_probabilities]
        probabilities_sum = sum(round_probabilities)
        normalized_probabilities = [prob / probabilities_sum for prob in round_probabilities]
        full_probabilities = sum(normalized_probabilities)
        while(full_probabilities != 1):
            normalized_probabilities[-1] -= full_probabilities - 1 
            full_probabilities = sum(normalized_probabilities)
        return normalized_probabilities

    def __calculateNodeProbabilityChoices(self, current_node):
        desirabilities = np.array([1 / self.G[current_node][node]['desirability'] for node in self.not_visited])  # Use 1 / desirability
        pheromones = np.array([self.G[current_node][node]['pheromone'] for node in self.not_visited])
    
        partial_probabilities = (pheromones ** self.ALPHA) * (desirabilities ** self.BETA)
        probability_sum = np.sum(partial_probabilities)
        normalized_probabilities = partial_probabilities / probability_sum

        return dict(zip(self.not_visited, normalized_probabilities))


class ACO():
    def __init__(self, ALPHA, BETA, dataset, cycles, ant_numbers, init_pheromone, pheromone_constant, min_pheromone, evaporation_rate, seed):
        self.ALPHA = ALPHA
        self.ant_numbers = ant_numbers
        self.BETA = BETA
        self.cycles = cycles
        self.pheromone_constant = pheromone_constant
        self.evaporation_rate = evaporation_rate
        self.seed = seed
        self.enviroment = Enviroment(dataset, init_pheromone, min_pheromone)
        self.time_of_executions = self.enviroment.getTimeOfExecutions()
        self.node_names = self.enviroment.getNodeNames()
        self.graph_edges = self.enviroment.getEdges()

    def releaseTheAnts(self):
        results_control = {}
        all_times = []
        for cycle_number in range(self.cycles):
            this_cycle_times = []
            this_cycle_Graph = self.enviroment.getGraph()
            this_cycle_edges_contributions = dict.fromkeys(self.graph_edges, 0) 
            for ant_number in range(self.ant_numbers):
                ant = Ant(this_cycle_Graph, self.node_names, self.ALPHA, self.BETA, self.seed, extended_seed=ant_number)
                ant_path = ant.walk()
                path_time = self.enviroment.calculateMakespanTime(ant_path)
                
                # Pheromones are updated using the inverse of makespan.
                for edge in ant_path:
                    this_cycle_edges_contributions[edge] += self.pheromone_constant / path_time
                
                this_cycle_times.append(path_time)
                all_times.append(path_time)

            self.enviroment.updatePheromone(self.evaporation_rate, this_cycle_edges_contributions)
            results_control.update({
                cycle_number: [min(this_cycle_times), mean(this_cycle_times), max(this_cycle_times)]
            })
        # Bring in the required libraries.

class ACOWithKalman(ACO):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)  # Make use of the superclass constructor.
        self.kalman_initial_covariance = np.eye(self.enviroment.num_jobs) * 100
        self.kalman_estimates = np.zeros(self.enviroment.num_jobs)
        self.min_pheromone = kwargs['min_pheromone']  # For inheriting min_pheromone, add this line.
        self.execution_times = []  # Set the constructor's execution_times list to zero.

    def runWithKalman(self):
        results_control = []
        all_times = []

        for cycle_number in range(self.cycles):
            this_cycle_times = []
            this_cycle_Graph = self.enviroment.getGraph()
            this_cycle_edges_contributions = dict.fromkeys(self.graph_edges, 0)

            # Note the beginning time.
            start_time = time.time()

            for ant_number in range(self.ant_numbers):
                ant = Ant(this_cycle_Graph, self.node_names, self.ALPHA, self.BETA, self.seed, extended_seed=ant_number)
                ant_path = ant.walk()
                path_time = self.enviroment.calculateMakespanTime(ant_path)
                self.updateKalmanEstimates(ant_path, path_time)

                for edge in ant_path:
                    this_cycle_edges_contributions[edge] += self.pheromone_constant / path_time
                this_cycle_times.append(path_time)
                all_times.append(path_time)

            # Note the finish time.
            end_time = time.time()

            # Determine the elapsed time.
            elapsed_time = end_time - start_time
            self.execution_times.append(elapsed_time)

            self.enviroment.updatePheromone(self.evaporation_rate, this_cycle_edges_contributions)
            self.updatePheromoneUsingKalmanEstimates()

            results_control.append({
                "cycle_number": cycle_number,
                "min_time": min(this_cycle_times),
                "mean_time": mean(this_cycle_times),
                "max_time": max(this_cycle_times),
                "elapsed_time": elapsed_time  # Add the elapsed time to the outcomes.
            })

        # Execution times can be printed as necessary or saved to a file.
        json.dump(self.execution_times, open("execution_times.json", 'w'))

        json.dump(results_control, open("ACO_cycles_results_with_kalman.json", 'w'))
        print("---------------------------------------------------")
        print("Mean: ", mean(all_times))
        print("Standard deviation: ", stdev(all_times))
        print("BEST PATH TIME: ", min(all_times), " seconds")
        print("---------------------------------------------------")

    def updatePheromoneUsingKalmanEstimates(self):
        for job in range(self.enviroment.num_jobs):
            new_pheromone = self.kalman_estimates[job]
            if new_pheromone > self.min_pheromone:
                self.enviroment.G[(-1, -1)][(job, 0)]['pheromone'] = new_pheromone
            else:
                self.enviroment.G[(-1, -1)][(job, 0)]['pheromone'] = self.min_pheromone

    def updateKalmanEstimates(self, ant_path, path_time):
        for edge in ant_path:
            this_job = edge[1][0]
            self.kalman_estimates[this_job] = self.updateKalmanFilter(this_job, path_time)

    def updateKalmanFilter(self, job, observed_time):
        dt = 1
        process_noise_variance = 0.01

        F = np.array([[1, dt], [0, 1]])
        H = np.array([[1, 0]])
        R = np.array([[0.1]])

        X = np.array([[self.kalman_estimates[job]], [0]])
        P = self.kalman_initial_covariance[job, job]

        X_predict = np.dot(F, X)
        P_predict = np.dot(np.dot(F, P), F.T) + process_noise_variance

        K = np.dot(np.dot(P_predict, H.T), np.linalg.inv(np.dot(np.dot(H, P_predict), H.T) + R))

        X[0, 0] = X_predict[0, 0] + np.dot(K, observed_time - np.dot(H, X_predict))[0, 0]
        self.kalman_estimates[job] = X[0, 0]
        self.kalman_initial_covariance[job, job] = P_predict[0, 0]

        return X[0, 0]


# Optimization parameters
parameters = {
    "seed": 0,
    "ALPHA": 1,
    "BETA": 1,
    "init_pheromone": 0.999,
    "pheromone_constant": 1,
    "min_pheromone": 0.001,
    "evaporation_rate": 0.91,
    "ant_numbers": 20,
    "cycles": 20,
    "dataset": 'la41.txt'
}

# Note the start time.
start_time_total = time.time()

# Run the algorithm with the Kalman filter after instantiating ACOWithKalman with your settings.
colony_with_kalman = ACOWithKalman(
    ALPHA=parameters['ALPHA'],
    BETA=parameters['BETA'],
    dataset=parameters['dataset'],
    cycles=parameters['cycles'],
    ant_numbers=parameters['ant_numbers'],
    init_pheromone=parameters['init_pheromone'],
    pheromone_constant=parameters['pheromone_constant'],
    min_pheromone=parameters['min_pheromone'],  # Here, add the min_pheromone argument.
    evaporation_rate=parameters['evaporation_rate'],
    seed=parameters['seed'])


colony_with_kalman.runWithKalman()


# Note the finish time.
end_time_total = time.time()

# Determine the total elapsed time and print it.
total_elapsed_time = end_time_total - start_time_total
print(f"Total Execution Time: {total_elapsed_time} seconds")