import bpy
import mathutils
import numpy as np
from . import createTransform
from . import add_driver
from . import delete_children
from ..cAPI import SimulationAPI
from ..properties import object_properties
from .. import ADDON_VERSION
# both grab_constraint and simulation_object are evaluated
def create(grab_constraint, simulation_object, weight, name, operator = None):
    n_verts = simulation_object.data.attributes.domain_size('POINT')
    positions_data = np.zeros((n_verts, 3), dtype=np.float32)
    simulation_object.data.attributes['position'].data.foreach_get('vector', np.ravel(positions_data))
    simulation_transform = createTransform(simulation_object.matrix_world);
    constraint_transform = createTransform(grab_constraint.matrix_world);
    points = SimulationAPI.CreateVector3Array(n_verts)
    SimulationAPI.CopyVector3Array(points, positions_data)
    weights = SimulationAPI.CreateFloatArray(n_verts)
    if (not grab_constraint.hi_phy.is_active or
        grab_constraint.hi_phy.object_type != "CONSTRAINT" or
        grab_constraint.hi_phy.constraint_type != SimulationAPI.ConstraintType.GrabConstraint):
        if operator:
            operator.report({"ERROR"}, grab_constraint.name + " is not a HiPhy grab constraint")
        return False
    # the center is always zero in the constraint space
    center = SimulationAPI.Vector3()
    SimulationAPI.Utils.ComputeSphericalMapWeights(weights, points,
                                                   grab_constraint.hi_phy.grab_constraint_inner_radius,
                                                   grab_constraint.hi_phy.grab_constraint_outer_radius,
                                                   center,
                                                   constraint_transform, simulation_transform)

    # We need to take only the vertices with non-zero weights
    verts = [p for p in range(n_verts) if weights[p] > 0]

    transform = grab_constraint.matrix_world.inverted() @ simulation_object.matrix_world;
    # this is in the grab constraint space
    vs = [transform @ mathutils.Vector((v[0], v[1], v[2])) for v in positions_data[verts]]
    
    n_points = len(verts)
    if (n_points == 0):
        if operator:
            operator.report({"ERROR"}, "Nothing was bond")
        return False

    pointcloud = bpy.data.meshes.new(name + "Mesh")
    pointcloud.from_pydata(vs, [], []) # Create a mesh with only verts

    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 = simulation_object.original
    obj_new.hi_phy.constraint_object1 = simulation_object.original
    obj_new.hi_phy.constraint_type = SimulationAPI.ConstraintType.PointConstraint
    obj_new.hi_phy.solver = simulation_object.hi_phy.solver
    obj_new.hi_phy.constraintIsSoft = True;
    obj_new.hi_phy.constraintStretchStiffness_use_attribute = True;
    obj_new.hi_phy.constraintStretchStiffness_attribute = "hi_phy_constraint_stiffness";
    object_properties.updt_bool_obj_type(obj_new.hi_phy)

    bpy.context.collection.objects.link(obj_new)
    obj_new.parent = grab_constraint.original

    weight_data = [1, 0, 0] * n_points
    zero_data = [0] * n_points
    zero_data_3 = [0, 0, 0] * n_points
    # weight_data_flat = [v for w in weight_data for v in w]

    stiffness_scale = [weights[v] * weight for v in verts]
    # 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', verts);

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

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

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

    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);

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

    rest_length_attr = pointcloud.attributes.new("hi_phy_constraint_rest_length", 'FLOAT', 'POINT')
    rest_length_attr.data.foreach_set('value', zero_data);

    positions_np = np.array(vs, dtype=np.float32)
    positions0_attr = pointcloud.attributes.new("hi_phy_constraint_point_positions_0", 'FLOAT_VECTOR', 'POINT')
    positions0_attr.data.foreach_set('vector', np.ravel(positions_np));

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

    # we are going to use it to set a map for scaling the extension stiffness
    # we need to set up a geometry node to scale this map with the parent's stretch stiffness
    stiffness_scale_attr = pointcloud.attributes.new("hi_phy_constraint_stiffness_scale", 'FLOAT', 'POINT')
    stiffness_scale_attr.data.foreach_set('value', stiffness_scale);

    pointcloud.update()

    # Create the geometry node to connect everything
    modifier = obj_new.modifiers.new(name="GN_GrabConstraint", type='NODES')

    gn_name = "GN_GrabsConstraintStiffness" + ADDON_VERSION
    node_tree = bpy.data.node_groups.get(gn_name)
    if not node_tree:
        # Create node tree
        node_tree = bpy.data.node_groups.new(gn_name, 'GeometryNodeTree')        
        # Clear default nodes
        nodes = node_tree.nodes
        links = node_tree.links
        nodes.clear()
        
        # Add Group Input & Output
        group_input = nodes.new("NodeGroupInput")
        group_output = nodes.new("NodeGroupOutput")

        # Add geometry sockets (required)
        node_tree.interface.new_socket(
            name="Geometry",
            in_out='INPUT',
            socket_type='NodeSocketGeometry'
        )

        node_tree.interface.new_socket(
            name="Stiffness",
            in_out='INPUT',
            socket_type='NodeSocketFloat'
        )

        node_tree.interface.new_socket(
            name="Geometry",
            in_out='OUTPUT',
            socket_type='NodeSocketGeometry'
        )

        # Add a Transform node
        input_attribute = nodes.new("GeometryNodeInputNamedAttribute")
        input_attribute.inputs["Name"].default_value = "hi_phy_constraint_stiffness_scale"

        multiply = nodes.new("ShaderNodeMath")
        multiply.operation = 'MULTIPLY'

        links.new(input_attribute.outputs["Attribute"], multiply.inputs[0])
        links.new(group_input.outputs["Stiffness"], multiply.inputs[1])

        output_attribute = nodes.new("GeometryNodeStoreNamedAttribute")
        output_attribute.data_type = 'FLOAT'
        output_attribute.domain = 'POINT'
        output_attribute.inputs["Name"].default_value = "hi_phy_constraint_stiffness"
        
        links.new(multiply.outputs[0], output_attribute.inputs["Value"])
        links.new(group_input.outputs["Geometry"], output_attribute.inputs["Geometry"])
        links.new(output_attribute.outputs["Geometry"], group_output.inputs["Geometry"])

    # Assign the geometry node
    modifier.node_group = node_tree
    # creating drivers
    add_driver(obj_new, grab_constraint.original, "hi_phy.constraintEnabled", "hi_phy.constraintEnabled")
    add_driver(modifier, grab_constraint.original, '["Socket_1"]', "hi_phy.constraintStretchStiffness")
    add_driver(obj_new, grab_constraint.original, "hi_phy.constraintDamping", "hi_phy.constraintDamping")
    add_driver(obj_new, grab_constraint.original, "hi_phy.constraintMaxStiffnessDistance", "hi_phy.constraintMaxStiffnessDistance")
    return True

def release(grab_constraint, operator = None):
    delete_children(grab_constraint.original)
    return True
