import pygame
import numpy as np
from board import Board
from constant import width, height, rows, cols
from tractor import Tractor
from kolejka import Stan, Odwiedzone
from queue import Queue
import pandas as pd
from neuralnetwork import load_model
from decisiontree import train_decision_tree

data = pd.read_csv('dane.csv')




model_path = 'model.pth'
neuralnetwork_model = load_model(model_path)

model, feature_columns = train_decision_tree(data)

fps = 15
WIN = pygame.display.set_mode((width, height))
pygame.display.set_caption('Inteligentny Traktor')

def goal_test(elem, board):
    return board.is_dirt(elem.row, elem.col)

def actions(elem, istate):
    akcje = []
    while elem.row != istate.row or elem.col != istate.col or elem.direction != istate.direction:
        akcje.append(elem.a)
        elem = elem.p[0]
    return akcje[::-1]

def graphsearch(istate, board):
    explored = Odwiedzone()
    fringe = Queue()
    fringe.put(istate)
    moves = ["up", "left", "right"]
    while not fringe.empty():
        elem = fringe.get()
        if goal_test(elem, board):
            return actions(elem, istate), elem
        explored.dodaj_stan(elem)
        for action in moves:
            stan = elem.succ(action, board)
            if stan is not None:
                if not fringe_check(fringe, stan) and not explored.check(stan):
                    stan.parrent(elem, action)
                    fringe.put(stan)
    return [], None

def fringe_check(fringe, stan):
    with fringe.mutex:
        for item in fringe.queue:
            if stan.direction == item.direction and stan.col == item.col and stan.row == item.row:
                return True
    return False



def main():
    initial_state = Stan(9, 1, "down")
    run = True
    clock = pygame.time.Clock()

    board = Board(load_from_file=True, filename='generated_board.npy')


    tractor = Tractor(9, 1, model, feature_columns, neuralnetwork_model)

    while run:
        clock.tick(fps)

        if all(not board.is_dirt(row, col) for row in range(rows) for col in range(cols)):
            print("Traktor odwiedził wszystkie pola.")
            break

        akcje, nowy_stan = graphsearch(initial_state, board)

        if not akcje:
            print("Nie znaleziono ścieżki do najbliższego pola dirt.")
            board = Board(load_from_file=True, filename='generated_board.npy')
            initial_state = Stan(0, 1, "down")
            tractor = Tractor(0, 1, model, feature_columns, neuralnetwork_model)
            while board.is_rock(initial_state.row, initial_state.col):
                board = Board(load_from_file=True, filename='generated_board.npy')
            continue

        print("akcje: >", akcje)

        while akcje:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    run = False

            akcja = akcje.pop(0)
            if akcja == "left":
                tractor.turn_left()
            elif akcja == "right":
                tractor.turn_right()
            elif akcja == "up":
                tractor.move_forward(board)

            board.draw_cubes(WIN)
            tractor.draw(WIN)
            pygame.display.update()


        initial_state = nowy_stan
        initial_state.direction = tractor.direction

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return

    pygame.quit()

main()