import numpy as np
import random
import pygame
from constant import width, height, size, rows, cols
from board import Board
from tractor import Tractor


routes_num = 20  # Ilość ścieżek, które będziemy generować
board = Board()
dirt_positions = board.get_dirt_positions()
dirt_count = len(dirt_positions)

def manhattan(a, b):
    return abs(a[0] - b[0]) + abs(a[1] - b[1])

def find_routes(routes_num):
    population_set = []  # zapisujemy trasy - losowe ułóżenia
    for i in range(routes_num):
        # losowo wygenerowane kolejności na trasie
        single_route = np.random.choice(list(range(dirt_count)), dirt_count, replace=False)
        population_set.append(single_route)
    return np.array(population_set)


def sum_up_for_route(route_indices):
    sum = 0
    for i in range(len(route_indices) - 1):
        current_dirt = dirt_positions[route_indices[i]]
        next_dirt = dirt_positions[route_indices[i + 1]]
        sum += manhattan(current_dirt, next_dirt)
    return sum


def routes_sum(population_set):  # zapisujemy na liście finalne sumy odległości dla każdej z opcji tras
    list_of_sums = np.zeros(routes_num)

    for i in range(routes_num):
        list_of_sums[i] = sum_up_for_route(population_set[i])  # wywołujemy dla każdej trasy na liście

    return list_of_sums


def calculate_fitness(distances):
    # Odwrotność odległości jako fitness
    # Dodajemy małą wartość (np. 1) aby uniknąć dzielenia przez zero
    return 1 / (distances + 1)


def selection(population_set, list_of_sums):
    # Oblicz wartości fitness dla każdej trasy
    fitness_values = calculate_fitness(list_of_sums)
    # Normalizuj wartości fitness, aby sumowały się do 1 (wymagane dla np.random.choice)
    probabilities = fitness_values / fitness_values.sum()
    # Wybierz rodziców na podstawie prawdopodobieństw (wartości fitness)
    progenitor_indices_a = np.random.choice(range(len(population_set)), len(population_set), p=probabilities,  replace=True)
    progenitor_indices_b = np.random.choice(range(len(population_set)), len(population_set), p=probabilities, replace=True)
    # Wybierz rzeczywiste trasy
    progenitor_a = population_set[progenitor_indices_a]
    progenitor_b = population_set[progenitor_indices_b]

    return np.array([progenitor_a, progenitor_b])

def one_point_crossover(parent_a, parent_b): #krzyzowanie jednopunktowe
    crossover_point = np.random.randint(1, len(parent_a))
    child = np.concatenate((parent_a[:crossover_point], [x for x in parent_b if x not in parent_a[:crossover_point]]))
    return child

def population_mating(progenitor_list):
    new_population_set = []
    for i in range(len(progenitor_list[0])):
        progenitor_a, progenitor_b = progenitor_list[0][i], progenitor_list[1][i]
        child = one_point_crossover(progenitor_a, progenitor_b)
        new_population_set.append(child)
    return new_population_set

def mutation_of_child(child):
    for i in range(dirt_count):  # dla każdego elementu dajemy losową szansę zamiany int *rate
        x = np.random.randint(0, dirt_count)
        y = np.random.randint(0, dirt_count)

        child[x], child[y] = child[y], child[x]  # zamiana miejscami

    return child

'''def mutation_of_child(child, mutation_rate=0.1):#procent moze pomoc w niezaklucaniu trasy gdy jesy duza trasa ale idk
    num_mutations = int(len(child) * mutation_rate)
    for _ in range(num_mutations):
        x = np.random.randint(0, len(child))
        y = np.random.randint(0, len(child))
        child[x], child[y] = child[y], child[x]
    return child'''


def mutate_population(new_population_set):
    final_mutated_population = []
    for child in new_population_set:
        final_mutated_population.append(mutation_of_child(child))  # dodajemy zmutowane dziecko do finalnej listy
    return final_mutated_population


if __name__ == '__main__':

    population_set = find_routes(routes_num)
    list_of_sums = routes_sum(population_set)
    progenitor_list = selection(population_set, list_of_sums)
    new_population_set = population_mating(progenitor_list)
    final_mutated_population = mutate_population(new_population_set)
    final_route = [-1, np.inf, np.array([])]  # format listy
    for i in range(20):
        list_of_sums = routes_sum(final_mutated_population)
        # zapisujemy najlepsze rozwiązanie
        if list_of_sums.min() < final_route[1]:
            final_route[0] = i
            final_route[1] = list_of_sums.min()
            final_route[2] = np.array(final_mutated_population)[list_of_sums.min() == list_of_sums]


        progenitor_list = selection(population_set, list_of_sums)
        new_population_set = population_mating(progenitor_list)
        final_mutated_population = mutate_population(new_population_set)

    print(f"Najlepsza trasa znaleziona w iteracji: {final_route[0]}")
    print(f"Minimalna suma odległości: {final_route[1]}")
    print(f"Kolejne pola: {final_route[2]}")