import bpy
import math

import numpy as np

from math import tan
from mathutils import Vector, Matrix
from bpy_extras.object_utils import world_to_camera_view
from mathutils.geometry import (barycentric_transform, intersect_point_tri, 
    intersect_point_line, intersect_line_plane, tessellate_polygon)

# --- Vector

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 get_bone_head_tail_2d(posebone, scene=None, cam=None) -> tuple[Vector, Vector]:
    '''Get 2D vectors in camera view of bone head and tails
    return tuple of 2d vectors (head_2d and tail_2d) 
    '''
    scene = scene or bpy.context.scene
    cam = cam or scene.camera

    arm = posebone.id_data

    # Get 3D locations of head and tail
    head_3d = arm.matrix_world @ posebone.head
    tail_3d = arm.matrix_world @ posebone.tail

    # Convert 3D locations to 2D
    head_2d = world_to_camera_view(scene, cam, head_3d)
    tail_2d = world_to_camera_view(scene, cam, tail_3d)
    
    ratio = scene.render.resolution_y / scene.render.resolution_x
    head_2d.y *= ratio
    tail_2d.y *= ratio
    
    return Vector((head_2d.x, head_2d.y)), Vector((tail_2d.x, tail_2d.y))


def rotate_matrix_around_pivot(matrix, angle, pivot, axis):
    '''Rotate a given matrix by a CW angle around pivot on a given axis
    matrix (Matrix): the matrix to rotate
    angle (Float, Radians): the angle in radians
    pivot (Vector3): the pivot 3D coordinate
    axis (Vector3): the vector axis of rotation
    '''

    # Convert angle to radians ?
    # angle = math.radians(angle)

    # Create a rotation matrix
    rot_matrix = Matrix.Rotation(angle, 4, axis)

    # Create translation matrices
    translate_to_origin = Matrix.Translation(-pivot)
    translate_back = Matrix.Translation(pivot)

    # Combine the transformations : The order of multiplication is important
    new_matrix = translate_back @ rot_matrix @ translate_to_origin @ matrix

    return new_matrix


# --- GREASE PENCIL

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