import joblib
from sklearn.calibration import LabelEncoder
from agentActionType import AgentActionType
import time
from garbage import Garbage, GarbageType, RecognizedGarbage
from garbageCan import GarbageCan
from turnCar import turn_left_orientation, turn_right_orientation
from garbageTruck import GarbageTruck
from typing import Tuple, Dict
from gridCellType import GridCellType
from gameContext import GameContext
from agentOrientation import AgentOrientation
import pygame
from bfs import find_path_to_nearest_can
from agentState import AgentState

import tensorflow as tf
from keras.models import load_model
import keras.utils as image
from keras.optimizers import Adam
import numpy as np


def collect_garbage(game_context: GameContext) -> None:
    while True:
        start_agent_state = AgentState(game_context.dust_car.position, game_context.dust_car.orientation)
        path = find_path_to_nearest_can(start_agent_state, game_context.grid, game_context.city)
        if path == None or len(path) == 0:
            break
        move_dust_car(path, game_context)
        next_position = calculate_next_position(game_context.dust_car)
        game_context.grid[next_position] = GridCellType.VISITED_GARBAGE_CAN
        can = game_context.city.cans_dict[next_position]
        can.is_visited = True
        _recognize_garbage(game_context.dust_car, can)
    pass

def _recognize_garbage(dust_car: GarbageTruck, can: GarbageCan) -> None:
    tree_model = joblib.load('machine_learning/model.pkl')
    optimizer = Adam(learning_rate=0.001)
    neural_model = load_model('machine_learning/neuralModel.h5', compile=False)
    neural_model.compile(optimizer=optimizer)
    for garbage in can.garbage:
        predicted_class = predict_class(garbage, tree_model, neural_model)
        garbage_type: GarbageType = None
        if predicted_class == 'PAPER':
            garbage_type = GarbageType.PAPER
        elif predicted_class == 'PLASTIC_AND_METAL':
            garbage_type = GarbageType.PLASTIC_AND_METAL
        elif garbage_type == 'GLASS':
            garbage_type = GarbageType.GLASS
        elif predicted_class == 'BIO' :
            garbage_type = GarbageType.BIO
        elif predicted_class == 'MIXED':
            garbage_type = GarbageType.MIXED
        print(predicted_class)
        recognized_garbage = RecognizedGarbage(garbage, garbage_type)
        dust_car.sort_garbage(recognized_garbage)

def predict_class(garbage: Garbage, tree_model, neural_model) -> str:
    if garbage.img is None:
        return predict_class_from_tree(garbage, tree_model)
    return predict_class_from_neural_model(garbage, neural_model)

def predict_class_from_tree(garbage: Garbage, tree_model) -> str:
    attributes = [garbage.shape, garbage.flexibility, garbage.does_smell, garbage.weight, garbage.size, garbage.color, garbage.softness, garbage.does_din]
    encoded = attributes_to_floats(attributes)
    return tree_model.predict([encoded])[0]

def predict_class_from_neural_model(garbage: Garbage, neural_model) -> str:
    img = image.load_img(garbage.img, target_size=(150, 150))
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array /= 255.

    predictions = neural_model.predict(img_array)
    prediction = np.argmax(predictions[0])
    if prediction == 0:
        return "BIO"
    if prediction == 1:
        return "GLASS"
    if prediction == 2:
        return "MIXED"
    if prediction == 3:
        return "PAPER"
    if prediction == 4:
        return "PLASTIC_AND_METAL"

def attributes_to_floats(attributes: list[str]) -> list[float]:
    output: list[float] = []
    if attributes[0] == 'Longitiudonal':
        output.append(0)
    elif attributes[0] == 'Round':
        output.append(1)
    elif attributes[0] == 'Flat':
        output.append(2)
    elif attributes[0] == 'Irregular':
        output.append(3)
    
    if attributes[1] == 'Low':
        output.append(0)
    elif attributes[1] == 'Medium':
        output.append(1)
    elif attributes[1] == 'High':
        output.append(2)
    

    if attributes[2] == "Yes":
        output.append(0)
    else:
        output.append(1)
    
    if attributes[3] == 'Low':
        output.append(0)
    elif attributes[3] == 'Medium':
        output.append(1)
    elif attributes[3] == 'High':
        output.append(2)

    if attributes[4] == 'Low':
        output.append(0)
    elif attributes[4] == 'Medium':
        output.append(1)
    elif attributes[4] == 'High':
        output.append(2)

    if attributes[5] == 'Transparent':
        output.append(0)
    elif attributes[5] == 'Light':
        output.append(1)
    elif attributes[5] == 'Dark':
        output.append(2)
    elif attributes[5] == "Colorful":
        output.append(3)

    if attributes[6] == 'Low':
        output.append(0)
    elif attributes[6] == 'Medium':
        output.append(1)
    elif attributes[6] == 'High':
        output.append(2)

    if attributes[7] == "Yes":
        output.append(0)
    else:
        output.append(1)
    return output


def move_dust_car(actions: list[AgentActionType], game_context: GameContext) -> None:
    for action in actions:
        street_position = game_context.dust_car.position
        has_to_render_street = False
        if action == AgentActionType.TURN_LEFT:
            game_context.dust_car.orientation = turn_left_orientation(game_context.dust_car.orientation)
        elif action == AgentActionType.TURN_RIGHT:
            game_context.dust_car.orientation = turn_right_orientation(game_context.dust_car.orientation)
        elif action == AgentActionType.MOVE_FORWARD:
            game_context.dust_car.position = calculate_next_position(game_context.dust_car)
            has_to_render_street = True
        game_context.dust_car.render(game_context)
        if has_to_render_street:
            if game_context.grid[street_position] == GridCellType.STREET_HORIZONTAL:
                game_context.render_in_cell(street_position, "imgs/street_horizontal.png")
            elif game_context.grid[street_position] == GridCellType.STREET_VERTICAL:
                game_context.render_in_cell(street_position, "imgs/street_vertical.png")
            elif game_context.grid[street_position] == GridCellType.SPEED_BUMP:
                game_context.render_in_cell(street_position, "imgs/speed_bump.png")
        pygame.display.update()
        time.sleep(0.15)


def calculate_next_position(car: GarbageTruck) -> Tuple[int, int]:
    if car.orientation == AgentOrientation.UP:
        if car.position[1] - 1 < 1:
            return None
        return (car.position[0], car.position[1] - 1)
    if car.orientation == AgentOrientation.DOWN:
        if car.position[1] + 1 > 27:
            return None
        return (car.position[0], car.position[1] + 1)
    if car.orientation == AgentOrientation.LEFT:
        if car.position[0] - 1 < 1:
            return None
        return (car.position[0] - 1, car.position[1])
    if car.position[0] + 1 > 27:
        return None
    return (car.position[0] + 1, car.position[1])