from typing import Callable, List, Dict, Tuple

import mathutils

from blenderproc.python.utility.CollisionUtility import CollisionUtility
from blenderproc.python.types.EntityUtility import Entity
from blenderproc.python.types.MeshObjectUtility import MeshObject, get_all_mesh_objects


def sample_poses(objects_to_sample: List[MeshObject], sample_pose_func: Callable[[MeshObject], None],
                 objects_to_check_collisions: List[MeshObject] = None, max_tries: int = 1000,
                 mode_on_failure: str = "last_pose") -> Dict[Entity, Tuple[int, bool]]:
    """ Samples positions and rotations of selected object inside the sampling volume while performing mesh and bounding box collision checks.

    :param objects_to_sample: A list of mesh objects whose poses are sampled based on the given function.
    :param sample_pose_func: The function to use for sampling the pose of a given object.
    :param objects_to_check_collisions: A list of mesh objects who should not be considered when checking for collisions.
    :param max_tries: Amount of tries before giving up on an object and moving to the next one.
    :param mode_on_failure: Define final state of objects that could not be placed without collisions within max_tries attempts. Options: 'last_pose', 'initial_pose'

    :return: A dict with the objects to sample as keys and a Tuple with the number of executed attempts to place the object as first element, and a bool whether it has been succesfully placed without collisions.
    """
    # Check if mode on failure is allowed
    allowed_modes_on_failure = ["last_pose", "initial_pose"]
    if mode_on_failure not in allowed_modes_on_failure:
        raise ValueError("{} is not an allowed mode_on_failure.")

    # After this many tries we give up on current object and continue with the rest
    if objects_to_check_collisions is None:
        objects_to_check_collisions = get_all_mesh_objects()

    # Among objects_to_sample only check collisions against already placed objects
    cur_objects_to_check_collisions = list(set(objects_to_check_collisions) - set(objects_to_sample))

    if max_tries <= 0:
        raise ValueError("The value of max_tries must be greater than zero: {}".format(max_tries))

    if not objects_to_sample:
        raise Exception("The list of objects_to_sample can not be empty!")

    # cache to fasten collision detection
    bvh_cache: Dict[str, mathutils.bvhtree.BVHTree] = {}

    sample_results: Dict[Entity, Tuple[int, bool]] = {}

    # for every selected object
    for obj in objects_to_sample:

        # Store the obejct's initial pose in case we need to place it back
        if mode_on_failure == 'initial_pose':
            initial_location = obj.get_location()
            initial_rotation = obj.get_rotation()
        
        no_collision = True

        amount_of_tries_done = -1

        # Try max_iter amount of times
        for i in range(max_tries):
            
            # Put the top object in queue at the sampled point in space
            sample_pose_func(obj)

            # Remove bvh cache, as object has changed
            if obj.get_name() in bvh_cache:
                del bvh_cache[obj.get_name()]

            no_collision = CollisionUtility.check_intersections(obj, bvh_cache, cur_objects_to_check_collisions, [])

            # If no collision then keep the position
            if no_collision:
                amount_of_tries_done = i
                break

        # After placing an object, we will check collisions with it
        cur_objects_to_check_collisions.append(obj)

        if no_collision:
            print("It took " + str(amount_of_tries_done + 1) + " tries to place " + obj.get_name())
        else:
            amount_of_tries_done = max_tries
            print("Could not place " + obj.get_name() + " without a collision.") 

            if mode_on_failure == 'initial_pose':
                obj.set_location(initial_location)
                obj.get_rotation(initial_rotation)          

        sample_results[obj] = (amount_of_tries_done, no_collision)
    
    return sample_results
