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

def createBindConstraintObject(name, sim_obj_1, sim_obj_2, indices1, indices2, binding_positions1, binding_positions2, weights1, weights2, zero_length_binding):
    n_points = int(np.array(binding_positions1).size / 3)
    edge_mesh = bpy.data.meshes.new(name + "Mesh")
    edges = [(i, i + n_points) for i in range(n_points)]
    edge_mesh.from_pydata(np.concatenate((binding_positions1, binding_positions2)), edges, [])

    obj_new = bpy.data.objects.new(name, edge_mesh)
    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_1.original
    obj_new.hi_phy.constraint_object1 = sim_obj_2.original
    obj_new.hi_phy.constraint_type = SimulationAPI.ConstraintType.BindingConstraint
    obj_new.hi_phy.solver = sim_obj_1.hi_phy.solver
    bpy.context.collection.objects.link(obj_new)

    # Do not create multiple attributes at the same time
    # It will make the previous handle invalid
    index0_0_attr = edge_mesh.attributes.new("hi_phy_constraint_point_indices0_0", 'INT', 'EDGE')
    index0_0_attr.data.foreach_set('value', indices1[0]);

    index0_1_attr = edge_mesh.attributes.new("hi_phy_constraint_point_indices0_1", 'INT', 'EDGE')
    index0_1_attr.data.foreach_set('value', indices1[1]);

    index0_2_attr = edge_mesh.attributes.new("hi_phy_constraint_point_indices0_2", 'INT', 'EDGE')
    index0_2_attr.data.foreach_set('value', indices1[2]);

    index1_0_attr = edge_mesh.attributes.new("hi_phy_constraint_point_indices1_0", 'INT', 'EDGE')
    index1_0_attr.data.foreach_set('value', indices2[0]);

    index1_1_attr = edge_mesh.attributes.new("hi_phy_constraint_point_indices1_1", 'INT', 'EDGE')
    index1_1_attr.data.foreach_set('value', indices2[1]);

    index1_2_attr = edge_mesh.attributes.new("hi_phy_constraint_point_indices1_2", 'INT', 'EDGE')
    index1_2_attr.data.foreach_set('value', indices2[2]);

    positions0_attr = edge_mesh.attributes.new("hi_phy_constraint_point_positions_0", 'FLOAT_VECTOR', 'EDGE')
    positions0_attr.data.foreach_set('vector', np.ravel(np.array(binding_positions1)));

    positions1_attr = edge_mesh.attributes.new("hi_phy_constraint_point_positions_1", 'FLOAT_VECTOR', 'EDGE')
    positions1_attr.data.foreach_set('vector', np.ravel(np.array(binding_positions2)));

    weight0_attr = edge_mesh.attributes.new("hi_phy_constraint_weight0", 'FLOAT_VECTOR', 'EDGE')
    weight0_attr.data.foreach_set('vector', np.ravel(np.array(weights1)));

    weight1_attr = edge_mesh.attributes.new("hi_phy_constraint_weight1", 'FLOAT_VECTOR', 'EDGE')
    weight1_attr.data.foreach_set('vector', np.ravel(np.array(weights2)));

    # note the rest length is in world space
    rest_length = [np.linalg.norm(np.array(binding_positions1[p]) - np.array(binding_positions2[p])) for p in range(n_points)] if not zero_length_binding else [0] * n_points
    rest_length_attr = edge_mesh.attributes.new("hi_phy_constraint_rest_length", 'FLOAT', 'EDGE')
    rest_length_attr.data.foreach_set('value', rest_length);

    edge_mesh.update()
    return edge_mesh

# we assume the object passed in has already been evaluated
def createMeshMesh(sim_obj_1, sim_obj_2, binding_radius, name, zero_length_binding, operator = None):
    if not sim_obj_1.hi_phy:
        if operator:
            operator.report({"ERROR"}, sim_obj_1.name + " is not a valid hi phy object")
        return False
    if not sim_obj_2.hi_phy:
        if operator:
            operator.report({"ERROR"}, sim_obj_2.name + " is not a valid hi phy object")
        return False

    # get mesh topology
    if (not sim_obj_1.type == 'MESH') :
        if operator:
            operator.report({"ERROR"}, "Binding mesh " + sim_obj_1.name + " is not a triangle mesh")
        return False
    if (not sim_obj_2.type == 'MESH') :
        if operator:
            operator.report({"ERROR"}, "Binding mesh " + sim_obj_2.name + " is not a triangle mesh")
        return False

    # build both BVH in world's space, as radius is in world space unit
    radius_1 = [binding_radius] * len(sim_obj_1.data.vertices)
    radius_2 = [binding_radius] * len(sim_obj_2.data.vertices)
    sim_obj_1_bvh = createMeshBVH(sim_obj_1, sim_obj_1.matrix_world, radius_1, operator)
    
    if (sim_obj_1.name != sim_obj_2.name):
        sim_obj_2_bvh = createMeshBVH(sim_obj_2, sim_obj_2.matrix_world, radius_2, operator)
    else:
        sim_obj_2_bvh = sim_obj_1_bvh

    if (not sim_obj_1_bvh) :
        if operator:
            operator.report({"ERROR"}, "Failed to build BVH for " + sim_obj_1.name)
        return False

    if (not sim_obj_2_bvh) :
        if operator:
            operator.report({"ERROR"}, "Failed to build BVH for " + sim_obj_1.name)
        return False

    points1 = SimulationAPI.CreateVector3Array()
    points2 = SimulationAPI.CreateVector3Array()
    verts1 = SimulationAPI.CreateInt3Array()
    verts2 = SimulationAPI.CreateInt3Array()
    barycentric_weights1 = SimulationAPI.CreateVector3Array()
    barycentric_weights2 = SimulationAPI.CreateVector3Array()
    rest_distance = SimulationAPI.CreateFloatArray()

    SimulationAPI.Utils.MeshMeshProximity(sim_obj_1_bvh, sim_obj_2_bvh,
                                          points1, points2,
                                          verts1, verts2,
                                          barycentric_weights1, barycentric_weights2,
                                          rest_distance);

    # bond_points_1 and 2 should have same points
    n_points = len(points1)
    if (n_points == 0):
        if operator:
            operator.report({"ERROR"}, "Nothing was bond")
            return False

    vs_world_1 = np.zeros((n_points, 3), dtype=np.float32)
    vs_world_2 = np.zeros((n_points, 3), dtype=np.float32)
    SimulationAPI.CopyVector3ArrayToPython(vs_world_1, points1)
    SimulationAPI.CopyVector3ArrayToPython(vs_world_2, points2)

    weights_1 = np.zeros((n_points, 3), dtype=np.float32)
    weights_2 = np.zeros((n_points, 3), dtype=np.float32)
    SimulationAPI.CopyVector3ArrayToPython(weights_1, barycentric_weights1)
    SimulationAPI.CopyVector3ArrayToPython(weights_2, barycentric_weights2)

    indices0_0 = [v[0] for v in verts1]
    indices0_1 = [v[1] for v in verts1]
    indices0_2 = [v[2] for v in verts1]
    indices1_0 = [v[0] for v in verts2]
    indices1_1 = [v[1] for v in verts2]
    indices1_2 = [v[2] for v in verts2]

    createBindConstraintObject(name, sim_obj_1, sim_obj_2,
                               [indices0_0, indices0_1, indices0_2],
                               [indices1_0, indices1_1, indices1_2],
                               vs_world_1, vs_world_2,
                               weights_1, weights_2, zero_length_binding)
    return True

# we assume the object passed in has already been evaluated
def createCurvesCurves(sim_obj_1, sim_obj_2, binding_radius, name, zero_length_binding, curve_self_binding, operator = None):
    if not sim_obj_1.hi_phy:
        if operator:
            operator.report({"ERROR"}, sim_obj_1.name + " is not a valid hi phy object")
        return False
    if not sim_obj_2.hi_phy:
        if operator:
            operator.report({"ERROR"}, sim_obj_2.name + " is not a valid hi phy object")
        return False

    # get mesh topology
    if (not sim_obj_1.type == 'CURVES') :
        if operator:
            operator.report({"ERROR"}, "Binding curves " + sim_obj_1.name + " is not a curves object")
        return False
    if (not sim_obj_2.type == 'CURVES') :
        if operator:
            operator.report({"ERROR"}, "Binding curves " + sim_obj_2.name + " is not a curves object")
        return False

    # build both BVH in world space (radius is in world space)
    n_cvs1 = sim_obj_1.data.attributes.domain_size('POINT')
    n_cvs2 = sim_obj_2.data.attributes.domain_size('POINT')
    radius_1 = [binding_radius] * n_cvs1
    radius_2 = [binding_radius] * n_cvs2
    sim_obj_1_bvh = createCurvesBVH(sim_obj_1, sim_obj_1.matrix_world, radius_1, operator)
    if (not sim_obj_1_bvh):
        if operator:
            operator.report({"ERROR"}, "Failed to create BVH for " + sim_obj_1.name)
        return False
    if (sim_obj_1.name != sim_obj_2.name):
        sim_obj_2_bvh = createCurvesBVH(sim_obj_2, sim_obj_2.matrix_world, radius_2, operator)
    else:
        sim_obj_2_bvh = sim_obj_1_bvh
    if (not sim_obj_2_bvh):
        if operator:
            operator.report({"ERROR"}, "Failed to create BVH for " + sim_obj_2.name)
        return False
    points1 = SimulationAPI.CreateVector3Array()
    points2 = SimulationAPI.CreateVector3Array()
    cv1 = SimulationAPI.CreateIntArray()
    cv2 = SimulationAPI.CreateIntArray()
    k1 = SimulationAPI.CreateFloatArray()
    k2 = SimulationAPI.CreateFloatArray()
    rest_distance = SimulationAPI.CreateFloatArray()

    SimulationAPI.Utils.CurvesCurvesProximity(sim_obj_1_bvh, sim_obj_2_bvh,
                                              points1, points2,
                                              cv1, cv2, k1, k2,
                                              rest_distance, not curve_self_binding);

    # bond_points_1 and 2 should have same points
    n_points = len(points1)
    if (n_points == 0):
        if operator:
            operator.report({"ERROR"}, "Nothing was bond")
            return False
    vs_world_1 = np.zeros((n_points, 3), dtype=np.float32)
    vs_world_2 = np.zeros((n_points, 3), dtype=np.float32)
    SimulationAPI.CopyVector3ArrayToPython(vs_world_1, points1)
    SimulationAPI.CopyVector3ArrayToPython(vs_world_2, points2)

    indices0_0 = [cv for cv in cv1]
    indices0_1 = [cv + 1 for cv in cv1]
    indices0_2 = [cv for cv in cv1]
    indices1_0 = [cv for cv in cv2]
    indices1_1 = [cv + 1 for cv in cv2]
    indices1_2 = [cv for cv in cv2]
    weight0 = [w for k in k1 for w in [k, 1 - k, 0]]
    weight1 = [w for k in k2 for w in [k, 1 - k, 0]]
    createBindConstraintObject(name, sim_obj_1, sim_obj_2,
                               [indices0_0, indices0_1, indices0_2],
                               [indices1_0, indices1_1, indices1_2],
                               vs_world_1, vs_world_2,
                               weight0, weight1, zero_length_binding)

    return True

# we assume the object passed in has already been evaluated
def createMeshCurve(sim_obj_1, sim_obj_2, binding_radius, name, zero_length_binding, operator = None):
    if not sim_obj_1.hi_phy:
        if operator:
            operator.report({"ERROR"}, sim_obj_1.name + " is not a valid hi phy object")
        return False
    if not sim_obj_2.hi_phy:
        if operator:
            operator.report({"ERROR"}, sim_obj_2.name + " is not a valid hi phy object")
        return False

    # get mesh topology
    if (not sim_obj_1.type == 'MESH') :
        if operator:
            operator.report({"ERROR"}, "Binding mesh " + sim_obj_1.name + " is not a triangle mesh")
        return False
    if (not sim_obj_2.type == 'CURVES') :
        if operator:
            operator.report({"ERROR"}, "Binding object " + sim_obj_2.name + " curves object")
        return False

    # Copy transforms
    # build both BVH in world space, we have to because radius is in world space unit
    radius_1 = [binding_radius] * sim_obj_1.data.attributes.domain_size('POINT')
    radius_2 = [binding_radius] * sim_obj_2.data.attributes.domain_size('POINT')
    sim_obj_1_bvh = createMeshBVH(sim_obj_1, sim_obj_1.matrix_world, radius_1, operator)
    if (not sim_obj_1_bvh):
        if operator:
            operator.report({"ERROR"}, "Failed to create BVH for " + sim_obj_1.name)
        return False
    sim_obj_2_bvh = createCurvesBVH(sim_obj_2, sim_obj_2.matrix_world, radius_2, operator)
    if (not sim_obj_2_bvh):
        if operator:
            operator.report({"ERROR"}, "Failed to create BVH for " + sim_obj_2.name)
        return False

    points1 = SimulationAPI.CreateVector3Array()
    points2 = SimulationAPI.CreateVector3Array()
    verts1 = SimulationAPI.CreateInt3Array()
    cv2 = SimulationAPI.CreateIntArray()
    barycentric_weights1 = SimulationAPI.CreateVector3Array()
    k2 = SimulationAPI.CreateFloatArray()
    rest_distance = SimulationAPI.CreateFloatArray()

    SimulationAPI.Utils.MeshCurvesProximity(sim_obj_1_bvh, sim_obj_2_bvh,
                                            points1, points2,
                                            verts1, cv2,
                                            barycentric_weights1, k2,
                                            rest_distance);

    # bond_points_1 and 2 should have same points
    n_points = len(points1)
    if (n_points == 0):
        if operator:
            operator.report({"ERROR"}, "Nothing was bond")
            return False

    vs_world_1 = np.zeros((n_points, 3), dtype=np.float32)
    vs_world_2 = np.zeros((n_points, 3), dtype=np.float32)
    SimulationAPI.CopyVector3ArrayToPython(vs_world_1, points1)
    SimulationAPI.CopyVector3ArrayToPython(vs_world_2, points2)

    indices0_0 = [v[0] for v in verts1]
    indices0_1 = [v[1] for v in verts1]
    indices0_2 = [v[2] for v in verts1]
    indices1_0 = [cv for cv in cv2]
    indices1_1 = [cv + 1 for cv in cv2]
    indices1_2 = [cv for cv in cv2]

    weights_1 = np.zeros((n_points, 3), dtype=np.float32)
    SimulationAPI.CopyVector3ArrayToPython(weights_1, barycentric_weights1)
    weights_2 = [w for k in k2 for w in [k, 1 - k, 0]]

    createBindConstraintObject(name, sim_obj_1, sim_obj_2,
                               [indices0_0, indices0_1, indices0_2],
                               [indices1_0, indices1_1, indices1_2],
                               vs_world_1, vs_world_2,
                               weights_1, weights_2, zero_length_binding)
    return True

