import random
import numpy as np
import pygame

# Wymiary planszy
rows, cols = 10, 10
size = 64
# Klasa reprezentująca planszę
class Board:
    def __init__(self):
        self.board = []
        self.vegetables = []
        self.vegetable_names = []
        self.load_images()
        self.generate_board()

    # Metoda do ładowania obrazów
    def load_images(self):
        try:
            self.grass = pygame.image.load("board/grass.png")  # Załaduj obraz trawy
            self.dirt = pygame.image.load("board/dirt.png")  # Załaduj obraz ziemi
            self.rock = pygame.image.load("board/rock.png")  # Załaduj obraz kamienia
        except pygame.error as e:
            print(f"Failed to load image: {e}")
            self.grass = pygame.Surface((size, size))
            self.grass.fill((0, 255, 0))  # Zastępczy kolor zielony (trawa)
            self.dirt = pygame.Surface((size, size))
            self.dirt.fill((139, 69, 19))  # Zastępczy kolor brązowy (ziemia)
            self.rock = pygame.Surface((size, size))
            self.rock.fill((128, 128, 128))  # Zastępczy kolor szary (kamień)

        # Tworzenie powierzchni dla różnych warzyw
        self.warzywa_images = {
            "marchewka": [pygame.image.load(f"warzywa/Carrot/{i}.jpg") for i in range(1, 10)],
            "ziemniak": [pygame.image.load(f"warzywa/Potato/{i}.jpg") for i in range(1, 10)],
            "pomidor": [pygame.image.load(f"warzywa/tomato/{i}.jpg") for i in range(1, 10)],
            "fasola": [pygame.image.load(f"warzywa/Bean/{i}.jpg") for i in range(1, 10)],
            "dynia": [pygame.image.load(f"warzywa/Pumpkin/{i}.jpg") for i in range(1, 10)],
            "rzodkiewka": [pygame.image.load(f"warzywa/Radish/{i}.jpg") for i in range(1, 10)],
            "ogorek": [pygame.image.load(f"warzywa/Cucumber/{i}.jpg") for i in range(1, 10)],
            "kalafior": [pygame.image.load(f"warzywa/Cauliflower/{i}.jpg") for i in range(1, 10)],
            "kapusta": [pygame.image.load(f"warzywa/Cabbage/{i}.jpg") for i in range(1, 10)],
            "brokul": [pygame.image.load(f"warzywa/Broccoli/{i}.jpg") for i in range(1, 10)]
        }

        # Typy warzyw przypisane do liczb
        self.vegetable_types = {
            "marchewka": 2,
            "ziemniak": 3,
            "pomidor": 4,
            "fasola": 5,
            "dynia": 6,
            "rzodkiewka": 7,
            "ogorek": 8,
            "kalafior": 9,
            "kapusta": 10,
            "brokul": 11
        }

    def count_vegetable_type(self, vegetable_type):
        # Funkcja zliczająca ilość pól z określonym typem warzywa
        count = 0
        for row in range(rows):
            for col in range(cols):
                if self.vegetable_names[row][col] == None:
                    count += 1
        return count
    
    # Metoda do generowania początkowej planszy
    def generate_board(self):
        # Zwiększenie prawdopodobieństwa generowania trawy i kamieni
        self.board = [[random.choice([0, 0, 0, 0, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) for _ in range(cols)] for _ in range(rows)]
        self.vegetables = [[None for _ in range(cols)] for _ in range(rows)]
        self.vegetable_names = [[None for _ in range(cols)] for _ in range(rows)]

        # Losowe przypisanie warzyw do planszy
        for row in range(rows):
            for col in range(cols):
                if self.board[row][col] in self.vegetable_types.values():
                    vegetable_type = list(self.warzywa_images.keys())[self.board[row][col] - 2]
                    vegetable_image = random.choice(self.warzywa_images[vegetable_type])
                    self.vegetables[row][col] = vegetable_image
                    self.vegetable_names[row][col] = vegetable_type


    # Metoda do generowania cech gleby


    # Metoda oceniająca jakość planszy
    def evaluate(self):
        score = 0
        directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]  # Kierunki: prawo, dół, lewo, góra
        type_0_count = self.count_vegetable_type(0)
        # Sprawdzanie sąsiednich pól dla każdego warzywa
        for row in range(rows):
            for col in range(cols):
                if self.is_vegetable(row, col):
                    for dr, dc in directions:
                        new_row, new_col = row + dr, col + dc
                        if 0 <= new_row < rows and 0 <= new_col < cols:
                            if self.vegetable_names[row][col] == self.vegetable_names[new_row][new_col]:
                                score += 1
        
        if type_0_count < 35:
            score -= (35 - type_0_count) * 2

        return score

    # Metoda mutująca planszę
    def mutate(self, mutation_rate):
        for row in range(rows):
            for col in range(cols):
                if random.random() < mutation_rate:
                    self.board[row][col] = random.choice([0, 0, 0, 0, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
                    if self.board[row][col] in self.vegetable_types.values():
                        vegetable_type = list(self.warzywa_images.keys())[self.board[row][col] - 2]
                        vegetable_image = random.choice(self.warzywa_images[vegetable_type])
                        self.vegetables[row][col] = vegetable_image
                        self.vegetable_names[row][col] = vegetable_type

    # Metoda krzyżująca planszę z inną planszą
    def crossover(self, other_board):
        crossover_point = random.randint(0, rows * cols - 1)
        for i in range(rows):
            for j in range(cols):
                if i * cols + j > crossover_point:
                    self.board[i][j], other_board.board[i][j] = other_board.board[i][j], self.board[i][j]
                    self.vegetables[i][j], other_board.vegetables[i][j] = other_board.vegetables[i][j], self.vegetables[i][j]
                    self.vegetable_names[i][j], other_board.vegetable_names[i][j] = other_board.vegetable_names[i][j], self.vegetable_names[i][j]

    # Metoda kopiująca planszę
    def copy(self):
        new_board = Board()
        new_board.board = [row[:] for row in self.board]
        new_board.vegetables = [row[:] for row in self.vegetables]
        new_board.vegetable_names = [row[:] for row in self.vegetable_names]
        return new_board

    # Sprawdza, czy dane pole zawiera warzywo
    def is_vegetable(self, row, col):
        return self.board[row][col] in self.vegetable_types.values()

    # Rysowanie planszy na oknie pygame
    def draw_cubes(self, win):
        for row in range(rows):
            for col in range(cols):
                cube_rect = pygame.Rect(col * size, row * size, size, size)
                cube = self.board[row][col]
                if cube == 0:
                    win.blit(self.grass, cube_rect)  # Użyj obrazu trawy
                elif cube == 1:
                    rock_scale = pygame.transform.scale(self.rock, (size, size))
                    win.blit(self.dirt, cube_rect)
                    win.blit(rock_scale, cube_rect)
                else:
                    if self.vegetables[row][col]:
                        vegetable_image = pygame.transform.scale(self.vegetables[row][col], (size, size))
                        win.blit(vegetable_image, cube_rect)

# Funkcja oceniająca planszę (wywołuje metodę evaluate)
def evaluate(board):
    return board.evaluate()

# Generowanie początkowej populacji plansz
def generate_population(size):
    return [Board() for _ in range(size)]

# Selekcja metodą ruletki
def roulette_wheel_selection(population, fitnesses):
    total_fitness = sum(fitnesses)
    selection_probs = [f / total_fitness for f in fitnesses]
    return population[random.choices(range(len(population)), weights=selection_probs, k=1)[0]]

# Krzyżowanie jednopunktowe
def crossover(parent1, parent2):
    child1, child2 = parent1.copy(), parent2.copy()
    child1.crossover(child2)
    return child1, child2

# Mutacja planszy
def mutate(board, mutation_rate):
    board.mutate(mutation_rate)

# Algorytm genetyczny
def genetic_algorithm(pop_size, generations, mutation_rate):
    population = generate_population(pop_size)
    for generation in range(generations):
        fitnesses = [evaluate(board) for board in population]
        new_population = [max(population, key=evaluate).copy()]  # Elityzm - zachowanie najlepszego osobnika
        while len(new_population) < pop_size:
            parent1 = roulette_wheel_selection(population, fitnesses)
            parent2 = roulette_wheel_selection(population, fitnesses)
            offspring1, offspring2 = crossover(parent1, parent2)
            mutate(offspring1, mutation_rate)
            mutate(offspring2, mutation_rate)
            new_population.extend([offspring1, offspring2])
        population = new_population[:pop_size]
        print(f"Generation {generation}: Best Fitness = {max(fitnesses)}")
    best_board = max(population, key=evaluate)
    return best_board

# Funkcja do zapisywania planszy jako obraz PNG
def save_board_as_image(board, filename):
    surface = pygame.Surface((cols * size, rows * size))
    board.draw_cubes(surface)
    pygame.image.save(surface, filename)


pop_size = 50
generations = 300
mutation_rate = 0.03

# Inicjalizacja Pygame
pygame.init()
win = pygame.display.set_mode((cols * size, rows * size))
pygame.display.set_caption('Generated Board')

# Generowanie najlepszej planszy
best_board = genetic_algorithm(pop_size, generations, mutation_rate)

# Rysowanie najlepszej planszy
best_board.draw_cubes(win)
pygame.display.update()

# Zapisywanie planszy do pliku
np.save('generated_board.npy', best_board.board)
save_board_as_image(best_board, 'generated_board.png')

# Utrzymanie okna otwartego do zamknięcia przez użytkownika
run = True
needs_redraw = True
while run:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            run = False
        elif event.type == pygame.MOUSEBUTTONDOWN or event.type == pygame.KEYDOWN:
            needs_redraw = True

    if needs_redraw:
        best_board.draw_cubes(win)
        pygame.display.flip()
        needs_redraw = False

pygame.quit()