import bpy
import bpy_extras.object_utils
import mathutils
import numpy as np
from ..cAPI import SimulationAPI
from . import createTransform
from . import createMeshBVH

# we assume the verts are already part of the object
# and the object passed in has already been evaluated
def create(sim_obj, mesh_obj, name, operator = None):
    if not sim_obj.hi_phy or sim_obj.hi_phy.object_type != "ELASTIC_ROD":
        if operator:
            operator.report({"ERROR"}, sim_obj.name + " is not a valid hi phy elastic rod object")
        return False

    # get mesh topology
    if (not mesh_obj.type == 'MESH') :
        if operator:
            operator.report({"ERROR"}, "Binding mesh " + mesh_obj.name + " is not a triangle mesh")
        return False
    # build bvh in world space (because radius is in world unit)
    mesh_bvh = createMeshBVH(mesh_obj, mesh_obj.matrix_world, None, operator)
    if (not mesh_bvh) :
        if operator:
            operator.report({"ERROR"}, "Failed to build BVH for " + mesh_obj.name)
        return False
    # create the testing points
    if ((not sim_obj.type == 'CURVES') and (not sim_obj.type == 'CURVE')) :
        return False
    if (sim_obj.type == 'CURVES'):
        n_curves = len(sim_obj.data.curves)
        n_points = sim_obj.data.attributes.domain_size('POINT')
        points_data = np.zeros((n_points, 3), dtype=np.float32)
        sim_obj.data.attributes['position'].data.foreach_get('vector', np.ravel(points_data))
        start_cvs = [sim_obj.data.curve_offset_data[i].value for i in range(n_curves)]
        base_normals_attr = sim_obj.data.attributes['hi_phy_base_normal']
        base_normals_data = np.zeros((n_curves, 3), dtype=np.float32)
        base_normals_attr.data.foreach_get('vector', np.ravel(base_normals_data))
        cvs = [p for c in range(n_curves) for p in [points_data[start_cvs[c]], points_data[start_cvs[c] + 1], base_normals_data[c]]]
        for c in range(n_curves):
            # orthogonalize the base_normals and make it the same size as the first edge
            cv0 = np.array(cvs[c * 3 + 0])
            cv1 = np.array(cvs[c * 3 + 1])
            bn = np.array(cvs[c * 3 + 2])
            l = np.linalg.norm(cv1 - cv0)
            e = (cv1 - cv0) / l
            n = bn - e * np.dot(e, bn)
            ln = np.linalg.norm(n)
            cvs[c * 3 + 2] = n / ln * l + cv0
    if (sim_obj.type == 'CURVE'):
        # TODO
        return False

    # For root constraints, we always constraint the first two cvs
    # The test points composed of the first two cvs and the base normal
    points = SimulationAPI.CreateVector3Array(n_curves * 3)
    SimulationAPI.CopyVector3Array(points, cvs)
    # Allocate data
    found_proximity = SimulationAPI.CreateBoolArray(n_curves)
    bind_points_on_mesh = SimulationAPI.CreateVector3Array(n_curves)
    bind_face_ids = SimulationAPI.CreateIntArray(n_curves)
    bind_barycentric_weights_on_mesh = SimulationAPI.CreateVector3Array(n_curves)
    bind_local_coord = SimulationAPI.CreateVector3Array(n_curves * 3)
    # Copy transforms
    mesh_transform = createTransform(mesh_obj.matrix_world)
    points_transform = createTransform(sim_obj.matrix_world)
    # Binding
    binding_radius = 1e10 # set to a large number so it always binds
    if not SimulationAPI.Utils.BindRootToMesh(points, mesh_bvh, found_proximity, bind_points_on_mesh, bind_face_ids, bind_barycentric_weights_on_mesh, bind_local_coord, points_transform, mesh_transform, binding_radius):
        if operator:
            operator.report({"ERROR"}, "Binding has failed")
            return False
    bond_indices = [index for index, value in enumerate(found_proximity) if value == 1]
    if (len(bond_indices) != n_curves):
        if operator:
            operator.report({"ERROR"}, "Binding has failed, some points has failed to bind")
            return False

    pointcloud = bpy.data.meshes.new(name + "Mesh")
    # The bond points are in sim object space
    matrix_world = sim_obj.matrix_world;
    vs_world = [matrix_world @ mathutils.Vector((v[0], v[1], v[2])) for v in points]
    edges = [e for c in range(n_curves) for e in [(c * 3 + 0, c * 3 + 1), (c * 3 + 0, c * 3 + 2)]]
    pointcloud.from_pydata(vs_world, edges, []) # Create a mesh with only verts and edges

    obj_new = bpy.data.objects.new(name, pointcloud)
    obj_new.hi_phy.is_active = True
    obj_new.hi_phy.object_type = "CONSTRAINT"
    # Use the original object. We will evaluate it when we need it
    obj_new.hi_phy.constraint_object0 = sim_obj.original
    obj_new.hi_phy.constraint_object1 = mesh_obj.original
    obj_new.hi_phy.constraint_type = SimulationAPI.ConstraintType.RootConstraint
    obj_new.hi_phy.solver = sim_obj.hi_phy.solver

    bpy.context.collection.objects.link(obj_new)

    weight_data = [1, 0, 0] * (n_curves * 3)
    zero_data = [0] * (n_curves * 3)
    zero_data_3 = [0, 0 , 0] * (n_curves * 3)

    bond_indices = [i for cv in start_cvs for i in [cv, cv + 1, 0]]
    # Do not create multiple attributes at the same time
    # It will make the previous handle invalid
    index0_0_attr = pointcloud.attributes.new("hi_phy_constraint_point_indices0_0", 'INT', 'POINT')
    index0_0_attr.data.foreach_set('value', bond_indices);

    index0_1_attr = pointcloud.attributes.new("hi_phy_constraint_point_indices0_1", 'INT', 'POINT')
    index0_1_attr.data.foreach_set('value', bond_indices);

    index0_2_attr = pointcloud.attributes.new("hi_phy_constraint_point_indices0_2", 'INT', 'POINT')
    index0_2_attr.data.foreach_set('value', bond_indices);

    # Use this to store face id
    index1_0_attr = pointcloud.attributes.new("hi_phy_constraint_point_indices1_0", 'INT', 'POINT')
    bond_face_ids = np.zeros(n_curves, dtype=np.int32)
    SimulationAPI.CopyIntArrayToPython(bond_face_ids, bind_face_ids)
    index1_0_attr.data.foreach_set('value', np.repeat(bond_face_ids, 3));

    index1_1_attr = pointcloud.attributes.new("hi_phy_constraint_point_indices1_1", 'INT', 'POINT')
    index1_1_attr.data.foreach_set('value', zero_data);

    index1_2_attr = pointcloud.attributes.new("hi_phy_constraint_point_indices1_2", 'INT', 'POINT')
    index1_2_attr.data.foreach_set('value', zero_data);

    weight0_attr = pointcloud.attributes.new("hi_phy_constraint_weight0", 'FLOAT_VECTOR', 'POINT')
    weight0_attr.data.foreach_set('vector', weight_data);

    # Use this to store barycentric weights
    weight1_attr = pointcloud.attributes.new("hi_phy_constraint_weight1", 'FLOAT_VECTOR', 'POINT')
    bond_barycentric_weights_on_mesh = np.zeros((n_curves, 3), dtype=np.float32)
    SimulationAPI.CopyVector3ArrayToPython(bond_barycentric_weights_on_mesh, bind_barycentric_weights_on_mesh)
    tmp = np.repeat(bond_barycentric_weights_on_mesh, 3, axis=0)
    weight1_attr.data.foreach_set('vector', np.ravel(tmp));

    # rest length is still zero (zero distance to the bond point)
    rest_length_attr = pointcloud.attributes.new("hi_phy_constraint_rest_length", 'FLOAT', 'POINT')
    rest_length_attr.data.foreach_set('value', zero_data);

    local_coord_attr = pointcloud.attributes.new("hi_phy_constraint_local_coord", 'FLOAT_VECTOR', 'POINT')
    bond_local_coord = np.zeros((n_curves * 3, 3), dtype=np.float32)
    SimulationAPI.CopyVector3ArrayToPython(bond_local_coord, bind_local_coord)
    local_coord_attr.data.foreach_set('vector', np.ravel(bond_local_coord));

    positions0_attr = pointcloud.attributes.new("hi_phy_constraint_point_positions_0", 'FLOAT_VECTOR', 'POINT')
    positions0_attr.data.foreach_set('vector', zero_data_3);

    positions1_attr = pointcloud.attributes.new("hi_phy_constraint_point_positions_1", 'FLOAT_VECTOR', 'POINT')
    positions1_attr.data.foreach_set('vector', zero_data_3);

    pointcloud.update()
    return True
