import bpy
import numpy as np

from math import tan, acos, degrees
from time import perf_counter
from mathutils import Vector, Matrix

from gp_interpolate import utils

from mathutils.geometry import (barycentric_transform, intersect_point_tri, 
    intersect_point_line, intersect_line_plane, tessellate_polygon)


def get_gp_draw_plane(obj=None):
    ''' return tuple with plane coordinate and normal
    of the curent drawing according to geometry'''

    if obj is None:
        obj = bpy.context.object

    settings = bpy.context.scene.tool_settings
    orient = settings.gpencil_sculpt.lock_axis #'VIEW', 'AXIS_Y', 'AXIS_X', 'AXIS_Z', 'CURSOR'
    loc = settings.gpencil_stroke_placement_view3d #'ORIGIN', 'CURSOR', 'SURFACE', 'STROKE'
    
    mat = obj.matrix_world
    plane_no = Vector((0.0, 0.0, 1.0))
    plane_co = mat.to_translation()

    # -> orientation
    if orient == 'VIEW':
        mat = bpy.context.scene.camera.matrix_world 
    # -> placement
    if loc == "CURSOR":
        plane_co = bpy.context.scene.cursor.location
        mat = bpy.context.scene.cursor.matrix

    elif orient == 'AXIS_Y':#front (X-Z)
        plane_no = Vector((0,1,0))

    elif orient == 'AXIS_X':#side (Y-Z)
        plane_no = Vector((1,0,0))

    elif orient == 'AXIS_Z':#top (X-Y)
        plane_no = Vector((0,0,1))

    plane_no.rotate(mat)
    
    return plane_co, plane_no

def triangle_normal(a, b, c):
    x = a[1] * b[2] - a[2] * b[1]
    y = a[2] * b[0] - a[0] * b[2]
    z = a[0] * b[1] - a[1] * b[0]
    
    return np.array([x, y, z], dtype='float64')

def plane_coords(size=1):
    v = size * 0.5
    return np.array([(-v, v, 0), (v, v, 0), (v, -v, 0), (-v, -v, 0)], dtype='float64')

def matrix_transform(coords, matrix):
    coords_4d = np.column_stack((coords, np.ones(len(coords), dtype='float64')))
    return np.einsum('ij,aj->ai', matrix, coords_4d)[:, :-1]

def vector_normalized(vec):
    return vec / np.sqrt(np.sum(vec**2))

def vector_magnitude(vec):
    return np.sqrt(vec.dot(vec))

def search_square(point, factor=0.05, cam=None):
    if cam is None:
        cam = bpy.context.scene.camera
        
    plane = plane_coords()
    mat = cam.matrix_world.copy()
    mat.translation = point
    depth = vector_magnitude(point - cam.matrix_world.to_translation())
    mat_scale = Matrix.Scale(tan(cam.data.angle*0.5)*depth*factor, 4)
    
    return matrix_transform(plane, mat @ mat_scale)

def ray_cast_point(point, origin, depsgraph):
    ray = (point - origin)#.normalized()
    hit, hit_location, normal, face_index, object_hit, matrix = bpy.context.scene.ray_cast(depsgraph, origin, ray)
    
    if not hit:
        return None, None, None, None
    
    eval_ob = object_hit.evaluated_get(depsgraph)
    
    face = eval_ob.data.polygons[face_index]
    vertices = [eval_ob.data.vertices[i] for i in face.vertices]
    face_co = matrix_transform([v.co for v in vertices], eval_ob.matrix_world)
    
    tri = None
    for tri_idx in tessellate_polygon([face_co]):
        tri = [face_co[i] for i in tri_idx]
        tri_indices = [vertices[i].index for i in tri_idx]
        if intersect_point_tri(hit_location, *tri):
            break
    
    return object_hit, np.array(hit_location), tri, tri_indices


def plane_on_bone(bone, arm=None, cam=None, set_rotation=True):
    if cam is None:
        cam = bpy.context.scene.camera
    
    if arm is None:
        arm = bone.id_data

    plane = plane_coords()
    mat = cam.matrix_world.copy()


    if set_rotation:
        head_world_coord = arm.matrix_world @ bone.head
        mat.translation = head_world_coord

        ## Apply 2d bone rotation facing camera

        # Get 2d camera space coords (NDC: normalized device coordinate, 0,0 is bottom-left)
        head_2d, tail_2d = utils.get_bone_head_tail_2d(bone, cam=cam)
        vec_from_corner_2d = (tail_2d - head_2d).normalized()
        up_vec_2d = Vector((0,1))
        # angle = acos(up_vec_2d.dot(vec_from_corner_2d)) ## equivalent but not signed!
        angle = up_vec_2d.angle_signed(vec_from_corner_2d)

        ## Axis camera aim (seem slightly off)
        # rot_axis = Vector((0, 0, -1))
        # rot_axis.rotate(cam.matrix_world)

        ## Axis camera origin -> pivot
        rot_axis = head_world_coord - cam.matrix_world.translation
        
        mat = utils.rotate_matrix_around_pivot(mat, angle, head_world_coord, rot_axis)

    else:
        ## Use mid bone to better follow movement
        mat.translation = arm.matrix_world @ ((bone.tail + bone.head) / 2) # Mid bone

    mat_scale = Matrix.Scale(10, 4)
    return matrix_transform(plane, mat @ mat_scale)

def intersect_with_tesselated_plane(point, origin, face_co):
    '''
    face_co: World face coordinates
    '''
    
    tri = None
    for tri_idx in tessellate_polygon([face_co]):
        tri = [face_co[i] for i in tri_idx]
        tri_indices = [i for i in tri_idx]
        hit_location = intersect_line_plane(origin, point, sum((Vector(v) for v in tri), Vector()) / 3, triangle_normal(*tri))  
        if intersect_point_tri(hit_location, *tri):
            break
    
    return np.array(hit_location), tri, tri_indices

def following_key(forward=True):
    direction = 1 if forward else -1
    cur_frame = bpy.context.scene.frame_current
    settings = bpy.context.scene.gp_interpo_settings
    
    if settings.mode == 'FRAME':
        return cur_frame + (settings.padding * direction)
    
    elif settings.mode == 'GPKEY':
        layers = bpy.context.object.data.layers
        frames = [f.frame_number for l in layers for f in l.frames]
    
    elif settings.mode == 'RIGKEY':
        col = settings.target_collection
        if not col:
            col = bpy.context.scene.collection            
        for arm in [o for o in col.all_objects if o.type == 'ARMATURE']:
            if not arm.animation_data or not arm.animation_data.action:
                continue
            frames = [k.co.x for fc in arm.animation_data.action.fcurves for k in fc.keyframe_points]
    
    if not frames:
        return

    frames.sort()
    if forward:
        new = next((f for f in frames if f > cur_frame), None)
    else:
        below = [f for f in frames if f < cur_frame]
        if not below:
            return
        new = below[-1]

    return int(new)

class GP_OT_interpolate_stroke(bpy.types.Operator):
    bl_idname = "gp.interpolate_stroke"
    bl_label = "Interpolate Stroke"
    bl_description = 'Interpolate Stroke'
    bl_options = {'REGISTER', 'UNDO'} 

    @classmethod
    def poll(cls, context):
        if context.active_object and context.object.type == 'GPENCIL'\
            and context.mode in ('EDIT_GPENCIL', 'SCULPT_GPENCIL', 'PAINT_GPENCIL'):
            return True
        cls.poll_message_set("Need a Grease pencil object in Edit or Sculpt mode")
        return False

    @classmethod
    def description(cls, context, properties):
        if properties.next:
            return f"Interpolate Stroke Forward"
        else:
            return f"Interpolate Stroke Backward"

    next : bpy.props.BoolProperty(name='Next', default=True, options={'SKIP_SAVE'})

    def execute(self, context):
        settings = context.scene.gp_interpo_settings

        auto_key_status = context.tool_settings.use_keyframe_insert_auto
        context.tool_settings.use_keyframe_insert_auto = True

        ## Determine on what key to jump
        # frame_to_jump = following_key(forward=self.next)
        frame_to_jump = following_key(forward=self.next)
        if frame_to_jump is None:
            self.report({'WARNING'}, 'No keyframe available in this direction')
            return {'CANCELLED'}
        
        # print('frame_to_jump: ', frame_to_jump)

        gp = context.object

        scn = bpy.context.scene
        dg = bpy.context.evaluated_depsgraph_get()
        matrix = np.array(gp.matrix_world, dtype='float64')#.inverted()
        col = settings.target_collection
        if not col:
            col = scn.collection

        origin = np.array(scn.camera.matrix_world.to_translation(), 'float64')
        # print('----')

        tgt_strokes = [s for s in gp.data.layers.active.active_frame.strokes if s.select]
        
        ## If nothing selected in sculpt/paint, Select all before triggering
        if not tgt_strokes and context.mode in ('SCULPT_GPENCIL', 'PAINT_GPENCIL'):
            for s in gp.data.layers.active.active_frame.strokes:
                s.select = True
            tgt_strokes = gp.data.layers.active.active_frame.strokes

        if not tgt_strokes:            
            self.report({'ERROR'}, 'No stroke selected !')
            return {'CANCELLED'}
        
        strokes_data = []
        if settings.method == 'BONE':
            ## Follow Bone method (Full WIP)
            if not settings.target_rig or not settings.target_bone:
                self.report({'ERROR'}, 'No Bone selected')
                return {'CANCELLED'}

            bone_plane = plane_on_bone(settings.target_rig.pose.bones.get(settings.target_bone),
                                        arm=settings.target_rig,
                                        set_rotation=settings.use_bone_rotation)

            strokes_data = []
            for stroke in tgt_strokes:
                nb_points = len(stroke.points)
                local_co = np.empty(nb_points * 3, dtype='float64')
                stroke.points.foreach_get('co', local_co)
                # local_co_3d = local_co.reshape((nb_points, 3))
                world_co_3d = matrix_transform(local_co.reshape((nb_points, 3)), matrix)

                stroke_data = []

                for i, point in enumerate(stroke.points):
                    point_co_world = world_co_3d[i]
                    hit_location, tri, tri_indices = intersect_with_tesselated_plane(point_co_world, origin, bone_plane)
                    ## probably easier to just generate a single vast triangle and use it
                    
                    ## Store same as other method (without object hit)
                    stroke_data.append((stroke, point_co_world, hit_location, tri, tri_indices))
                
                strokes_data.append(stroke_data)

            bpy.ops.gpencil.copy()

            scn.frame_set(frame_to_jump)

            dg = bpy.context.evaluated_depsgraph_get()

            plan_co, plane_no = get_gp_draw_plane(gp)
            
            bpy.ops.gpencil.paste()


            matrix_inv = np.array(gp.matrix_world.inverted(), dtype='float64')#.inverted()
            new_strokes = gp.data.layers.active.active_frame.strokes[-len(strokes_data):]

            bone_plane = plane_on_bone(settings.target_rig.pose.bones.get(settings.target_bone),
                                        arm=settings.target_rig,
                                        set_rotation=settings.use_bone_rotation)

            for new_stroke, stroke_data in zip(new_strokes, strokes_data):
                world_co_3d = [] # np.array(len()dtype='float64')#np.
                for stroke, point_co, hit_location, tri_a, tri_indices in stroke_data:
                    tri_b = [bone_plane[i] for i in tri_indices]
                    # tri_b = matrix_transform(tri_b, settings.target_rig.matrix_world)
                    ## rotate tri_b by bone differential angle camera's aim axis ?

                    new_loc = barycentric_transform(hit_location, *tri_a, *tri_b)
                    world_co_3d.append(new_loc)

    
                # Reproject on plane
                new_world_co_3d = [intersect_line_plane(origin, p, plan_co, plane_no) for p in world_co_3d]        
                new_local_co_3d = matrix_transform(new_world_co_3d, matrix_inv)
                
                nb_points = len(new_stroke.points)
                new_stroke.points.foreach_set('co', new_local_co_3d.reshape(nb_points*3))
                new_stroke.points.update()


        else:
            ## Geometry method
            for stroke in tgt_strokes:
                nb_points = len(stroke.points)

                local_co = np.empty(nb_points * 3, dtype='float64')
                stroke.points.foreach_get('co', local_co)
                # local_co_3d = local_co.reshape((nb_points, 3))
                world_co_3d = matrix_transform(local_co.reshape((nb_points, 3)), matrix)

                stroke_data = []
                for i, point in enumerate(stroke.points):
                    point_co_world = world_co_3d[i]
                    
                    object_hit, hit_location, tri, tri_indices = ray_cast_point(point_co_world, origin, dg)
                    if not object_hit or object_hit not in col.all_objects[:]:
                        for square_co in search_square(point_co_world, factor=settings.search_range):
                            object_hit, hit_location, tri, tri_indices = ray_cast_point(square_co, origin, dg)
                            if object_hit and object_hit in col.all_objects[:]:
                                
                                hit_location = intersect_line_plane(origin, point_co_world, tri[0], triangle_normal(*tri))
                                
                                break
                    
                    stroke_data.append((stroke, point_co_world, object_hit, hit_location, tri, tri_indices))
                
                strokes_data.append(stroke_data)


            bpy.ops.gpencil.copy()

            scn.frame_set(frame_to_jump)

            dg = bpy.context.evaluated_depsgraph_get()

            plan_co, plane_no = get_gp_draw_plane(gp)
            
            bpy.ops.gpencil.paste()


            matrix_inv = np.array(gp.matrix_world.inverted(), dtype='float64')#.inverted()
            new_strokes = gp.data.layers.active.active_frame.strokes[-len(strokes_data):]

            for new_stroke, stroke_data in zip(new_strokes, strokes_data):
                world_co_3d = [] # np.array(len()dtype='float64')#np.
                for stroke, point_co, object_hit, hit_location, tri_a, tri_indices in stroke_data:
                    eval_ob = object_hit.evaluated_get(dg)
                    tri_b = [eval_ob.data.vertices[i].co for i in tri_indices]
                    tri_b = matrix_transform(tri_b, eval_ob.matrix_world)
                    
                    new_loc = barycentric_transform(hit_location, *tri_a, *tri_b)
                    world_co_3d.append(new_loc)

                # Reproject on plane
                new_world_co_3d = [intersect_line_plane(origin, p, plan_co, plane_no) for p in world_co_3d]        
                new_local_co_3d = matrix_transform(new_world_co_3d, matrix_inv)
                
                nb_points = len(new_stroke.points)
                new_stroke.points.foreach_set('co', new_local_co_3d.reshape(nb_points*3))
                new_stroke.points.update()

        context.tool_settings.use_keyframe_insert_auto = auto_key_status

        return {'FINISHED'}


classes = (
    GP_OT_interpolate_stroke,
)

def register():
    for c in classes:
        bpy.utils.register_class(c)
    
        
def unregister():
    for c in reversed(classes):
        bpy.utils.unregister_class(c)