from matplotlib.pylab import randint
from numpy import ones, zeros, random
from sklearn.metrics import mean_squared_error as MSE


random.seed(123)

OBJECTS = {
    "book": {"dims": [4, 4, 2], "stack": True},
    "penstand": {"dims": [2, 2, 4], "stack": True},
    "bottle": {"dims": [2, 2, 6], "stack": False},
    "pen": {"dims": [1, 1, 4], "stack": False},
    "pencil": {"dims": [1, 1, 6], "stack": False},
    "eraser": {"dims": [2, 1, 1], "stack": False},
    "powerbank": {"dims": [4, 2, 1], "stack": False},
    "mobile": {"dims": [4, 2, 1], "stack": False},
    "laptop": {"dims": [6, 4, 1], "stack": True},
    "monitor": {"dims": [6, 4, 2], "stack": False},
    "keyboard": {"dims": [6, 2, 1], "stack": False},
    "mouse": {"dims": [4, 2, 1], "stack": False},
    "headphones": {"dims": [4, 4, 2], "stack": False},
    "charger": {"dims": [2, 2, 1], "stack": False},
    "notebook": {"dims": [4, 4, 1], "stack": True},
    "folder": {"dims": [4, 4, 1], "stack": True},
    "backpack": {"dims": [6, 4, 2], "stack": False},
    "pouch": {"dims": [4, 4, 2], "stack": False},
}

OBJECT_NAMES = [
    "book",
    "penstand",
    "bottle",
    "pen",
    "pencil",
    "eraser",
    "powerbank",
    "mobile",
    "laptop",
    "monitor",
    "keyboard",
    "mouse",
    "headphones",
    "charger",
    "notebook",
    "folder",
    "backpack",
    "pouch",
]

ACTION_CONFIG = {
    "RIGHT": [1, 0, 0],
    "LEFT": [-1, 0, 0],
    "UP": [0, 1, 0],
    "DOWN": [0, -1, 0],
    "FORWARD": [0, 0, 1],
    "BACKWARD": [0, 0, -1],
    "ROTATE": [0, 0, 0],
}


def appendRewardFeedback(state, choice, feedback, reward):
    state.rewardList.append(reward)
    state.rewardFeedback.append(feedback)
    if choice == "segment":
        state.rewardFeedbackSegment.append(feedback)
        state.rewardListSegment.append(reward)
    elif choice == "place":
        state.rewardFeedbackPlace.append(feedback)
        state.rewardListPlace.append(reward)
    elif choice == "adjust":
        state.rewardFeedbackAdjust.append(feedback)
        state.rewardListAdjust.append(reward)


def initDimentions(obj):
    dims = obj.get("dims")
    if dims is None:
        return []

    return ones(dims, dtype=int).tolist()


def initGrid():
    sizeX, sizeY, sizeZ = randint(8, 12), randint(8, 12), randint(8, 12)
    grid = zeros((sizeX, sizeY, sizeZ), dtype=int).tolist()

    numObjs = randint(3, len(OBJECT_NAMES) + 1)
    chosenNames = random.choice(OBJECT_NAMES, size=numObjs, replace=False)

    placed = {}

    for name in chosenNames:
        obj = OBJECTS.get(name)

        dimX, dimY, dimZ = obj["dims"]

        if dimX > sizeX or dimY > sizeY or dimZ > sizeZ:
            continue

        isPlaced = False
        tryPlaced = 0

        while not isPlaced and tryPlaced < 100:
            posX = randint(0, sizeX - dimX + 1)
            posY = randint(0, sizeY - dimY + 1)
            posZ = 0

            canPlace = True
            for i in range(dimX):
                for j in range(dimY):
                    for k in range(dimZ):
                        if (
                            grid[posX + i][posY + j][posZ + k] != 0
                            and obj["stack"] == False
                        ):
                            canPlace = False
                            break
                        else:
                            canPlace = True
                    if not canPlace:
                        break
                if not canPlace:
                    break

            if canPlace:
                for i in range(dimX):
                    for j in range(dimY):
                        for k in range(dimZ):
                            if (
                                obj["stack"]
                                and grid[posX + i][posY + j][posZ + k] > 0
                                and posZ + k + 1 < sizeZ
                            ):
                                grid[posX + i][posY + j][posZ + k + 1] += 1
                            else:
                                grid[posX + i][posY + j][posZ + k] += 1

            placed[name] = (posX, posY, posZ, obj["stack"])
            isPlaced = True

    return (grid, placed)


def initWeightedGrid(shape=None):
    if shape is None:
        shape = (randint(8, 12), randint(8, 12), randint(8, 12))

    grid = random.uniform(0, 1, shape)

    x_mid = grid.shape[0] // 2
    x_span = grid.shape[0] // 4
    y_front = grid.shape[1] // 3

    grid[x_mid - x_span : x_mid + x_span, :y_front, :] *= 0.2

    return grid


def _get_weight_value(weight, x, y, z):
    if not weight or not weight[0] or not weight[0][0]:
        return 0.0

    if (
        x < 0
        or y < 0
        or z < 0
        or x >= len(weight)
        or y >= len(weight[0])
        or z >= len(weight[0][0])
    ):
        return 0.0

    return weight[x][y][z]


def place(segment, objects, state):
    dims = state.currentGrid
    weight = state.weightedGrid
    objsPresent = state.ObjectsPresent

    reward = 0.0
    totalObjs = len(objects)
    reward_per_obj_placed = 45.0 / totalObjs

    placement_failed = False

    if segment:
        appendRewardFeedback(
            state, "place", "Placing objects with segmentation is not allowed.", -60.0
        )
        return (-60.0, True)

    for obj_name, pos in objects.items():

        obj = OBJECTS.get(obj_name)
        if obj is None:
            appendRewardFeedback(
                state,
                "place",
                f"Object '{obj_name}' is not recognized.",
                -reward_per_obj_placed,
            )
            reward -= reward_per_obj_placed
            continue

        objGrid = initDimentions(obj)
        for i in range(len(objGrid)):
            for j in range(len(objGrid[0])):
                for k in range(len(objGrid[0][0])):
                    if (
                        pos[0] + i >= len(dims)
                        or pos[1] + j >= len(dims[0])
                        or pos[2] + k >= len(dims[0][0])
                    ):
                        reward -= reward_per_obj_placed
                        appendRewardFeedback(
                            state,
                            "place",
                            f"Object '{obj_name}' placement is out of bounds.",
                            -reward_per_obj_placed,
                        )
                        placement_failed = True
                        break

                    if dims[pos[0] + i][pos[1] + j][pos[2] + k] > 0 and pos[3] == False:
                        reward -= reward_per_obj_placed
                        appendRewardFeedback(
                            state,
                            "place",
                            f"Object '{obj_name}' placement overlaps with another object and stacking is not allowed.",
                            -reward_per_obj_placed,
                        )
                        placement_failed = True
                        break

                    elif (
                        dims[pos[0] + i][pos[1] + j][pos[2] + k] > 0 and pos[3] == True
                    ):
                        if pos[2] + k + 1 < len(dims[0][0]):
                            dims[pos[0] + i][pos[1] + j][pos[2] + k + 1] += 1
                            bonus = (
                                _get_weight_value(
                                    weight,
                                    pos[0] + i,
                                    pos[1] + j,
                                    pos[2] + k + 1,
                                )
                                * reward_per_obj_placed
                            )
                            reward += bonus
                            appendRewardFeedback(
                                state,
                                "place",
                                f"Object '{obj_name}' placed with stacking. Bonus: {bonus:.2f}",
                                bonus,
                            )
                        else:
                            reward -= reward_per_obj_placed
                            appendRewardFeedback(
                                state,
                                "place",
                                f"Object '{obj_name}' placement failed. No space for stacking.",
                                -reward_per_obj_placed,
                            )
                            placement_failed = True

                        break

                    else:
                        dims[pos[0] + i][pos[1] + j][pos[2] + k] = 1
                        bonus = reward_per_obj_placed * _get_weight_value(
                            weight, pos[0] + i, pos[1] + j, pos[2] + k
                        )
                        reward += bonus
                        appendRewardFeedback(
                            state,
                            "place",
                            f"Object '{obj_name}' placed successfully. Bonus: {bonus:.2f}",
                            bonus,
                        )
                if placement_failed:
                    break
            if placement_failed:
                break

        if not placement_failed:
            state.ObjectsPlaced[obj_name] = pos
            state.numberPlaced += 1
            try:
                if objsPresent[obj_name] == state.ObjectsPlaced[obj_name]:
                    reward -= 45.0 / totalObjs
                    appendRewardFeedback(
                        state,
                        "place",
                        f"Object '{obj_name}' is being placed in the same location",
                        -reward_per_obj_placed,
                    )
            except KeyError:
                reward -= reward_per_obj_placed
                appendRewardFeedback(
                    state,
                    "place",
                    f"Object '{obj_name}' is present in the environment, but is placed in same location as originally found.",
                    -reward_per_obj_placed,
                )

                continue

    return (reward, placement_failed)


def findobject(segment, objects, state):

    if not segment or segment is None:
        appendRewardFeedback(
            state,
            "segment",
            "Finding objects without segmentation is not allowed.",
            -60.0,
        )
        return -60.0

    if set(state.objectsFound) == set(state.ObjectsPresent.keys()):
        appendRewardFeedback(
            state,
            "segment",
            "No point in finding more objects as all are already found. Make the isSegement attribute false and execute the place method.",
            -60.0,
        )
        return -60.0

    reward = 0.0
    glMetric = 45.0 / len(state.ObjectsPresent)
    objs = []
    for obj_found, pos_found in objects.items():
        pos_real = state.ObjectsPresent.get(obj_found)
        if pos_real is None:
            reward -= glMetric
            appendRewardFeedback(
                state,
                "segment",
                f"Object '{obj_found}' not found in the environment.",
                -glMetric,
            )
            continue

        if pos_found == pos_real:
            reward += glMetric
            appendRewardFeedback(
                state,
                "segment",
                f"Object '{obj_found}' found with correct position and stacking.",
                glMetric,
            )
            objs.append(obj_found)
        else:
            mse = MSE(pos_real[:3], pos_found[:3])
            reward -= mse
            appendRewardFeedback(
                state,
                "segment",
                f"Object '{obj_found}' found with incorrect position. MSE: {mse:.2f}",
                -mse,
            )

        if pos_found[3] != pos_real[3]:
            reward -= glMetric / 4.0
            appendRewardFeedback(
                state,
                "segment",
                f"Object '{obj_found}' found with incorrect stacking. Penalty: {glMetric / 4.0}",
                -glMetric / 4.0,
            )
        else:
            reward += glMetric / 4.0
            appendRewardFeedback(
                state,
                "segment",
                f"Object '{obj_found}' found with correct stacking. Bonus: {glMetric / 4.0}",
                glMetric / 4.0,
            )

    for obj in objs:
        if obj in state.objectsLeft:
            state.objectsLeft.remove(obj)
        if obj not in state.objectsFound:
            state.objectsFound.append(obj)

    return reward


def _remove_object(state, obj_name):
    reward = 0
    try:
        pos = state.ObjectsPlaced.pop(obj_name)
    except KeyError:
        reward -= 45.0 / len(state.ObjectsPresent)
        appendRewardFeedback(
            state,
            "adjust",
            f"Object '{obj_name}' is not placed in the environment.",
            -reward,
        )
        return reward

    state.numberPlaced -= 1
    dims = state.currentGrid
    obj = OBJECTS.get(obj_name)
    objGrid = initDimentions(obj)

    for i in range(len(objGrid)):
        for j in range(len(objGrid[0])):
            for k in range(len(objGrid[0][0])):
                if dims[pos[0] + i][pos[1] + j][pos[2] + k] > 0:
                    dims[pos[0] + i][pos[1] + j][pos[2] + k] -= 1


def _adjustment_helper(state, name, pos, change, direction):
    _remove_object(state, name)

    if direction == "ROTATE":
        newPos = (pos[1], pos[0], pos[2], pos[3])
    else:
        newPos = (pos[0] + change[0], pos[1] + change[1], pos[2] + change[2], pos[3])

    reward, isNotPlaced = place(False, {name: newPos}, state)

    if isNotPlaced:
        dummyReward = place(False, {name: pos}, state)[0]
        appendRewardFeedback(
            state,
            "adjust",
            f"Failed to adjust object '{name}' in direction {direction}. Reverting to original position.",
            -dummyReward,
        )
        return -dummyReward

    appendRewardFeedback(
        state,
        "adjust",
        f"Object '{name}' moved {direction} successfully.",
        reward,
    )
    return reward


def adjustment(segment, action, state):
    objsPlaced = state.ObjectsPlaced

    if segment:
        appendRewardFeedback(
            state, "adjust", "Placing objects with segmentation is not allowed.", -60.0
        )
        return -60.0

    try:
        initPos = objsPlaced[action[0]]
        name = action[0]
    except KeyError:
        reward_per_obj_placed = 45.0 / len(state.ObjectsPresent)
        appendRewardFeedback(
            state,
            "adjust",
            f"Object '{action[0]}' is not placed in the environment, so it cannot be adjusted.",
            -reward_per_obj_placed,
        )
        return -reward_per_obj_placed

    if action[1] in ACTION_CONFIG:
        reward = _adjustment_helper(
            state, name, initPos, ACTION_CONFIG.get(action[1]), action[1]
        )
        return reward
    else:
        reward_per_obj_placed = 45.0 / len(state.ObjectsPresent)
        appendRewardFeedback(
            state,
            "adjust",
            f"Invalid adjustment direction '{action[1]}'. Valid directions are RIGHT, LEFT, UP, DOWN, FORWARD, BACKWARD, ROTATE.",
            -reward_per_obj_placed,
        )
        return -reward_per_obj_placed