"""Utility functions for making arrangments / piles of objects.

Main functions are :func:`make_object_arrangement` and
:func:`make_object_pile`.
"""

from klampt import *
from klampt.math import vectorops,so3,se3
from klampt.model import collide
import random
import math
from typing import Union,List,Tuple,Sequence,Callable,Any

def _get_bound(objects):
    """Obtains the tight outer bounds of the object(s) at their current transforms, in world coordinates."""
    if hasattr(objects,'__iter__'):
        bbs = [_get_bound(o) for o in objects]
        if len(bbs) == 1:
            return bbs[0]
        return collide.bb_union(*bbs)
    else:
        return objects.geometry().getBBTight()

def xy_randomize(obj,bmin,bmax):
    """Randomizes the xy position and z orientation of an object inside of a bounding box bmin->bmax.
    Assumes the bounding box is large enough to hold the object in any orientation.
    """
    R,t = obj.getTransform()
    R = so3.mul(so3.rotation([0,0,1],random.uniform(0,math.pi*2)),R)
    t[0] = 0
    t[1] = 0
    obj.setTransform(R,t)
    obmin,obmax = obj.geometry().getBB()
    t[0] = random.uniform(bmin[0]-obmin[0],bmax[0]-obmax[0])
    t[1] = random.uniform(bmin[1]-obmin[1],bmax[1]-obmax[1])
    obj.setTransform(R,t)

def xy_jiggle(world,objects,fixed_objects,bmin,bmax,iters,randomize=True,
    verbose=0):
    """Jiggles the objects' x-y positions within the range bmin - bmax, and randomizes orientation about the z
    axis until the objects are collision free.  A list of fixed objects (fixed_objects) may be given as well.

    Objects for which collision-free resolutions are not found are returned.
    """
    if randomize:
        for obj in objects:
            xy_randomize(obj,bmin,bmax)
    object_geometries = [o.geometry() for o in objects]
    fixed_geometries = [o.geometry() for o in fixed_objects]
    inner_iters = 10
    while iters > 0:
        numConflicts = [0]*len(objects)
        for (i,j) in collide.self_collision_iter(object_geometries):
            numConflicts[i] += 1
            numConflicts[j] += 1
        for (i,j) in collide.group_collision_iter(object_geometries,fixed_geometries):
            numConflicts[i] += 1
        
        amax = max((c,i) for (i,c) in enumerate(numConflicts))[1]
        cmax = numConflicts[amax]
        if cmax == 0:
            #conflict free
            return []
        if verbose: 
            print(cmax,"conflicts with object",objects[amax].getName())
        other_geoms = [o.geometry() for o in objects[:amax]+objects[amax+1:]+fixed_objects]
        nc = 0
        for it in range(inner_iters):
            xy_randomize(objects[amax],bmin,bmax)
            nc = sum([1 for p in collide.group_collision_iter([objects[amax].geometry()],other_geoms)])
            if nc < cmax:
                break
            iters-=1
        if verbose: 
            print("Now",nc,"conflicts with object",objects[amax].getName())

    numConflicts = [0]*len(objects)
    for (i,j) in collide.self_collision_iter(object_geometries):
        numConflicts[i] += 1
        numConflicts[j] += 1
    for (i,j) in collide.group_collision_iter(object_geometries,fixed_geometries):
        numConflicts[i] += 1
    removed = []