text:
from mpl_toolkits.mplot3d import Axes3D

from fabrik_chain_3d import Bone as Bone, Joint as Joint, Mat as Mat, Utils as Util
from fabrik_chain_3d.output_writer import *
import numpy as np
import matplotlib.pyplot as plt
import math

class Chain3d:
    def __init__(self, is_base_bone_fixed, base_address="./test/output/", joints_file_address="joints-position.txt"):
        self.is_base_bone_fixed = is_base_bone_fixed
        self.target_position = [1, 0, 0]
        self.target_orientation = [1, 0, 0, 0]
        self.chain = []
        self.solve_distance_threshold = 0.001
        self.max_iteration_attempts = 20
        self.min_iteration_change = 0.01
        self.fixed_base_location = [0, 0, 0]
        self.fixed_base_location_2 = [0, 0, 0]
        self.chain_length = 0
        self.base_bone_constraint_type = "None"
        self.base_bone_constraint_uv = [0, 0, 0]
        self.current_solve_distance = 100000000
        self.bone_twist_limit = 2.8973*180/math.pi
        self.deg =[0]*4
        self. rotations = [0]*4 # the last one belongs to the base bone which is fixed!
        self.fixed_base_orientation = [1,0,0,0]
        self.base_address = base_address
        self.joints_file_address = joints_file_address
    def update_chain_length(self):
        self.chain_length = 0
        for i in self.chain:
            self.chain_length += 1

    def add_bone(self, bone):
        self.chain.append(bone)
        # if it is the base bone:
        if len(self.chain) == 1:
            self.fixed_base_location = bone.get_start_point_position()
            if self.is_base_bone_fixed == 1:
                self.fixed_base_location_2 = bone.get_end_point_position()
            self.base_bone_constraint_uv = bone.get_direction_uv()

        self.update_chain_length()

    def add_consecutive_bone(self, bone):
        direction = bone.get_direction_uv()
        length = bone.get_length()
        # If we have at least one bone already in the chain...
        if self.get_chain_length() != 0:
            prev_bone_end = self.get_bone(self.chain_length - 1).get_end_point_position()
            bone.set_start_point_position(prev_bone_end)
            scale = np.dot(direction, length)
            end_point = [x + y for x, y in zip(prev_bone_end, scale)]
            bone.set_end_point_position(end_point)
            self.add_bone(bone)
        else:
            raise Exception(
                "You cannot add the base bone to a chain using this method as it does not provide a start location.")

    def add_consecutive_freely_rotating_hinged_bone(self, bone_direction_uv, bone_length, joint_type,
                                                    hinge_rotation_axis, is_fixed,bone_orientation):
        self.add_consecutive_hinged_bone(bone_direction_uv, bone_length, joint_type, hinge_rotation_axis, 180, 180,
                                         Util.Utils().gen_perpendicular_vector_quick(hinge_rotation_axis), is_fixed,bone_orientation)

    def add_consecutive_hinged_bone(self, bone_direction_uv, bone_length, joint_type, hinge_rotation_axis,
                                    clockwise_degs,
                                    anticlockwise_deg, hinge_reference_axis, is_bone_fixed, bone_orientation):
        Util.Utils().validate_direction(bone_direction_uv)
        Util.Utils().validate_direction(hinge_rotation_axis)
        Util.Utils().validate_length(bone_length)

        if self.chain_length == 0:
            raise Exception("You must add a base bone before adding a consectutive bone.")

        bone_direction_uv = (bone_direction_uv)
        hinge_rotation_axis = Util.Utils().normalization(hinge_rotation_axis)
        prev_bone_end = self.get_bone(self.chain_length - 1).get_end_point_position()
        scale_direction = [i * bone_length for i in bone_direction_uv]
        bone_end_location = [x + y for x, y in zip(prev_bone_end, scale_direction)]
        m_bone = Bone.Bone3D(prev_bone_end, bone_end_location, bone_direction_uv, bone_length, is_bone_fixed,
                             bone_orientation)
        m_joint = Joint.Joint3D()
        if joint_type == "GLOBAL_HINGE":
            m_joint.set_as_global_hinge(hinge_rotation_axis, clockwise_degs, anticlockwise_deg, hinge_reference_axis)
        elif joint_type == "LOCAL_HINGE":
            m_joint.set_as_local_hinge(hinge_rotation_axis, clockwise_degs, anticlockwise_deg, hinge_reference_axis)
        else:
            raise Exception("Hinge joint types may be only JointType.GLOBAL_HINGE or JointType.LOCAL_HINGE.")

        m_bone.set_joint(m_joint)
        self.add_bone(m_bone)

    def add_consecutive_rotor_constrained_bone(self, bone_direction_uv, bone_length, constraint_angle_degs, is_fixed,bone_orientation):
        if self.get_chain_length() == 0:
            raise Exception("Add a basebone before attempting to add consectuive bones.")
        prev_bone_end =self.get_bone(self.get_chain_length() - 1).get_end_point_position()
        scale_direction = [i * bone_length for i in bone_direction_uv]
        bone_end_location = [x + y for x, y in zip(prev_bone_end, scale_direction)]

        m_bone = Bone.Bone3D(prev_bone_end, bone_end_location, bone_direction_uv, bone_length, is_fixed, bone_orientation)
        rotor_joint = Joint.Joint3D()
        rotor_joint.set_as_ball_joint(constraint_angle_degs)
        rotor_joint.set_as_ball_joint(constraint_angle_degs)
        m_bone.set_joint(rotor_joint)
        m_bone.set_ball_joint_constraint_degs(constraint_angle_degs)
        self.add_bone(m_bone)

    def get_base_bone_constraint_uv(self):
        return self.base_bone_constraint_uv

    def get_base_bone_constraint_type(self):
        if self.base_bone_constraint_type is None:
            raise Exception('Cannot return the basebone constraint when the basebone constraint type is NONE.')
        return self.base_bone_constraint_type

    def get_base_location(self):
        return self.get_bone(0).get_start_point_position()

    def get_bone(self, bone_number):
        return self.chain[bone_number]

    def get_chain(self):
        return self.chain

    def get_chain_length(self):
        return self.chain_length

    def get_end_effector_location(self):
        return self.get_bone(self.chain_length - 1).get_end_point_position()

    def remove_bone(self, bone_number):
        if bone_number < self.chain_length:
            self.chain.remove(self.get_bone(bone_number))
            self.update_chain_length()
        else:
            raise Exception(
                "Bone " + bone_number + " does not exist to be removed from the chain. Bones are zero indexed.")

    def set_rotor_base_bone_constraint(self, rotor_type, constraint_axis, angle_degs):
        # Sanity checking
        if len(self.chain) == 0:
            raise Exception("Chain must contain a basebone before we can specify the basebone constraint type.")

        if len(constraint_axis) <= 0.0:
            raise Exception("Constraint axis cannot be zero.")

        if angle_degs < 0.0:
            angle_degs = 0.0

        if angle_degs > 180.0:
            angle_degs = 180.0

        # if rotor_type != "GLOBAL_ROTOR" or rotor_type != "LOCAL_ROTOR":
        #     raise Exception("The only valid rotor types for this method are GLOBAL_ROTOR and LOCAL_ROTOR.")

        #  Set the constraint type, axis and angle
        self.base_bone_constraint_type = rotor_type
        self.base_bone_constraint_uv = Util.Utils().normalization(constraint_axis)
        rotor_joint = Joint.Joint3D()
        rotor_joint.set_as_ball_joint(angle_degs)
        self.get_bone(0).set_joint(rotor_joint)

    def set_hinge_base_bone_constraint(self, hinge_type, hinge_rotation_axis, cw_constraint_degs, acw_constraint_axis,
                                       hinge_reference_axis, ):
        # Sanity checking
        if len(self.chain) == 0:
            raise Exception("Chain must contain a base bone before we can specify the base bone constraint type.")

        if Util.length_calc(hinge_rotation_axis) <= 0.0:
            raise Exception("Constraint axis cannot be zero.")

        if Util.length_calc(hinge_reference_axis) <= 0.0:
            raise Exception("Constraint axis cannot be zero.")

        # if np.inner(hinge_rotation_axis, hinge_reference_axis) == 0:
        #     # raise Exception(
            #     "The hinge reference-axis must be in the plane of the hinge rotation axis, i.e"
            #     ", they must not be perpendicular")

        # if hinge_type != "GLOBAL_HINGE" and hinge_type != "LOCAL_HINGE":
        #     raise Exception("The only valid rotor types for this method are GLOBAL_HINGE and LOCAL_HINGE.")

        #  Set the constraint type, axis and angle
        self.base_bone_constraint_type = hinge_type
        self.base_bone_constraint_uv = Util.normalization(hinge_rotation_axis)

        hinge_joint = Joint.Joint3D()

        if hinge_type == "GLOBAL_HINGE":
            hinge_joint.set_hinge("GLOBAL_HINGE", hinge_rotation_axis, cw_constraint_degs, acw_constraint_axis,
                                  hinge_reference_axis)

        else:
            hinge_joint.set_hinge("LOCAL_HINGE", hinge_rotation_axis, cw_constraint_degs, acw_constraint_axis,
                                  hinge_reference_axis)

        self.get_bone(0).set_joint(hinge_joint)

    def set_freely_rotating_global_hinged_base_bone(self, hinge_rotation_axis):
        self.set_hinge_base_bone_constraint("GLOBAL_HINGE", hinge_rotation_axis, 180, 180,
                                            Util.Utils().gen_perpendicular_vector_quick(hinge_rotation_axis))

    def set_freely_rotating_local_hinged_base_bone(self, hinge_rotation_axis):
        self.set_hinge_base_bone_constraint("LOCAL_HINGE", hinge_rotation_axis, 180, 180,
                                            Util.Utils().gen_perpendicular_vector_quick(hinge_rotation_axis))

    def set_local_hinged_base_bone(self, hinge_rotation_axis, cw_constraint_degs, acw_constraint_degs,
                                   hinge_reference_axis):
        self.set_hinge_base_bone_constraint("LOCAL_HINGE", hinge_rotation_axis, cw_constraint_degs, acw_constraint_degs,
                                            hinge_reference_axis)

    def set_global_hinged_base_bone(self, hinge_rotation_axis, cw_constraint_degs, acw_constraint_degs,
                                    hinge_reference_axis):
        self.set_hinge_base_bone_constraint("GLOBAL_HINGE", hinge_rotation_axis, cw_constraint_degs,
                                            acw_constraint_degs,
                                            hinge_reference_axis)

    def set_base_bone_constraint_uv(self, constraint_uv):
        if len(constraint_uv) == 0:
            raise Exception("direction unit vector cannot be zero")
        self.base_bone_constraint_uv = Util.Utils().normalization(constraint_uv)

    def set_base_location(self, base_location):
        self.fixed_base_location = base_location

    def set_max_iteration_attempt(self, max_iteration):
        if max_iteration < 1:
            raise Exception("Minimum number of iteration is 1")
        self.max_iteration_attempts = max_iteration

    def get_max_iteration_attempt(self):
        return self.max_iteration_attempts

    def set_min_iteration_change(self, min_iteration_change):
        if min_iteration_change < 0:
            raise Exception("Minimum iteration change greater than zero")
        self.min_iteration_change = min_iteration_change

    def get_min_iteration_change(self):
        return self.min_iteration_change

    def set_solve_distance_threshold(self, solve_distance):
        if solve_distance < 0:
            raise Exception("The solve distance threshold must be greater than or equal to zero.")
        self.solve_distance_threshold = solve_distance

    def get_solve_distance_threshold(self):
        return self.solve_distance_threshold

    def set_target(self, target_position, target_orientation):
        self.target_position = target_position
        self.target_orientation = target_orientation

    def solve_for_orientation(self, outer_joint_orientation, inner_joint_orientation,loop):
        q1 = outer_joint_orientation
        q2 = inner_joint_orientation
        # finding the rotor that express rotation between two orientational frame(between outer and inner joint)
        rotor = Util.Utils().find_rotation_quaternion(q1, q2)
        if rotor[0]>1:
            rotor[0] =0.99
        if rotor[0] < -1:
            rotor[0] = -0.99
        needed_rotation = math.acos(rotor[0]) * 2 * (180 / np.pi)
        self.rotations[loop] = needed_rotation*(np.pi/180)
        if needed_rotation <= self.bone_twist_limit:
            # if the rotation is inside the limited
            return Mat.Mat().multiply_two_quaternion(rotor, outer_joint_orientation)
        else:
            # the maximum allowed rotation angle
            theta = (self.bone_twist_limit) * (np.pi / 180)
            self.rotations[loop] = theta
            # the rotation axis
            if abs(rotor[0]) ==1:
                return rotor
            v1 = np.dot(rotor[1:], (1 / math.sqrt(1 - rotor[0] ** 2)))
            w = math.cos(theta / 2)
            x = v1[0] * math.sin(theta / 2)
            y = v1[1] * math.sin(theta / 2)
            z = v1[2] * math.sin(theta / 2)
            return [w, x, y, z]

    def forward(self, target_position, target_orientation):
        for loop in range(self.chain_length - 1, -1, -1):
            #  Get the length of the bone we're working on
            this_bone = self.get_bone(loop)
            this_bone_length = this_bone.get_length()
            this_bone_joint = this_bone.get_joint()
            this_bone_joint_type = this_bone_joint.get_joint_type()

            # 	If we are NOT working on the end effector bone
            if loop != (self.chain_length - 1):
                if this_bone.is_fix_bone() == 1:
                    this_bone_outer_to_inner_uv = Util.Utils().negated(this_bone.get_fixed_bone_direction_uv())
                else:
                    # Get the outer-to-inner unit vector of the bone further out
                    outer_bone_outer_to_inner_uv = Util.Utils().negated(self.get_bone(loop + 1).get_direction_uv())
                    # Get the outer-to-inner unit vector of this bone
                    this_bone_outer_to_inner_uv = Util.Utils().negated(self.get_bone(loop).get_direction_uv())
                    next_bone_orientation = self.get_bone(loop+1).get_bone_orientation()
                    this_bone_orientation = self.get_bone(loop).get_bone_orientation()
                    this_bone.set_bone_orientation(self.solve_for_orientation(next_bone_orientation,this_bone_orientation,loop))
                    # Get the joint type for this bone and handle constraints on thisBoneOuterToInnerUV
                    if this_bone_joint_type == "BALL":
                        # Constrain to relative angle between this bone and the outer bone if required
                        angle_between_degs = Util.Utils().get_angle_between_degs(outer_bone_outer_to_inner_uv,
                                                                         this_bone_outer_to_inner_uv)
                        constrain_angle_degs = this_bone_joint.get_ball_joint_constraint_degs()
                        if angle_between_degs > constrain_angle_degs:
                            this_bone_outer_to_inner_uv = Util.Utils().get_angle_limited_uv(this_bone_outer_to_inner_uv,
                                                                                outer_bone_outer_to_inner_uv,
                                                                                constrain_angle_degs)
                    elif this_bone_joint_type == "GLOBAL_HINGE":
                        # Project this bone outer-to-inner direction onto the hinge rotation axis
                        this_bone_outer_to_inner_uv = Util.Utils().project_on_to_plane(this_bone_outer_to_inner_uv,
                                                                               this_bone_joint.get_hinge_rotation_axis())
                    elif this_bone_joint_type == "LOCAL_HINGE":
                        # Not a base bone? Then construct a rotation matrix based on the previous bones
                        # inner-to-outer direction...
                        if loop > 0:
                            m = Util.Utils().create_rotation_matrix(self.get_bone(loop - 1).get_direction_uv())
                            relative_hinge_rotation_axis = Util.Utils().normalization(
                                Util.Utils().times(m, this_bone_joint.get_hinge_rotation_axis()))

                            # transform the hinge rotation axis into the previous bones frame of reference.
                            # Project this bone's outer-to-inner direction onto the plane described by the relative hinge rotation axis
                            this_bone_outer_to_inner_uv = Util.Utils().project_on_to_plane(this_bone_outer_to_inner_uv,
                                                                                relative_hinge_rotation_axis)
                        else:
                            raise Exception("The base bone joint can't be LOCAL HINGE")

                scale = [i * this_bone_length for i in this_bone_outer_to_inner_uv]
                end_location = this_bone.get_end_point_position()
                new_start_location = [x + y for x, y in zip(end_location, scale)]

                this_bone.set_start_point_position(new_start_location)

                # If we are not working on the basebone, then we also set the end joint location of
                # the previous bone in the chain
                if loop > 0:
                    self.get_bone(loop - 1).set_end_point_position(new_start_location)
            # If we ARE working on the end effector bone..
            else:
                # put end effector end location to the target
                this_bone.set_end_point_position(target_position)
                this_bone.set_bone_orientation(
                    self.solve_for_orientation(target_orientation, this_bone.get_bone_orientation(), loop))
                if this_bone.is_fix_bone() == 1:
                    this_bone_outer_to_inner_uv = Util.Utils().negated(this_bone.get_fixed_bone_direction_uv())
                else:
                    this_bone_outer_to_inner_uv = Util.Utils().negated(this_bone.get_direction_uv())

                    if this_bone_joint_type == "BALL":
                        i=0
                    elif this_bone_joint_type == "GLOBAL_HINGE":
                        this_bone_outer_to_inner_uv = Util.Utils().project_on_to_plane(this_bone_outer_to_inner_uv,
                                                                               this_bone_joint.get_hinge_rotation_axis())
                    elif this_bone_joint_type == "LOCAL_HINGE":
                        m = Util.Utils().create_rotation_matrix(self.get_bone(loop - 1).get_direction_uv())
                        relative_hinge_rotation_axis = Util.Utils().normalization(
                            Util.Utils().times(m, this_bone_joint.get_hinge_rotation_axis()))
                        # Project this bone's outer-to-inner direction onto the plane described by the relative hinge
                        # rotation axis
                        this_bone_outer_to_inner_uv = Util.Utils().project_on_to_plane(this_bone_outer_to_inner_uv,
                                                                               relative_hinge_rotation_axis)

                scale = [i * this_bone_length for i in this_bone_outer_to_inner_uv]
                end_location = this_bone.get_end_point_position()
                new_start_location = [x + y for x, y in zip(end_location, scale)]

                this_bone.set_start_point_position(new_start_location)
                # If we are not working on the base bone, then we also set the end joint location of
                # the previous bone in the chain
                if loop > 0:
                    self.get_bone(loop - 1).set_end_point_position(new_start_location)
        # self.draw_chain()

    def backward(self):
        for loop in range(self.chain_length):
            this_bone = self.get_bone(loop)
            this_bone_length = self.get_bone(loop).get_length()
            # If we are not working on the base bone
            if loop != 0:
                if this_bone.is_fix_bone() == 1:
                    this_bone_inner_to_outer_uv = this_bone.get_fixed_bone_direction_uv()
                else:
                    this_bone_inner_to_outer_uv = this_bone.get_direction_uv()
                    prev_bone_inner_to_outer_uv = self.get_bone(loop - 1).get_direction_uv()
                    this_bone_joint = this_bone.get_joint()
                    this_bone_joint_type = this_bone_joint.get_joint_type()
                    if this_bone_joint_type == "BALL":
                        angle_between_degs = Util.Utils().get_angle_between_degs(prev_bone_inner_to_outer_uv,
                                                                         this_bone_inner_to_outer_uv)
                        constraint_angle_degs = this_bone_joint.get_ball_joint_constraint_degs()
                        self.deg[loop] =angle_between_degs
                        if angle_between_degs > constraint_angle_degs:
                            this_bone_inner_to_outer_uv = Util.Utils().get_angle_limited_uv(this_bone_inner_to_outer_uv,
                                                                                prev_bone_inner_to_outer_uv,
                                                                                constraint_angle_degs)
                            self.deg[loop] =constraint_angle_degs
                    elif this_bone_joint_type == "GLOBAL_HINGE":
                        # Get the hinge rotation axis and project our inner-to-outer UV onto it
                        this_bone_inner_to_outer_uv = Util.Utils().project_on_to_plane(this_bone_inner_to_outer_uv,
                                                                               this_bone_joint.get_hinge_rotation_axis())
                        # If there are joint constraints, then we must honour them...
                        cw_constraint_degs = -this_bone_joint.get_hinge_clockwise_constraint_degs()
                        acw_constraint_degs = this_bone_joint.get_hinge_anticlockwise_constraint_degs()
                        if not Util.Utils().approximately_equal(cw_constraint_degs,
                                                        -this_bone_joint.get_MAX_CONSTRAINT_ANGLE_DEGS(),
                                                        0.001) and not Util.Utils().approximately_equal(acw_constraint_degs,
                                                                                this_bone_joint.get_MAX_CONSTRAINT_ANGLE_DEGS(),
                                                                                0.001):
                            hinge_reference_axis = this_bone_joint.get_reference_axis()
                            hinge_rotation_axis = this_bone_joint.get_hinge_rotation_axis()
                            # Get the signed angle (about the hinge rotation axis) between the hinge reference axis and the hinge-rotation aligned bone UV
                            signed_angle_degs = Util.Utils().get_signed_angle_between_degs(hinge_reference_axis,
                                                                                this_bone_inner_to_outer_uv,
                                                                                hinge_rotation_axis)
                            self.deg[loop] = signed_angle_degs*math.pi/180
                            # Make our bone inner-to-outer UV the hinge reference axis rotated by its maximum clockwise or anticlockwise rotation as required
                            if signed_angle_degs > acw_constraint_degs:
                                this_bone_inner_to_outer_uv = Util.Utils().normalization(
                                    Mat.Mat().rotate_about_axis(hinge_reference_axis, acw_constraint_degs,
                                                          hinge_rotation_axis))
                                self.deg[loop] = acw_constraint_degs*math.pi/180
                            elif signed_angle_degs < cw_constraint_degs:
                                this_bone_inner_to_outer_uv = Util.Utils().normalization(
                                    Mat.Mat().rotate_about_axis(hinge_reference_axis, cw_constraint_degs,
                                                          hinge_rotation_axis))
                                self.deg[loop] = cw_constraint_degs*math.pi/180
                    elif this_bone_joint_type == "LOCAL_HINGE":
                        # Transform the hinge rotation axis to be relative to the previous bone in the chain
                        hinge_rotation_axis = this_bone_joint.get_hinge_rotation_axis()
                        m = Util.Utils().create_rotation_matrix(prev_bone_inner_to_outer_uv)
                        relative_hinge_rotation_axis = Util.Utils().normalization(Util.Utils().times(m, hinge_rotation_axis))
                        this_bone_inner_to_outer_uv = Util.Utils().project_on_to_plane(this_bone_inner_to_outer_uv,
                                                                               relative_hinge_rotation_axis)
                        # Constrain rotation about reference axis if required
                        cw_constraint_degs = -this_bone_joint.get_hinge_clockwise_constraint_degs()
                        acw_constraint_degs = this_bone_joint.get_hinge_anticlockwise_constraint_degs()
                        if not Util.Utils().approximately_equal(cw_constraint_degs,
                                                        -this_bone_joint.get_MAX_CONSTRAINT_ANGLE_DEGS(),
                                                        0.001) and not Util.Utils().approximately_equal(acw_constraint_degs,
                                                                                this_bone_joint.get_MAX_CONSTRAINT_ANGLE_DEGS(),
                                                                                0.001):
                            relative_hinge_reference_axis = Util.Utils().normalization(
                                Util.Utils().times(m, this_bone_joint.get_reference_axis()))

                            signed_angle_degs = Util.Utils().get_signed_angle_between_degs(relative_hinge_reference_axis,
                                                                                this_bone_inner_to_outer_uv,
                                                                                relative_hinge_rotation_axis)
                            self.deg[loop] = signed_angle_degs*math.pi/180
                            if signed_angle_degs > acw_constraint_degs:
                                this_bone_inner_to_outer_uv = Util.Utils().normalization(
                                    Mat.Mat().rotate_about_axis(relative_hinge_reference_axis, acw_constraint_degs,
                                                          relative_hinge_rotation_axis))
                                self.deg[loop] = acw_constraint_degs*math.pi/180
                            elif signed_angle_degs < cw_constraint_degs:
                                this_bone_inner_to_outer_uv = Util.Utils().normalization(
                                    Mat.Mat().rotate_about_axis(relative_hinge_reference_axis, cw_constraint_degs,
                                                          relative_hinge_rotation_axis))
                                self.deg[loop] = cw_constraint_degs*math.pi/180
                        # twisted = np.cross(this_bone_inner_to_outer_uv,relative_hinge_rotation_axis)
                        # print("bone"+str(loop))
                        # print(twisted)

                scale = [i * this_bone_length for i in this_bone_inner_to_outer_uv]
                start_location = this_bone.get_start_point_position()
                new_end_location = [x + y for x, y in zip(start_location, scale)]
                this_bone.set_end_point_position(new_end_location)
                if loop < self.chain_length - 1:
                    self.get_bone(loop + 1).set_start_point_position(new_end_location)

            #  If we ARE working on the basebone...
            else:
                self.get_bone(0).set_start_point_position(self.fixed_base_location)
                if self.is_base_bone_fixed == 1:
                    self.get_bone(0).set_end_point_position(self.fixed_base_location_2)
                    if self.chain_length > 1:
                        self.get_bone(1).set_start_point_position(self.fixed_base_location_2)
                else:
                    this_bone_joint = this_bone.get_joint()
                    this_bone_joint_type = this_bone_joint.get_joint_type()
                    if this_bone_joint_type == "GLOBAL_HINGE":
                        hinge_rotation_axis = this_bone_joint.get_hinge_rotation_axis()
                        cw_constraint_degs = -this_bone_joint.get_hinge_clockwise_constraint_degs()
                        acw_constraint_degs = this_bone_joint.get_hinge_anticlockwise_constraint_degs()
                        this_bone_inner_to_outer_uv = Util.Utils().project_on_to_plane(this_bone.get_direction_uv(),
                                                                               hinge_rotation_axis)
                        # If we have a global hinge which is not freely rotating then we must constrain about the reference axis
                        if not Util.Utils().approximately_equal(cw_constraint_degs,
                                                        -this_bone_joint.get_MAX_CONSTRAINT_ANGLE_DEGS(),
                                                        0.001) and not Util.Utils().approximately_equal(acw_constraint_degs,
                                                                                this_bone_joint.get_MAX_CONSTRAINT_ANGLE_DEGS(),
                                                                                0.001):
                            hinge_reference_axis = this_bone_joint.get_reference_axis()
                            signed_angle_degs = Util.Utils().get_signed_angle_between_degs(hinge_reference_axis,
                                                                                this_bone_inner_to_outer_uv,
                                                                                hinge_rotation_axis)
                            self.deg[loop] = signed_angle_degs * math.pi / 180
                            if signed_angle_degs > acw_constraint_degs:
                                this_bone_inner_to_outer_uv = Util.Utils().normalization(
                                    Mat.Mat().rotate_about_axis(hinge_reference_axis, acw_constraint_degs,
                                                          hinge_rotation_axis))
                                self.deg[loop] = acw_constraint_degs * math.pi / 180
                            elif signed_angle_degs < cw_constraint_degs:
                                this_bone_inner_to_outer_uv = Util.Utils().normalization(
                                    Mat.Mat().rotate_about_axis(hinge_reference_axis, cw_constraint_degs,
                                                          hinge_rotation_axis))
                                self.deg[loop] = cw_constraint_degs * math.pi / 180

                        # twisted = np.cross(this_bone_inner_to_outer_uv, hinge_rotation_axis)
                        # print("bone" + str(loop))
                        # print(twisted)
                        scale = [i * this_bone_length for i in this_bone_inner_to_outer_uv]
                        start_location = this_bone.get_start_point_position()
                        new_end_location = [x + y for x, y in zip(start_location, scale)]
                        this_bone.set_end_point_position(new_end_location)
                        if self.chain_length > 1:
                            self.get_bone(1).set_start_point_position(new_end_location)
                    if this_bone_joint_type=="BALL":
                        this_bone_inner_to_outer_uv = this_bone.get_direction_uv()
                        angle_between_degs = Util.Utils().get_angle_between_degs(self.base_bone_constraint_uv,this_bone_inner_to_outer_uv)
                        constraint_angle_degs = this_bone.get_ball_joint_constraint_degs()
                        self.deg[loop] = angle_between_degs * math.pi / 180
                        if angle_between_degs>constraint_angle_degs:
                            this_bone_inner_to_outer_uv = Util.Utils().get_angle_limited_uv(this_bone_inner_to_outer_uv,self.base_bone_constraint_uv,
                                                                                constraint_angle_degs)
                            self.deg[loop] = constraint_angle_degs * math.pi / 180

                        scale = [i * this_bone_length for i in this_bone_inner_to_outer_uv]
                        start_location = this_bone.get_start_point_position()
                        new_end_location = [x + y for x, y in zip(start_location, scale)]
                        this_bone.set_end_point_position(new_end_location)
                        if self.chain_length > 1:
                            self.get_bone(1).set_start_point_position(new_end_location)
                    else:
                        this_bone_inner_to_outer_uv = this_bone.get_direction_uv()
                        scale = [i * this_bone_length for i in this_bone_inner_to_outer_uv]
                        start_location = this_bone.get_start_point_position()
                        new_end_location = [x + y for x, y in zip(start_location, scale)]
                        this_bone.set_end_point_position(new_end_location)
                        if self.chain_length > 1:
                            self.get_bone(1).set_start_point_position(new_end_location)
        # self.draw_chain()

    def solve_fabrik_ik(self):
        dist_base_to_target = Util.Utils().get_distance_between(self.get_bone(0).get_start_point_position(),
                                                        self.target_position)
        total_chain_length = 0
        for i in range(0, self.chain_length):
            total_chain_length += self.get_bone(i).get_length()
        # print("The initial joint angles and diagram. close the figure to see the final configuration\n")
        # self.draw_chain()
        if dist_base_to_target <= total_chain_length:
            if self.get_chain_length == 0:
                raise Exception("It makes no sense to solve an IK chain with zero bones.")

            dist_to_target = Util.Utils().get_distance_between(self.get_bone(self.chain_length - 1).get_end_point_position(),
                                                       self.target_position)

            counter = 0
            while dist_to_target > self.get_solve_distance_threshold() and counter < 10000:
                self.forward(self.target_position,self.target_orientation)
                self.backward()

                dist_to_target = Util.Utils().get_distance_between(
                    self.get_bone(self.get_chain_length() - 1).get_end_point_position(),
                    self.target_position)
                counter += 1
                # self.draw_chain()

            # after finding these joint position we can do anything with them.
            # Here I calculate the joints_angle:
            # print("Final joint angles:\n")
            self.draw_chain()

            # self.output_joint_angles()
        else:
            print("Target is so far! can't be reached")
            return

    def angles(self):
        angles =[]
        # for base bone twist
        base_bone = self.get_bone(0)
        base_bone_orientation = base_bone.get_bone_orientation()
        # number 3 belongs to the rotation matrix which is 0 for bone 3 rotation, 1 for bone 5 rotation,
        # 2 for bone 7 rotation, 3 for bone 0 rotation
        self.solve_for_orientation(base_bone_orientation,self.fixed_base_orientation, 3)
        angles.append(self.rotations[3])
        # for bone 2(in schematic)
        angles.append(self.deg[0])
        angles.append(self.rotations[0])
        # for bone 3(in schematic)
        angles.append(self.deg[1])
        angles.append(self.rotations[1])
        # for bone 4(in schematic)
        angles.append(self.deg[2])
        angles.append(self.rotations[2])

        s = ','.join([str(n) for n in angles])
        print(s)

    def draw_chain(self):
        self.angles()
        coordinate = self.fill_array(self.points_retrieval())
        x_prime = coordinate[0]
        y_prime = coordinate[1]
        z_prime = coordinate[2]

        f = OutputWriter(self.base_address, self.joints_file_address).joint_writer()
        for i in range(0, len(x_prime)):
            f.write(str(x_prime[i]))
            f.write(' ')
            f.write(str(y_prime[i]))
            f.write(' ')
            f.write(str(z_prime[i]))

            f.write("\n")
        f.close()
        fig = plt.figure()
        # ax = fig.gca(projection='3d')
        ax = Axes3D(fig)
        ax.plot3D(x_prime, y_prime, z_prime, color='red')
        ax.scatter3D(self.target_position[0], self.target_position[1], self.target_position[2])
        ax.scatter3D(x_prime, y_prime, z_prime)

        # for i in range(len(x_prime) - 1):
        #     print('length bone ' + str(i + 1))
        #     x = x_prime[i] - x_prime[i + 1]
        #     y = y_prime[i] - y_prime[i + 1]
        #     z = z_prime[i] - z_prime[i + 1]
        #     print("%.2f" % round(math.sqrt(x * x + y * y + z * z), 2))
        # plt.show()

    def set_axes_radius(self, ax, origin, radius):
        ax.set_xlim3d([origin[0] - radius, origin[0] + radius])
        ax.set_ylim3d([origin[1] - radius, origin[1] + radius])
        ax.set_zlim3d([origin[2] - radius, origin[2] + radius])

    def set_axes_equal(self, ax):
        limits = np.array([
            ax.get_xlim3d(),
            ax.get_ylim3d(),
            ax.get_zlim3d(),
        ])
        origin = np.mean(limits, axis=1)
        radius = 0.5 * np.max(np.abs(limits[:, 1] - limits[:, 0]))
        self.set_axes_radius(ax, origin, radius)

    def fill_array(self, start_locations):
        w, h = 3, len(start_locations)
        coordinate = [[0 for x in range(w)] for y in range(h)]
        # coordinate =[][len(body_part_index)]
        x = [0]
        y = [0]
        z = [0]
        length = [0.316, 0.088,0.088]
        for i in range(len(start_locations)):
            if i>=1:
                uv = [(start_locations[i][0]-start_locations[i-1][0]),
                      (start_locations[i][1]-start_locations[i-1][1]),
                      (start_locations[i][2]-start_locations[i-1][2])]
                uv = Mat.Mat().normalization(uv)
                scale = np.dot(uv, length[i-1])
                middle_points = [a + b for a, b in zip(start_locations[i-1], scale)]
                x.append(middle_points[0])
                y.append(middle_points[1])
                z.append(middle_points[2])

            x.append(start_locations[i][0])
            y.append(start_locations[i][1])
            z.append(start_locations[i][2])

        coordinate[0][:] = x
        coordinate[1][:] = y
        coordinate[2][:] = z
        return coordinate

    def points_retrieval(self):
        start_locations = []
        end_locations = []
        for i in range(0, self.get_chain_length()):
            start_loc = self.get_bone(i).get_start_point_position()
            # end_loc = self.chain.get_bone(i).set_end_point()
            start_locations.append(start_loc)
            # end_locations.append(end_loc)
        end_effector_bone = self.chain[self.get_chain_length() - 1].end_point
        start_locations.append(end_effector_bone)
        return start_locations

max_stars_repo_name: Atiehmerikh/PaintingWithFabrik

max_issues_repo_name: Atiehmerikh/PaintingWithFabrik

max_forks_repo_name: Atiehmerikh/PaintingWithFabrik


text:
from math import sqrt

import numpy as np
from sklearn.metrics import mean_squared_error, precision_recall_curve, auc

def rmse(y_actual: float, y_pred: float) -> float:
    """Root mean squared error

    Args:
        y_actual: Array of actual y values
        y_pred: Array of predicted y values

    Returns:
        Root mean squared error
    """
    return sqrt(mean_squared_error(y_actual, y_pred))

def mape(y_actual: float, y_pred: float) -> float:
    """Mean average precision error

    Args:
        y_actual: Array of actual y values
        y_pred: Array of predicted y values

    Returns:
        Mean average precision error
    """
    mask = y_actual != 0
    return (np.fabs(y_actual - y_pred) / y_actual)[mask].mean()

def pr_auc_score(y_true: np.ndarray, y_score: np.ndarray):
    """
    Area under Curve for Precision Recall Curve
    Args:
        y_true: Array of actual y values
        y_score: Array of predicted probability for all y values

    Returns:
        Area under Curve for Precision Recall Curve
    """
    assert y_true.shape == y_score.shape
    precision, recall, _ = precision_recall_curve(y_true, y_score)
    return auc(recall, precision)

ext: py

max_stars_repo_name: AdityaSidharta/bamboos

max_issues_repo_name: AdityaSidharta/bamboos

max_forks_repo_name: AdityaSidharta/bamboos

reason: import numpy


text:
# Python code to demonstrate
# converting string to json
# using json.loads
import json

##### inititialising json object
ini_string = {'nikhil': 1, 'akash' : 5,
              'manjeet' : 10, 'akshat' : 15}
####
##### printing initial json
####ini_string = json.dumps(ini_string)
####print ("initial 1st dictionary", ini_string)
####print ("type of ini_object", type(ini_string))
####
##### converting string to json
####final_dictionary = json.loads(ini_string)
####
##### printing final result
####print ("final dictionary", str(final_dictionary))
####print ("type of final_dictionary", type(final_dictionary))

ini_string =[{IDUser:1,IDCarts:1,NameProduct:Laptop inspiron 15,Producer:DELL,Price:1000,Quantity:3},{IDUser:1,IDCarts:2,NameProduct:P30 lite,Producer:HUAWEI,Price:200,Quantity:2},{IDUser:1,IDCarts:1,NameProduct:Camara,Producer:Canon,Price:300.0,Quantity:1},{IDUser:1,IDCarts:1,NameProduct:Camara,Producer:Canon,Price:300.0,Quantity:1}]

print(ini_string)

####from MySQLEngine import *
####import configparser
####SQLEngine = MySQLEngine()
####SQLEngine.start()
####import json

#import time

#app = Flask(__name__)
#app.secret_key = 'my_secret_key'
#csrf = CsrfProtect(app)

from flask import Flask, render_template, make_response
from flask import redirect, request, jsonify, url_for

import io
import os
import uuid

import numpy as np

import json
import sys

def contraseñaExi(p):
    dict = {"val1":"this is x", "val2":True}
    return json.dumps('usuario exitoso')
def contraseñaFall(p):
    dict = {"val1":"this is x", "val2":True}
    return json.dumps('no encontado')
def ListProductP():
    dict = [{
        "Producto": "Camara",
        "Fabricante": "Cannon",
        "color": ["white", "red", "Black"],
        "Precio": 5500},{
        "Producto": "Camara",
        "Fabricante": "Cannon",
        "color": ["white", "red", "Black"],
        "Precio": 5500}]
    return json.dumps(dict)
def DetailProduct():
    #sys.argv[2] es el IDProducto
    dict = {
        "Producto": "Camara",
        "Fabricante": "Cannon",
        "color": ["white", "red", "Black"],
        "Precio": 5500}
    return json.dumps(dict)
def Cart(IdCl):
    IdCliente = IdCl
    return json.dumps(IdCliente)
def invoices_list(Invo):
    invoice = Invo
    return json.dumps(invoice)
def Packages_list(Packa):
    Packages = Packa
    return json.dumps("Packages")
def PayPackageProduct(Packa):
    Packages = Packa
    return json.dumps("pago paquetes Packages")
def UpProductNew(Name, Producer, Price, Discount, Quantity, Category, Description, Specification):
    NameProducts = Name
    ProducerProducts = Producer
    Phy_Price = Price
    Phy_Discount = Discount
    Phy_Quantity = Quantity
    Phy_Category = Category
    Phy_Description = Description
    Phy_Specification = Specification
    return json.dumps("Subir producto nuevo")
def ShoresBuyList(IdCliente,IdShores):
    print("LLenado de compras ")
def ShoresStocktakingList(IdCliente,IdShores):
    print("LLenado de inventario ")

if sys.argv[1] == 'Login':
    print( contraseñaFall(sys.argv[2]))

#print (len(sys.argv),sys.argv)

if sys.argv[1] == 'ListProd':
    print(ListProductP())
if sys.argv[1] == 'DetailProduct':
    print(DetailProduct())
if sys.argv[1] == 'Cart':
    print(Cart(sys.argv[2]))
if sys.argv[1] == 'invoices':
    print(invoices_list(sys.argv[2]))
if sys.argv[1] == 'Packages':
    print(Packages_list(sys.argv[2]))
if sys.argv[1] == 'PayPackageProducts':
    print(PayPackageProduct(sys.argv[2]))
if sys.argv[1] == 'UpProduct':
    print( UpProductNew(sys.argv[2],sys.argv[3],sys.argv[4],sys.argv[5],sys.argv[6],sys.argv[7],sys.argv[8],sys.argv[9]))
if sys.argv[1] == 'ShoresBuyList':
    print( ShoresBuyList(sys.argv[2],sys.argv[3]))
if sys.argv[1] == 'ShoresStocktakingList':
    print( ShoresStocktakingList(sys.argv[2],sys.argv[3]))

#INSERT INTO Customer (Name_customer) VALUES ('1llllllllllllllll'),('1-2-afc'),('2-asdasdasdasd2-afc');


text:
from numpy import ndindex, savetxt

def save_array(
    file, data, fmt="%7.2f", delimiter=",", header="data", slice="slice", sep=":"
):
    """Function to save numpy nD arrays. Therefore the array is sliced except for
    the last 2 dimenions.

    Parameters
    ----------

    file: str
        the save filename

    header: str
        some file header string

    sep: str
        seperator that delimits the shape information in the header

    delimiter: str
        data delimiter

    fmt: str
        string to define the output number format that is passed to numpy.savetext

    Returns
    -------

    None
    """
    # Write the array to disk
    with open(file, "w") as outfile:
        # writing a header to get the shape while loading
        outfile.write(f"#{header}{sep}{data.shape}\n")

        # iterating through ndarray except and write slices of the last 2 dims
        if len(data.shape) > 2:
            d = len(data.shape) - 2
            for i in ndindex(data.shape[:d]):
                # writing a break to indicate different slices...
                outfile.write(f"#{slice}{sep}{i}\n")
                savetxt(outfile, data[i], delimiter=delimiter, fmt=fmt)
        else:
            savetxt(outfile, data, delimiter=delimiter, fmt=fmt)

max_issues_repo_name: ecs-kev/pyleecan

max_forks_repo_name: ecs-kev/pyleecan

reason: from numpy


text:
import json
import os
import torch
import shutil
import numpy as np
from torch.utils.data import Dataset, DataLoader

def my_collate(batch):
    '''
        Custom collate function for the dataloader in
        order to handle the case where the adjacency
        matrices have different sizes (no padding).
    '''
    enc_sent = [item[0] for item in batch]
    atte_mask = [item[1] for item in batch]
    ne_tags = [item[2] for item in batch]

    return [enc_sent, atte_mask, ne_tags]

class ADE_Dataset(Dataset):
    def __init__(self, files, path_in_general):
        super().__init__()
        '''
            Args:
                files (list): list with paths of the json files
        '''
        self.files = files
        self.path_in_general = path_in_general

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        # Load json file
        with open(self.path_in_general + self.files[idx]) as json_file:
            data = json.load(json_file)

        # Load the padded encoded sentence and the attention mask
        # which are the inputs of the text encoder.
        enc_sent = torch.LongTensor(data['padded encoded sentence'])
        atte_mask = torch.LongTensor(data['attention mask'])

        # Load the relation pairs, the NE tags and the word pieces
        # which will be used to apply CL in the sentence level.
        ne_tags = data['ne tags']

        return enc_sent, atte_mask, ne_tags


text:
import numpy as np
from django.contrib.auth import login, authenticate
from django.contrib.auth.models import User
from django.http import JsonResponse
from django.shortcuts import render
from django.views.decorators.csrf import csrf_exempt
from rest_framework.authtoken.models import Token

from Account.forms import SignUpForm
from Api.forms import MessageForm
from Api.models import Message
from Data.forms import UserDataForm
from Data.models import UserData, Version, Profile

def data_delete(request):
    # POST request should contain list of all the "data_id"s to be deleted separated by commas

    if request.method == 'POST':
        item_id = request.POST.get('data_id')

        try:
            item = UserData.objects.get(id=item_id)
        except:
            return JsonResponse({"success": False, "error": "Invalid id"})

        # There must be a simple verification so check if the user who sent the request is the owner of the data
        if request.user == item.user:
            item.delete()
            return JsonResponse({"success": True})

        return JsonResponse({"success": False, "error": "Not authorized"})

    else:
        return render(request, '404.html', status=404)

def reset_token(request):
    # POST request should contain

    if request.method == 'POST':
        token = request.POST.get('token')

        try:
            token = Token.objects.get(key=token)
        except:
            return JsonResponse({'success': False})

        # There must be a simple verification so check if the user who sent the request is the owner of the data
        if request.user == token.user:
            token.delete()
            token = Token.objects.create(user=request.user)

            return JsonResponse({'token': token.key, 'success': True})

        return JsonResponse({'success': False})

    else:
        return render(request, '404.html', status=404)

@csrf_exempt
def validate_token(request):
    """
    View to allow client to login at the startup of the desktop-app

    CLient should send POST request as:
    headers={"Authorization": f"Token {token}"}, data={'username': f'{username}'}
    """
    if request.method == 'POST':

        if 'HTTP_AUTHORIZATION' in request.META:
            try:
                key = request.META['HTTP_AUTHORIZATION'].split(" ")[1]

                token = Token.objects.get(key=key)

                return JsonResponse({"valid-token": True, 'username': str(token.user)})

            except:
                return JsonResponse({"valid-token": False})

        else:
            return JsonResponse({"valid-token": False, 'error': 'No auth token or username in POST request'})

    else:
        return render(request, '404.html', status=404)

@csrf_exempt
def api_create_account(request):
    if request.method == 'POST':

        form = SignUpForm(request.POST)

        if form.is_valid():

            form.save()
            username = form.cleaned_data.get('username')
            raw_password = form.cleaned_data.get('password1')
            user = authenticate(username=username, password=raw_password)

            login(request, user)
            token = Token.objects.get(user=user).key

            return JsonResponse({'success': True, 'token': token})

        else:
            return JsonResponse({'success': False, 'errors': form.errors.as_json()})

    else:
        return render(request, '404.html', status=404)

@csrf_exempt
def score_check(request):
    user_total_score = 0
    if request.method == 'POST':

        if 'username' in request.POST and 'HTTP_AUTHORIZATION' in request.META:
            try:
                username = request.POST['username']
                key = request.META['HTTP_AUTHORIZATION'].split(" ")[1]

                user = User.objects.get(username=username)
                token = Token.objects.get(key=key)

                if token.user == user:
                    user_total_score = Profile.objects.get(user=user).score

                else:
                    user_total_score = 0

            except:
                user_total_score = 0

    return JsonResponse({'score': user_total_score})

def version_control(request):
    version_control = [obj.as_dict() for obj in Version.objects.all()]
    return JsonResponse({"Version Control": version_control})

@csrf_exempt
def bug_report(request):
    if request.method == 'POST':

        # Gathering all info sent by the client
        try:
            key = request.META['HTTP_AUTHORIZATION'].split(" ")[1]
            username = request.POST['user']
            msg = request.POST['message']
            version = float(request.POST['version'])

        except:
            return JsonResponse({"success": False, "error": "Wrong formatting or missing information"})

        # Validating info
        try:

            user = User.objects.get(username__exact=username)
            version = Version.objects.get(version__exact=version)
            token = Token.objects.get(key=key)

        except:
            return JsonResponse({"success": False, "error": "User or version not found on database"})

        # Validanting form
        form = MessageForm({
            'message': msg,
            'user': user.id,
            'version': version.id,
        })

        # Checking authorization and form validity
        if form.is_valid() and token.user == user:

            message = Message(message=msg, user=user, version=version, read=False)
            message.save()

            return JsonResponse({"success": True})

        else:
            return JsonResponse({'success': False, 'error': 'Wrong formatting or missing information'})

    else:
        return render(request, '404.html', status=404)

@csrf_exempt
def data_upload(request):
    if request.method == 'POST':
        if request.FILES:

            try:
                username = request.POST['username']
                key = request.META['HTTP_AUTHORIZATION'].split(" ")[1]

                user = User.objects.get(username=username)
                token = Token.objects.get(key=key)

                if token.user == user:

                    form = UserDataForm(request.POST, request.FILES)

                    file = request.FILES['file']
                    version = request.POST['version']

                    version = Version.objects.get(version__exact=version)

                    if form.is_valid() and (file.name == 'training_data.npy'):  # and (file.size > 140000):

                        score = np.load(file).shape[0] # len(list(np.load(file))) ---- .shape[0] is faster but untested!
                        file = UserData(file=request.FILES['file'], user=user, score=score, processed=False,
                                        version=version)
                        file.save()

                        return JsonResponse({'success': True})

                    else:
                        return JsonResponse({
                            'success': False,
                            'valid form': form.is_valid(),
                            'file name': file.name,
                            'file size': file.size
                        })
                else:
                    return JsonResponse({"success": False, "error": "Not authorized"})

            except Exception as e:
                return JsonResponse({"success": False, "error": f"{e}"})

    else:
        return render(request, '404.html', status=404)

# def download(request):
#     file_path = 'Bot-Latest-Version/Globals.py'
#     if os.path.exists(file_path):
#         with open(file_path, 'rb') as fh:
#             response = HttpResponse(fh.read(), content_type="application/x-python")
#             response['Content-Disposition'] = 'inline; filename=' + os.path.basename(file_path)
#             return response


text:
import os
import cv2
import numpy as np
import matplotlib.pyplot as plt

class Detections:
    """
    Detection and recognition class for inference results
    Keyword Arguments:
        :param img: {Image.Image}
        :param comps: {tuple} -> (predicted_names, cropped_faces)
        :param outputs: {tuple} -> (locations, probs, landmarks if landmarks True)
        :param landmarks: {bool} -> True or False
    """
    def __init__(self, img, comps, outputs, landmarks):
        self.img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
        self.comps = comps
        self.outputs = outputs
        self.landmarks = landmarks

    def display(self, pprint=False, loc=False, crop=False, show=False, plot=False, save=False, save_dir=''):
        if self.comps is not None:
            names, crops = self.comps
            if pprint:
                # showing names and probs of detected face
                print('Names:{}, Probs:{}'.format(names, self.outputs[1]))

            if loc:
                # location of detected faces
                return dict(Boxes=self.outputs[0], landMarks=self.outputs[2] if self.landmarks else None)

            if crop:
                # cropped_face image array and save it if u want
                if save:
                    os.makedirs(f'{save_dir}/crops', exist_ok=True)
                    for i in range(len(crops)):
                        file = f'{save_dir}/crops/{names[i]}.jpg'
                        cv2.imwrite(file, self.to_numpy(crops[i]) * 255)
                        print('Cropped Faces saved!')
                return crops

            if show:
                # draw bounding box names and aligned point on the face
                for i, box in enumerate(self.outputs[0]):
                    x1, y1, x2, y2 = list(map(lambda x: int(x), box))
                    scale = round(((x2 - x1) + 78) / 75)
                    (w, h), _ = cv2.getTextSize(names[i], cv2.FONT_HERSHEY_PLAIN, scale, 2)
                    cv2.rectangle(self.img, (x1, y2 + h + 2), (x1 + w, y2), (255, 0, 0), -1)
                    cv2.rectangle(self.img, (x1, y1), (x2, y2), (255, 0, 0), thickness=5)
                    cv2.putText(self.img, names[i], (x1, y2 + h), cv2.FONT_HERSHEY_PLAIN, scale, (255, 255, 255), 2)

                    if self.landmarks:
                        for point in self.outputs[2][i]:
                            x, y = int(point[0]), int(point[1])
                            cv2.circle(self.img, (x, y), 5, (0, 0, 255), -1)
                return self.img

            if plot:
                # plot cropped faces
                faces = self.cropped_face()
                plt.figure(figsize=(15, 15))
                for i in range(len(faces)):
                    plt.subplot(3, 3, i + 1)
                    plt.imshow(faces[i])
                    plt.axis('off')
                    plt.title(names[i])
                plt.show()

            if save:
                # saving draw image
                file = save_dir + '/saved_pred.jpg'
                os.makedirs(save_dir, exist_ok=True)
                cv2.imwrite(file, self.show())
                print('Image Saved!')

        else:
            if pprint or loc or crop or plot or save:
                print('No Face Detected!')
            else:
                return self.img

    def print(self):
        self.display(pprint=True)

    def locations(self):
        return self.display(loc=True)

    def cropped_face(self, save=False, save_dir='results'):
        return self.display(crop=True, save=save, save_dir=save_dir)

    def show(self):
        return self.display(show=True)

    def plot(self):
        return self.display(plot=True)

    def save(self, save_dir='results'):
        self.display(save=True, save_dir=save_dir)

    @staticmethod
    def to_numpy(img):
        return img.numpy().transpose(1, 2, 0)


text:
import numpy as np
from scipy.sparse.csgraph import minimum_spanning_tree
from .nns import k_nearest_neighbors
from .operator import gather, square_distance

# Note:
# (1) code is same code below:
# normals = np.zeros((len(v), 3), dtype=np.float32) # (N, 3)
# for i, idx in enumerate(w_min_idxs):
#     normals[i] = v[i, :, idx]
def normal_estimation(coords:np.ndarray, k:int=10) -> np.ndarray:
    """Estimate normals each point with eigenvector of covariance matrix.
    This function use k nearest neighbor (kNN) to get covariance matrixs.

    Args:
        coords: coordinates of points, (N, 3)
        k: number of neighbor for kNN

    Returns
        normals: estimated normals (N, 3)
    """
    # Get neighbor points. (TODO: add radius)
    idxs, _  = k_nearest_neighbors(coords, coords, k)
    knn_points = gather(coords, idxs)

    # Get covariance matrix of each point.
    knn_mean_points = np.mean(knn_points, axis=-2, keepdims=True)
    deviation = knn_points - knn_mean_points
    deviation_t = deviation.transpose(0, 2, 1)
    covariance_matrixs = np.matmul(deviation_t, deviation) / k # (N, 3, 3)

    # Get eigenvector and eigenvalue of each point
    w, v = np.linalg.eig(covariance_matrixs.transpose(0, 2, 1))
    # w, v = np.linalg.eig(covariance_matrixs)
    w_min_idxs = np.argmin(w, axis=1)

    # Get normal of each point (1)
    normals = np.take_along_axis(v, w_min_idxs[:, np.newaxis, np.newaxis], axis=2)
    normals = normals.squeeze(2) # (N, normal_vector, 1) -> (N, normal_vector)

    return normals

def normal_orientation(coords:np.ndarray, normals:np.ndarray):
    """Normal orientation.

    Args:
        coords: (N, 3)
        normals: (N, 3)

    Returns:
        normals: (N, 3)
    """
    N, _ = normals.shape

    def cosine_similarity(vector1:np.ndarray, vector2:np.ndarray):
        """Cosine similarity

        Args:
            vector1: (C)
            vector2: (C)

        Return:
            similarity: value range (-1 ~ 1)

        Note:
            C: number of vector dimension
        """
        return np.dot(vector1, vector2)/(np.linalg.norm(vector1)*np.linalg.norm(vector2))

    def tree_recursion(_normals:np.ndarray, adjacency_matrix:np.ndarray, idx:int, rest_idxs:np.ndarray=None):
        """tree recursion for normal orientation

        Args:
            _normals: (N, 3)
            adjacency_matrix: adjacency matrix between points, (N, N)
            idx: first index for tree (first node)
            rest_idxs: rest node indexs for recursion processing, (N)

        Returns:
            oriented_normals: fixed normals(N, 3)
        """
        if rest_idxs is None:
            rest_idxs = np.full(N, fill_value=True,dtype=np.bool8)
        mask_row = adjacency_matrix[idx] > 0
        mask_col = adjacency_matrix[:, idx] > 0
        mask = mask_row + mask_col
        connection = np.arange(N)[rest_idxs & mask]

        for c_idx in connection:
            rest_idxs[c_idx] = False
            parent = _normals[idx]
            child = _normals[c_idx]

            cos_sim = cosine_similarity(parent, child)
            cos_sim_minus = cosine_similarity(parent, -child)

            if cos_sim_minus > cos_sim:
                _normals[c_idx] = child * -1

            _normals = tree_recursion(_normals, adjacency_matrix, c_idx, rest_idxs)

        return _normals

    _normals = np.copy(normals)
    weighted_adjacency_matrix = square_distance(coords, coords)
    weighted_adjacency_matrix[np.identity(N, dtype=np.bool8)] = 0.0
    tcsr = minimum_spanning_tree(weighted_adjacency_matrix)
    tree_weighted_adjacency_matrix =  tcsr.toarray()

    _normals = tree_recursion(_normals, tree_weighted_adjacency_matrix, 0)

    return _normals

reason: import numpy,from scipy


text:
# -*- coding: utf-8 -*-
"""
Created on Mon Sep 14 17:41:37 2020

@author: Lucas & Caro
"""

import streamlit as st
import base64

from pathlib import Path
from sklearn.metrics import roc_curve, precision_recall_curve, average_precision_score
from sklearn.model_selection import train_test_split
from rdkit.ML.Scoring import Scoring
from sklearn.metrics import confusion_matrix
import math
import pandas as pd
import statistics
# import plotly as py
from plotly.subplots import make_subplots
import plotly.graph_objects as go
import numpy as np

from sklearn.metrics import accuracy_score
from sklearn.metrics import balanced_accuracy_score
from sklearn.metrics import f1_score
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score
from sklearn.metrics import matthews_corrcoef
#---------------------------------#
# Page layout
## Page expands to full width
st.set_page_config(page_title='LIDEB Tools - Metrics', page_icon="📏", layout='wide')

######
# Function to put a picture as header
def img_to_bytes(img_path):
    img_bytes = Path(img_path).read_bytes()
    encoded = base64.b64encode(img_bytes).decode()
    return encoded

from PIL import Image
image = Image.open('cropped-header.png')
st.image(image)
#####
st.write("&nbsp[![Website](https://img.shields.io/badge/website-LIDeB-blue)](https://lideb.biol.unlp.edu.ar)&nbsp[![Twitter Follow](https://img.shields.io/twitter/follow/LIDeB_UNLP?style=social)](https://twitter.com/intent/follow?screen_name=LIDeB_UNLP)")
st.markdown("###### Made in 🐍 and [![this is an image link](https://i.imgur.com/iIOA6kU.png)](https://www.streamlit.io/)&nbsp with :heart: by [Lucas Alberca](https://twitter.com/capigol) and [Caro Bellera](https://twitter.com/carobellera)")
#####
st.subheader(":pushpin:" "About Us")
st.markdown("We are a drug discovery team with an interest in the development of publicly available open-source customizable cheminformatics tools to be used in computer-assisted drug discovery. We belong to the Laboratory of Bioactive Research and Development (LIDeB) of the National University of La Plata (UNLP), Argentina. Our research group is focused on computer-guided drug repurposing and rational discovery of new drug candidates to treat epilepsy and neglected tropical diseases.")

# Introduction
#---------------------------------#

st.write("""
# LIDeB Tools - Metrics
**Metrics to evaluate the performance of the classificatory models **

Web App to evaluate the perfomance of classificatory models by calculation of AUCROC, BEDROC, Accuracy, F-measure, PR and EF.

Additionaly, using the generated **PPV surfaces** you can select an adequate score threshold for prospective virtual screenings.

The tool uses the following packages [RDKIT](https://www.rdkit.org/docs/index.html), [Scikit-learn](https://scikit-learn.org/stable/), [Plotly](https://plotly.com/python/)

""")

# image = Image.open('clustering_workflow.png')
# st.image(image, caption='Clustering Workflow')

st.sidebar.header('Upload your File')

uploaded_file_1 = st.sidebar.file_uploader("Upload a TXT file with the scores", type=["txt"])
st.sidebar.markdown("""
[Example TXT input file](https://raw.githubusercontent.com/cbellera/metrics/main/example_file_metrics.txt)
""")

st.sidebar.header('SCORE THRESHOLD')
threshold_cm = st.sidebar.number_input('Threshold to split between actives and inactives', value=0.5)

st.sidebar.header('BEDROC')
alpha_bedroc = st.sidebar.slider('Alpha BEDROC', 5, 100, 20, 5)

st.sidebar.header('ENRICHMENT FACTOR')
fraccion_enriquesimiento = st.sidebar.slider('Fraction enrichment', 0.0, 0.95, 0.01 , 0.01)

st.sidebar.header('STANDARD DEVIATION')
repeticiones = st.sidebar.slider('Number of iterations', 100, 1000, 100 , 50)
particion = st.sidebar.slider('Fraction of evaluation', 0.5, 0.95, 0.85 , 0.05)

st.sidebar.title(":speech_balloon: Contact Us")
st.sidebar.info(
"""
If you are looking to contact us, please
[:e-mail:](mailto:lideb@biol.unlp.edu.ar) or [Twitter](https://twitter.com/LIDeB_UNLP)
""")

def metrics_calculation(uploaded_file_1):

    # Leemos el archivo cargado
    scores1 = uploaded_file_1
    # Lo ordenamos
    scores1.sort_values(by=["score"],axis = 0, ascending = False,inplace = True, na_position ='last')

    y_modelo = list(scores1["score"])
    y_real = list(scores1["class"])

    clases_predichas = []
    for value in y_modelo:
        if value > threshold_cm:
            clases_predichas.append(1)
        else:
            clases_predichas.append(0)

    clases_predichas_serie = pd.Series(clases_predichas)
    # result = scores1["class"].eq(clases_predichas_serie)
    # proporcion_bc = round(sum(result) / len(result),4) # Numero de compuestos bien clasificados / numero de compuestos del dataset

    st.markdown(":point_down: **Here you can see the perfomance of the model for the selected threshold**", unsafe_allow_html=True)

    metric1 = round(accuracy_score(y_real, clases_predichas_serie),4)
    metric2 = round(balanced_accuracy_score(y_real, clases_predichas_serie),4)
    metric4 = round(f1_score(y_real, clases_predichas_serie),4)
    metric5 = round(precision_score(y_real, clases_predichas_serie),4)
    metric6 = round(recall_score(y_real, clases_predichas_serie),4)
    metric7 = round(matthews_corrcoef(y_real, clases_predichas_serie),4)

    results = [metric1,metric2,metric4,metric5,metric6,metric7]
    metric_name = ["Accuracy","Balanced accuracy","F-measure","Precision","Recall", "MCC"]
    table_metrics = pd.Series(results,index=metric_name,name="Value")
    st.write(table_metrics)

    st.markdown(":point_down: **Here you can find the confusion matrix for the selected threshold**", unsafe_allow_html=True)
    matrix = pd.DataFrame(confusion_matrix(y_real, clases_predichas))
    # matrix=matrix.rename(index={0: 'Predicted Positive', 1:"Predicted Negative"},columns={0:"Real Positive",1:"Real Negative"})
    matrix=matrix.rename(index={0: 'Real Negative', 1:"Real Positive"},columns={0:"Predicted Negative",1:"Predicted Positive"})

    st.write(matrix)

    # Hacemos los calculos para X repeticiones
    i=0
    aucs=[]
    bedrocs=[]
    efs=[]
    prs=[]

    # aupr_max= 1

    while i <= repeticiones:    # Para hacer el cross validation y poder sacar la desviacion del auc

        scores_modelo, X_test, clases_real, y_test = train_test_split(y_modelo, y_real, test_size= 1 - particion, random_state=i,stratify = y_real)

        array = tuple(zip(scores_modelo,clases_real))
        array_1 = pd.DataFrame(array)
        array_1.sort_values(0, axis = 0, ascending = False,inplace = True, na_position ='last')
        array_metricas = tuple(zip(array_1[0],array_1[1]))

        auc = Scoring.CalcAUC(scores= array_metricas, col= 1)
        aucs.append(auc)
        bedroc = Scoring.CalcBEDROC(scores= array_metricas, col= 1,alpha=alpha_bedroc)
        bedrocs.append(bedroc)
        enrichment = Scoring.CalcEnrichment(scores= array_metricas, col= 1,fractions=[fraccion_enriquesimiento])
        efs.append(enrichment[0])
        auc_PR = average_precision_score(clases_real, scores_modelo)                        # Calculo el AUPR
        prs.append(auc_PR)

        i = i + 1

    # # Esto ese para el EFmax #
    n_total = len(y_real)           # TOT(tot)
    fraccion_ef = math.ceil(n_total*fraccion_enriquesimiento) # devuelve el numero entero mayor (TOT(1%))
    n_activos = sum(y_real)     # ACT(TOT)
    if fraccion_ef > n_activos:
        efmax = (n_activos/fraccion_ef)/(n_activos/n_total)
    else:
        efmax = 1/(n_activos/n_total)

    # Sacamos las medias y las desviaciones
    auc_ok= round(statistics.mean(aucs),4)
    sd_auc = round(statistics.stdev(aucs),4)

    bedroc_ok= round(statistics.mean(bedrocs),4)
    sd_bedroc = round(statistics.stdev(bedrocs), 4)

    efs1 = [x/efmax for x in efs]
    ef_ok= round(statistics.mean(efs1),4)
    if ef_ok > 1:
        ef_ok = 1.000
    sd_ef = round(statistics.stdev(efs1),4)

    pr_ok= round(statistics.mean(prs),4)
    sd_pr = round(statistics.stdev(prs),4)

    resultado = [auc_ok,bedroc_ok,ef_ok, pr_ok]
    desviaciones = [sd_auc,sd_bedroc, sd_ef, sd_pr]

    final = pd.DataFrame([resultado,desviaciones]).T

    final=final.rename(index={0: 'AUC', 1:"BEDROC_" + str(alpha_bedroc), 2: 'EFrel_' + str(fraccion_enriquesimiento), 3:'AUC_PR'},columns={0:"Mean",1:"SD"})

    st.markdown(":point_down: **Here you can see the calculated metrics**", unsafe_allow_html=True)
    st.write(final)
    # st.markdown(":point_down: **Here you can download this metrics**", unsafe_allow_html=True)
    st.markdown(filedownload_1(final), unsafe_allow_html=True)

    return scores1, auc_ok, sd_auc,pr_ok, sd_pr, y_modelo, y_real

def roc_aupr_curves(auc_ok, sd_auc, pr_ok, sd_pr , y_modelo, y_real):

    title1= "ROC CURVE " + " " + "(" + "AUCROC = " + str(format(auc_ok, '.4f')) + " ± " + str(format(sd_auc, '.4f')) + ")"
    title2="PR CURVE " + " " + "(" + "AUPR = " + str(format(pr_ok, '.4f')) + " ± " + str(format(sd_pr, '.4f')) + ")"
    fig = make_subplots(rows=1, cols=2, subplot_titles=(title1,title2))

    # # CURVA ROC
    fpr, tpr, thresholds = roc_curve(y_score= y_modelo,y_true= y_real)
    fig.add_trace(go.Scatter(x=fpr, y=tpr, fill='tozeroy', showlegend=False),row=1, col=1)
    fig.add_shape(type='line', line=dict(dash='dash'), x0=0, x1=1, y0=0, y1=1,row=1, col=1)
    fig.update_xaxes(title_text="False Positive Rate (FPR)", row=1, col=1)
    fig.update_yaxes(title_text="True Positive Rate (TPR)", row=1, col=1)

    # # CURVA PRECISION RECALL
    precision, recall, thresholds = precision_recall_curve(y_true= y_real, probas_pred = y_modelo)
    fig.add_trace(go.Scatter(x=recall, y=precision, fill='tozeroy', showlegend=False),row=1, col=2)
    fig.add_shape(type='line', line=dict(dash='dash'),x0=0, x1=1, y0=1, y1=0,row=1, col=2)

    fig.update_layout(margin = dict(t=60,r=20,b=20,l=20), autosize = True)
    fig.update_xaxes(title_text="Recall", row=1, col=2)
    fig.update_yaxes(title_text="Precision", row=1, col=2)
    st.plotly_chart(fig)

def ppv_plot(scores1):
# importamos los paquetes necesarios:
    y_modelo = list(scores1["score"])
    y_real = list(scores1["class"])

    # Determino la sensibilidad, especificidad y los valores de corte
    fpr, tpr, thresholds = roc_curve(y_score= y_modelo, y_true= y_real,drop_intermediate=False)
    se = tpr
    sp = 1 - fpr
    se_sp = se/sp

    tabla_final = pd.DataFrame(list(zip(thresholds,fpr,tpr,se,sp,se_sp)),
                   columns =['Cutoff', 'fpr',"tpr","Se","Sp","Se/Sp"])

    # Calculos para agregar el valor de PPV a dos Ya distintos a la tabla

    prev1 = [0.001,0.005,0.01]

    se1=tabla_final["Se"].tolist()
    sp1=tabla_final["Sp"].tolist()

    tamanio = list(range(1, len(se1), 1))

    ppv_ok1 =[]

    for prev in prev1:
        ppv1a=[]
        for i in tamanio:
            ppv1 = (se1[i]*prev)/(se1[i]*prev + (1 - sp1[i])*(1-prev))
            ppv1a.append(ppv1)
        ppv_ok1.append(ppv1a)

    # ppv_para_dos_ya = pd.DataFrame(ppv_ok1).T

    tabla_final_1 = pd.DataFrame(list(zip(thresholds,se,sp,se_sp,ppv_ok1[0],ppv_ok1[1],ppv_ok1[2])),
                    columns =['Cutoff', "Se","Sp","Se/Sp","PPV (Ya = 0.001)","PPV (Ya = 0.005)","PPV (Ya = 0.01)"])

    # preparacion de datos para el grafico de  ppv

    prevalencia = pd.DataFrame(np.arange(0.00000000001, 0.012, 0.001))
    ya=prevalencia[0].tolist()

    se1=tabla_final["Se"].tolist()
    sp1=tabla_final["Sp"].tolist()

    tamanio = list(range(1, len(se1), 1))

    ppv_ok =[]

    for x in ya:
        ppv1=[]
        for i in tamanio:
            ppv = (se1[i]*x)/(se1[i]*x + (1 - sp1[i])*(1-x))
            ppv1.append(ppv)
        ppv_ok.append(ppv1)

    ppv_tabla = pd.DataFrame(ppv_ok).T

    # GRAFICO PPV

    fig = go.Figure(go.Surface(opacity=1,
        contours = {
            "x": {"show": False, "start": 0, "end": 0.01, "size": 0.04, "color":"white"},
            "y": {"show": False, "start": 0, "end": 2, "size": 0.04, "color":"white"},
            "z": {"show": False, "start": 0, "end": 1, "size": 0.2,"color":"white"}
        },
        x = ya,
        y = se_sp,
        z = ppv_tabla))

    fig.update_layout(
          scene = dict(
            xaxis = dict(title="Ya", nticks=10, range=[0,0.011],gridwidth = 10,linewidth = 10,
                         backgroundcolor="rgb(200, 200, 230)", gridcolor="white", showbackground=True, zerolinecolor="white"),
            yaxis = dict(title="Se/Sp",nticks=10, range=[0,2],gridwidth = 10,linewidth = 10,
                         backgroundcolor="rgb(230, 200,230)",
                            gridcolor="white",
                            showbackground=True,
                            zerolinecolor="white"),
            zaxis = dict(title="PPV",nticks=10, range=[0,1],gridwidth = 10,linewidth = 10,
                         backgroundcolor="rgb(230, 230,200)",
                            gridcolor="white",
                            showbackground=True,
                            zerolinecolor="white"),),
        margin=dict(r=20, l=10, b=10, t=10),
        font=dict(
            size=14,
            color="black")
        )
    st.plotly_chart(fig)

    st.markdown(":point_down: **Here you can see the table with cutoffs, Se, Sp and PPVs**", unsafe_allow_html=True)
    st.write(tabla_final_1)
    return tabla_final_1

#%%
### Exporting files ###

def filedownload(df):
    csv = df.to_csv(index=False,header=True)
    b64 = base64.b64encode(csv.encode()).decode()  # strings <-> bytes conversions
    href = f'Download CSV File with the previous table'
    return href

def filedownload_1(df):
    csv = df.to_csv(index=True,header=True)
    b64 = base64.b64encode(csv.encode()).decode()  # strings <-> bytes conversions
    href = f'Download CSV File with the metrics'
    return href

if uploaded_file_1 is not None:
    scores = pd.read_csv(uploaded_file_1, sep='\t', delimiter=None, header='infer', names=None)

    scores1, auc_ok, sd_auc, pr_ok, sd_pr , y_modelo, y_real = metrics_calculation(scores)
    st.write("")
    st.markdown(":point_down: **Here you can see the ROC and AUPR curves**", unsafe_allow_html=True)
    roc_aupr_curves(auc_ok, sd_auc, pr_ok, sd_pr , y_modelo, y_real)

    st.markdown(":point_down: **Here you can see the PPV surface**", unsafe_allow_html=True)
    tabla_final_1 = ppv_plot(scores1)

    st.markdown(filedownload(tabla_final_1), unsafe_allow_html=True)

else:
    st.info('Awaiting for TXT file to be uploaded.')
    if st.button('Press to use Example File'):
        st.write("A file with 1132 predictions (only 22 real actives) have been loaded as example")
        scores = pd.read_csv("example_file_metrics.txt",sep="\t",header="infer")
        scores1, auc_ok, sd_auc, pr_ok, sd_pr , y_modelo, y_real = metrics_calculation(scores)
        st.write("")
        st.markdown(":point_down: **Here you can see the ROC and AUPR curves**", unsafe_allow_html=True)
        roc_aupr_curves(auc_ok, sd_auc, pr_ok, sd_pr , y_modelo, y_real)

        st.markdown(":point_down: **Here you can see the PPV surface**", unsafe_allow_html=True)
        tabla_final_1 = ppv_plot(scores1)

        st.markdown(filedownload(tabla_final_1), unsafe_allow_html=True)

max_stars_repo_name: cbellera/Metrics

max_issues_repo_name: cbellera/Metrics

max_forks_repo_name: cbellera/Metrics


text:
# Copyright 2020 IBM Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
from .matlab_utils import find
from .matrix_constructor import matrix_constructor
from .lp_tools.pivot import pivot_mn

#'#@profile
def extract_rates_from_basis(dct, problem_dims):
    klist1 = find(dct.prim_name > 0)
    jlist2 = find(dct.dual_name < 0)
    kn1 =  dct.prim_name[klist1]
    jn2 = -dct.dual_name[jlist2]
    if problem_dims.KK < problem_dims.totalK:
        kn2 =  dct.dual_name[dct.dual_name > 0]
        kord = np.argsort(np.argsort(np.hstack((kn1, kn2))))[:len(kn1)]
        dx = (dct.simplex_dict[klist1+1,0], kord)
    else:
        dx = (dct.simplex_dict[klist1+1,0], kn1-1)
    if problem_dims.JJ < problem_dims.totalJ:
        jn1 = -dct.prim_name[dct.prim_name < 0]
        jord = np.argsort(np.argsort(np.hstack((jn1, jn2))))[len(jn1):]
        dq = (dct.simplex_dict[0,jlist2+1], jord)
    else:
        dq = (dct.simplex_dict[0,jlist2+1], jn2-1)
    return dx, dq

def extract_rates_from_partial(prim_vars, dual_vars, prim_name, dual_name, problem_dims):
    klist1 = find(prim_name > 0)
    jlist2 = find(dual_name < 0)
    kn1 =  prim_name[klist1]
    jn2 = -dual_name[jlist2]
    if problem_dims.KK < problem_dims.totalK:
        kn2 =  dual_name[dual_name > 0]
        kord = np.argsort(np.argsort(np.hstack((kn1, kn2))))[:len(kn1)]
        dx = (prim_vars[klist1+1], kord)
    else:
        dx = (prim_vars[klist1+1], kn1-1)
    if problem_dims.JJ < problem_dims.totalJ:
        jn1 = -prim_name[prim_name < 0]
        jord = np.argsort(np.argsort(np.hstack((jn1, jn2))))[len(jn1):]
        dq = (dual_vars[jlist2+1], jord)
    else:
        dq = (dual_vars[jlist2+1], jn2-1)
    return dx, dq

def extract_rates_from_subproblem(pivots, AAN1, AAN2, problem_dims):
    # Warning this based on assumption that first basis in new_base_sequence is equal to the AAN1 and/or last basis is equal to the AAN2
    if len(pivots) > 0:
        if AAN1 is not None:
            AAN1 = AAN1.copy()
            if AAN2 is not None:
                ran = enumerate(pivots[:-1])
                dx = matrix_constructor(None, None, problem_dims.KK, -1, len(pivots))
                dq = matrix_constructor(None, None, problem_dims.JJ, -1, len(pivots))
            else:
                ran = enumerate(pivots)
                dx = matrix_constructor(None, None, problem_dims.KK, -1, len(pivots)+1)
                dq = matrix_constructor(None, None, problem_dims.JJ, -1, len(pivots)+1)
            for i, piv1 in ran:
                AAN1 = pivot_mn(AAN1, piv1[0], piv1[1])
                ndx, ndq = extract_rates_from_basis(AAN1, problem_dims)
                dx.append(ndx)
                dq.append(ndq)
        else:
            dx = matrix_constructor(None, None, problem_dims.KK, 1, len(pivots) + 1)
            dq = matrix_constructor(None, None, problem_dims.JJ, 1, len(pivots) + 1)
            AAN2 = AAN2.copy()
            for i, piv1 in enumerate(reversed(pivots)):
                AAN2 = pivot_mn(AAN2, piv1[1], piv1[0])
                ndx, ndq = extract_rates_from_basis(AAN2, problem_dims)
                dx.prepend(ndx)
                dq.prepend(ndq)
    else:
        dx = matrix_constructor(None, None, problem_dims.KK, -1, 1)
        dq = matrix_constructor(None, None, problem_dims.JJ, -1, 1)
    return dx.get_matrix(), dq.get_matrix()

max_stars_repo_name: haroldship/SCLPsolver

max_issues_repo_name: haroldship/SCLPsolver

max_forks_repo_name: haroldship/SCLPsolver


text:
from pyrfuniverse.envs import RFUniverseGymGoalWrapper
from pyrfuniverse.utils import RFUniverseController
import numpy as np
from gym import spaces
from gym.utils import seeding
import copy

class FrankaClothFoldEnv(RFUniverseGymGoalWrapper):
    metadata = {'render.modes': ['human']}
    height_offset = 0.01

    def __init__(
            self,
            asset_bundle_file,
            executable_file=None,
            reward_type='sparse',
            tolerance=0.05,
            object_xz_range=0.055,
            object_rotation_range=180,
    ):
        super().__init__(
            executable_file=executable_file,
            camera_channel_id=None,
            articulation_channel_id="581ec41b-1241-11ec-a092-18c04d443e7d",
            obi_cloth_with_grasping_channel_id="581ec41a-1241-11ec-a092-18c04d443e7d",
            game_object_channel_id=None,
        )
        self.reward_type = reward_type
        self.tolerance = tolerance
        self.object_xz_range = object_xz_range
        self.object_rotation_range = object_rotation_range
        self.object_range_low = np.array([-object_xz_range, self.height_offset, -object_xz_range])
        self.object_range_high = np.array([object_xz_range, self.height_offset, object_xz_range])
        self.asset_bundle_file = asset_bundle_file
        self.ik_controller = RFUniverseController('franka', base_pos=np.array([-0.6, 0, 0]))
        self.t = 0

        self.has_loaded_object = False
        self.goal = None
        self.seed()
        self._load_object()

        self.action_space = spaces.Box(
            low=-1, high=1, shape=(4,), dtype=np.float32
        )
        obs = self._get_obs()
        self.observation_space = spaces.Dict({
            'observation': spaces.Box(-np.inf, np.inf, shape=obs['observation'].shape, dtype=np.float32),
            'desired_goal': spaces.Box(-np.inf, np.inf, shape=obs['desired_goal'].shape, dtype=np.float32),
            'achieved_goal': spaces.Box(-np.inf, np.inf, shape=obs['achieved_goal'].shape, dtype=np.float32)
        })

    def step(self, action: np.ndarray):
        pos_ctrl = action[:3] * 0.05
        curr_pos = np.array(self.articulation_channel.data[1]['positions'][3])
        pos_ctrl = curr_pos + pos_ctrl
        joint_positions = self.ik_controller.calculate_ik_recursive(pos_ctrl)

        curr_gripper_width = self._get_gripper_width()
        target_gripper_width = curr_gripper_width + action[3] * 0.2
        target_gripper_width = np.clip(target_gripper_width, 0, 0.08)
        joint_positions.append(target_gripper_width)

        self._set_franka_joints(np.array(joint_positions))
        self.t += 1

        obs = self._get_obs()
        done = False
        info = {
            'is_success': self._check_success(obs)
        }
        reward = self.compute_reward(obs['achieved_goal'], obs['desired_goal'], info)

        return obs, reward, done, info

    def reset(self):
        super().reset()
        self._destroy_object()
        self.env.reset()
        self.ik_controller.reset()
        self.t = 0
        self._load_object()

        grasp_position = self._get_grasp_position()
        joint_positions = self.ik_controller.calculate_ik_recursive(grasp_position)
        self.articulation_channel.set_action(
            'SetJointPositionDirectly',
            index=0,
            joint_positions=list(joint_positions),
        )
        self._step()

        return self._get_obs()

    def seed(self, seed=1234):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def render(self, mode='human'):
        self._step()

    def compute_reward(self, achieved_goal, desired_goal, info):
        distance = self._compute_goal_distance(achieved_goal, desired_goal)
        if self.reward_type == 'sparse':
            return -(distance > self.tolerance).astype(np.float32)
        else:
            return -distance

    def _get_obs(self):
        gripper_position = np.array(self.articulation_channel.data[1]['positions'][3])
        gripper_velocity = np.array(self.articulation_channel.data[1]['velocities'][3])
        gripper_width = self._get_gripper_width()

        panda_obs = np.concatenate((gripper_position, gripper_velocity, [gripper_width]))

        grasp_position = np.array(self.obi_cloth_with_grasping_channel.data[0]['grasp_position'])
        grasp_rotation = np.array(self.obi_cloth_with_grasping_channel.data[0]['grasp_rotation'])
        grasp_velocity = np.array(self.obi_cloth_with_grasping_channel.data[0]['grasp_velocity'])
        grasp_angular_vel = np.array(self.obi_cloth_with_grasping_channel.data[0]['grasp_angular_vel'])
        achieved_goal = grasp_position.copy()

        object_obs = np.concatenate((grasp_position, grasp_rotation, grasp_velocity, grasp_angular_vel))

        obs = np.concatenate((panda_obs, object_obs))

        return {
            'observation': obs.copy(),
            'achieved_goal': achieved_goal.copy(),
            'desired_goal': self.goal.copy()
        }

    def _generate_random_float(self, min: float, max: float) -> float:
        assert min < max, \
            'Min value is {}, while max value is {}.'.format(min, max)
        random_float = np.random.rand()
        random_float = random_float * (max - min) + min

        return random_float

    def _set_franka_joints(self, a: np.ndarray):
        self.articulation_channel.set_action(
            'SetJointPosition',
            index=0,
            joint_positions=list(a[0:7]),
        )
        self._step()

        a[7] = -1 * a[7] / 2
        self.articulation_channel.set_action(
            'SetJointPosition',
            index=1,
            joint_positions=[a[7], a[7]],
        )
        self._step()

    def _get_gripper_width(self):
        gripper_joint_positions = copy.deepcopy(self.articulation_channel.data[1]['joint_positions'])
        return -1 * (gripper_joint_positions[0] + gripper_joint_positions[1])

    def _check_success(self, obs):
        achieved_goal = obs['achieved_goal']
        desired_goal = obs['desired_goal']
        distance = self._compute_goal_distance(achieved_goal, desired_goal)

        return (distance < self.tolerance).astype(np.float32)

    def _compute_goal_distance(self, goal_a, goal_b):
        assert goal_a.shape == goal_b.shape
        return np.linalg.norm(goal_a - goal_b, axis=-1)

    def _load_object(self):
        assert self.asset_bundle_file is not None, \
            'There must be an asset bundle file to load.'

        object_name = 'Obi Solver For Grasp'
        position = self.np_random.uniform(self.object_range_low, self.object_range_high)
        rotation = [0, self._generate_random_float(-self.object_rotation_range, self.object_rotation_range), 0]

        self.obi_cloth_with_grasping_channel.set_action(
            'LoadObiClothWithGrasping',
            filename=self.asset_bundle_file,
            name=object_name,
            position=position,
            rotation=rotation
        )
        self._step()
        self.has_loaded_object = True
        self.goal = np.array(self.obi_cloth_with_grasping_channel.data[0]['target_position'])

    def _destroy_object(self):
        self.obi_cloth_with_grasping_channel.set_action(
            'Destroy',
            index=0
        )
        self._step()
        self.has_loaded_object = False
        self.goal = None

    def _get_grasp_position(self):
        assert self.has_loaded_object, \
            'No object loaded.'
        grasp_position = np.array(self.obi_cloth_with_grasping_channel.data[0]['grasp_position'])

        return grasp_position

max_stars_repo_name: happyCoderJDFJJ/pyrfuniverse

max_issues_repo_name: happyCoderJDFJJ/pyrfuniverse

max_forks_repo_name: happyCoderJDFJJ/pyrfuniverse