import pygame
from treelearn import treelearn
import loadmodel
from astar import astar
from state import State
import time
from garbage_truck import GarbageTruck
from heuristicfn import heuristicfn
from map import randomize_map
from heuristicfn import heuristicfn
import pygame as pg
import random
from request import Request

def determine_fitness(requests_list):
    distances = []
    for i in range(len(requests_list)+1): #from: request_list[i].x_pos and .y_pos
        temp = []
        for j in range(len(requests_list)+1):
            if j<i:
                temp.append('-')
            elif j==i:
                temp.append(0)
            elif j>i:
                if i==0:
                    dist = heuristicfn(0, 0, requests_list[j-1].x_pos, requests_list[j-1].y_pos)
                    temp.append(dist)
                else:
                    dist = heuristicfn(requests_list[i-1].x_pos, requests_list[i-1].y_pos, requests_list[j-1].x_pos, requests_list[j-1].y_pos)
                    temp.append(dist)
        distances.append(temp)
    return(distances)


def perform_permutation(obj_list, perm_list):
    result = [None] * len(obj_list)

    for i, index in enumerate(perm_list):
        result[int(index)-1] = obj_list[i-1]
    return result

def apply_genetic(request_list):
    print("Genetic algorithm started")

    distances = determine_fitness(request_list)
    population_size = 12
    num_generations = 8
    mutation_rate = 0.3
    NUM = len(distances)

    def initialize_population():
        population = []
        for _ in range(population_size):
            chromosome = ['0']
            while True:
                if len(chromosome) == NUM:
                    chromosome.append('0')
                    break

                temp = random.randint(1, NUM-1)
                temp_str = str(temp)
                if temp_str not in chromosome:
                    chromosome.append(temp_str)
            population.append(chromosome)
        return population
    
    def calculate_route_length(route):
        length = 0
        for i in range(len(route)-1):
            p = int(route[i])
            q = int(route[i + 1])
            length += distances[int(min(p,q))][int(max(p,q))]
        return length

    def calculate_fitness(population):
        fitness_scores = []
        for chromosome in population:
            fitness_scores.append(1 / calculate_route_length(chromosome))
        return fitness_scores

    def parents_selection(population, fitness_scores):
        selected_parents = []
        for _ in range(len(population)):
            candidates = random.sample(range(len(population)), 2)
            fitness1 = fitness_scores[candidates[0]]
            fitness2 = fitness_scores[candidates[1]]
            selected_parent = population[candidates[0]] if fitness1 > fitness2 else population[candidates[1]]
            selected_parents.append(selected_parent)
        return selected_parents



    def david_crossover(parent1, parent2):
        start_index = random.randint(1, len(parent1)-3)
        end_index = random.randint(start_index+1, len(parent1)-2)
        parent1_chain = parent1[start_index:end_index+1]
        parent2_letters = []
        for trash in parent2[1:-1]:
            if trash not in parent1_chain:
                parent2_letters.append(trash)
        child = [parent2[0]]+parent2_letters[0:start_index] + parent1_chain + parent2_letters[start_index:]+[parent2[-1]]
        """         print('PARENTS: ')
        print(parent1)
        print(parent2)
        print('CHILDS:')
        print(child) """
        return child

    def mutation(chromosome):
        index1 = random.randint(1, len(chromosome)-2)
        index2 = random.randint(1, len(chromosome)-2)
        chromosome[index1], chromosome[index2] = chromosome[index2], chromosome[index1]
        return chromosome

    def genetic_algorithm():
        population = initialize_population()
        for _ in range(num_generations):
            fitness_scores = calculate_fitness(population)
            parents = parents_selection(population, fitness_scores)
            offspring = []
            for i in range(0, len(parents), 2):
                parent1 = parents[i]
                parent2 = parents[i+1]
                child1 = david_crossover(parent1, parent2)
                child2 = david_crossover(parent2, parent1)
                offspring.extend([child1, child2])
            population = offspring
            for i in range(len(population)):
                if random.random() < mutation_rate:
                    population[i] = mutation(population[i])
        return population


    best_route = None
    best_length = float('inf')
    population = genetic_algorithm()
    for chromosome in population:
        length = calculate_route_length(chromosome)
        if length < best_length:
            best_length = length
            best_route = chromosome

    print("Permutation chosen: ", best_route)
    print("Its length:", best_length)
    permuted_list = perform_permutation(request_list, best_route[1:-1])