import bpy
import bmesh
from mathutils import Vector  


def raycast_from_loc_to_obj(src, tgt, direction=None, dg=None):
    ''' Raycast from a world space source on a target ovject using a direction Vector

    :origin: a world coordinate location as source of the ray
    :tgt: an object onto project the ray
    :direction: a direction vector. If not passed, point bottom : Vector((0,0,-1))
    return world coordiante of the hit if any
    '''

    if direction is None:
        direction = Vector((0,0,-1))

    mw = tgt.matrix_world
    origin = mw.inverted() @ src

    hit, loc, _norm, _face = tgt.ray_cast(origin, direction, depsgraph=dg)

    if hit:
        # print("Hit at ", loc, " (local)")
        world_loc = mw @ loc
        # bpy.ops.object.empty_add(location = world_loc) # test
        return world_loc

    return False

def worldspace_move_posebone(b, vec, is_target=False):
    ''' Move or snap a posebone using a vector in worldspace
    :b: posebone
    :vec: Vector to move posebone worldspace (if not attached to parent)
    :is_target: if True the posebone snap to the vector, else the vector is added
    '''

    a = b.id_data

    if is_target:
        target_vec = vec
    else:
        target_vec = (a.matrix_world @ b.matrix).translation + vec

    mw = a.convert_space(pose_bone=b, 
            matrix=b.matrix, 
            from_space='POSE', 
            to_space='WORLD')
    
    mw.translation = target_vec
    
    b.matrix = a.convert_space(pose_bone=b, 
            matrix=mw, 
            from_space='WORLD', 
            to_space='POSE')
    return target_vec

def snap_foot(pb, gnd):
    '''Get posebone and ground to touch'''

    # arm = bpy.context.object# bpy.context.scene.objects.get('Armature')
    arm = pb.id_data
    print('arm: ', arm)

    # find tip bone :
    tip = [p for p in pb.children_recursive if p.name.startswith('DEF')][-1]
    print('tip: ', tip)
    # get deformed object VG (find skinned mesh)

    ob = None

    for o in arm.proxy_collection.instance_collection.all_objects:
        if o.type != 'MESH':
            continue
        for m in o.modifiers:
            if m.type == 'ARMATURE':
                # print(o.name, m.object)
                if m.object == arm.proxy: # if point to orignal rig
                    ## here we want body, not body_deform
                    if not 'body' in o.name:
                        continue
                    if '_deform' in o.name:
                        continue
                    ob = o
                    break
    if not ob:
        print('ERROR', 'no skinned mesh found')
        return

    print('check skinning of', ob.name)
    ### MESH baking
    #-# Get Vertices position for a specific vertex group if over weight limit

    # me0 = simple_to_mesh(ob) # if no need to apply modifier just make ob.data.copy()
    # # generate new
    # bm = 
    # bm.from_mesh(me0)
    # bm.verts.ensure_lookup_table()
    # bm.edges.ensure_lookup_table()
    # bm.faces.ensure_lookup_table()

    # # store weight values
    # weight = []
    # ob_tmp = bpy.data.objects.new("temp", me0)
    # for g in ob.vertex_groups:
    #     ob_tmp.vertex_groups.new(name=g.name)
    # for v in me0.vertices:
    #     try:
    #         weight.append(ob_tmp.vertex_groups[tip.name].weight(v.index))
    #     except:
    #         weight.append(0)
    
    # verts = [vert for vid, vert in enumerate(bake_mesh.vertices) \
    #     if ob_tmp.vertex_groups[tip.name].index in [i.group for i in vert.groups] \
    #     and weight[vid] > 0.5]

    #-# /

    #-# Get Vertices position for a specific vertex group if over weight limit
    #-# (Does not work if a subdivision modifier is on)

    for m in ob.modifiers:
        # if m.type in ('SUBSURF', 'TRIANGULATE'):
        if m.type == 'SUBSURF':
            m.show_viewport = False

    dg = bpy.context.evaluated_depsgraph_get()
    obeval = ob.evaluated_get(dg) #.copy()
    print('object: ', ob.name)
    
    ## bpy.context.object.proxy_collection.instance_collection.all_objects['body_deform']
    ## bpy.context.object.proxy_collection.instance_collection.all_objects['body']

    ## Hide modifier 
    # for m in obeval.modifiers:
    #     if m.type == 'SUBSURF':
    #         m.show_viewport = False # m.levels = 0

    bake_mesh = obeval.to_mesh(preserve_all_data_layers=True, depsgraph=dg)

    ct = 0
    vg = obeval.vertex_groups[tip.name]
    world_co = []
    for idx, vert in enumerate(bake_mesh.vertices):
        grp_indexes = [i.group for i in vert.groups]
        if vg.index in grp_indexes and vg.weight(idx) > 0.5:
            ct +=1
            world_co.append(ob.matrix_world @ vert.co)

    if not ct:
        print('ERROR', 'No vertices found')
        return
    
    #-# # list comprehension
    # verts = [vert for vid, vert in enumerate(bake_mesh.vertices) \
    #     if obeval.vertex_groups[tip.name].index in [i.group for i in vert.groups] \
    #     and obeval.vertex_groups[tip.name].weight(vid) > 0.5]   
    #world_co = [ob.matrix_world @ v.co for v in verts]
    #-# /
    
    
    print(len(world_co), 'vertices')

    # sort by height
    world_co.sort(key=lambda x: x[2])
    
    ### Raycast and find lowest distance
    up_check = True
    updists = []
    dists = []

    for co in world_co: # [:6] (no neede to get all)
        contact = raycast_from_loc_to_obj(co, gnd, Vector((0,0,-1)), dg=dg)
        if contact:
            dists.append((co - contact).length)
        
        if not contact and up_check:
            contact = raycast_from_loc_to_obj(co, gnd, Vector((0,0,1)), dg=dg)
            if contact:
                updists.append((co - contact).length)

        if not contact:
            continue
        # empty_at(contact, size=0.2)
    
    if not dists and not updists:
        print('ERROR', 'raycast could not found contact')
        return

    # move bones by the minimal amount.
    if updists:
        move = max(updists)
        vec = Vector((0,0, move))
        worldspace_move_posebone(pb, vec)
        print('INFO', f'move up by {move}')
    
    else:
        move = min(dists)
        vec = Vector((0,0, -move))
        worldspace_move_posebone(pb, vec)
        print('INFO', f'move down by {move}')
    

    ## restore

    for m in ob.modifiers:
        if m.type == 'SUBSURF':
        # if m.type in ('SUBSURF', 'TRIANGULATE'):
            m.show_viewport = True
    
    obeval.to_mesh_clear()


def snap_feet():

    ## add undo push if launched from shelf (TODO need test !!!)
    # bpy.ops.ed.undo_push(message='Snap to ground')

    # if bpy.context.object.type != 'ARMATURE':
    #     print('ERROR', 'Selection is not an armature')
    #     return

    if bpy.context.mode != 'POSE':
        print('ERROR', 'Not in pose mode')
        return
    
    
    gnd = bpy.context.scene.anim_cycle_settings.gnd
    print('ground: ', gnd.name)
    if not gnd:
        return ('ERROR', 'Need to point ground object in "ground" field')
    # Snap all selected feets, posebone to ground

    for pb in bpy.context.selected_pose_bones:
        ## find the foot bones.
        if '_shoe' in pb.name:
            # get pb lowest surface deformed point
            snap_foot(pb, gnd)

class UAC_OT_contact_to_ground(bpy.types.Operator):
    bl_idname = "anim.contact_to_ground"
    bl_label = "Ground Feet"
    bl_description = "Ground selected feets"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'ARMATURE'

    def execute(self, context):
        # context.scene.anim_cycle_settings.expand_on_selected_bones
        err = snap_feet()
        if err:
            self.report({err[0]}, err[1])
            if err[0] == 'ERROR':
                return {"CANCELLED"}

        return {"FINISHED"}


classes=(
UAC_OT_contact_to_ground,
)

def register():
    for cls in classes:
        bpy.utils.register_class(cls)

def unregister():
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)