from enum import Enum
import random
import pygame
import sys
from animals import create_animals
from agent import Agent
from enclosure import Enclosure, create_enclosures
from spawner import Spawner
from state_space_search import graphsearch
from terrain_obstacle import Terrain_Obstacle

BLACK = (0, 0, 0)
RED = (255, 0, 0)

GRID_SIZE = 50
GRID_WIDTH = 30
GRID_HEIGHT = 15

pygame.init()

WINDOW_SIZE = (GRID_WIDTH * GRID_SIZE, GRID_HEIGHT * GRID_SIZE)
screen = pygame.display.set_mode(WINDOW_SIZE)
pygame.display.set_caption("Mini Zoo")



background_image = pygame.image.load('images/tło.jpg')
background_image = pygame.transform.scale(background_image, WINDOW_SIZE)
fenceH = pygame.image.load('images/fenceHor.png')
fenceV = pygame.image.load('images/fenceVer.png')
gate = pygame.image.load('images/gate.png')
puddle_image = pygame.image.load('images/puddle.png')
bush_image = pygame.image.load('images/bush.png')
obstacles = set()
animals_position = set()
terrain_obstacles_position = set()

Animals = create_animals()
Enclosures = create_enclosures()

puddle1 = Terrain_Obstacle(0,0,'puddle', puddle_image)
puddle2 = Terrain_Obstacle(0,0,'puddle', puddle_image)
puddle3 = Terrain_Obstacle(0,0,'puddle', puddle_image)
puddle4 = Terrain_Obstacle(0,0,'puddle', puddle_image)
puddle5 = Terrain_Obstacle(0,0,'puddle', puddle_image)
puddle6 = Terrain_Obstacle(0,0,'puddle', puddle_image)
puddle7 = Terrain_Obstacle(0,0,'puddle', puddle_image)
bush1 = Terrain_Obstacle(0,0,'bush', bush_image)
bush2 = Terrain_Obstacle(0,0,'bush', bush_image)
bush3 = Terrain_Obstacle(0,0,'bush', bush_image)
bush4 = Terrain_Obstacle(0,0,'bush', bush_image)
bush5 = Terrain_Obstacle(0,0,'bush', bush_image)
Terrain_Obstacles = [puddle1, puddle2, puddle3, puddle4,  puddle5, puddle6, puddle7, bush1, bush2, bush3, bush4, bush5]

def draw_grid():
    for y in range(0, GRID_HEIGHT * GRID_SIZE, GRID_SIZE):
        for x in range(0, GRID_WIDTH * GRID_SIZE, GRID_SIZE):
            rect = pygame.Rect(x, y, GRID_SIZE, GRID_SIZE)
            pygame.draw.rect(screen, BLACK, rect, 1)

def draw_enclosures():
    for enclosure in Enclosures:
        enclosure.draw(screen, GRID_SIZE)

def draw_gates():
    for enclosure in Enclosures:
        enclosure.gatebuild(screen, GRID_SIZE)

def draw_Animals():
    for Animal in Animals:
        Animal.draw(screen, GRID_SIZE)
        if Animal.feed() == 'True':
           Animal.draw_exclamation(screen, GRID_SIZE, Animal.x, Animal.y)
        else: 
            Animal.draw_food(screen,GRID_SIZE,Animal.x,Animal.y)

def spawn_all_animals():
    for Animal in Animals:
        spawner1 = Spawner(Animal)
        spawner1.spawn_animal(obstacles, animals_position, Enclosures)

def draw_Terrain_Obstacles():
    for terrain_obstacle in Terrain_Obstacles:
        terrain_obstacle.draw(screen, GRID_SIZE)

def spawn_obstacles():
    for terrain_obstacle in Terrain_Obstacles:
        spawner2= Spawner(terrain_obstacle)
        spawner2.spawn_terrain_obstacles(obstacles,animals_position, terrain_obstacles_position, GRID_WIDTH, GRID_HEIGHT)


obstacles = set()
def generate_obstacles():

    for en in Enclosures:
        # Pobierz współrzędne bramy
        gate_x, gate_y = en.gate
        gate_x -= 1
        gate_y -= 1
        
        # Dodaj lewy brzeg prostokąta
        for y in range(en.y1, en.y2 + 1):
            if (en.x1, y) != (gate_x, gate_y):
                obstacles.add((en.x1, y))

        # Dodaj prawy brzeg prostokąta
        for y in range(en.y1, en.y2 + 1):
            if (en.x2, y) != (gate_x, gate_y):
                obstacles.add((en.x2, y))

        # Dodaj górny brzeg prostokąta
        for x in range(en.x1+1, en.x2):
            if (x, en.y1) != (gate_x, gate_y):
                obstacles.add((x, en.y1))

        # Dodaj dolny brzeg prostokąta
        for x in range(en.x1+1, en.x2):
            if (x, en.y2) != (gate_x, gate_y):
                obstacles.add((x, en.y2))

    return obstacles

cost_map = {}
def generate_cost_map():
    adult_animal_cost = 10
    baby_animal_cost = 5
    puddle_cost = 50
    bush_cost = 20
    wall_cost = 1000
    for animal in Animals:
        if animal.adult:
            cost_map[(animal.x + 1, animal.y + 1)] = baby_animal_cost
            cost_map[(animal.x + 1, animal.y)] = baby_animal_cost
            cost_map[(animal.x, animal.y + 1)] = baby_animal_cost
            cost_map[(animal.x, animal.y)] = adult_animal_cost
        else:
            cost_map[(animal.x, animal.y)] = baby_animal_cost

    for terrain_obstacle in Terrain_Obstacles:
        if terrain_obstacle.type == 'puddle':
            cost_map[(terrain_obstacle.x , terrain_obstacle.y )] = puddle_cost
        else:
            cost_map[(terrain_obstacle.x , terrain_obstacle.y )] = bush_cost

    for wall in Walls:
        cost_map[wall[0], wall[1]] = wall_cost
    # Inne pola z różnym kosztem
    # cost_map[(x, y)] = cost_value

# region Fields Tests
available_fields_small = set()
available_fields_large = set()

WHITE = (255,255,255)
GREEN = (0, 255, 0)
YELLOW = (255, 255, 0)
BLACK = (0, 0, 0)

def generate_available_fields():
    for enclosure in Enclosures:
        for x in range(enclosure.x1 + 1, enclosure.x2):
            for y in range(enclosure.y1 + 1, enclosure.y2):
                field = (x, y)
                if field not in obstacles:
                    available_fields_small.add(field)
                    if x < enclosure.x2 - 1 and y < enclosure.y2 - 1:
                        available_fields_large.add(field)

def draw_fields(fields, color):
    for field in fields:
        x, y = field
        pygame.draw.rect(screen, color, (x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE, GRID_SIZE))

def available_fields_tests():
    obstacles = generate_obstacles()

    while True:
        screen.fill(WHITE)

        draw_grid()
        draw_fields(obstacles, BLACK)

        generate_available_fields()

        draw_fields(available_fields_small, GREEN)
        draw_fields(available_fields_large, YELLOW)

        # Odświeżenie ekranu
        pygame.display.flip()
# endregion

# region Main Code
def main():
    initial_state = (0,0,'S')
    agent = Agent(initial_state, 'images/agent1.png', GRID_SIZE)
    
    obstacles = generate_obstacles()
    actions = []
    clock = pygame.time.Clock()
    
    spawned = False
    while True:
        # Manualne poruszanie agentem
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            agent.handle_event(event, GRID_WIDTH, GRID_HEIGHT, Animals, obstacles)

        screen.blit(background_image,(0,0))
        draw_grid()
        draw_enclosures()
        draw_gates()
        if not spawned:
            spawn_all_animals()
            spawn_obstacles()
            generate_cost_map()
            for animal in Animals:
                animal._feed = 2 # Ustawienie aby zwierzę było głodne
            spawned = True
        draw_Animals()
        draw_Terrain_Obstacles()
        agent.draw(screen, GRID_SIZE)
        pygame.display.flip()
        clock.tick(10)

        if actions:
            action = actions.pop(0)
            agent.move(action, GRID_WIDTH, GRID_HEIGHT, obstacles, Animals, goal)
            pygame.time.wait(200)
        else:
            animal = random.choice(Animals)
            goal = (animal.x, animal.y)
 
            # --- Zaznaczenie celu ---
            pygame.draw.rect(screen, RED, (animal.x * GRID_SIZE, animal.y * GRID_SIZE, GRID_SIZE, GRID_SIZE))
            pygame.display.flip()
            pygame.time.delay(2000) 
            # ------------------------

            actions = graphsearch(agent.istate, goal, GRID_WIDTH, GRID_HEIGHT, obstacles, cost_map)
# endregion

Walls = []
# region A* Test
from elephant import Elephant
puddle1 = Terrain_Obstacle(15,8,'puddle', puddle_image)
bush1 = Terrain_Obstacle(15,6,'bush', bush_image)
animal = Elephant(15, 10)
animal1 = Elephant(14, 10)
animal2 = Elephant(13, 10)
animal3 = Elephant(12, 10)
animal4 = Elephant(11, 10)
Animals = [animal, animal1, animal2, animal3, animal4]
Terrain_Obstacles = [puddle1, bush1]

empty_rows = [5, 7, 9]

def generate_test_walls():
    for x in range(4,26):
        for y in range(0, 15):
            if y not in empty_rows:
                Walls.append((x, y))
    return Walls

def draw_test_walls():
    for wall in generate_test_walls():
        pygame.draw.rect(screen, BLACK, (wall[0] * GRID_SIZE, wall[1] * GRID_SIZE, GRID_SIZE, GRID_SIZE))

def a_star_testing():
    initial_state = (0, 7, 'E')
    agent = Agent(initial_state, 'images/agent1.png', GRID_SIZE)
    goal = (29, 7)

    actions = []
    clock = pygame.time.Clock()

    generated = False
    while True:
        screen.fill(WHITE)
        draw_grid()
        draw_test_walls() 
        draw_Terrain_Obstacles()
        draw_Animals()
        if not generated:
            generate_cost_map()

        agent.draw(screen, GRID_SIZE)
        pygame.draw.rect(screen, RED, (goal[0] * GRID_SIZE, goal[1] * GRID_SIZE, GRID_SIZE, GRID_SIZE))
        
        pygame.display.flip()
        clock.tick(10)

        if actions:
            action = actions.pop(0)
            agent.move(action, GRID_WIDTH, GRID_HEIGHT, obstacles, Animals, goal)
            pygame.time.wait(100)
        else:
            actions = graphsearch(agent.istate, goal, GRID_WIDTH, GRID_HEIGHT, obstacles, cost_map)
# endregion

class DebugMode(Enum):
    MAIN = 1
    AVAILABLE_FIELDS = 2
    A_STAR_TESTING = 3

if __name__ == "__main__":
    debug_mode = DebugMode.A_STAR_TESTING

    if debug_mode == DebugMode.MAIN:
        main()
    elif debug_mode == DebugMode.AVAILABLE_FIELDS:
        available_fields_tests()
    elif debug_mode == DebugMode.A_STAR_TESTING:
        a_star_testing()